CUTLASS 2.6 (#298)

CUTLASS 2.6

include/cutlass/gemm/kernel/gemm_splitk_parallel.h

@@ -63,6 +63,7 @@ struct GemmSplitKParallel {
   struct Params {
     cutlass::gemm::GemmCoord problem_size;
     cutlass::gemm::GemmCoord grid_tiled_shape;
+    int swizzle_log_tile;
     typename Mma::IteratorA::Params params_A;
     typename Mma::IteratorA::TensorRef ref_A;
     typename Mma::IteratorB::Params params_B;
@@ -78,7 +79,7 @@ struct GemmSplitKParallel {
     //
 
     CUTLASS_HOST_DEVICE
-    Params() { }
+    Params(): swizzle_log_tile(0) { }
 
     CUTLASS_HOST_DEVICE
     Params(
@@ -92,6 +93,7 @@ struct GemmSplitKParallel {
     ):
       problem_size(problem_size),
       grid_tiled_shape(grid_tiled_shape),
+      swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)),
       params_A(ref_A.layout()),
       ref_A(ref_A),
       params_B(ref_B.layout()),
@@ -129,7 +131,7 @@ struct GemmSplitKParallel {
     ThreadblockSwizzle threadblock_swizzle;
 
     cutlass::gemm::GemmCoord threadblock_tile_offset =
-        threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
     // Early exit if CTA is out of range
     if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
@@ -207,7 +209,7 @@ struct GemmSplitKParallel {
     //
 
     threadblock_tile_offset =
-        threadblock_swizzle.get_tile_offset(params.grid_tiled_shape);
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
     //assume identity swizzle
     MatrixCoord threadblock_offset(
@@ -243,4 +245,3 @@ struct GemmSplitKParallel {
 } // namespace kernel
 } // namespace gemm
 } // namespace cutlass
-