Iris Class

Factory Function

Prefer using the convenience factory over calling the constructor directly:

iris(heap_size=1073741824, allocator_type='torch', engines=())

Create and return an Iris instance with the specified heap size.

Parameters:

• heap_size (int) – Size of the heap in bytes. Defaults to 1GB.

• allocator_type (str) – Type of allocator to use. Options: “torch” (default), “vmem”. Can be overridden with IRIS_ALLOCATOR environment variable.

• engines (str or tuple of str, optional) – See Iris.

Returns:

An initialized Iris instance.

Return type:

Iris

Example

>>> import iris
>>> iris_ctx = iris.iris(2**30)  # 1GB heap with default (torch) allocator
>>> tensor = iris_ctx.zeros(1024, 1024)

>>> # Use VMem allocator
>>> iris_ctx = iris.iris(2**30, allocator_type="vmem")
>>> tensor = iris_ctx.zeros(1024, 1024)

Optional SDMA Copy Engine

The engines parameter defaults to (). In this mode, Iris skips SDMA initialization and queue creation, and get_copy_engine_ctx() returns None. Regular Triton and Gluon memory operations and collectives work without SDMA.

Pass engines=("sdma",) to initialize both host- and device-initiated SDMA queues. This needs the sdma extra (Installation):

ctx = iris.iris(heap_size=1 << 30, engines=("sdma",))
copy_engine_ctx = ctx.get_copy_engine_ctx()

This is required for the host methods put, put_tile, put_tiles, and quiet, and for device operations using use_copy_engine=True. Host SDMA methods raise RuntimeError if the copy engine is disabled. Pass the returned context to device operations through their copy_engine_ctx argument.

Iris.get_copy_engine_ctx()

Get the device-side copy engine context for SDMA operations.

The returned tensor holds the per-rank SDMA queue, read and write pointer, and doorbell addresses that device-side kernels need in order to enqueue copy packets. Pass it to iris.put as copy_engine_ctx when driving the copy engine from inside a kernel.

Returns:

The copy engine device context, or None when Iris was not initialized with engines=("sdma",).

Return type:

torch.Tensor

Core Methods

Iris.get_heap_bases()

Return the tensor of symmetric heap base addresses for all ranks.

Returns:

A 1D tensor of uint64 heap base addresses of size num_ranks on the Iris device. Pass this to device-side Triton kernels that require heap translation.

Return type:

torch.Tensor

Example

>>> ctx = iris.iris(1 << 20)
>>> heap_bases = ctx.get_heap_bases()
>>> print(heap_bases.shape)  # torch.Size([num_ranks])

Iris.get_device_context()

Get the device context tensor for DeviceContext initialization.

Returns a tensor encoding: [cur_rank, world_size, heap_base_0, heap_base_1, …] If tracing is enabled, also includes: [trace_enabled, max_events, trace_counter_ptr, trace_buffer_ptrs…]

This opaque format allows future extension without breaking the API.

Returns:

Encoded context data as int64 tensor on device

Return type:

torch.Tensor

Example

>>> import iris
>>> from iris import DeviceContext
>>> import triton
>>> import triton.language as tl
>>>
>>> ctx = iris.iris()
>>> context_tensor = ctx.get_device_context()
>>>
>>> @triton.jit
>>> def my_kernel(context_tensor, rank: tl.constexpr, world_size: tl.constexpr, ...):
>>>     ctx = DeviceContext.initialize(context_tensor, rank, world_size)
>>>     data = ctx.load(buffer, from_rank=1)

Iris.barrier(stream=None, group=None)

Synchronize ranks within the specified group and their CUDA devices.

This first calls torch.cuda.synchronize() or stream.synchronize() to ensure the local GPU has finished all queued work, then performs a distributed barrier so that all ranks in the group reach the same point before proceeding.

Parameters:

• stream – If stream is given: wait only for that stream before barrier. If stream is None: legacy behavior (device-wide sync).

• group (ProcessGroup, optional) – The process group to synchronize. If None, uses the default process group (all ranks).

Example

>>> ctx = iris.iris(1 << 20)
>>> ctx.barrier()  # Synchronize all ranks
>>> ctx.barrier(group=my_group)  # Synchronize only ranks in my_group

Iris.get_device()

