From 3784329b68a1900b99f08e771af6f413475cf1e7 Mon Sep 17 00:00:00 2001
From: coderfeli <coderfeli@163.com>
Date: Thu, 26 Dec 2024 13:01:07 +0000
Subject: [PATCH] can run

---
 .../blockwise_gemm_pipeline_xdlops_v3.hpp     | 20 +++++++++----------
 ...hread_group_tensor_slice_transfer_v4r1.hpp |  6 ++++++
 .../threadwise_tensor_slice_transfer_v3r1.hpp |  7 +++++++
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3.hpp
index a1bd1fd1fe..77f548d5ae 100644
--- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3.hpp
+++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3.hpp
@@ -363,15 +363,15 @@ struct BlockwiseGemmXdlops_pipeline_v3<BlockGemmPipelineScheduler::Intrawave,
                     static_for<0, MRepeat, 1>{}([&](auto m0) {
                         static_for<0, NRepeat, 1>{}([&](auto n0) {
                             vector_type<ComputeDataType, KPack> a_thread_vec;
-                            vector_type<ComputeDataType, KPack> b_thread_vec;
-
+                            vector_type<ComputeDataType, KPack> b_thread_vec = 
+                                b_blockwise_copy.template GetSrcThreadScratchIdx<Sequence<0, k0, 0>>();
                             static_for<0, KPack, 1>{}([&](auto ik) {
                                 a_thread_vec.template AsType<ComputeDataType>()(ik) =
                                     a_thread_buf[Number<a_thread_desc_.CalculateOffset(
                                         make_tuple(m0, I0, k0, ik))>{}];
-                                b_thread_vec.template AsType<ComputeDataType>()(ik) =
-                                    b_thread_buf[Number<b_thread_desc_.CalculateOffset(
-                                        make_tuple(n0, I0, k0, ik))>{}];
+                                // b_thread_vec.template AsType<ComputeDataType>()(ik) = 
+                                //     b_thread_buf[Number<b_thread_desc_.CalculateOffset(
+                                //         make_tuple(n0, I0, k0, ik))>{}];
                             });
 
                             using mfma_input_type =
@@ -423,15 +423,15 @@ struct BlockwiseGemmXdlops_pipeline_v3<BlockGemmPipelineScheduler::Intrawave,
                 static_for<0, MRepeat, 1>{}([&](auto m0) {
                     static_for<0, NRepeat, 1>{}([&](auto n0) {
                         vector_type<ComputeDataType, KPack> a_thread_vec;
-                        vector_type<ComputeDataType, KPack> b_thread_vec;
-
+                        vector_type<ComputeDataType, KPack> b_thread_vec = 
+                            b_blockwise_copy.template GetSrcThreadScratchIdx<Sequence<0, k0, 0>>();
                         static_for<0, KPack, 1>{}([&](auto ik) {
                             a_thread_vec.template AsType<ComputeDataType>()(ik) =
                                 a_thread_buf[Number<a_thread_desc_.CalculateOffset(
                                     make_tuple(m0, I0, k0, ik))>{}];
-                            b_thread_vec.template AsType<ComputeDataType>()(ik) =
-                                b_thread_buf[Number<b_thread_desc_.CalculateOffset(
-                                    make_tuple(n0, I0, k0, ik))>{}];
+                            // b_thread_vec.template AsType<ComputeDataType>()(ik) =
+                            //     b_thread_buf[Number<b_thread_desc_.CalculateOffset(
+                            //         make_tuple(n0, I0, k0, ik))>{}];
                         });
 
                         using mfma_input_type =
diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp
index 2c5fbc3937..218e409543 100644
--- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp
+++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp
@@ -109,6 +109,12 @@ struct ThreadGroupTensorSliceTransfer_v4r1
         }
     }
 
+    template <typename SeqIdx>
+    __device__ auto GetSrcThreadScratchIdx()
+    {
+       return threadwise_transfer_.template GetSrcThreadScratchIdx<SeqIdx>();
+    }
+
     template <typename SrcBuffer, index_t ThreadScratchId = 0>
     __device__ void RunRead(const SrcDesc& src_desc,
                             const SrcBuffer& src_buf,
diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp
index 96ea04c8fa..4a73f77eb6 100644
--- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp
+++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v3r1.hpp
@@ -267,6 +267,13 @@ struct ThreadwiseTensorSliceTransfer_v3r1
         }
     }
 
+    template <typename SeqIdx, index_t ThreadScratchId = 0>
+    __device__ auto GetSrcThreadScratchIdx(Number<ThreadScratchId> thread_scratch_id = Number<ThreadScratchId>{})
+    {
+        using vector_t = typename vector_type_maker<SrcData, SrcScalarPerVector>::type::type;
+        return src_thread_scratch_tuple_(thread_scratch_id).template GetAsType<vector_t>(SeqIdx{});
+    }
+    
     template <index_t ThreadScratchId>
     __device__ void
     TransferDataFromSrcThreadScratchToDstThreadScratch(Number<ThreadScratchId> thread_scratch_id)