int64 indexing and max grid size

[iclementine]

Indxing & Task partitioning in Triton jit functions

In triton kernels, we compose indices & strides to compute the offset of elements in a tensor, then add the offset to the data pointer of the tensor to get pointers to these elements. Almost everything starts from tl.program_id, tl.arange and tl.num_programs.

As to the task partitioning, the most-common example is to use a ceil-div pattern. Each CTA takes a tile of data, and the number of CTAs to process the whole Tensor is the cdiv(size, TILE_SIZE).

Almost everything tutorial tells us to write triton jit functions like this, but in practice you have to take care of a lot of other stuffs.

@triton.jit
def add_kernel(a_ptr, b_ptr, o_ptr, N, TILE_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * TILE_SIZE + tl.arange(0, TILE_SIZE)
    mask = offsets < N

    a = tl.load(a_ptr + offsets, mask=mask)
    b = tl.load(b_ptr + offsets, mask=mask)
    o = (a + b) * 2

    tl.store(o_ptr + offsets, o, mask=mask)

def add(a, b):
    o = torch.empty_like(a)
    N = a.numel()
    TILE_SIZE = 32
    grid = (triton.cdiv(N, TILE_SIZE), 1, 1)
    add_kernel[grid](a, b, o, N, TILE_SIZE)
    return o

Q1: Will offsets exceed the limit of int32 indexing? Yes, it would. Actually int32 has the max value 2147483647, which is 2 * 1024 * 1024 * 1024 - 1, which means a contiguuos tensor with more than 2 * 1024 * 1024 * 1024 elements would make offsets exceed the limit and overflow to some negative value and cause incorrect memory address to access. If this is a tensor with arbitrary stride, things would be worse.

But how to avoid this?

1. Carefully add parenthesis into the code of offset computing. After all the base pointer is a 64bit integer, so base_ptr + offset1 * stride1 + offset2 * stride2 + ... would avoid this, since the first addition promotes the second argument to 64-bit, and so do the consecutive additions.
2. Determine whether there would be a least a tensor that has a larger offset than 2 * 1024 * 1024 * 1024 - 1. If none of them has such a large offset, use int32 index as shown above, otherwise, cast the result of tl.program_id, tl.arange and tl.num_programs to int64 from the beginning.
3. Finer grained control. Do as many as possible offset computation in i32, until it is impossible to do so. But this fine grained control is hard to do.

Q2: Will triton.cdiv(N, TILE_SIZE) exceed the limit of int32 indexing? Yes. In nvidia gpus, the max grid size in the x dimension is also 2147483647. So you should also take care not to make a super-large grid.

Answer:

1. use larger tile size to make sure that the grid sizes doe not exceed the limits;
2. use grid-stride-loop style task partitioning.

We want the best-practice to be follow by triton developers, but taking care of indexing code in triton jit functions and task partitioning is sometimes tedious and not everyone follows them. Do you have any good ideas?