[CK TILE] Grouped conv fwd split image

example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp

@@ -13,7 +13,6 @@
 #include "grouped_convolution_utils.hpp"
 #include "grouped_convolution_forward_invoker.hpp"
 #include "run_grouped_convolution_fwd_example.inc"
-
 template <typename GemmWarpConfig>
 int run_grouped_conv_fwd_example(int argc, char* argv[])
 {

example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp

@@ -21,6 +21,10 @@ struct GroupedConvolutionForwardInvoker
     static float grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args,
                                   const ck_tile::stream_config& s)
     {
+        if(s.log_level_ > 0)
+        {
+            std::cout << "[INVOKER] grouped_conv_fwd called, NDimSpatial=" << NDimSpatial << "\n";
+        }
         constexpr int kBlockPerCu = 1;
 
         constexpr ck_tile::index_t M_Tile = 64;
@@ -99,32 +103,40 @@ struct GroupedConvolutionForwardInvoker
                                                                     TilePartitioner,
                                                                     CodegenPipeline,
                                                                     ConvEpilogue>;
-            auto kargs   = Kernel::MakeKernelArgs(args);
 
-            const dim3 grids  = Kernel::GridSize(kargs);
-            const dim3 blocks = Kernel::BlockSize();
+            float ave_time = 0.0f;
 
-            if(!Kernel::IsSupportedArgument(kargs))
-            {
-                throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
-            }
+            // Create kargs and check if split-image is needed
+            auto kargs = Kernel::MakeKernelArgs(args);
+
+            // Check if split-image is needed (uses unified threshold internally)
+            auto split_info = kargs.GetSplitImageInfo();
 
-            if(s.log_level_ > 0)
+            if(!split_info.should_split)
             {
-                std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                          << "shape: " << CodegenShape::GetName() << '\n'
-                          << "problem: " << CodegenPipelineProblem::GetName() << '\n'
-                          << "pipeline: " << CodegenPipeline::GetName() << '\n'
-                          << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
-                          << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
-                          << "}" << '\n'
-                          << "Vector size A: " << CodegenPipeline::GetVectorSizeA()
-                          << ", Vector size B: " << CodegenPipeline::GetVectorSizeB()
-                          << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
+                // No split-image needed - use kargs directly (may have Split-N)
+                if(s.log_level_ > 0)
+                {
+                    std::cout << "[INVOKER] No split-image needed - launching with kargs"
+                              << std::endl;
+                }
+                const dim3 grids  = Kernel::GridSize(kargs);
+                const dim3 blocks = Kernel::BlockSize();
+
+                if(!Kernel::IsSupportedArgument(kargs))
+                {
+                    throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
+                }
+
+                ave_time = ck_tile::launch_kernel(
+                    s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+                return ave_time;
             }
 
-            float ave_time = ck_tile::launch_kernel(
-                s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+            // RECURSIVE split-image path - delegate to transformer helper
+            ave_time = decltype(kargs.transformer_)::template LaunchWithRecursiveSplit<Kernel,
+                                                                                       kBlockPerCu>(
+                args, s, kargs);
 
             return ave_time;
         };

include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp

@@ -73,23 +73,21 @@ struct GroupedConvFwdKernelArgs
         }
         out_ptr = args.out_ptr;
 
-        ConvToGemmFwdTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
-                                                          wei_g_k_c_xs_lengths,
-                                                          out_g_n_k_wos_lengths,
-                                                          conv_filter_strides,
-                                                          conv_filter_dilations,
-                                                          input_left_pads,
-                                                          input_right_pads};
+        // Create and STORE transformer (for split-image support)
+        transformer_ = ConvToGemmFwdTransformer{in_g_n_c_wis_lengths,
+                                                wei_g_k_c_xs_lengths,
+                                                out_g_n_k_wos_lengths,
+                                                conv_filter_strides,
+                                                conv_filter_dilations,
+                                                input_left_pads,
+                                                input_right_pads};
 
         a_grid_desc_m_k =
