Device Context

Device Context#

Device-side methods provided by Iris for remote memory operations, atomics and collectives. These methods are part of the DeviceContext aggregate used inside @triton.jit kernels, and are reached through the context tensor passed in from the host:

from iris import DeviceContext

@triton.jit
def producer(context_tensor, target_buffer, ..., rank: tl.constexpr, world_size: tl.constexpr):
    ctx = DeviceContext.initialize(context_tensor, rank, world_size)
    ...
    ctx.store(target_buffer + offsets, values, to_rank=consumer_rank, mask=mask)
    ctx.atomic_cas(flag + pid, 0, 1, to_rank=consumer_rank, sem="release", scope="sys")

Taken from examples/06_message_passing/message_passing_device_context.py.

Note

The collectives below are tile-level – they run inside your kernel, on the tile a program is already holding, and fuse with whatever compute surrounds them. They are a different surface from the host collectives in Collective Communication Operations, which launch their own kernels and are called from Python.

Note

Every operation here also exists as a free function – see Device Functions. The free form is explicit about both ends of the transfer and the heap: load(pointer, to_rank, from_rank, heap_bases). The bound form below knows the current rank and carries the heap bases itself, so the same call is ctx.load(pointer, from_rank).

Initialization#

initialize#

static Context.initialize(context_tensor, rank, world_size, tracing=False)[source]

Initialize Context from the encoded context tensor.

The context tensor has the format:

[cur_rank, num_ranks, heap_base_0, ..., heap_base_N, trace_info...]

If tracing=True, trace buffer pointers are also extracted from it.

Parameters:

  • context_tensor – Pointer to encoded context data (from Iris.get_device_context())

  • rank – Current rank (must be constexpr in kernel signature)

  • world_size – Total number of ranks (must be constexpr in kernel signature)

  • tracing (constexpr) – Enable event tracing (constexpr, default: False)

Returns:

Initialized device context

Return type:

Context

Example

>>> import iris
>>> from iris import DeviceContext
>>>
>>> ctx = iris.iris()
>>> ctx.tracing.enable(max_events=1_000_000)
>>> context_tensor = ctx.get_device_context()
>>>
>>> @triton.jit
>>> def kernel(context_tensor, rank: tl.constexpr, world_size: tl.constexpr, ...):
>>>     # Without tracing
>>>     ctx = DeviceContext.initialize(context_tensor, rank, world_size)
>>>
>>>     # With tracing
>>>     ctx = DeviceContext.initialize(context_tensor, rank, world_size, tracing=True)
>>>     mask = tl.full([64], True, dtype=tl.int1)  # Example mask
>>>     ctx.tracing.record_event_start(event_id=TraceEvent().put, target_rank=1, address=ptr, pid_m=0, pid_n=0, mask=mask)

Memory transfer operations#

load#

Context.load(pointer, from_rank, mask=None, other=None, cache_modifier=None, volatile=False, hint=None)[source]

Loads a value from the specified rank’s memory location.

This method performs a memory read operation by translating the pointer from the current rank’s address space to the from_rank’s address space and loading data from the target memory location. If the current rank and from_rank are the same, this performs a local load operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that will be translated to the from_rank’s address space.

  • from_rank (int) – The rank ID from which to read the data.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not load the data at address pointer[idx]. Defaults to None.

  • other (Block, optional) – Value to return for masked-out elements. If not provided, the result for masked-out elements is undefined. Defaults to None.

  • cache_modifier (str, optional) –

    Controls cache behavior of the load.

    Supported values:

    • None: (default) — Same as “.ca”. Uses cache at all levels (CU, L2, LLC) with LRU policy.

    • ”.ca”: Cache at all levels (CU, L2, LLC) with LRU policy

    • ”.cg”: Bypasses the CU (L1) cache, streams through L2, and may hit in LLC but the line is not retained or inserted.

    • ”.cv”: Bypasses all GPU caches (CU and L2) and fetches directly from system memory. If data exists in the LLC, it may hit, but is not retained or inserted.

      Ensures global coherence by invalidating stale GPU cache lines.

  • volatile (bool, optional) – If True, disables compiler optimizations that could reorder or eliminate the load. Defaults to False.

  • hint (int or tuple, optional) – Vectorization hint for the translated pointer. Defaults to None.

