From bbfe4501fa9f45edd38f4134b344fc231e087167 Mon Sep 17 00:00:00 2001
From: Mohsen Saffari <mohsen.saffari@amd.com>
Date: Mon, 20 Oct 2025 14:43:32 +0000
Subject: [PATCH] Add complete multi-dimensional stride support via descriptors

---
 .../batched_contraction.cpp                   |   3 +-
 .../kernel/batched_contraction_kernel.hpp     | 111 ++++++++++--------
 2 files changed, 64 insertions(+), 50 deletions(-)

diff --git a/example/ck_tile/41_batched_contraction/batched_contraction.cpp b/example/ck_tile/41_batched_contraction/batched_contraction.cpp
index ea78f09dff..20b25fc99a 100644
--- a/example/ck_tile/41_batched_contraction/batched_contraction.cpp
+++ b/example/ck_tile/41_batched_contraction/batched_contraction.cpp
@@ -219,9 +219,8 @@ float batched_contraction(const ck_tile::BatchedContractionHostArgs<DsDataType::
     HANDLE_CASE(2, 1, 1, 1);
     HANDLE_CASE(2, 2, 2, 1);
     HANDLE_CASE(1, 2, 1, 1);
-    HANDLE_CASE(1, 1, 1, 2);
+    HANDLE_CASE(2, 1, 1, 1);
     HANDLE_CASE(2, 2, 2, 2);
-    HANDLE_CASE(4, 4, 4, 4);
 
     throw std::runtime_error(
         "Unsupported dimension combination: G=" + std::to_string(num_g_dims) +
diff --git a/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp b/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp
index 0d2ec7a624..a93c6f5064 100644
--- a/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp
+++ b/include/ck_tile/ops/batched_contraction/kernel/batched_contraction_kernel.hpp
@@ -232,7 +232,8 @@ struct BatchedContractionKernelArgs
     AGridDesc_M_K_ a_grid_desc_m_k; ///< Tensor descriptor for A[M, K] with actual strides
     BGridDesc_N_K_ b_grid_desc_n_k; ///< Tensor descriptor for B[N, K] with actual strides
     EGridDesc_M_N_ e_grid_desc_m_n; ///< Tensor descriptor for E[M, N] with actual strides
-    // TODO: Add D descriptors array when needed
+    std::array<EGridDesc_M_N_, NumDTensor>
+        ds_grid_desc_m_n; ///< Descriptors for D tensors (same shape as E, independent strides)
 };
 
 /// @brief GPU kernel for batched tensor contraction operations.
@@ -356,15 +357,14 @@ struct BatchedContractionKernel
     /// non-contiguous support
     /// @details This function creates tensor views from descriptors and runs GEMM pipeline,
     ///          similar to UniversalGemmKernel::RunGemm but with descriptor-based tensor views
-    CK_TILE_DEVICE static void
-    RunGemm(const ADataType* a_ptr,
-            const BDataType* b_ptr,
-            [[maybe_unused]] const std::array<const void*, NumDTensor>& ds_ptr,
-            EDataType* e_ptr,
-            void* smem_ptr,
-            const KernelArgs& kargs,
-            const index_t i_m,
-            const index_t i_n)
+    CK_TILE_DEVICE static void RunGemm(const ADataType* a_ptr,
+                                       const BDataType* b_ptr,
+                                       const std::array<const void*, NumDTensor>& ds_ptr,
+                                       EDataType* e_ptr,
+                                       void* smem_ptr,
+                                       const KernelArgs& kargs,
+                                       const index_t i_m,
+                                       const index_t i_n)
     {
         // Create tensor views from descriptors (handles multi-dimensional strides)
         auto a_tensor_view =
@@ -401,11 +401,23 @@ struct BatchedContractionKernel
         const auto& c_block_tile = GemmPipeline{}(
             a_block_window, AElementWise{}, b_block_window, BElementWise{}, num_loop, smem_ptr);
 
-        // Create empty D windows tuple (for NumDTensor=0 case)
-        // TODO: Create actual D windows from descriptors when NumDTensor > 0
-        auto ds_block_windows = make_tuple();
+        // Create D windows from descriptors (for each D tensor)
+        auto ds_block_windows = generate_tuple(
+            [&](auto i) {
+                using DDataType        = remove_cvref_t<std::tuple_element_t<i.value, DsDataType>>;
+                const DDataType* d_ptr = static_cast<const DDataType*>(ds_ptr[i]);
 
-        // Run Epilogue Pipeline (same as UniversalGemmKernel)
+                auto d_tensor_view =
+                    make_tensor_view<address_space_enum::global>(d_ptr, kargs.ds_grid_desc_m_n[i]);
+
+                return make_tile_window(d_tensor_view,
+                                        make_tuple(number<TilePartitioner::MPerBlock>{},
+                                                   number<TilePartitioner::NPerBlock>{}),
+                                        {i_m, i_n});
+            },
+            number<NumDTensor>{});
+
+        // Run Epilogue Pipeline with descriptor-based D windows
         EpiloguePipeline{}(e_block_window, c_block_tile, ds_block_windows, smem_ptr);
     }
 
@@ -526,6 +538,13 @@ struct BatchedContractionKernel
         kargs.e_grid_desc_m_n =
             DescriptorUtils::Make_E_GridDescriptor_M_N(host_args.E_dims, host_args.E_strides);
 
+        // Create D descriptors with their own strides (same shape as E, independent strides)
+        for(ck_tile::index_t d = 0; d < NumDTensor; ++d)
+        {
+            kargs.ds_grid_desc_m_n[d] = DescriptorUtils::Make_E_GridDescriptor_M_N(
+                host_args.Ds_dims[d], host_args.Ds_strides[d]);
+        }
+
         // Keep simple strides for backward compatibility
         kargs.stride_A = kargs.K_total;
         kargs.stride_B = kargs.K_total;
@@ -582,43 +601,39 @@ struct BatchedContractionKernel
         // Allocate shared memory
         __shared__ char smem_ptr[GetSmemSize()];
 
-        if constexpr(NumDTensor > 0)
-        {
-            // OLD PATH: Use UniversalGemmKernel for cases with D tensors
-            typename UniversalGemmKernel::KernelArgs gemm_kargs{{a_ptr},
-                                                                {b_ptr},
-                                                                ds_batch_ptr,
-                                                                e_ptr,
-                                                                kargs.M_total,
-                                                                kargs.N_total,
-                                                                kargs.K_total,
-                                                                {kargs.stride_A},
-                                                                {kargs.stride_B},
-                                                                kargs.stride_Ds,
-                                                                kargs.stride_E,
-                                                                kargs.k_batch};
+#if 0 // OLD PATH: UniversalGemmKernel (kept for reference, can be deleted after full validation)
+        typename UniversalGemmKernel::KernelArgs gemm_kargs{{a_ptr},
+                                                            {b_ptr},
+                                                            ds_batch_ptr,
+                                                            e_ptr,
+                                                            kargs.M_total,
+                                                            kargs.N_total,
+                                                            kargs.K_total,
+                                                            {kargs.stride_A},
+                                                            {kargs.stride_B},
+                                                            kargs.stride_Ds,
+                                                            kargs.stride_E,
+                                                            kargs.k_batch};
 
-            const typename UniversalGemmKernel::SplitKBatchOffset splitk_batch_offset(gemm_kargs,
-                                                                                      i_splitk);
+        const typename UniversalGemmKernel::SplitKBatchOffset splitk_batch_offset(gemm_kargs,
+                                                                                  i_splitk);
 
-            const ADataType* a_ptr_final = a_ptr + splitk_batch_offset.as_k_split_offset[0];
-            const BDataType* b_ptr_final = b_ptr + splitk_batch_offset.bs_k_split_offset[0];
+        const ADataType* a_ptr_final = a_ptr + splitk_batch_offset.as_k_split_offset[0];
+        const BDataType* b_ptr_final = b_ptr + splitk_batch_offset.bs_k_split_offset[0];
 
-            UniversalGemmKernel::RunGemm({a_ptr_final},
-                                         {b_ptr_final},
-                                         ds_batch_ptr,
-                                         e_ptr,
-                                         smem_ptr,
-                                         gemm_kargs,
-                                         splitk_batch_offset,
-                                         i_m,
-                                         i_n);
-        }
-        else
-        {
-            // NEW PATH: Use descriptor-based RunGemm for num_d=0
-            RunGemm(a_ptr, b_ptr, ds_batch_ptr, e_ptr, smem_ptr, kargs, i_m, i_n);
-        }
+        UniversalGemmKernel::RunGemm({a_ptr_final},
+                                     {b_ptr_final},
+                                     ds_batch_ptr,
+                                     e_ptr,
+                                     smem_ptr,
+                                     gemm_kargs,
+                                     splitk_batch_offset,
+                                     i_m,
+                                     i_n);
+#else // NEW PATH: Descriptor-based RunGemm
+      // custom descriptor-based RunGemm with full multi-dimensional stride support
+        RunGemm(a_ptr, b_ptr, ds_batch_ptr, e_ptr, smem_ptr, kargs, i_m, i_n);
+#endif
     }
 };