Merge pull request #34 from ashvardanian/wgmma-stats

Counting Tensor Ops Correctly 🤯 

Measuring Tensor-Core throughput is tricky! Many families of matrix-multiplications instructions exist. Practically every Nvidia GPU generation brings new tiles, new numeric types, mixed-precision schemes, and "structured sparsity" models. All of those together form [some of the longest PTX IR instructions](https://ashvardanian.com/posts/longest-ptx-instruction/). To make things worse, across generations, Tensor Core scheduling and collective execution scale are different!

- Before Volta and Tensor Cores, each GPU thread would execute its own scalar Fused-Multiply-Add — easy-peasy, as long as you know how to choose the optimal grid size for your GPU model.
- On Volta, with new `mma.*` instructions and `wmma::` intrinsics, 8 threads would execute every tiled Mat-Mul together. This scale of collaboration was creatively called by Nvidia engineers ~~a octet~~ a "quadpair", of course 🤦‍♂️ 
- On Ampere, with new `wmma.mma.*` instructions, all of the 32 threads in a single "warp" would work together. This abstraction makes sense to people familiar with CUDA C++ and how scheduling works on the GPU. Great!
On Hopper, things changed again, of course, with `wgmma.mma_async.sync.*`, which supports basic asynchronous primitives at the hardware level. It has 128 threads across 4 consecutive "warps" forming a "warp group."
- On Blackwell, you would be wise to expect a new change, and it came with a broader set of functionality refactored into an all-new `tcgen05.*` namespace of instructions 🧠 🔫 

This new PR addresses this by explicitly marking the collaboration "scale" and counting TOPS differently for each family of instructions.

---

This way, the compiler will see that I'm trying to export the accumulated value and will not remove our `mma_sync` call, even if the target address is a NULL pointer. Another approach I'd often use in PTX is to define dummy global variables and export a few values there:

```cuda
template <typename input_type_, typename output_type_, int m_, int n_, int k_, int repetitions_ = 128>
__device__ inline void tops_tc_cuda_kernel() {
    using namespace nvcuda;
    wmma::fragment<wmma::matrix_a, m_, n_, k_, input_type_, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, m_, n_, k_, input_type_, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, m_, n_, k_, output_type_> c_frag;
    for (int i = 0; i != repetitions_; ++i) wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    if (threadIdx.x == 2147483647) wmma::store_matrix_sync(nullptr, c_frag, 16, wmma::mem_row_major);
}
```

This way, the compiler will see that I'm trying to export the accumulated value and will not remove our `mma_sync` call, even if the target address is a NULL pointer. Another approach I'd often use in PTX is to define dummy global variables and export a few values there:

```ptx
.visible .global .align 4 .s32 dummy_sink_s32[32];
.visible .global .align 4 .f32 dummy_sink_f32[32];
.visible .entry tops_f16f32_sm90tc_m64n256k16_loop128_ptx_kernel() {
    ...
loop_exit:
    // Zero argument means - wait for all committed WGMMAs to complete.
    wgmma.wait_group.sync.aligned 0;

    // Use volatile stores to force the accumulator values to be written out.
    // This dummy write (to a global variable) makes the work observable and 
    // prevents the multiplication  pipeline from being optimized out.
    st.global.volatile.f32 [dummy_sink_f32],      accum0;
    st.global.volatile.f32 [dummy_sink_f32+4],    accum1;
    ret;
}
```

But with WGMMA, the `PTXAS` tool will optimize our multiplications if the shared-memory tile descriptors aren't valid. Even if it's just for a benchmark. So this PR shows how to assemble valid descriptors 🤗 

---

This PR fixes those issues and adds more PTX kernels to highlight the different aspects of GPGPU development 🤗

Author: ashvardanian

.gitignore

@@ -7,6 +7,8 @@ build_release/
 
 # Temporary binaries
 /tmp/
-less_slow_from_ptx.cubin
 less_slow_from_cu.cubin
 less_slow_from_cu.ptx
+less_slow_sm70_from_ptx.cubin
+less_slow_sm80_from_ptx.cubin
+less_slow_sm90a_from_ptx.cubin

CMakeLists.txt

@@ -334,7 +334,7 @@ endif()
 # https://cmake.org/cmake/help/latest/variable/CMAKE_LANG_COMPILER_ID.html
 if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA" OR CMAKE_CUDA_COMPILER_ID STREQUAL "NVHPC")
   set_target_properties(less_slow PROPERTIES POSITION_INDEPENDENT_CODE ON)
-  set_target_properties(less_slow PROPERTIES CUDA_ARCHITECTURES "70;75;80;89;90")
+  set_target_properties(less_slow PROPERTIES CUDA_ARCHITECTURES "70;75;80;89;90a")
   target_compile_options(less_slow PRIVATE
     -Wfatal-errors # Stop on first error
     -fopenmp # OpenMP support, also requires linking
@@ -434,7 +434,7 @@ if(USE_NVIDIA_CCCL)
   # target_link_libraries(less_slow PRIVATE nvidia::cutlass::cutlass)
 
   # List the PTX files you want to copy
-  set(PTX_FILES less_slow_sm70.ptx less_slow_sm90a.ptx)
+  set(PTX_FILES less_slow_sm70.ptx less_slow_sm80.ptx less_slow_sm90a.ptx)
 
   # Loop over each PTX file and add a custom command to copy it
   foreach(PTX ${PTX_FILES})

less_slow.cpp

@@ -18,8 +18,16 @@
  *  You can validate this file by asking the Nvidia PTX Assembler to compile it
  *  to `.cubin` for some target architecture:
  * 
- *  $ ptxas -o less_slow_from_ptx.cubin -arch=sm_70 less_slow_sm70.ptx
- *  $ cuobjdump -sass less_slow_from_ptx.cubin | grep -i mma
+ *  $ ptxas -o less_slow_sm70_from_ptx.cubin -arch=sm_70 less_slow_sm70.ptx
+ *  $ cuobjdump -sass less_slow_sm70_from_ptx.cubin | grep -i mma
+ *
+ *  Assuming how aggressively NVCC unrolls loops and the number of kernels in
+ *  this file, you may want to deduplicate them:
+ *
+ *  $ cuobjdump -sass less_slow_sm70_from_ptx.cubin | grep -i mma | \
+ *  $   sed -r 's/\/\*[^*]+\*\///g' | \
+ *  $   sed -r 's/^[[:space:]]+//; s/[[:space:]]+$//' | \
+ *  $   sort -u
  * 
  *  @section Register File
  * 
@@ -37,18 +45,95 @@
 .target sm_70            // Target architecture (SM 7.0 - Volta GPUs)
 .address_size 64         // 64-bit addressing
 
-.visible .entry tops_f16f16_sm70tc_16x16x16_loop128_ptx_kernel()
+.visible .entry tops_f16f16_sm70mma_8x8x4_loop128_ptx_kernel()
 {
     // Accumulator registers used for both input and output of the MMA operation
     .reg .b32 accum_0, accum_1, accum_2, accum_3;
 
-    // Registers to hold packed 16-bit data for matrix a (8 registers)
-    .reg .b32 matrix_a_0, matrix_a_1, matrix_a_2, matrix_a_3,
-              matrix_a_4, matrix_a_5, matrix_a_6, matrix_a_7;
+    // Registers to hold packed pairs of 16-bit data for matrix a (2 registers)
+    .reg .b32 matrix_a_0, matrix_a_1;
 
-    // Registers to hold packed 16-bit data for matrix b (8 registers)
-    .reg .b32 matrix_b_0, matrix_b_1, matrix_b_2, matrix_b_3,
-              matrix_b_4, matrix_b_5, matrix_b_6, matrix_b_7;
+    // Registers to hold packed pairs of 16-bit data for matrix b (2 registers)
+    .reg .b32 matrix_b_0, matrix_b_1;
+
+    // General-purpose registers for loop control and constant values
+    .reg .b32 loop_counter, loop_limit, packed_const;
+
+    // Predicate register for conditional branching (loop exit)
+    .reg .pred exit_predicate;
+
+    // Set up loop counter and loop limit
+    mov.u32 loop_counter, 0;
+    mov.u32 loop_limit, 128;
+
+    // Zero-initialize the accumulator registers
+    mov.f32 accum_0, 0.0;
+    mov.f32 accum_1, 0.0;
+    mov.f32 accum_2, 0.0;
+    mov.f32 accum_3, 0.0;
+
+    // Initialize constant for packed matrix data (placeholder)
+    mov.b32 packed_const, 0x00010001;
+
+    // Initialize matrix a registers with the packed constant
+    mov.b32 matrix_a_0, packed_const;
+    mov.b32 matrix_a_1, packed_const;
+
+    // Initialize matrix b registers with the packed constant
+    mov.b32 matrix_b_0, packed_const;
+    mov.b32 matrix_b_1, packed_const;
+
+    // The main loop will repeat for 128 iterations
+loop_start:
+    setp.ge.u32 exit_predicate, loop_counter, loop_limit;
+    @exit_predicate bra loop_end;
+
+    mma.sync.aligned.m8n8k4.row.col.f16.f16.f16.f16 
+         { accum_0, accum_1, accum_2, accum_3 },
+         { matrix_a_0, matrix_a_1 },
+         { matrix_b_0, matrix_b_1 },
+         { accum_0, accum_1, accum_2, accum_3 };
+
+    // Increment the loop counter
+    add.u32 loop_counter, loop_counter, 1;
+
+    // Branch back to the beginning of the loop
+    bra loop_start;
+
+loop_end:
+    // If we simply exit, the computation will be optimized out!
+    // Instead, let's check for an impossible condition, like if the thread ID
+    // is equal to `UINT_MAX`, and if so - write accumulators to global memory
+    // NULL address.
+    .reg .u32 tid;
+    .reg .pred impossible_predicate;
+    mov.u32 tid, %tid.x; //? Special system registers start with `%`
+    setp.ne.u32 impossible_predicate, tid, 0xFFFFFFFF;
+    @impossible_predicate bra loop_exit;
+
+    // Write into memory:
+    .reg .u64 store_ptr;
+    mov.u64 store_ptr, 0;
+    st.global.f32 [store_ptr],      accum_0;
+    st.global.f32 [store_ptr+4],    accum_1;
+    st.global.f32 [store_ptr+8],    accum_2;
+    st.global.f32 [store_ptr+12],   accum_3;
+
+loop_exit:
+    ret;
+}
+
+.visible .entry tops_f16f32_sm70mma_8x8x4_loop128_ptx_kernel()
+{
+    // Accumulator registers used for both input and output of the MMA operation
+    .reg .b32 accum_0, accum_1, accum_2, accum_3,
+              accum_4, accum_5, accum_6, accum_7;
+
+    // Registers to hold packed 16-bit data for matrix a (4 registers)
+    .reg .b32 matrix_a_0, matrix_a_1, matrix_a_2, matrix_a_3;
+
+    // Registers to hold packed 16-bit data for matrix b (4 registers)
+    .reg .b32 matrix_b_0, matrix_b_1, matrix_b_2, matrix_b_3;
 
     // General-purpose registers for loop control and constant values
     .reg .b32 loop_counter, loop_limit, packed_const;
@@ -74,33 +159,25 @@
     mov.b32 matrix_a_1, packed_const;
     mov.b32 matrix_a_2, packed_const;
     mov.b32 matrix_a_3, packed_const;
-    mov.b32 matrix_a_4, packed_const;
-    mov.b32 matrix_a_5, packed_const;
-    mov.b32 matrix_a_6, packed_const;
-    mov.b32 matrix_a_7, packed_const;
 
     // Initialize matrix b registers with the packed constant
     mov.b32 matrix_b_0, packed_const;
     mov.b32 matrix_b_1, packed_const;
     mov.b32 matrix_b_2, packed_const;
     mov.b32 matrix_b_3, packed_const;
-    mov.b32 matrix_b_4, packed_const;
-    mov.b32 matrix_b_5, packed_const;
-    mov.b32 matrix_b_6, packed_const;
-    mov.b32 matrix_b_7, packed_const;
 
     // The main loop will repeat for 128 iterations
 loop_start:
     setp.ge.u32 exit_predicate, loop_counter, loop_limit;
     @exit_predicate bra loop_end;
 
-    wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 
-         { accum_0, accum_1, accum_2, accum_3 },
-         { matrix_a_0, matrix_a_1, matrix_a_2, matrix_a_3,
-           matrix_a_4, matrix_a_5, matrix_a_6, matrix_a_7 },
-         { matrix_b_0, matrix_b_1, matrix_b_2, matrix_b_3,
-           matrix_b_4, matrix_b_5, matrix_b_6, matrix_b_7 },
-         { accum_0, accum_1, accum_2, accum_3 };
+    mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32
+         { accum_0, accum_1, accum_2, accum_3,
+           accum_4, accum_5, accum_6, accum_7 },
+         { matrix_a_0, matrix_a_1 },
+         { matrix_b_0, matrix_b_1 },
+         { accum_0, accum_1, accum_2, accum_3,
+           accum_4, accum_5, accum_6, accum_7 };
 
     // Increment the loop counter
     add.u32 loop_counter, loop_counter, 1;
@@ -147,7 +224,8 @@ loop_exit:
  *    with both arguments in shared memory!
  *
  *  Because only one `.version` directive can be placed in each file, for newer
- *  kernels, go to `less_slow_sm90a.ptx`.
+ *  kernels, go to `less_slow_sm80.ptx` for Ampere and `less_slow_sm90a.ptx`
+ *  for Hopper.
  *
  *  @see PTX module-level directives:
  *  https://docs.nvidia.com/cuda/parallel-thread-execution/#ptx-module-directives