Returns:

The loaded value from the target memory location.

Return type:

Block

Example

>>> data = ctx.load(buffer + offsets, from_rank=1, mask=mask)

store#

Context.store(pointer, value, to_rank, mask=None, cache_modifier=None, hint=None)[source]

Writes data to the specified rank’s memory location.

This method performs a memory write operation by translating the pointer from the current rank’s address space to the to_rank’s address space and storing the provided data to the target memory location. If the current rank and to_rank are the same, this performs a local store operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that will be translated to the to_rank’s address space.

  • value (Block) – The tensor of elements to be stored.

  • to_rank (int) – The rank ID to which the data will be written.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not store the data at address pointer[idx]. Defaults to None.

  • cache_modifier (str, optional) –

    Controls cache behavior of the store. Supported values are:

    • None: (default) — Same as “.wb”. Uses write-back caching at all levels (CU, L2, LLC) with LRU policy.

    • ”.wb”: Write-back. Write-allocate on L1 miss, inserted into caches and written back later.

    • ”.cg”: Cache Global. Equivalent to “.wb” — stored through L1 → L2 → LLC under LRU.

    • ”.cs”: Cache Streaming. Bypasses L1, streamed through L2, not retained in LLC.

    • ”.wt”: Write-Through. Bypasses L1 and L2 (coherent cache bypass), may hit in LLC with LRU.

  • hint (constexpr)

Returns:

None

Example

>>> ctx.store(buffer + offsets, values, to_rank=1, mask=mask)

get#

Context.get(from_ptr, to_ptr, from_rank, mask=None, other=None, load_cache_modifier=None, store_cache_modifier=None, hint=None)[source]

Copies data from the specified rank’s memory into current rank’s local memory.

This method performs a remote load operation by translating from_ptr from the current rank’s address space to the from_rank’s address space, loading the data, and storing it to to_ptr in the current rank’s local memory. If the current rank and from_rank are the same, this performs a local copy operation.

Parameters:

  • from_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that references memory in from_rank.

  • to_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer to local memory in current rank where the data will be written.

  • from_rank (int) – The rank ID from which to read the data.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not load from from_ptr[idx] and do not store to to_ptr[idx]. Defaults to None.

  • other (Block, optional) – Value to return for masked-out elements during the load operation. If not provided, the result for masked-out elements is undefined. Defaults to None.

  • load_cache_modifier (str, optional) – Controls cache behavior of the load. Supported values are: - None: (default) — Same as “.ca”. Uses cache at all levels (CU, L2, LLC) with LRU policy. - “.ca”: Cache at all levels (CU, L2, LLC) with LRU policy. - “.cg”: Bypasses the CU (L1) cache, streams through L2, and may hit in LLC but the line is not retained or inserted. - “.cv”: Bypasses all GPU caches (CU and L2) and fetches directly from system memory. If data exists in the LLC, it may hit, but is not retained or inserted.

  • store_cache_modifier (str, optional) – Controls cache behavior of the store. Supported values are: - None: (default) — Same as “.wb”. Uses write-back caching at all levels (CU, L2, LLC) with LRU policy. - “.wb”: Write-back. Write-allocate on L1 miss, inserted into caches and written back later. - “.cg”: Cache Global. Equivalent to “.wb” — stored through L1 → L2 → LLC under LRU. - “.cs”: Cache Streaming. Bypasses L1, streamed through L2, not retained in LLC. - “.wt”: Write-Through. Bypasses L1 and L2 (coherent cache bypass), may hit in LLC with LRU.

  • hint (constexpr)

Returns:

None

Example

>>> ctx.get(remote_ptr + offsets, local_ptr + offsets, from_rank=1, mask=mask)

put#

Context.put(from_ptr, to_ptr, to_rank, mask=None, other=None, load_cache_modifier=None, store_cache_modifier=None, hint=None)[source]

Copies data from current rank’s local memory to the specified rank’s memory.

This method performs a remote store operation by loading data from from_ptr in the current rank’s local memory, translating to_ptr from the current rank’s address space to the to_rank’s address space, and storing the data to the target memory location. If the current rank and to_rank are the same, this performs a local copy operation.

