4 ms·

I tried this (and backprop) in GPGPU compute shaders once. One issue that really tripped me up is matrix multiplication along dimensions that are not powers of

I tried this (and backprop) in GPGPU compute shaders once. One issue that really tripped me up is matrix multiplication along dimensions that are not powers of 2. It will take me awhile to learn to read CUDA code -- is this just an easier thing to do in CUDA or did you figure out a good trick for it?

ryao 2y ago
I know a trick, although it relies on the closed source cuBLAS library. The matrix multiplication in non-batched forward passes is really matrix-vector multiplication, but you can shoehorn that into matrix-matrix multiplication by setting the right dimension to 1. cuBLAS uses column major ordering to match Fortran instead of the row major ordering of C/C++, although you can swap transa and transb, m and n, a and b, lda and ldb; and if you are using cublasGemmEx() to use tensor cores, atype and btype. This relies on the property of mathematics where A * B = C and B’ * A’ = C’ for matrices. You can experiment with this on a CPU using a BLAS library like OpenBLAS or the Intel MKL before trying to do this on a GPU. That said, to do matrix vector multiplication (GEMV), you just need to compute dot products on all of the rows of the matrix with the vector to get your output vector. You could probably just handle each dot product in power of 2 chunks on the GPU to compute partial sums. When you get to the last chunk that is not a power of 2, just have the threads that go pass the end of the dot product contribute 0 to the partial sums. Then you would have the threads handle the horizontal sum of the partial sums to finish the computation of an entry in the output vector. As for matrix-matrix multiplication (GEMM), that is much harder to do in a performant manner. If you do it by treating the multiplication as a series of matrix-vector multiplications, you will be bandwidth limited rather than compute limited. You would need to implement tiling to get good performance, but there are many tricks needed and it is very hard to outperform cuBLAS on Nvidia hardware. Here is one person’s attempts to do it: https://siboehm.com/articles/22/CUDA-MMM https://siboehm.com/articles/22/CUDA-MMM Here is a fairly cryptic description of how to get close to cuBLAS that mentions a few key techniques: https://accu.org/journals/overload/32/181/schuetze/ https://accu.org/journals/overload/32/181/schuetze/ Here is another guy’s attempts that actually managed to outperform cuBLAS on the H100: https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog https://cudaforfun.substack.com/p/outperforming-cublas-on-h1...
- upghost 2y ago
  Wow, great answer. That deserves more upvotes than it's probably going to get. So what that means for compute shaders is I will need to partition the vectors being multiplied into power of 2 sized partitions and repeatedly call the kernels, but with offsets appropriate to the partition size, and then have the overrun contribute zero as you said. Obviously not as efficient as doing it all GPU side, still faster than doing it on the CPU. Thank you!
  - ryao 2y ago
    You don’t need to repeatedly call the kernels. You just need to loop inside the kernels to have each process every nth multiplication according to their thread id, adding it to the thread’s partial sum and have a branch in the loop skip the addition (or add 0) when you exceed the range. Afterward, you would do a horizontal sum of the partial sums to get the value for the corresponding field in the output vector, and you would have a “grid” do this for all members of the output vector. The original poster’s mathmul kernel does this. See the example under “we can perform a warp sum reduction to combine results from all threads”: https://andrewkchan.dev/posts/yalm.html#section-3.2 https://andrewkchan.dev/posts/yalm.html#section-3.2 It does “if (i >= d) return;“ to handle non-multiples in one dimension and skips the addition in its dot product, matmul_row(), to handle non-multiples in the other dimension. In CUDA, you have threads that are organized into 32 thread warps, that are organized into blocks of warps, that are organized into grids. Threads within a block can share memory, but blocks within a grid cannot. The hardware will effectively loop over blocks in the grid until it finishes. The dimension for the output vector does not require sharing memory, so we can set that to the grid size. The hardware might actually round the grid up, so “if (i >= d) return;” is a check for the case that we are executing past the end of the output vector in case the hardware schedules pass the end of the vector. As for the other dimension, we only have a limited number of threads (in this example 32, which is a single warp), so we need each to do partial sums of the multiplications in the dot product. We essentially break the dot product into thread count (32 here) chunks. Each thread will know which multiplication in the thread count chunk to compute based on its thread id, and we keep adding the thread count (32 here) to that to get the next one until we have finished. The for loop’s loop condition protects us from going past the end of the dot product dimension. Then we have thread count number of partial sums, so we need to do a horizontal sum across the threads in a block to get the final sum. This part is a bit magical and is done via a warp sum reduction. Since the threads share memory, you basically have them start summing their neighbors’ values into their own. Basically, all threads find pairs and one of the pair does the sum, cutting the threads doing the sum in half. Then the ones that did the sums pair up again. This again cuts the number in half. This continues until only 1 thread remains. Then only the remaining thread is allowed to write to “out[i] = rowSum” due to “if (threadIdx.x == 0)”. The warp itself is somewhat special since it can do a really fast horizontal sum via shuffles that cannot be done when using thread blocks ghat utilize multiple warps. That is why the author chose to have the block size be exactly 1 warp rather than a multiple of warps. It helps to think about the threads in the block as a spreadsheet. The columns are the threads and the rows represent instructions. Every thread executes the same instruction, and when a branch says not to do something for only some threads, the ones excluded do noops until execution reaches a point where they all do the same thing. An if else similarly has noops in threads not executing true, followed by noops in the threads not executing false. This is called branch divergence. There are diagrams showing a spreadsheet like view of this, such as the following, which shows a cutdown example: https://image1.slideserve.com/2427671/branch-divergence-l.jpg https://image1.slideserve.com/2427671/branch-divergence-l.jp... I have mentioned a horizontal sum multiple times. That just means summing across the columns inside rows of the “spreadsheet”. It is called horizontal since rows are horizontal. Coincidentally, the partial sum calculation is called a vertical sum since the columns are vertical. Now we need to consider the grid, which is yet another dimension and is the dimension along which we populate the output vector. Basically, you can view the grid as doing multiple spreadsheets in parallel, but not all at once. One comical way to imagine it is to imagine having printed the same spreadsheet hundreds of times an and then feeding them into a paper shredder head first. A normal paper shredder can only do a few at a time, and similarly, the GPU can only do so many at a time too. You can imagine the processing to proceed along the rows in all of the spreadsheets being shredded at the same time.