Get the underlying device where the Iris symmetric heap resides.

Returns:

The CUDA device of Iris-managed memory.

Return type:

torch.device

Example

>>> ctx = iris.iris(1 << 20)
>>> device = ctx.get_device()
>>> print(device)  # cuda:0

Iris.get_cu_count()

Get the number of compute units (CUs) for the current GPU.

Returns:

Number of compute units on this rank’s GPU.

Return type:

int

Example

>>> ctx = iris.iris(1 << 20)
>>> cu_count = ctx.get_cu_count()
>>> print(f"GPU has {cu_count} CUs")  # GPU has 304 CUs

Iris.get_rank()

Get this process’s rank id in the distributed communicator.

Returns:

Zero-based rank id of the current process.

Return type:

int

Example

>>> ctx = iris.iris(1 << 20)
>>> rank = ctx.get_rank()
>>> print(f"This is rank {rank}")  # This is rank 0

Iris.get_num_ranks()

Get the total number of ranks in the distributed communicator.

Returns:

World size (number of ranks).

Return type:

int

Example

>>> ctx = iris.iris(1 << 20)
>>> num_ranks = ctx.get_num_ranks()
>>> print(f"Total ranks: {num_ranks}")  # Total ranks: 1

Logging Helpers

Use Iris-aware logging that automatically annotates each message with the current rank and world size. This is helpful when debugging multi-rank programs.

set_logger_level(level)

Set the logging level for the iris logger.

Parameters:

level – Logging level (iris.DEBUG, iris.INFO, iris.WARNING, iris.ERROR)

Example

>>> ctx = iris.iris()
>>> iris.set_logger_level(iris.DEBUG)
>>> ctx.debug("This will now be visible")  # [Iris] [0/1] This will now be visible

Iris.debug(message)

Log a debug message with rank information.

Parameters:

message (str) – Human-readable message to log at debug level.

Notes

The log record is enriched with iris_rank and iris_num_ranks so formatters can display the originating rank and world size.

Example

>>> ctx = iris.iris()
>>> iris.set_logger_level(iris.DEBUG)
>>> ctx.debug("Allocating buffers")  # [Iris] [0/1] Allocating buffers

Iris.info(message)

Log an info message with rank information.

Parameters:

message (str) – Human-readable message to log at info level.

Example

>>> ctx = iris.iris()
>>> ctx.info("Starting iteration 0")  # [Iris] [0/1] Starting iteration 0

Iris.warning(message)

Log a warning message with rank information.

Parameters:

message (str) – Human-readable message to log at warning level.

Example

>>> ctx = iris.iris()
>>> ctx.warning("Memory usage is high")  # [Iris] [0/1] Memory usage is high

Iris.error(message)

Log an error message with rank information.

Parameters:

message (str) – Human-readable message to log at error level.

Example

>>> ctx = iris.iris()
>>> ctx.error("Failed to allocate memory")  # [Iris] [0/1] Failed to allocate memory

Utility Functions

do_bench(fn, barrier_fn=<function <lambda>>, preamble_fn=<function <lambda>>, n_warmup=25, n_repeat=100, quantiles=None, return_mode='mean')

Benchmark a function by timing its execution.

Parameters:

• fn (callable) – Function to benchmark.

• barrier_fn (callable, optional) – Function to call for synchronization. Default: no-op.

• preamble_fn (callable, optional) – Function to call before each execution. Default: no-op.

• n_warmup (int, optional) – Number of warmup iterations. Default: 25.

• n_repeat (int, optional) – Number of timing iterations. Default: 100.

• quantiles (list, optional) – Quantiles to return instead of summary statistic. Default: None.

• return_mode (str, optional) – Summary statistic to return (“mean”, “min”, “max”, “median”, “all”). Default: “mean”.

Returns:

Timing result(s) in milliseconds.

Return type:

float or list

Example

>>> import iris
>>> iris_ctx = iris.iris(1 << 20)
>>> def test_fn():
>>>     tensor = iris_ctx.zeros(1000, 1000)
>>> time_ms = iris.do_bench(test_fn, barrier_fn=iris_ctx.barrier)

Broadcast Helper

Broadcast data from a source rank to all ranks. This method automatically detects whether the value is a tensor/array or a scalar and uses the appropriate broadcast mechanism.