Parameters:

  • from_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer to local memory in current rank from which to read data.

  • to_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that references memory in to_rank.

  • to_rank (int) – The rank ID to which the data will be written.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not load from from_ptr[idx] and do not store to to_ptr[idx]. Defaults to None.

  • other (Block, optional) – Value to return for masked-out elements during the load operation. If not provided, the result for masked-out elements is undefined. Defaults to None.

  • load_cache_modifier (str, optional) – Controls cache behavior of the load. Supported values are: - None: (default) — Same as “.ca”. Uses cache at all levels (CU, L2, LLC) with LRU policy. - “.ca”: Cache at all levels (CU, L2, LLC) with LRU policy. - “.cg”: Bypasses the CU (L1) cache, streams through L2, and may hit in LLC but the line is not retained or inserted. - “.cv”: Bypasses all GPU caches (CU and L2) and fetches directly from system memory. If data exists in the LLC, it may hit, but is not retained or inserted.

  • store_cache_modifier (str, optional) – Controls cache behavior of the store. Supported values are: - None: (default) — Same as “.wb”. Uses write-back caching at all levels (CU, L2, LLC) with LRU policy. - “.wb”: Write-back. Write-allocate on L1 miss, inserted into caches and written back later. - “.cg”: Cache Global. Equivalent to “.wb” — stored through L1 → L2 → LLC under LRU. - “.cs”: Cache Streaming. Bypasses L1, streamed through L2, not retained in LLC. - “.wt”: Write-Through. Bypasses L1 and L2 (coherent cache bypass), may hit in LLC with LRU.

  • hint (constexpr)

Returns:

None

Example

>>> ctx.put(local_ptr + offsets, remote_ptr + offsets, to_rank=1, mask=mask)

copy#

Context.copy(src_ptr, dst_ptr, from_rank, to_rank, mask=None, other=None, load_cache_modifier=None, store_cache_modifier=None, hint=None)[source]

Copies data from one rank’s memory to another rank’s memory.

This method performs a data transfer by translating src_ptr from the current rank’s address space to the from_rank’s address space, performing a masked load from the translated source, translating dst_ptr to the to_rank’s address space, and storing the loaded data to the target memory location. If from_rank and to_rank are the same, this performs a local copy operation. It is undefined behaviour if the current rank is neither from_rank nor to_rank.

