Introduce cache modifiers

[mawad-amd]

Motivation

This PR introduces support for cache modifiers and volatile flags in Iris' distributed memory operations (load, store, get, put).
The goal is to provide fine-grained control over caching behavior at different cache levels. This allows developers to have fine grained control over caching behavoir.

Technical Details

Changes include:

• Extended load with cache_modifier and volatile options.
• Extended store with cache_modifier.
• Extended get and put with both load_cache_modifier and store_cache_modifier.
• Docstrings updated to document supported values and behavior.
• Fixed error propagation in tests/run_tests_distributed.py to correctly forward worker exit codes.
• Added new unit tests covering all cache modifier combinations and volatile option.

Test Plan

I added comprehensive tests under tests/unittests:

• test_load_cache_modifiers.py – verifies iris.load with all cache modifiers and volatile settings.
• test_store_cache_modifiers.py – verifies iris.store with all store cache modifiers across ranks.
• test_get_cache_modifiers.py – verifies iris.get with load/store cache modifier combinations.
• test_put_cache_modifiers.py – verifies iris.put with load/store cache modifier combinations.

Each test checks correctness by comparing against expected tensors across distributed ranks.

Test Result

See CI

Submission Checklist

• Looked over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull Request Overview

This PR adds cache modifier support to the Iris framework's load and store operations. The changes enable fine-grained control over GPU cache behavior during memory operations, allowing developers to optimize performance by specifying how data should be cached across the CU (L1), L2, and LLC levels.

Key changes:

• Enhanced load() and store() functions with cache_modifier and volatile parameters
• Added comprehensive test coverage for different cache behaviors (default, writeback, cache global, cache streaming, write-through, and volatile)
• Updated function documentation to explain cache modifier options and their effects on GPU cache hierarchy

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
iris/iris.py	Enhanced load() and store() functions with cache_modifier and volatile parameters, plus detailed documentation
tests/unittests/test_load_cache_modifiers.py	New comprehensive test suite covering all cache modifier behaviors and volatile flag

[Copilot]

Pull Request Overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

[Copilot]

Pull Request Overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

[mawad-amd]

@copilot fix ci errors

[Copilot]

@mawad-amd I've opened a new pull request, #310, to work on those changes. Once the pull request is ready, I'll request review from you.

[neoblizz]

My high-level comment is that the tests/unittests should simply just check if the correct amdgcn is being generated for the load/stores.

[neoblizz]

@copilot fix the following tests;

 =========================== short test summary info ============================
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[8-dtype0]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[8-dtype1]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[8-dtype2]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[16-dtype0]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[16-dtype1]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[16-dtype2]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[32-dtype0]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[32-dtype1]
  FAILED tests/unittests/test_get_other_triton.py::test_get_other_api[32-dtype2]

[Copilot]

@neoblizz I've opened a new pull request, #357, to work on those changes. Once the pull request is ready, I'll request review from you.

[mawad-amd]

I don't remember if I fixed this or not, but cache modifiers do not work on remote accesses. Tests should not test that. Doc strings need to be updated. I don't think we can assert for that.

[neoblizz]

I don't remember if I fixed this or not, but cache modifiers do not work on remote accesses. Tests should not test that. Doc strings need to be updated. I don't think we can assert for that.

It still is an issue because if you do curr_rank -> curr_rank using iris.load, we should be able to use cache-modifiers there?

[mawad-amd]

I don't remember if I fixed this or not, but cache modifiers do not work on remote accesses. Tests should not test that. Doc strings need to be updated. I don't think we can assert for that.

It still is an issue because if you do curr_rank -> curr_rank using iris.load, we should be able to use cache-modifiers there?

Yes, correct. Local ops should work fine. Remote ops are not allowed to have cache modifiers. Something like this won't work and I was initially testing for it:
iris.load(to_rank=local_rank, from_rank=remote_rank, modifer="anything")