Fix local_tile for version 3.3 and newer, NVIDIA/cutlass#1201

cloudhan · Dec 12, 2023 · a26f6ad · a26f6ad
1 parent f804b4c
commit a26f6ad
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/gemm/cuda_cute/04_smem/matmul_smem_0.cu b/gemm/cuda_cute/04_smem/matmul_smem_0.cu
@@ -133,8 +133,8 @@ MATMUL_KERNEL_SIGNATURE(matmul_kernel_smem_0) {
     store_smem_load_global_b<NumThreads, SmemShapeN, SmemShapeK>(sB, stripe_gB(_, _, _0{}, block_p), stripe_cB(_, _, _0{}, block_p), n, k);
     __syncthreads();
 
-    const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), threadIdx.x % (CtaShapeM / ThreadShapeM));
-    const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), threadIdx.x / (CtaShapeM / ThreadShapeM));
+    const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x % (CtaShapeM / ThreadShapeM)));
+    const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x / (CtaShapeM / ThreadShapeM)));
 
 #pragma unroll
     for (int smem_AB_thread_p = 0; smem_AB_thread_p < SmemShapeK; smem_AB_thread_p++) {

diff --git a/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_0_naive.cu b/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_0_naive.cu
@@ -166,8 +166,8 @@ __launch_bounds__(NumThreads, 2) MATMUL_KERNEL_SIGNATURE(matmul_smem_and_registe
   load_global_a<NumThreads, SmemShapeM, SmemShapeK, SmemALoadStoreVec>(staging_a, stripe_gA(_, _, _0{}, 0), stripe_cA(_, _, _0{}, 0), m, k);
   load_global_b<NumThreads, SmemShapeN, SmemShapeK, SmemBLoadStoreVec>(staging_b, stripe_gB(_, _, _0{}, 0), stripe_cB(_, _, _0{}, 0), n, k);
 
-  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), threadIdx.x % (CtaShapeM / ThreadShapeM));
-  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), threadIdx.x / (CtaShapeM / ThreadShapeM));
+  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x % (CtaShapeM / ThreadShapeM)));
+  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x / (CtaShapeM / ThreadShapeM)));
 
   const auto num_smem_block = size<3>(stripe_gA);
 #pragma unroll 1  // no unroll

diff --git a/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_1_cutlass.cu b/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_1_cutlass.cu
@@ -173,8 +173,8 @@ MATMUL_KERNEL_SIGNATURE(matmul_smem_and_register_pipelining_1) {
   store_smem_b<NumThreads, SmemShapeN, SmemShapeK, SmemBLoadStoreVec>(sB(_, _, _0{}), staging_b, n, k);
   __syncthreads();
 
-  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), threadIdx.x % (CtaShapeM / ThreadShapeM));
-  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), threadIdx.x / (CtaShapeM / ThreadShapeM));
+  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x % (CtaShapeM / ThreadShapeM)));
+  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x / (CtaShapeM / ThreadShapeM)));
 
   copy(stripe_sA(_, 0, _0{}, _0{}), fragA[0]);  // load_fragment a
   copy(stripe_sB(_, 0, _0{}, _0{}), fragB[0]);  // load_fragment b

diff --git a/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_2_single_buffer.cu b/gemm/cuda_cute/05_pipelining/matmul_smem_and_register_pipelining_2_single_buffer.cu
@@ -171,8 +171,8 @@ MATMUL_KERNEL_SIGNATURE(matmul_smem_and_register_pipelining_2) {
   store_smem_b<NumThreads, SmemShapeN, SmemShapeK, SmemBLoadStoreVec>(sB, staging_b, n, k);
   __syncthreads();
 
-  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), threadIdx.x % (CtaShapeM / ThreadShapeM));
-  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), threadIdx.x / (CtaShapeM / ThreadShapeM));
+  const auto stripe_sA = local_tile(sA, make_tile(Int<ThreadShapeM>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x % (CtaShapeM / ThreadShapeM)));
+  const auto stripe_sB = local_tile(sB, make_tile(Int<ThreadShapeN>{}, Int<SmemShapeK>{}), make_coord(threadIdx.x / (CtaShapeM / ThreadShapeM)));
 
   copy(stripe_sA(_, 0, _0{}), fragA[0]);  // load_fragment a
   copy(stripe_sB(_, 0, _0{}), fragB[0]);  // load_fragment b