-            conv_to_gemm_transformer
-                .template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
+            transformer_.template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
         b_grid_desc_n_k =
-            conv_to_gemm_transformer
-                .template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
+            transformer_.template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
         c_grid_desc_m_n =
-            conv_to_gemm_transformer
-                .template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
+            transformer_.template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
 
         group_stride_a = args.C_;
         group_stride_b = args.K_ * args.C_ *
@@ -101,13 +99,15 @@ struct GroupedConvFwdKernelArgs
 
         // Initialize Split-N support fields for 1D convolution (NWGC layout)
         // Get the actual split N from transformer
-        n_per_split = conv_to_gemm_transformer.GetN();
-        original_n  = conv_to_gemm_transformer.GetOriginalN();
+        n_per_split = transformer_.GetN();
+        original_n  = transformer_.GetOriginalN();
         n_splits    = ck_tile::integer_divide_ceil(original_n, n_per_split);
 
-        // Calculate batch strides for NWGC layout
-        input_batch_stride  = args.C_ * args.input_spatial_lengths_[0];
-        output_batch_stride = args.K_ * args.output_spatial_lengths_[0];
+        // FIX: Calculate batch strides using args dimensions
+        // These are the ORIGINAL dimensions passed to constructor, not modified by invoker yet
+        // (invoker modifies args AFTER calling MakeKernelArgs)
+        input_batch_stride  = args.G_ * args.C_ * args.input_spatial_lengths_[0];
+        output_batch_stride = args.G_ * args.K_ * args.output_spatial_lengths_[0];
 
         // Update GemmM to use split N (not original N)
         GemmM = n_per_split * args.output_spatial_lengths_[0];
@@ -163,23 +163,21 @@ struct GroupedConvFwdKernelArgs
         }
         out_ptr = args.out_ptr;
 
-        ConvToGemmFwdTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
-                                                          wei_g_k_c_xs_lengths,
-                                                          out_g_n_k_wos_lengths,
-                                                          conv_filter_strides,
-                                                          conv_filter_dilations,
-                                                          input_left_pads,
-                                                          input_right_pads};
+        // Create and STORE transformer (for split-image support)
+        transformer_ = ConvToGemmFwdTransformer{in_g_n_c_wis_lengths,
+                                                wei_g_k_c_xs_lengths,
+                                                out_g_n_k_wos_lengths,
+                                                conv_filter_strides,
+                                                conv_filter_dilations,
+                                                input_left_pads,
+                                                input_right_pads};
 
         a_grid_desc_m_k =
-            conv_to_gemm_transformer
-                .template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
+            transformer_.template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
         b_grid_desc_n_k =
-            conv_to_gemm_transformer
-                .template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
+            transformer_.template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
         c_grid_desc_m_n =
-            conv_to_gemm_transformer
-                .template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
+            transformer_.template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
 
         group_stride_a = args.C_;
         group_stride_b = args.K_ * args.C_ *
@@ -191,15 +189,16 @@ struct GroupedConvFwdKernelArgs
 
         // Initialize Split-N support fields for 2D convolution (NHWGC layout)
         // Get the actual split N from transformer
-        n_per_split = conv_to_gemm_transformer.GetN();
-        original_n  = conv_to_gemm_transformer.GetOriginalN();
+        n_per_split = transformer_.GetN();
+        original_n  = transformer_.GetOriginalN();
         n_splits    = ck_tile::integer_divide_ceil(original_n, n_per_split);
 
         // Calculate batch strides for NHWGC layout
+        // Need to account for G dimension when moving between batches
         input_batch_stride =
-            args.C_ * args.input_spatial_lengths_[0] * args.input_spatial_lengths_[1];
+            args.G_ * args.C_ * args.input_spatial_lengths_[0] * args.input_spatial_lengths_[1];
         output_batch_stride =
