NVIDIA/Fuser

Opportunistically encourage uniform register usage for some integer scalars

Open

#2,323 opened on 2024年5月31日

GitHub で見る
 (7 comments) (0 reactions) (0 assignees)C++ (80 forks)github user discovery
good first issue

Repository metrics

Stars
 (392 stars)
PR merge metrics
 (PR metrics pending)

説明

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:

  1. Analyze fusion to determine whether any of the block dimensions is constant. This can be done by finding parallelized dimensions with constant extents.
  2. If the condition above holds, then we know that threadIdx.x / 32 is warp-constant. In that case, we should look through all scalar expressions to try and find that expression or one equivalent to it such as threadIdx.x >> n for n>4.
  3. If we find that expression, we should insert a new expression that performs this warp broadcast. This can be done with a new cacheAfter CacheOp or with a new kir node like kir::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 threadIdx.x will probably hit all known opportunities for this optimization. Generally we should find the first out of 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.

コントリビューターガイド