ryao 2y ago

I know a trick, although it relies on the closed source cuBLAS library. The matrix multiplication in non-batched forward passes is really matrix-vector multiplication, but you can shoehorn that into matrix-matrix multiplication by setting the right dimension to 1. cuBLAS uses column major ordering to match Fortran instead of the row major ordering of C/C++, although you can swap transa and transb, m and n, a and b, lda and ldb; and if you are using cublasGemmEx() to use tensor cores, atype and btype. This relies on the property of mathematics where A * B = C and B’ * A’ = C’ for matrices. You can experiment with this on a CPU using a BLAS library like OpenBLAS or the Intel MKL before trying to do this on a GPU. That said, to do matrix vector multiplication (GEMV), you just need to compute dot products on all of the rows of the matrix with the vector to get your output vector. You could probably just handle each dot product in power of 2 chunks on the GPU to compute partial sums. When you get to the last chunk that is not a power of 2, just have the threads that go pass the end of the dot product contribute 0 to the partial sums. Then you would have the threads handle the horizontal sum of the partial sums to finish the computation of an entry in the output vector. As for matrix-matrix multiplication (GEMM), that is much harder to do in a performant manner. If you do it by treating the multiplication as a series of matrix-vector multiplications, you will be bandwidth limited rather than compute limited. You would need to implement tiling to get good performance, but there are many tricks needed and it is very hard to outperform cuBLAS on Nvidia hardware. Here is one person’s attempts to do it: https://siboehm.com/articles/22/CUDA-MMM https://siboehm.com/articles/22/CUDA-MMM Here is a fairly cryptic description of how to get close to cuBLAS that mentions a few key techniques: https://accu.org/journals/overload/32/181/schuetze/ https://accu.org/journals/overload/32/181/schuetze/ Here is another guy’s attempts that actually managed to outperform cuBLAS on the H100: https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog https://cudaforfun.substack.com/p/outperforming-cublas-on-h1...

