7 / 11 Kernels (CUDA / Triton / Pallas)

Which Tile Am I

For pl.BlockSpec(block_shape, index_map), the index_map function is called with the current grid indices and must return, for each array dimension, a value the Pallas runtime treats as:

A

An already-scaled element offset, i.e. index_map must do the multiplication by block_shape itself

B

The flattened linear program_id collapsed across all grid axes

C

A block index, which is multiplied by that dimension's block_shape to compute the starting element offset

D

A raw byte offset directly into the backing HBM buffer

E

The total number of blocks along that axis