iris.x: Device-side communication + .x APIs.

[neoblizz]

Introduces iris.x APIs.

Submission Checklist

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

[Copilot]

Pull request overview

This PR introduces iris.x, a new module providing device-side tile-level primitives for fine-grained collective operations. Unlike iris.ccl which handles full tensors with internal tiling, iris.x provides composable functions that users can call from their own kernels to manage tile iteration themselves.

Key Changes:

• New iris.x module with tile-level communication primitives (all-reduce, all-gather, all-to-all, reduce-scatter)
• Fused GEMM+Communication operations requiring tritonBLAS (gemm_all_reduce, gemm_all_gather, etc.)
• Comprehensive test suite for new primitives in tests/x/
• CI/CD modernization with unified workflow replacing 3 separate workflows
• Documentation updates and benchmark enhancements

Reviewed changes

Copilot reviewed 33 out of 33 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
iris/x/__init__.py	Module initialization exposing all tile-level primitives with optional GEMM operations
iris/x/all_reduce.py	Five all-reduce variants (atomic, one-shot, two-shot, spinlock, ring) for different use cases
iris/x/all_gather.py	Tile-level all-gather primitive for gathering data from all ranks
iris/x/all_to_all.py	Tile-level all-to-all primitive for bidirectional data exchange
iris/x/reduce_scatter.py	Tile-level reduce-scatter that reduces and scatters to assigned ranks
iris/x/gemm_all_reduce.py	Fused GEMM + all-reduce using tritonBLAS stages
iris/x/gemm_all_gather.py	Fused GEMM + all-gather combining computation and communication
iris/x/gemm_reduce_scatter.py	Fused GEMM + reduce-scatter for column-parallel workloads
iris/x/all_gather_gemm.py	Fused all-gather + GEMM for tensor-parallel workloads
iris/x/common.py	Shared utilities for tile indexing and offset computation
tests/x/test_*.py	Comprehensive test suite validating all primitives against PyTorch references
.github/workflows/iris-tests.yml	New unified test workflow supporting multiple test directories and install methods
.github/scripts/run_tests.sh	Updated test runner with tritonBLAS installation for iris.x tests
tests/ccl/test_all_reduce.py	Modified to add explicit preamble calls for better test isolation
pyproject.toml	Added optional gemm dependency group for tritonBLAS
docs/reference/examples.md	Updated documentation with new example references
benchmark/ccl/all_to_all/benchmark.py	Added RCCL comparison benchmarking option

[Copilot]

Inconsistent accumulator dtype logic compared to other GEMM kernels.

In gemm_all_reduce.py and gemm_reduce_scatter.py, the logic is:

acc_dtype = tl.int32 if C.type.element_ty != tl.int8 else tl.float32

But here it's:

acc_dtype = tl.int32 if C.type.element_ty == tl.int8 else tl.float32

This appears to be inverted. For int8 output types, you'd want int32 accumulator (correct in other files), but for floating point types, you'd want float32 accumulator. The condition should be != tl.int8 for consistency.

Suggested change

	acc_dtype = tl.int32 if C.type.element_ty == tl.int8 else tl.float32
	acc_dtype = tl.int32 if C.type.element_ty != tl.int8 else tl.float32

[Copilot]

Same accumulator dtype logic inconsistency as in gemm_all_gather.py. The condition should be != tl.int8 instead of == tl.int8 to match the logic in gemm_all_reduce.py and gemm_reduce_scatter.py.

Suggested change

	acc_dtype = tl.int32 if C.type.element_ty != tl.int8 else tl.float32
	acc_dtype = tl.int32 if C.type.element_ty == tl.int8 else tl.float32

[Copilot]

Inconsistent bias vector broadcasting. In gemm_all_reduce.py line 172, the code uses:

acc = add_vector(acc, bias_vector, QUANTIZED=False)

But here it uses:

acc = add_vector(acc, bias_vector[:, None], QUANTIZED=False)

The [:, None] adds an extra dimension. For consistency across GEMM kernels, the bias handling should be uniform. Based on typical GEMM bias usage where bias is (M,) and acc is (M, N), the [:, None] expansion may be needed, but all GEMM kernels should handle it the same way.

[Copilot]

Variable K_local is not used.

Suggested change

K_local = K // world_size

[Copilot]

Variable total_tiles is not used.

Suggested change

total_tiles = num_pid_m * num_pid_n

[Copilot]

Variable tiles_per_rank is not used.

[Copilot]

Import of 'torch' is not used.

Suggested change

import torch

[Copilot]

Import of 'torch' is not used.

Suggested change

import torch

[Copilot]

Import of 'torch' is not used.

Suggested change

import torch

[Copilot]

Import of 'torch' is not used.

Suggested change

import torch