Opportunistically encourage uniform register usage for some integer scalars
#2 323 ouverte le 31 mai 2024
Métriques du dépôt
- Stars
- (392 stars)
- Métriques de merge PR
- (Métriques PR en attente)
Description
Integer arithmetic can follow the uniform data path if all the values match among thread within a warp, which improves performance by reducing interruptions to floating point ops and non-uniform instructions. The cuda compiler is good at determining when this optimization can be used. However this is not always possible. For example, consider
nvfuser_index_t i3 = threadIdx.x / 32;
Whenever blockDim.x % 32 == 0 || (blockDim.y == 1 && blockDim.z == 1) then i3 can be stored in a uniform register. However, nvcc cannot guarantee that the kernel will only be launched with block size satisfying that condition, so it will not use a uniform register in this case.
In order to encourage uniform register use we can do the following:
nvfuser_index_t i3 = threadIdx.x / 32;
i3 = __shfl_sync(0xffffffff, i3, 0);
See this code in cutlass https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cutlass/gemm/kernel/gemm.h#L259 and related issue https://github.com/NVIDIA/cutlass/issues/971 for an example in the wild.
In order to enable this type of optimization in nvFuser, I suggest the following plan:
- Analyze fusion to determine whether any of the block dimensions is constant. This can be done by finding parallelized dimensions with constant extents.
- If the condition above holds, then we know that
threadIdx.x / 32is warp-constant. In that case, we should look through all scalar expressions to try and find that expression or one equivalent to it such asthreadIdx.x >> nfor n>4. - If we find that expression, we should insert a new expression that performs this warp broadcast. This can be done with a new
cacheAfterCacheOp or with a new kir node likekir::WarpBroadcast.
We could get even fancier: for example we could try and handle cases where only (threadIdx.x * threadIdx.y) % 32 == 0 or (threadIdx.x * threadIdx.y * threadIdx.z) % 32 == 0. However I think just handling Generally we should find the first out of threadIdx.x will probably hit all known opportunities for this optimization.TIDx, TIDx*TIDy, or TIDx*TIDy*TIDz that is a multiple of 32 scanned in that order. Then we should ensure that value over 32 as well as the higher TIDs themselves are warp broadcasted.