Predicates and Registers

[leimao]

To implement some generic kernels that support arbitrary problem sizes, sometimes we will have to rely on using predicates to decide whether certain operation has to be performed in each thread.

Typically we will first construct an identity tensor whose size seems to be large in the thread, and then we partition the tensor for each thread based on the thread id.

Tensor identity = make_identity_tensor(shape(mC));
// Partition. Got smaller coordinate tensor for each thread.
// ...

My question is, suppose the identity tensor is usually large, will it be allocated in local memory instead of registers? If so, how does the compiler optimize this? If we have 256 threads, are there 256 identity tensors or just 1 identity tensors allocated in local memory (I think it is 256)?

Because accessing local memory has the same performance as accessing global memory, for some memory bound kernels, it will reduce the performance to some extent.

[ccecka]

These are implicit Tensors that have no storage. There will be very few URs or Registers to possibly represent the offset, but otherwise there is no physical storage and the elements of the tensor are generated on-the-fly.

https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/cute/0z_tma_tensors.md

[leimao]

Thanks @ccecka