Device Functions

Device Functions#

Device-side functions provided by Iris for remote memory operations and atomics.

Iris: Multi-GPU Communication and Memory Management Framework

Iris is a high-performance framework that enables seamless multi-GPU programming in Triton, enabling fine-grained communication and compute overlap natively in Triton across multiple GPUs with SHMEM-like Remote Memory Access (RMA) capabilities.

Key Features: - Symmetric heap management across multiple GPUs - High-performance atomic operations (add, cas, xchg, xor, and, or, min, max) - Efficient load/store operations with rank-to-rank communication - Memory allocation and deallocation utilities - Built-in logging with rank information - PyTorch distributed integration for distributed computing

Example

>>> import iris
>>> ctx = iris.iris(heap_size=2**30)  # 1GB heap
>>> tensor = ctx.zeros(1024, 1024, dtype=torch.float32)

class Iris(heap_size=1073741824)[source]

Main Iris class for multi-GPU communication and memory management.

This class provides a unified interface for distributed GPU operations including memory allocation, atomic operations, and inter-rank communication.

Parameters:: heap_size (int) – Size of the symmetric heap in bytes. Default: 1GB (2^30)

Example

>>> ctx = iris.iris(heap_size=2**31)  # 2GB heap
>>> print(f"Rank {ctx.cur_rank} of {ctx.num_ranks}") # Rank 0 of 1
>>> tensor = ctx.zeros(1000, 1000, dtype=torch.float32)

__init__(heap_size=1073741824)[source]

debug(message)[source]

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

info(message)[source]

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

warning(message)[source]

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

error(message)[source]

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

broadcast(value, source_rank)[source]

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

Parameters:

value (Any) – The value to broadcast. Can be a scalar, tensor, numpy array, or any picklable object. Only the source_rank value is used; other ranks should pass a placeholder (e.g., None).
source_rank (int) – Rank id 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, source_rank=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, source_rank=0)  # All ranks get the same array

zeros_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format)[source]

Returns a tensor filled with the scalar value 0, with the same size as input, allocated on the Iris symmetric heap.

Parameters:

input (Tensor) – the size of input will determine size of the output tensor.

Keyword Arguments:

dtype (torch.dtype, optional) – the desired data type of returned Tensor. Default: if None, defaults to the dtype of input.
layout (torch.layout, optional) – the desired layout of returned tensor. Default: if None, defaults to the layout of input. Note: Iris tensors are always contiguous (strided).
device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input. Must be compatible with this Iris instance.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.
memory_format (torch.memory_format, optional) – the desired memory format of returned Tensor. Default: torch.preserve_format.

Example

>>> ctx = iris.iris(1 << 20)
>>> input_tensor = ctx.ones(2, 3)
>>> zeros_tensor = ctx.zeros_like(input_tensor)
>>> print(zeros_tensor.shape)  # torch.Size([2, 3])

arange(start=0, end=None, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Returns a 1-D tensor of size ⌈(end - start) / step⌉ with values from the interval [start, end) taken with common difference step beginning from start. The tensor is allocated on the symmetric heap.

Note: When using floating-point dtypes (especially reduced precision types like bfloat16), the results may be affected by floating-point rounding behavior. Some values in the sequence might not be exactly representable in certain floating-point formats, which can lead to repeated values or unexpected rounding. For precise sequences, it is recommended to use integer dtypes instead of floating-point dtypes.

Note that non-integer step is subject to floating point rounding errors when comparing against end; to avoid inconsistency, we advise subtracting a small epsilon from end in such cases.

Parameters:

start (Number, optional) – the starting value for the set of points. Default: 0.
end (Number) – the ending value for the set of points
step (Number, optional) – the gap between each pair of adjacent points. Default: 1.
out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.get_default_dtype()). If dtype is not given, infer the data type from the other input arguments. If any of start, end, or step are floating-point, the dtype is inferred be the default dtype, see get_default_dtype(). Otherwise, the dtype is inferred to be torch.int64.
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.arange(0, 10, 2)  # [0, 2, 4, 6, 8]
>>> print(tensor.shape)  # torch.Size([5])