Parameters:

  • src_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that references from_rank’s local memory.

  • dst_ptr (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that references to_rank’s local memory.

  • from_rank (int) – The rank ID that owns src_ptr (source rank).

  • to_rank (int) – The rank ID that will receive the data (destination rank).

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not load from src_ptr[idx] and do not store to dst_ptr[idx]. Defaults to None.

  • other (Block, optional) – Value to return for masked-out elements during the load operation. If not provided, the result for masked-out elements is undefined. Defaults to None.

  • load_cache_modifier (str, optional) – Controls cache behavior of the load. Supported values are: - None: (default) — Same as “.ca”. Uses cache at all levels (CU, L2, LLC) with LRU policy. - “.ca”: Cache at all levels (CU, L2, LLC) with LRU policy. - “.cg”: Bypasses the CU (L1) cache, streams through L2, and may hit in LLC but the line is not retained or inserted. - “.cv”: Bypasses all GPU caches (CU and L2) and fetches directly from system memory. If data exists in the LLC, it may hit, but is not retained or inserted.

  • store_cache_modifier (str, optional) – Controls cache behavior of the store. Supported values are: - None: (default) — Same as “.wb”. Uses write-back caching at all levels (CU, L2, LLC) with LRU policy. - “.wb”: Write-back. Write-allocate on L1 miss, inserted into caches and written back later. - “.cg”: Cache Global. Equivalent to “.wb” — stored through L1 → L2 → LLC under LRU. - “.cs”: Cache Streaming. Bypasses L1, streamed through L2, not retained in LLC. - “.wt”: Write-Through. Bypasses L1 and L2 (coherent cache bypass), may hit in LLC with LRU.

  • hint (constexpr)

Returns:

None

Example

>>> ctx.copy(src_ptr + offsets, dst_ptr + offsets, from_rank=1, to_rank=0, mask=mask)

Atomic operations#

atomic_add#

Context.atomic_add(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic add at the specified rank’s memory location.

This method performs an atomic addition operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically adding the provided data to the to_rank memory location. If the current rank and to_rank are the same, this performs a local atomic addition operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block of dtype=pointer.dtype.element_ty) – The values with which to perform the atomic operation.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel” (stands for “ACQUIRE_RELEASE”), and “relaxed”. If not provided, the function defaults to using “acq_rel” semantics.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are “gpu” (default), “cta” (cooperative thread array, thread block), or “sys” (stands for “SYSTEM”). The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

Example

>>> old_val = ctx.atomic_add(counter, 1, to_rank=1)

atomic_sub#

Context.atomic_sub(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Atomically subtracts data from the specified rank’s memory location.

This method performs an atomic subtraction operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically subtracting the provided data from the to_rank memory location. If the current rank and to_rank are the same, this performs a local atomic subtraction operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – Pointer in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The tensor of elements to be subtracted atomically.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_cas#

Context.atomic_cas(pointer, cmp, val, to_rank, sem=None, scope=None, hint=None)[source]

Performs an atomic compare-and-swap at the specified rank’s memory location.

This method performs an atomic compare-and-swap operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically comparing the value at the memory location with cmp. If they match, it replaces the value with val. If the current rank and to_rank are the same, this performs a local atomic CAS operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory location in the current rank’s address space that will be translated to the to_rank’s address space.

  • cmp (Block) – The expected value to compare against.

  • val (Block) – The new value to store if comparison succeeds.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_xchg#

Context.atomic_xchg(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic exchange at the specified rank’s memory location.

This method performs an atomic exchange operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically swapping the value at the memory location with val. If the current rank and to_rank are the same, this performs a local atomic exchange operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The new values to store.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_xor#

Context.atomic_xor(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic XOR at the specified rank’s memory location.

This method performs an atomic bitwise XOR operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically XOR’ing the value at the memory location with val. If the current rank and to_rank are the same, this performs a local atomic XOR operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The values to XOR with.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_and#

Context.atomic_and(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic AND at the specified rank’s memory location.

This method performs an atomic bitwise AND operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically AND’ing the value at the memory location with val. If the current rank and to_rank are the same, this performs a local atomic AND operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The values to AND with.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_or#

Context.atomic_or(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic OR at the specified rank’s memory location.

This method performs an atomic bitwise OR operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically OR’ing the value at the memory location with val. If the current rank and to_rank are the same, this performs a local atomic OR operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The values to OR with.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_min#

Context.atomic_min(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic minimum at the specified rank’s memory location.

This method performs an atomic minimum operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically updating the memory location to the minimum of its current value and val. If the current rank and to_rank are the same, this performs a local atomic min operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The values to compare with.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

atomic_max#

Context.atomic_max(pointer, val, to_rank, mask=None, sem=None, scope=None, hint=None)[source]

Performs an atomic maximum at the specified rank’s memory location.

This method performs an atomic maximum operation by translating the pointer from the current rank’s address space to the to_rank’s address space and atomically updating the memory location to the maximum of its current value and val. If the current rank and to_rank are the same, this performs a local atomic max operation.

Parameters:

  • pointer (triton.PointerType, or block of dtype=triton.PointerType) – The memory locations in the current rank’s address space that will be translated to the to_rank’s address space.

  • val (Block) – The values to compare with.

  • to_rank (int) – The rank ID to which the atomic operation will be performed.

  • mask (Block of triton.int1, optional) – If mask[idx] is false, do not perform the atomic operation at address pointer[idx]. Defaults to None.

  • sem (str, optional) – Specifies the memory semantics for the operation. Acceptable values are “acquire”, “release”, “acq_rel”, and “relaxed”. Defaults to “acq_rel”.

  • scope (str, optional) – Defines the scope of threads that observe the synchronizing effect. Acceptable values are “gpu” (default), “cta”, or “sys”. The default value is “gpu”.

  • hint (constexpr)

Returns:

The data stored at pointer before the atomic operation.

Return type:

Block

Collective operations#

all_reduce_atomic#

Context.all_reduce_atomic(tile, dst_view)[source]

Tile-level all-reduce using atomic operations.

Atomically adds tile.data to the destination on all ranks.

Parameters:

  • tile (Tile) – Tile with position, dimensions, and data to reduce.

  • dst_view (TensorView) – TensorView for output tensor.

all_reduce_spinlock#

Context.all_reduce_spinlock(tile, dst_view, locks)[source]

Tile-level all-reduce using spinlock synchronization.

For each rank’s tile, acquires a lock, reads current value, adds local contribution, writes back, and releases the lock.

Parameters:

  • tile (Tile) – Tile with position, dimensions, and local data.

  • dst_view (TensorView) – TensorView for output tensor.

  • locks – Pointer to locks array (one lock per tile).

all_reduce_one_shot#

Context.all_reduce_one_shot(tile, src_view, dst_view, locks)[source]

Tile-level all-reduce using one-shot algorithm.

Each rank reads from all ranks and computes the reduction locally. Uses locks as ready flags (producer-consumer).

Parameters:

  • tile (Tile) – Tile with position, dimensions, and local data.

  • src_view (TensorView) – TensorView for source tensor (to load remote data).

  • dst_view (TensorView) – TensorView for output tensor.

  • locks – Pointer to lock array used as ready flags.

all_reduce_ring#

Context.all_reduce_ring(tile, src_view, dst_view)[source]

Tile-level all-reduce using ring algorithm.

Parameters:

  • tile (Tile) – Tile with position and dimensions.

  • src_view (TensorView) – TensorView for input tensor.

  • dst_view (TensorView) – TensorView for output tensor.

all_reduce_two_shot#

Context.all_reduce_two_shot(tile, src_view, dst_view, locks)[source]

Tile-level all-reduce using two-shot algorithm with work distribution.

Each rank reduces only its assigned tiles, then scatters the result. Uses interleaved distribution: rank handles tiles where tile_id % world_size == rank.

Parameters:

  • tile (Tile) – Tile with position, dimensions, and local data.

  • src_view (TensorView) – TensorView for source tensor.

  • dst_view (TensorView) – TensorView for output tensor.

  • locks – Pointer to lock array used as ready flags.

all_gather#

Context.all_gather(tile, dst_view, dim)[source]

Tile-level all-gather operation.

Scatters a pre-computed tile to all ranks at correct offsets.

Parameters:

  • tile (Tile) – Tile with position, dimensions, and computed data.

  • dst_view (TensorView) – TensorView for destination (full gathered size).

  • dim (constexpr) – Dimension to gather along (0 for rows, 1 for columns).

gather#

Context.gather(tile, src_view, source_rank, hint=None)[source]

Tile-level gather from a specific rank.

Loads a tile from source_rank’s memory and returns it directly.

Parameters:

  • tile (TileView) – Tile with position and dimensions.

  • src_view (TensorView) – TensorView for source tensor on source_rank.

  • source_rank (constexpr) – Specific rank to load from (constexpr).

  • hint (constexpr) – Vectorization hint passed to tl.multiple_of / tl.max_contiguous on the translated pointer. Use a scalar (e.g. 16) or a tuple (e.g. (1, 16)) to indicate alignment. Defaults to None (no hint).

Returns:

Loaded tile data as a tensor.

all_to_all#

Context.all_to_all(tile, src_view, dst_view, N_per_rank)[source]

Tile-level all-to-all communication.

Each rank sends portions of its data to every other rank and receives data from every other rank, organized by columns (N dimension).

Parameters:

  • tile (TileView) – Tile with position and dimensions.

  • src_view (TensorView) – TensorView for input tensor.

  • dst_view (TensorView) – TensorView for output tensor.

  • N_per_rank (constexpr) – Number of columns each rank sends/receives per rank.

reduce_scatter#

Context.reduce_scatter(tile, src_view, dst_view, locks)[source]

Tile-level reduce-scatter using contiguous work distribution.

Each rank reduces only its assigned contiguous block of tiles, then stores the result locally.

Parameters:

  • tile (Tile) – Tile with position, dimensions, and local data.

  • src_view (TensorView) – TensorView for source tensor.

  • dst_view (TensorView) – TensorView for output tensor.

  • locks – Pointer to lock array used as ready flags.

Contents