Use Kbatch when calculate KPadded.

b39cbc7f · Adam Osewski · 6f166ff7 · b39cbc7f
Commit b39cbc7f authored May 11, 2023 by Adam Osewski
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_splitk_c_shuffle.hpp ...tion/gpu/device/impl/device_gemm_xdl_splitk_c_shuffle.hpp +2 -2

No files found.
--- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_splitk_c_shuffle.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_splitk_c_shuffle.hpp
@@ -265,7 +265,7 @@ struct DeviceGemmXdlSplitKCShuffle : public DeviceGemmSplitK<ALayout,
                        StrideC,
                        GridwiseGemm::CalculateMPadded(M),
                        GridwiseGemm::CalculateNPadded(N),
-                        GridwiseGemm::CalculateKPadded(K),
+                        GridwiseGemm::CalculateKPadded(K, KBatch),
                        GridwiseGemm::CalculateK0(K, KBatch),
                        KBatch};
    }
@@ -298,7 +298,7 @@ struct DeviceGemmXdlSplitKCShuffle : public DeviceGemmSplitK<ALayout,
                                          StrideC,
                                          GridwiseGemm::CalculateMPadded(M),
                                          GridwiseGemm::CalculateNPadded(N),
-                                          GridwiseGemm::CalculateKPadded(K),
+                                          GridwiseGemm::CalculateKPadded(K, KBatch),
                                          GridwiseGemm::CalculateK0(K, KBatch),
                                          KBatch);
    }