Iris.broadcast(value, src=0)

Broadcast a value from one rank to all ranks.

This method automatically detects the type of value and uses the appropriate broadcast mechanism:

• For tensors and arrays: uses efficient PyTorch distributed tensor collectives

• For scalars and other objects: uses object broadcast

Matches torch.distributed.broadcast parameter naming.

Parameters:

• value (Any) – The value to broadcast. Can be a scalar, tensor, numpy array, or any picklable object. Only the src rank’s value is used; other ranks should pass a placeholder (e.g., None).

• src (int) – Source rank that holds the authoritative value.

Returns:

The value broadcast to all ranks. Tensors and arrays are returned as

numpy arrays; scalars and objects are returned in their original type.

Return type:

Any

Examples

>>> ctx = iris.iris()
>>> # Broadcasting a scalar
>>> value = 42 if ctx.cur_rank == 0 else None
>>> value = ctx.broadcast(value, src=0)  # All ranks get 42
>>>
>>> # Broadcasting a tensor
>>> if ctx.cur_rank == 0:
>>>     data = torch.randn(10, 10)
>>> else:
>>>     data = None
>>> data = ctx.broadcast(data, src=0)  # All ranks get the same array

allocate_symmetric

Iris.allocate_symmetric(*size, dtype=None)

Allocate a symmetric tensor and return it with its peer-base table.

Kernels take the pair as two ordinary arguments – a pointer and a tensor – and inline the address translation, so the same device code works for a tensor from any provider that returns this shape.

Parameters:

• *size (int...) – Shape of the tensor, as a sequence of integers or a single collection.

• dtype (torch.dtype, optional) – Element type. Defaults to the torch default dtype.

Returns:

The tensor, uninitialized, and an int64[world_size] device-resident table whose entry r is the address of this tensor on rank r. peer_bases[cur_rank] is the tensor’s own data_ptr(), which is the base device-side translation subtracts.

Return type:

tuple[torch.Tensor, torch.Tensor]

Note

Collective. All ranks must call this together, as with the other Iris allocation ops.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor, peer_bases = ctx.allocate_symmetric(1024, dtype=torch.float32)

device_barrier

Iris.device_barrier(group=None)

Device-side barrier that is CUDA graph capturable.

Unlike barrier() which uses host-side torch.distributed.barrier(), this uses device-side atomic operations on the symmetric heap to synchronize ranks. Stateless w.r.t. host-side epoch tracking: each rank’s flag on the heap serves as its own epoch counter, managed entirely by the GPU via atomic_add. A persistent per-group flags tensor is cached in _device_barrier_state.

Parameters:

group (ProcessGroup, optional) – The process group to synchronize. If None, uses all ranks in the shmem context.

Example

>>> ctx = iris.iris(1 << 20)
>>> ctx.device_barrier()  # Synchronize all ranks on device

get_device_id

Iris.get_device_id()

Get the device ID used by this Iris instance.

In simulation mode, this may differ from the local rank if multiple ranks share a single GPU. This is the device ID that was set during Iris initialization.

Returns:

The GPU device ID used by this Iris instance.

Return type:

int

Example

>>> ctx = iris.iris(1 << 20)
>>> device_id = ctx.get_device_id()
>>> print(f"Using GPU {device_id}")  # Using GPU 0

ops

property Iris.ops

Access fused GEMM+CCL operations.

This property provides a namespace for high-level fused operations that combine matrix multiplication with collective communication. Operations automatically infer dimensions, strides, and hardware parameters from input tensors.

Available operations:

• matmul_all_reduce: GEMM + All-Reduce

• all_gather_matmul: All-Gather + GEMM

• matmul_all_gather: GEMM + All-Gather

• matmul_reduce_scatter: GEMM + Reduce-Scatter

Returns:

Namespace with fused operation methods

Return type:

OpsNamespace

Raises:

ImportError – If tritonBLAS is not available

Example

>>> ctx = iris.iris()
>>> A = ctx.randn((1024, 512), dtype=torch.float16)
>>> B = ctx.randn((512, 2048), dtype=torch.float16)
>>> output = ctx.zeros((1024, 2048), dtype=torch.float16)
>>> ctx.ops.matmul_all_reduce(output, A, B, ctx)