Tile Types

Tile Types#

The tile abstractions the device APIs operate on. All of these are exported from the top level, so iris.Tile, iris.TensorView and so on are the import paths.

A TensorView wraps a pointer plus the tensor’s shape and strides, and hands out tile pointers and masks. A TileView carries a tile’s coordinates; a Tile carries its coordinates and the data. AllReduceConfig is the device-side configuration consumed by the all-reduce methods on Device Context.

TensorView#

class TensorView(ptr, M, N, stride_m, stride_n)[source]#

Bases: object

TensorView storing pointer and tensor metadata.

Parameters:

  • ptr (tensor)

  • M (tensor)

  • N (tensor)

  • stride_m (tensor)

  • stride_n (tensor)

ptr#

Pointer to the tensor’s data.

Type:

triton.language.tensor

M#

Number of rows.

Type:

triton.language.tensor

N#

Number of columns.

Type:

triton.language.tensor

stride_m#

Stride in the M dimension.

Type:

triton.language.tensor

stride_n#

Stride in the N dimension.

Type:

triton.language.tensor

Example

>>> view = make_tensor_view(ptr, M, N, stride_m, stride_n)
>>> tile_ptr, mask = view.tile_ptr(tile)
tile_ptr(tile)[source]#

Compute tile pointer and mask using stored dimensions/strides.

Parameters:

tile (Tile)

tile_ptr_from_indices(rm, rn, block_m, block_n)[source]#

Compute tile pointer and mask from custom row/column indices.

Parameters:

  • block_m (constexpr)

  • block_n (constexpr)

offset_tile_ptr(tile, offset_m=0, offset_n=0, src_mask=None)[source]#

Compute tile pointer with row/column offsets applied.

Parameters:

tile (Tile)

offset(offset_m=0, offset_n=0)[source]#

Create a new view with pointer offset applied.

number_of_tiles(block_m, block_n)[source]#

Compute the total number of tiles needed to cover the tensor.

TileView#

class TileView(pid_m, pid_n, block_m, block_n)[source]#

Bases: object

TileView storing BOTH runtime coordinates AND compile-time block sizes.

This class uses the @constexpr_function pattern:

  • Stores runtime coordinates (pid_m, pid_n) as tl.tensor (computed from tl.program_id)

  • Stores compile-time block sizes (block_m, block_n) as tl.constexpr

  • Constructor is decorated with @constexpr_function to execute at compile-time

Parameters:

  • pid_m (tensor)

  • pid_n (tensor)

  • block_m (constexpr)

  • block_n (constexpr)

pid_m#

Tile row index, computed from tl.program_id.

Type:

triton.language.tensor

pid_n#

Tile column index, computed from tl.program_id.

Type:

triton.language.tensor

block_m#

Tile height. Compile-time constant.

Type:

triton.language.constexpr

block_n#

Tile width. Compile-time constant.

Type:

triton.language.constexpr

Example

>>> pid = tl.program_id(0)
>>> pid_m = pid // num_tiles_n
>>> pid_n = pid % num_tiles_n
>>> tile = TileView(pid_m, pid_n, BLOCK_M, BLOCK_N)
>>> rm, rn, mask = tile.layout(M, N)
layout(M, N)[source]#

Compute memory layout using stored coordinates.

indices()[source]#

Compute base row and column indices for this tile (without bounds checking).

Tile#

class Tile(pid_m, pid_n, block_m, block_n, data)[source]#

Bases: object

Tile with embedded data for computed results.

Extends TileView concept to include a data field for storing computed tile data (e.g., GEMM results in registers).

Parameters:

  • pid_m (tensor)

  • pid_n (tensor)

  • block_m (constexpr)

  • block_n (constexpr)

  • data (tensor)

pid_m#

Tile row index, computed from tl.program_id.

Type:

triton.language.tensor

pid_n#

Tile column index, computed from tl.program_id.

Type:

triton.language.tensor

block_m#

Tile height. Compile-time constant.

Type:

triton.language.constexpr

block_n#

Tile width. Compile-time constant.

Type:

triton.language.constexpr

data#

The tile’s values, for example a GEMM result held in registers.

Type:

triton.language.tensor

Example

>>> c = tl.dot(a, b)
>>> tile = Tile(pid_m, pid_n, BLOCK_M, BLOCK_N, c)
>>> ctx.all_reduce_atomic(tile, dst_view)
layout(M, N)[source]#

Compute memory layout using stored coordinates.

indices()[source]#

Compute base row and column indices for this tile (without bounds checking).

AllReduceConfig#

class AllReduceConfig(variant_code, locks_ptr)[source]#

Bases: object

Device-side configuration for all_reduce collective.

Parameters:

  • variant_code (constexpr)

  • locks_ptr (tensor)

variant_code#

Integer code for algorithm variant (0=atomic, 1=ring, 2=one_shot, 3=two_shot, 4=spinlock)

Type:

triton.language.constexpr

locks_ptr#

Pointer to locks array (ALWAYS required - pass dummy tensor if variant doesn’t need locks)

Type:

triton.language.tensor

make_tensor_view#

make_tensor_view(ptr, M, N, stride_m, stride_n)[source]#

Factory function to create a TensorView inside a JIT context.

This wrapper is needed because @triton.constexpr_function constructors require a JIT context for proper semantic handling. It also converts int/constexpr values to tensors using the +0 trick.

Parameters:

  • ptr – Pointer to tensor data

  • M – Number of rows

  • N – Number of columns

  • stride_m – Stride in M dimension

  • stride_n – Stride in N dimension

Returns:

TensorView instance

Contents