-            args.K_ * args.output_spatial_lengths_[0] * args.output_spatial_lengths_[1];
+            args.G_ * args.K_ * args.output_spatial_lengths_[0] * args.output_spatial_lengths_[1];
 
         // Update GemmM to use split N (not original N)
         GemmM = n_per_split * args.output_spatial_lengths_[0] * args.output_spatial_lengths_[1];
@@ -263,23 +262,21 @@ struct GroupedConvFwdKernelArgs
         }
         out_ptr = args.out_ptr;
 
-        ConvToGemmFwdTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
-                                                          wei_g_k_c_xs_lengths,
-                                                          out_g_n_k_wos_lengths,
-                                                          conv_filter_strides,
-                                                          conv_filter_dilations,
-                                                          input_left_pads,
-                                                          input_right_pads};
+        // Create and STORE transformer (for split-image support)
+        transformer_ = ConvToGemmFwdTransformer{in_g_n_c_wis_lengths,
+                                                wei_g_k_c_xs_lengths,
+                                                out_g_n_k_wos_lengths,
+                                                conv_filter_strides,
+                                                conv_filter_dilations,
+                                                input_left_pads,
+                                                input_right_pads};
 
         a_grid_desc_m_k =
-            conv_to_gemm_transformer
-                .template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
+            transformer_.template MakeADescriptor_M_K<typename GroupedConvTraitsType_::InLayout>();
         b_grid_desc_n_k =
-            conv_to_gemm_transformer
-                .template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
+            transformer_.template MakeBDescriptor_N_K<typename GroupedConvTraitsType_::WeiLayout>();
         c_grid_desc_m_n =
-            conv_to_gemm_transformer
-                .template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
+            transformer_.template MakeCDescriptor_M_N<typename GroupedConvTraitsType_::OutLayout>();
 
         group_stride_a = args.C_;
         group_stride_b = args.K_ * args.C_ *
@@ -291,14 +288,15 @@ struct GroupedConvFwdKernelArgs
 
         // Initialize Split-N support fields for 3D convolution (NDHWGC layout)
         // Get the actual split N from transformer
-        n_per_split = conv_to_gemm_transformer.GetN();
-        original_n  = conv_to_gemm_transformer.GetOriginalN();
+        n_per_split = transformer_.GetN();
+        original_n  = transformer_.GetOriginalN();
         n_splits    = ck_tile::integer_divide_ceil(original_n, n_per_split);
 
         // Calculate batch strides for NDHWGC layout
-        input_batch_stride = args.C_ * args.input_spatial_lengths_[0] *
+        // Need to account for G dimension when moving between batches
+        input_batch_stride = args.G_ * args.C_ * args.input_spatial_lengths_[0] *
                              args.input_spatial_lengths_[1] * args.input_spatial_lengths_[2];
-        output_batch_stride = args.K_ * args.output_spatial_lengths_[0] *
+        output_batch_stride = args.G_ * args.K_ * args.output_spatial_lengths_[0] *
                               args.output_spatial_lengths_[1] * args.output_spatial_lengths_[2];
 
         // Update GemmM to use split N (not original N)
@@ -351,6 +349,19 @@ struct GroupedConvFwdKernelArgs
     index_t original_n          = 1; // Original batch size before splitting
     index_t input_batch_stride  = 0; // Stride to next batch in input tensor
     index_t output_batch_stride = 0; // Stride to next batch in output tensor
+
+    // Split-image support - spatial offsets (applied per-batch in operator())
+    long_index_t spatial_offset_in  = 0; // Spatial offset for input (e.g., W/2 for 1D split)
+    long_index_t spatial_offset_out = 0; // Spatial offset for output (e.g., W/2 for 1D split)
+
+    // Split-image support - transformer instance
+    // We store the transformer so invoker can call CalculateSplitImage()
+    // which uses N_ (after Split-N) for correct offset calculation
+    ConvToGemmFwdTransformer transformer_;
+
+    // Method to get split-image information from transformer
+    // Uses unified TwoGB threshold internally
+    CK_TILE_HOST auto GetSplitImageInfo() const { return transformer_.CalculateSplitImage(); }
 };
 
 /// @brief The Grouped Convolution Forward kernel template.
