So, I was just practicing triton and decided to try to create such offsets so I’m able to access blocks of memory in a single operation, contrary to straightforward looping with multiple load/store operations, and what often is generated by LLMs (for now). This I also came up with the help of Claude, but after some time of deliberation. Something like this can also be found in (Triton 2024) from Philippe Tillet on triton-lang.org, but just everything is jammed up there, so it’s harder to dissect. Meanwhile, I’m sharing simplest kernel possible, that just does block fill.
@triton.jit
def _blocky_fill_kernel_2(output_matrix, block_dim_x:tl.constexpr, block_dim_y:tl.constexpr, stride:tl.constexpr, len:tl.constexpr):
pid = tl.program_id(0)
block_y_pos = (pid * block_dim_x) // stride
block_x_pos = pid % (stride // block_dim_x)
offset = block_x_pos * block_dim_x + block_y_pos * stride * block_dim_y + tl.arange(0, block_dim_x)
# Create offsets for the entire block at once
block_base = block_x_pos * block_dim_x + block_y_pos * stride * block_dim_y
# Create row offsets (0, 1, 2, 3, 0, 1, 2, 3, ...) for each element in block
col_idx = tl.arange(0, block_dim_x * block_dim_y) % block_dim_x
# Create row numbers (0, 0, 0, 0, 1, 1, 1, 1, ...) for each element
row_idx = tl.arange(0, block_dim_x * block_dim_y) // block_dim_x
# Combine into final offsets
offsets = block_base + row_idx * stride + col_idx
mask = offsets < len
#if pid == 0: # Debug
tl.store(output_matrix + offsets, pid + 1, mask=mask)it works in this way: col_idx is a matrix of column indices for each cell in the block, and for 16x16 matrix with 4x4 block size will look like this, if we isolate block offsets only for pid==0:
tensor([[0., 1., 2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 1., 2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 1., 2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 1., 2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
device='cuda:0')`row_idx` is a matrix of row indices for each cell in the block, and for 16x16 matrix with 4x4 block size will look like this, if we isolate block offsets only for pid==0:tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[2., 2., 2., 2., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[3., 3., 3., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
device='cuda:0')So, after this, we can calculate their offsets:
# Combine into final offsets
offsets = block_base + row_idx * stride + col_idxWhich for our block with pid == 0 of 4x4 for 16x16 offset matrix will look like this:
tensor([[ 0., 1., 2., 3., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[16., 17., 18., 19., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[32., 33., 34., 35., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[48., 49., 50., 51., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
device='cuda:0')Really convenient, no loops, single load/store operation.
…
Or you could just use tl.make_block_ptr :)
References
Triton. 2024. “Matrix Multiplication Tutorial.” https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html.