zeros(*size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Returns a tensor filled with the scalar value 0, with the shape defined by the variable argument size. The tensor is allocated on the Iris symmetric heap.

Parameters:

*size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.

Keyword Arguments:

out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.zeros(2, 3)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([0., 0., 0.], device='cuda:0')

randn(*size, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False)[source]

Returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1 (also called the standard normal distribution). The tensor is allocated on the Iris symmetric heap.

\[\text{out}_i \sim \mathcal{N}(0, 1)\]

For complex dtypes, the tensor is i.i.d. sampled from a complex normal distribution with zero mean and unit variance as

\[\text{out}_i \sim \mathcal{CN}(0, 1)\]

This is equivalent to separately sampling the real \((\text{Re})\) and imaginary \((\text{Im})\) part of \(\text{out}_i\) as

\[\text{Re}(\text{out}_i) \sim \mathcal{N}(0, \frac{1}{2}), \quad \text{Im}(\text{out}_i) \sim \mathcal{N}(0, \frac{1}{2})\]

The shape of the tensor is defined by the variable argument size.

Parameters:

*size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.

Keyword Arguments:

generator (torch.Generator, optional) – a pseudorandom number generator for sampling
out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type (see torch.set_default_device()). device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.
pin_memory (bool, optional) – If set, returned tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.randn(2, 3)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([ 0.3982, -0.0059, -0.4365], device='cuda:0')

ones(*size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument size. The tensor is allocated on the Iris symmetric heap.

Parameters:

*size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.

Keyword Arguments:

out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.ones(2, 3)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([1., 1., 1.], device='cuda:0')

full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Creates a tensor of size size filled with fill_value. The tensor’s dtype is inferred from fill_value. The tensor is allocated on the Iris symmetric heap.

Parameters:

size (int...) – a list, tuple, or torch.Size of integers defining the shape of the output tensor.
fill_value (Scalar) – the value to fill the output tensor with.

Keyword Arguments:

out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.full((2, 3), 3.14)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([3.1400, 3.1400, 3.1400], device='cuda:0')

uniform(size, low=0.0, high=1.0, dtype=torch.float)[source]

Returns a tensor filled with random numbers from a uniform distribution, allocated on the Iris symmetric heap.

Parameters:

size (int or tuple of ints) – the size of the output tensor.
low (float, optional) – the lower bound of the uniform distribution. Default: 0.0.
high (float, optional) – the upper bound of the uniform distribution. Default: 1.0.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: torch.float.

Returns:

A tensor filled with random numbers from a uniform distribution.

Return type:

Tensor

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.uniform((2, 3), low=0.0, high=1.0)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([0.1234, 0.5678, 0.9012], device='cuda:0')

empty(*size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, memory_format=torch.contiguous_format)[source]

Returns a tensor filled with uninitialized data. The shape of the tensor is defined by the variable argument size. The tensor is allocated on the Iris symmetric heap.

Note

If torch.use_deterministic_algorithms() and torch.utils.deterministic.fill_uninitialized_memory are both set to True, the output tensor is initialized to prevent any possible nondeterministic behavior from using the data as an input to an operation. Floating point and complex tensors are filled with NaN, and integer tensors are filled with the maximum value.

Parameters:

*size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.

Keyword Arguments:

out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.
pin_memory (bool, optional) – If set, returned tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False. Note: Iris tensors are always on GPU.
memory_format (torch.memory_format, optional) – the desired memory format of returned Tensor. Default: torch.contiguous_format.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.empty(2, 3)
>>> print(tensor.shape)  # torch.Size([2, 3])

randint(*args, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive). The shape of the tensor is defined by the variable argument size. The tensor is allocated on the Iris symmetric heap.

Note

With the global dtype default (torch.float32), this function returns a tensor with dtype torch.int64.

Parameters:

low (int, optional) – Lowest integer to be drawn from the distribution. Default: 0.
high (int) – One above the highest integer to be drawn from the distribution.
size (tuple) – a tuple defining the shape of the output tensor.

Keyword Arguments:

generator (torch.Generator, optional) – a pseudorandom number generator for sampling.
out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – if None, this function returns a tensor with dtype torch.int64.
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.randint(0, 10, (2, 3))  # Random integers [0, 10)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([7, 2, 9], device='cuda:0')

linspace(start, end, steps, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)[source]

Creates a one-dimensional tensor of size steps whose values are evenly spaced from start to end, inclusive. The tensor is allocated on the Iris symmetric heap.

The values are: (start, start + (end-start)/(steps-1), …, start + (steps-2)*(end-start)/(steps-1), end)

Parameters:

start (float or Tensor) – the starting value for the set of points. If Tensor, it must be 0-dimensional.
end (float or Tensor) – the ending value for the set of points. If Tensor, it must be 0-dimensional.
steps (int) – size of the constructed tensor.

Keyword Arguments:

out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the data type to perform the computation in. Default: if None, uses the global default dtype when both start and end are real, and corresponding complex dtype when either is complex.
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.linspace(0, 10, 5)  # [0, 2.5, 5, 7.5, 10]
>>> print(tensor) # tensor([ 0.0000,  2.5000,  5.0000,  7.5000, 10.0000], device='cuda:0')

rand(*size, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False)[source]

Returns a tensor filled with random numbers from a uniform distribution on the interval [0, 1). The tensor is allocated on the Iris symmetric heap.

Parameters:

*size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.

Keyword Arguments:

generator (torch.Generator, optional) – a pseudorandom number generator for sampling.
out (Tensor, optional) – the output tensor.
dtype (torch.dtype, optional) – the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
layout (torch.layout, optional) – the desired layout of returned Tensor. Default: torch.strided. Note: Iris tensors always use torch.strided regardless of this parameter.
device (torch.device, optional) – the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default: False.
pin_memory (bool, optional) – If set, returned tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False. Note: Iris tensors are always on GPU.

Example

>>> ctx = iris.iris(1 << 20)
>>> tensor = ctx.rand(2, 3)  # Random values in [0, 1)
>>> print(tensor.shape)  # torch.Size([2, 3])
>>> print(tensor[0])  # tensor([0.1234, 0.5678, 0.9012], device='cuda:0')

get_heap_bases()[source]

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])

barrier(stream=None)[source]

Synchronize all ranks 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 global distributed barrier so that all ranks reach the same point before proceeding. :param stream: If stream is given: wait only for that stream before barrier. If stream is None: legacy behavior (device-wide sync).

Example

>>> ctx = iris.iris(1 << 20)
>>> ctx.barrier()  # Synchronize all ranks

get_device()[source]

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

get_cu_count()[source]

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

get_rank()[source]

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

get_num_ranks()[source]

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

iris(heap_size=1073741824)[source]

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

Parameters:: heap_size (int) – Size of the heap in bytes. Defaults to 1GB.
Returns:: An initialized Iris instance.
Return type:: Iris

Example

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

Memory transfer operations#

load#

load(*args, **kwargs)

MagicMock is a subclass of Mock with default implementations of most of the magic methods. You can use MagicMock without having to configure the magic methods yourself.