From 7c8e92facee443c45925100a80f2814370034b92 Mon Sep 17 00:00:00 2001
From: aska-0096 <haocwang@amd.com>
Date: Tue, 3 Sep 2024 07:53:04 +0000
Subject: [PATCH] tempsave

---
 example/01_gemm/gemm_xdl_bf16_v3.cpp                 |  6 +++---
 .../gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp       | 12 +++++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/example/01_gemm/gemm_xdl_bf16_v3.cpp b/example/01_gemm/gemm_xdl_bf16_v3.cpp
index f02456fbb6..934539bc15 100644
--- a/example/01_gemm/gemm_xdl_bf16_v3.cpp
+++ b/example/01_gemm/gemm_xdl_bf16_v3.cpp
@@ -28,13 +28,13 @@ using DeviceGemmV2Instance =
         ADataType,   BDataType,  CDataType,  AccDataType,  CShuffleDataType, 
         PassThrough, PassThrough, PassThrough, GemmDefault, 
         256,
-        128, 128, 
+        224, 256, 
         64, 8, 1,
         16,   16,
-        4,    4,
+        7,    8,
         S<8, 32, 1>,  S<1, 0, 2>,  S<1, 0, 2>, 
         2, 8, 8, 0,
-        S<16, 16, 1>,  S<0, 2, 1>,  S<0, 2, 1>, 
+        S<8, 32, 1>,  S<0, 2, 1>,  S<0, 2, 1>, 
         1, 8, 8, 1,
         1, 2, S<1, 32, 1, 8>, 8,
         ck::BlockGemmPipelineScheduler::Intrawave,ck::BlockGemmPipelineVersion::v3>;
diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
index ff37111c09..3a9eae16b0 100644
--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
@@ -752,12 +752,18 @@ struct GridwiseGemm_xdl_cshuffle_v3
     __device__ static constexpr auto GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()
     {
         // B matrix in LDS memory, dst of blockwise copy
-        if constexpr(BBlockLdsExtraN)
+        // if constexpr(BBlockLdsExtraN)
+        // {
+        //     return make_naive_tensor_descriptor(
+        //         make_tuple(BK0Number, Number<NPerBlock>{}, BK1Number),
+        //         make_tuple(BK1Number, Number<KPerBlock + BBlockLdsExtraN>{}, I1));
+        // }
+        // else
+        if constexpr(BBlockLdsExtraN && is_same<tensor_layout::gemm::RowMajor, BLayout>::value)
         {
             return make_naive_tensor_descriptor(
                 make_tuple(BK0Number, Number<NPerBlock>{}, BK1Number),
-                // make_tuple(BK1Number, Number<KPerBlock + BBlockLdsExtraN>{}, I1));
-                make_tuple(BK1Number, Number<KPerBlock>{}, I1));
+                make_tuple(BK1Number * Number<NPerBlock>{}, I1, Number<NPerBlock>{}));
         }
         else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, BLayout>::value)
         {