define cluster_ctarank

Author: XiaobingSuper

include/common/util.cuh

@@ -285,6 +285,21 @@ struct shared_allocator {
  */
 using tma_allocator = shared_allocator<1024>;
 using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
+
+/* Get CTA ID within a cluster */
+__device__ static inline int3 clusterIdx() {
+    int3 cluster_idx;
+    asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
+    asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
+    asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
+    return cluster_idx;
+}
+__device__ static inline int cluster_ctarank() {
+    uint32_t ctarank;
+    asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
+    return ctarank;
+}
+
 #endif
 
 } // namespace kittens

kernels/group_gemm/group_gemm.cu

@@ -127,7 +127,7 @@ struct matmul_template {
                 } else {
                     tma::expect(args.inputs_cluster_arrived, args.input.b);
                     tma::load_async(args.input.b, args.globals.B,
-                        {args.common.group_idx, args.common.block_n_idx, args.iter}, args.inputs_arrived);
+                        {args.common.group_idx, args.common.block_n_idx, args.iter}, args.inputs_cluster_arrived);
                 }
             }
         }