@@ -460,7 +471,8 @@ struct GroupedConvolutionForwardKernel
     CK_TILE_HOST static constexpr GroupedConvFwdKernelArgsSpecialized
     MakeKernelArgs(const GroupedConvFwdHostArgs& hostArgs)
     {
-        return GroupedConvFwdKernelArgsSpecialized(hostArgs);
+        auto kargs = GroupedConvFwdKernelArgsSpecialized(hostArgs);
+        return kargs;
     }
 
     CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
@@ -821,12 +833,8 @@ struct GroupedConvolutionForwardKernel
     CK_TILE_DEVICE void operator()(GroupedConvFwdKernelArgsSpecialized kargs) const
     {
         const auto blockIdX = amd_wave_read_first_lane(blockIdx.x);
-        const auto [iM, iN] =
-            TilePartitioner{kargs.GemmM, kargs.GemmN}.GetOutputTileIndex(blockIdX);
-        const index_t i_m = amd_wave_read_first_lane(iM * TilePartitioner::MPerBlock);
-        const index_t i_n = amd_wave_read_first_lane(iN * TilePartitioner::NPerBlock);
+        const auto blockIdY = amd_wave_read_first_lane(blockIdx.y);
 
-        const auto blockIdY       = amd_wave_read_first_lane(blockIdx.y);
         const auto group_offset_a = amd_wave_read_first_lane(kargs.group_stride_a * blockIdY);
         const auto group_offset_b = amd_wave_read_first_lane(kargs.group_stride_b * blockIdY);
         const auto group_offset_c = amd_wave_read_first_lane(kargs.group_stride_c * blockIdY);
@@ -844,13 +852,26 @@ struct GroupedConvolutionForwardKernel
             static_cast<long_index_t>(batch_offset) *
             static_cast<long_index_t>(kargs.output_batch_stride);
 
-        // Adjust pointers: combine group offset and batch offset
-        const InDataType* a_ptr =
-            static_cast<const InDataType*>(kargs.in_ptr) + group_offset_a + input_batch_offset;
+        // FIX: Adjust pointers with formula: base + group_offset + batch_offset + spatial_offset
+        // This ensures spatial offset is applied per-batch, not globally
+        const InDataType* base_a_ptr =
+            static_cast<const InDataType*>(kargs.in_ptr) + group_offset_a + input_batch_offset +
+            kargs.spatial_offset_in; // Add spatial offset from split-image
         const WeiDataType* b_ptr = static_cast<const WeiDataType*>(kargs.wei_ptr) +
                                    group_offset_b; // No batch offset for weights!
-        OutDataType* c_ptr =
-            static_cast<OutDataType*>(kargs.out_ptr) + group_offset_c + output_batch_offset;
+        OutDataType* base_c_ptr = static_cast<OutDataType*>(kargs.out_ptr) + group_offset_c +
+                                  output_batch_offset +
+                                  kargs.spatial_offset_out; // Add spatial offset from split-image
+
+        // Use base pointers directly
+        const InDataType* a_ptr = base_a_ptr;
+        OutDataType* c_ptr      = base_c_ptr;
+
+        // Tile partitioning
+        const auto [iM, iN] = TilePartitioner{kargs.GemmM, kargs.GemmN}.GetOutputTileIndex(
+            static_cast<index_t>(blockIdX));
+        const index_t i_m = amd_wave_read_first_lane(iM * TilePartitioner::MPerBlock);
+        const index_t i_n = amd_wave_read_first_lane(iN * TilePartitioner::NPerBlock);
 
         // allocate LDS
         __shared__ char smem_ptr_0[GetSmemSize()];