upghost 2y ago
Wow, great answer. That deserves more upvotes than it's probably going to get. So what that means for compute shaders is I will need to partition the vectors being multiplied into power of 2 sized partitions and repeatedly call the kernels, but with offsets appropriate to the partition size, and then have the overrun contribute zero as you said. Obviously not as efficient as doing it all GPU side, still faster than doing it on the CPU. Thank you!
- ryao 2y ago
  You don’t need to repeatedly call the kernels. You just need to loop inside the kernels to have each process every nth multiplication according to their thread id, adding it to the thread’s partial sum and have a branch in the loop skip the addition (or add 0) when you exceed the range. Afterward, you would do a horizontal sum of the partial sums to get the value for the corresponding field in the output vector, and you would have a “grid” do this for all members of the output vector. The original poster’s mathmul kernel does this. See the example under “we can perform a warp sum reduction to combine results from all threads”: https://andrewkchan.dev/posts/yalm.html#section-3.2 https://andrewkchan.dev/posts/yalm.html#section-3.2 It does “if (i >= d) return;“ to handle non-multiples in one dimension and skips the addition in its dot product, matmul_row(), to handle non-multiples in the other dimension. In CUDA, you have threads that are organized into 32 thread warps, that are organized into blocks of warps, that are organized into grids. Threads within a block can share memory, but blocks within a grid cannot. The hardware will effectively loop over blocks in the grid until it finishes. The dimension for the output vector does not require sharing memory, so we can set that to the grid size. The hardware might actually round the grid up, so “if (i >= d) return;” is a check for the case that we are executing past the end of the output vector in case the hardware schedules pass the end of the vector. As for the other dimension, we only have a limited number of threads (in this example 32, which is a single warp), so we need each to do partial sums of the multiplications in the dot product. We essentially break the dot product into thread count (32 here) chunks. Each thread will know which multiplication in the thread count chunk to compute based on its thread id, and we keep adding the thread count (32 here) to that to get the next one until we have finished. The for loop’s loop condition protects us from going past the end of the dot product dimension. Then we have thread count number of partial sums, so we need to do a horizontal sum across the threads in a block to get the final sum. This part is a bit magical and is done via a warp sum reduction. Since the threads share memory, you basically have them start summing their neighbors’ values into their own. Basically, all threads find pairs and one of the pair does the sum, cutting the threads doing the sum in half. Then the ones that did the sums pair up again. This again cuts the number in half. This continues until only 1 thread remains. Then only the remaining thread is allowed to write to “out[i] = rowSum” due to “if (threadIdx.x == 0)”. The warp itself is somewhat special since it can do a really fast horizontal sum via shuffles that cannot be done when using thread blocks ghat utilize multiple warps. That is why the author chose to have the block size be exactly 1 warp rather than a multiple of warps. It helps to think about the threads in the block as a spreadsheet. The columns are the threads and the rows represent instructions. Every thread executes the same instruction, and when a branch says not to do something for only some threads, the ones excluded do noops until execution reaches a point where they all do the same thing. An if else similarly has noops in threads not executing true, followed by noops in the threads not executing false. This is called branch divergence. There are diagrams showing a spreadsheet like view of this, such as the following, which shows a cutdown example: https://image1.slideserve.com/2427671/branch-divergence-l.jpg https://image1.slideserve.com/2427671/branch-divergence-l.jp... I have mentioned a horizontal sum multiple times. That just means summing across the columns inside rows of the “spreadsheet”. It is called horizontal since rows are horizontal. Coincidentally, the partial sum calculation is called a vertical sum since the columns are vertical. Now we need to consider the grid, which is yet another dimension and is the dimension along which we populate the output vector. Basically, you can view the grid as doing multiple spreadsheets in parallel, but not all at once. One comical way to imagine it is to imagine having printed the same spreadsheet hundreds of times an and then feeding them into a paper shredder head first. A normal paper shredder can only do a few at a time, and similarly, the GPU can only do so many at a time too. You can imagine the processing to proceed along the rows in all of the spreadsheets being shredded at the same time.