add example for preshuffle

example/65_gemm_multiply_multiply/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_example_executable(example_gemm_multiply_multiply_xdl_fp8 gemm_multiply_multiply_xdl_fp8.cpp)
 add_example_executable(example_gemm_multiply_multiply_xdl_fp8_ab_scale gemm_multiply_multiply_xdl_fp8_ab_scale.cpp)
 add_example_executable(example_gemm_multiply_multiply_xdl_fp8_bpreshuffle gemm_multiply_multiply_xdl_fp8_bpreshuffle.cpp)
+add_example_executable(example_gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp)
 add_example_executable(example_gemm_add_add_xdl_fp16 gemm_add_add_xdl_fp16.cpp)
 add_example_executable(example_gemm_multiply_multiply_xdl_int8 gemm_multiply_multiply_xdl_int8.cpp)
 add_example_executable(example_moe_gemm1_xdl_fp8 moe_gemm1_xdl_fp8.cpp)

example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp

@@ -0,0 +1,472 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
+
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/utility/literals.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
+#include "ck/library/utility/check_err.hpp"
+
+#include "ck/utility/blkgemmpipe_scheduler.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16  = ck::half_t;
+using BF16 = ck::bhalf_t;
+using FP8  = ck::f8_t;
+using F32  = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+using A0DataType       = FP8;
+using B0DataType       = FP8;
+using AccDataType      = F32;
+using CShuffleDataType = F32;
+using D0DataType       = F32;
+using D1DataType       = F32;
+using DsDataType       = ck::Tuple<D0DataType, D1DataType>;
+using EDataType        = F16;
+
+using A0Layout = Row;
+using B0Layout = Col;
+using D0Layout = Row;
+using D1Layout = Col;
+using DsLayout = ck::Tuple<D0Layout, D1Layout>;
+using ELayout  = Row;
+
+struct MultiplyMultiply
+{
+    template <typename E, typename C, typename D0, typename D1>
+    __host__ __device__ constexpr void
+    operator()(E& e, const C& c, const D0& d0, const D1& d1) const;
+
+    template <>
+    __host__ __device__ constexpr void operator()<F16, float, float, float>(F16& e,
+                                                                            const float& c,
+                                                                            const float& d0,
+                                                                            const float& d1) const
+    {
+        const float x0_f = c * d0 * d1;
+
+        e = ck::type_convert<F16>(x0_f);
+    }
+
+    template <>
+    __host__ __device__ constexpr void operator()<BF16, float, float, float>(BF16& e,
+                                                                             const float& c,
+                                                                             const float& d0,
+                                                                             const float& d1) const
+    {
+        const float x0_f = c * d0 * d1;
+
+        e = ck::type_convert<BF16>(x0_f);
+    }
+
+    template <>
+    __host__ __device__ constexpr void operator()<ck::half_t, int, float, float>(
+        ck::half_t& e, const int& c, const float& d0, const float& d1) const
+    {
+        const float x0_f =
+            ck::type_convert<float>(c) * ck::type_convert<float>(d0) * ck::type_convert<float>(d1);
+
+        e = ck::type_convert<ck::half_t>(x0_f);
+    }
+
+    template <>
+    __host__ __device__ constexpr void operator()<ck::bhalf_t, int, float, float>(
+        ck::bhalf_t& e, const int& c, const float& d0, const float& d1) const
+    {
+        const float x0_f =
+            ck::type_convert<float>(c) * ck::type_convert<float>(d0) * ck::type_convert<float>(d1);
+
+        e = ck::type_convert<ck::bhalf_t>(x0_f);
+    }
+};
+
+template <typename DataType>
+void preShuffleBuffer(const DataType* src, DataType* dst, int N, int K, int NXdl, int Knew)
+{
+    int KPack = 16;
+    int NLane = NXdl;
+    int KLane = 64 / NLane;
+
+    int K0 = Knew / (KLane * KPack);
+    // K -> K0 KLane KPack
+    // N -> N0 NLane
+    // N, K -> N0 K0 KLane NLane KPack
+    int tempk;
+    for(int n = 0; n < N; ++n)
+    {
+        for(int k = 0; k < K; ++k)
+        {
+            int n0 = n / NLane;
+            int n1 = n % NLane;
+
+            int k0 = k / (KLane * KPack);
+            tempk  = k % (KLane * KPack);
+            int k1 = tempk / KPack;
+            int k2 = tempk % KPack;
+
+            int outputIndex = n0 * KPack * NLane * KLane * K0 + k0 * KPack * NLane * KLane +
+                              k1 * KPack * NLane + n1 * KPack + k2;
+
+            dst[outputIndex] = src[n * K + k];
+        }
+    }
+}
+
+int GetKPreShufflePadded(int K)
+{
+    return (K + KShufflePadded - 1) / KShufflePadded * KShufflePadded;
+}
+int GetNPreShufflePadded(int N)
+{
+    return (N + NShufflePadded - 1) / NShufflePadded * NShufflePadded;
+}
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+using AElementOp   = PassThrough;
+using BElementOp   = PassThrough;
+using CDEElementOp = MultiplyMultiply;
+
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::NKPadding;
+
+using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
+    // clang-format off
+    // <   Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
+    //     AElementOp,  BElementOp, CDEElementOp, GemmSpec, 256,
+    //     256,   256,    128,
+    //     16,   16,
+    //     16,   16,
+    //     8,    8,
+    //     S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
+    //     S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
+    //     1,    2,   S<1, 32, 1, 8>, S<8, 8, 1>,
+    //     ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>;
+
+    //  <  Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
+    //     AElementOp,  BElementOp, CDEElementOp, GemmSpec, 256,
+    //     128,    128,   128,
+    //     16,  16,
+    //     32,  32,
+    //     4,   1,
+    //     S<8, 32, 1>,    S<1, 0, 2>,    S<1, 0, 2>,               2,             16,             16,          0,
+    //     S<8, 32, 1>,     S<1, 0, 2>,    S<1, 0, 2>,               2,             16,             16,          0,
+    //     1,           1,                   S<1, 32, 1, 8>,    S<8, 8, 1>,
+    //     ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>;
+
+    <  Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
+          AElementOp,  BElementOp, CDEElementOp, GemmSpec, 256,
+          64,    512,   128,
+          16,  16,
+          32,   32,
+          2,    4,
+          S<8, 32, 1>,     S<1, 0, 2>,    S<1, 0, 2>,               2,             16,             16,          0,
+          S<8, 32, 1>,     S<1, 0, 2>,    S<1, 0, 2>,               2,             16,             16,          0,
+          1,           1,                   S<1, 32, 1, 8>,     S<8, 8, 1>,
+    ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>;
+
+// clang-format on
+
+int main(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = false;
+
+    // GEMM shape
+    ck::index_t M = 3840;
+    ck::index_t N = 4096;
+    ck::index_t K = 4096;
+
+    ck::index_t StrideA = K;
+    ck::index_t StrideB = K;
+    ck::index_t StrideD = 0;
+    ck::index_t StrideE = N;
+
+    ck::index_t KBatch = 1;
+
+    ck::index_t Warmup = 50;
+    ck::index_t Repeat = 50;
+
+    if(argc == 1)
+    {
+        // use default case
+    }
+    else if(argc == 4)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc == 12)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        M = std::stoi(argv[4]);
+        N = std::stoi(argv[5]);
+        K = std::stoi(argv[6]);
+
+        StrideA = std::stoi(argv[7]);
+        StrideB = std::stoi(argv[8]);
+        StrideD = std::stoi(argv[9]);
+        StrideE = std::stoi(argv[10]);
+
+        KBatch = std::stoi(argv[11]);
+    }
+    else if(argc == 14)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        M = std::stoi(argv[4]);
+        N = std::stoi(argv[5]);
+        K = std::stoi(argv[6]);
+
+        StrideA = std::stoi(argv[7]);
+        StrideB = std::stoi(argv[8]);
+        StrideD = std::stoi(argv[9]);
+        StrideE = std::stoi(argv[10]);
+
+        KBatch = std::stoi(argv[11]);
+
+        Warmup = std::stoi(argv[12]);
+        Repeat = std::stoi(argv[13]);
+    }
+    else
+    {
+        printf("arg1: verification (0=no, 1=yes)\n");
+        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
+        printf("arg3: time kernel (0=no, 1=yes)\n");
+        printf(
+            "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE, KBatch\n");
+        printf("arg10 to 11: Warmup, Repeat\n");
+        exit(0);
+    }
+
+    auto f_host_tensor_descriptor =
+        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
+            using namespace ck::literals;
+
+            if(std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value)
+            {
+                return HostTensorDescriptor({row, col}, {stride, 1_uz});
+            }
+            else
+            {
+                return HostTensorDescriptor({row, col}, {1_uz, stride});
+            }
+        };
+
+    auto Knew       = GetKPreShufflePadded(K);
+    auto StrideBnew = Knew;
+    auto Nnew       = GetNPreShufflePadded(N);
+    std::cout << "Knew: " << Knew << " Nnew: " << Nnew << std::endl;
+    Tensor<A0DataType> a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{}));
+    Tensor<B0DataType> b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{}));
+    Tensor<B0DataType> b0_preshuffled(
+        f_host_tensor_descriptor(Knew, Nnew, StrideBnew, B0Layout{})); // use laout only for size
+    Tensor<D0DataType> d0_m_n(f_host_tensor_descriptor(M, N, StrideD, D0Layout{}));
+    Tensor<D1DataType> d1_m_n(f_host_tensor_descriptor(M, N, StrideD, D1Layout{}));
+    Tensor<EDataType> e_m_n_host_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{}));
+    Tensor<EDataType> e_m_n_device_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{}));
+
+    std::cout << "a0_m_k: " << a0_m_k.mDesc << std::endl;
+    std::cout << "b0_k_n: " << b0_k_n.mDesc << std::endl;
+    std::cout << "d1_m_n: " << d1_m_n.mDesc << std::endl;
+    std::cout << "d0_m_n: " << d0_m_n.mDesc << std::endl;
+    std::cout << "e_m_n: " << e_m_n_host_result.mDesc << std::endl;
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        a0_m_k.GenerateTensorValue(GeneratorTensor_2<A0DataType>{-5, 5});
+        b0_k_n.GenerateTensorValue(GeneratorTensor_2<B0DataType>{-5, 5});
+        d0_m_n.GenerateTensorValue(GeneratorTensor_2<D0DataType>{-2, 2});
+        d1_m_n.GenerateTensorValue(GeneratorTensor_2<D1DataType>{-2, 2});
+        break;
+    case 2:
+        a0_m_k.GenerateTensorValue(GeneratorTensor_1<A0DataType>{});
+        b0_k_n.GenerateTensorValue(GeneratorTensor_1<B0DataType>{});
+        d0_m_n.GenerateTensorValue(GeneratorTensor_1<D0DataType>{});
+        d1_m_n.GenerateTensorValue(GeneratorTensor_1<D1DataType>{});
+        break;
+    default:
+        a0_m_k.GenerateTensorValue(GeneratorTensor_3<A0DataType>{0.0, 1.0});
+        b0_k_n.GenerateTensorValue(GeneratorTensor_3<B0DataType>{-0.5, 0.5});
+        d0_m_n.GenerateTensorValue(GeneratorTensor_3<D0DataType>{0.0, 1.0});
+        d1_m_n.GenerateTensorValue(GeneratorTensor_3<D1DataType>{0.0, 1.0});
+    }
+    DeviceMem a0_device_buf(sizeof(A0DataType) * a0_m_k.mDesc.GetElementSpaceSize());
+    DeviceMem b0_device_buf(sizeof(B0DataType) * b0_preshuffled.mDesc.GetElementSpaceSize());
+    DeviceMem d0_device_buf(sizeof(D0DataType) * d0_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem d1_device_buf(sizeof(D1DataType) * d1_m_n.mDesc.GetElementSpaceSize());
+    DeviceMem e_device_buf(sizeof(EDataType) * e_m_n_device_result.mDesc.GetElementSpaceSize());
+
+    a0_device_buf.ToDevice(a0_m_k.mData.data());
+    d0_device_buf.ToDevice(d0_m_n.mData.data());
+    d1_device_buf.ToDevice(d1_m_n.mData.data());
+    e_device_buf.ToDevice(e_m_n_device_result.mData.data());
+
+    auto a_element_op   = AElementOp{};
+    auto b_element_op   = BElementOp{};
+    auto cde_element_op = CDEElementOp{};
+
+    constexpr ck::index_t NumDTensor = DsDataType::Size();
+
+    constexpr auto I0 = ck::Number<0>{};
+
+    // do GEMM
+    auto device_op = DeviceOpInstance{};
+    int NPerXdl    = device_op.GetPreShuffleParameters();
+
+#if 0
+{   //test shuffle result
+    auto ouput_matirx=[](auto data, ck::index_t N0, ck::index_t K0){
+        std::cout << std::endl;
+        int ii = 0;
+        for(int n = 0; n < N0; n++)
+        {
+            std::cout << ii++ << " line: ";
+            for(int k = 0; k < K0; k++)
+            {
+                std::cout << data(k,n) << " ";
+                // std::cout << ck::type_convert<float>(data.mData[n*K0 + k]) << " ";
+            }
+            std::cout << std::endl;
+        }
+    };
+
+    Tensor<int> b0_k_n2(f_host_tensor_descriptor(K, N, StrideB, B0Layout{}));
+    Tensor<int> b0_preshuffled2(
+        f_host_tensor_descriptor(Knew, N, StrideBnew, B0Layout{})); // use laout only for size
+    int nCount = 0;
+    for(int n = 0; n < N; n++)
+    {
+        for(int k = 0; k < K; k++)
+        {
+            b0_k_n2(k,n) =  nCount++;
+        }
+    }
+    ouput_matirx(b0_k_n2, N, K);
+    b0_preshuffled2.SetZero();
+    preShuffleBuffer(b0_k_n2.mData.data(), b0_preshuffled2.mData.data(), N, K, NPerXdl, Knew);
+    std::cout << "after shuffle" << std::endl;
+    ouput_matirx(b0_preshuffled2, N, Knew);
+}
+#endif
+
+    b0_preshuffled.SetZero(); // must set to zero
+    preShuffleBuffer(b0_k_n.mData.data(), b0_preshuffled.mData.data(), N, K, NPerXdl, Knew);
+
+    b0_device_buf.ToDevice(b0_preshuffled.mData.data());
+
+    auto invoker = device_op.MakeInvoker();
+    auto argument =
+        device_op.MakeArgument(a0_device_buf.GetDeviceBuffer(),
+                               b0_device_buf.GetDeviceBuffer(),
+                               std::array<const void*, NumDTensor>{d0_device_buf.GetDeviceBuffer(),
+                                                                   d1_device_buf.GetDeviceBuffer()},
+                               e_device_buf.GetDeviceBuffer(),
+                               M,
+                               N,
+                               K,
+                               StrideA,
+                               StrideB,
+                               std::array<ck::index_t, NumDTensor>{I0, I0},
+                               StrideE,
+                               KBatch,
+                               a_element_op,
+                               b_element_op,
+                               cde_element_op,
+                               Nnew,
+                               Knew);
+
+    if(!device_op.IsSupportedArgument(argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_gemm with the specified compilation parameters does "
+            "not support this GEMM problem");
+    }
+
+    size_t total_size =
+        (M * K * sizeof(A0DataType) + N * K * sizeof(B0DataType) + M * sizeof(D0DataType) +
+         N * sizeof(D1DataType) + M * N * sizeof(EDataType));
+    int rotate_buf_num =
+        ck::math::min(size_t(Repeat), ck::math::integer_divide_ceil(512 * 1024 * 1024, total_size));
+
+    float ave_time = invoker.Run(
+        argument, StreamConfig{nullptr, time_kernel, 0, Warmup, Repeat, true, rotate_buf_num});
+
+    std::size_t flop = std::size_t(2) * M * N * K;
+    std::size_t num_btype =
+        sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(EDataType) * M * N;
+
+    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
+              << std::endl;
+
+    if(do_verification)
+    {
+        invoker.Run(argument, StreamConfig{nullptr, false});
+
+        e_device_buf.FromDevice(e_m_n_device_result.mData.data());
+
+        Tensor<CShuffleDataType> c_m_n({M, N});
+
+        using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<A0DataType,
+                                                                                B0DataType,
+                                                                                CShuffleDataType,
+                                                                                AccDataType,
+                                                                                PassThrough,
+                                                                                PassThrough,
+                                                                                PassThrough>;
+        auto ref_gemm               = ReferenceGemmInstance{};
+        auto ref_invoker            = ref_gemm.MakeInvoker();
+
+        auto ref_argument = ref_gemm.MakeArgument(
+            a0_m_k, b0_k_n, c_m_n, PassThrough{}, PassThrough{}, PassThrough{});
+
+        ref_invoker.Run(ref_argument);
+
+        for(int m = 0; m < M; ++m)
+        {
+            for(int n = 0; n < N; ++n)
+            {
+                cde_element_op(e_m_n_host_result(m, n), c_m_n(m, n), d0_m_n(m, n), d1_m_n(m, n));
+            }
+        }
+
+        e_device_buf.FromDevice(e_m_n_device_result.mData.data());
+
+        return ck::utils::check_err(
+                   e_m_n_device_result, e_m_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2)
+                   ? 0
+                   : 1;
+    }
+
+    return 0;
+}

include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp

@@ -125,7 +125,6 @@ struct DeviceGemmMultipleDSplitKBPreShuffle : public BaseOperator
 {
     static constexpr index_t NumDTensor = DsDataType::Size();
 
-#ifndef CK_CODE_GEN_RTC
     virtual std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const void* p_a,
                         const void* p_b,
@@ -141,14 +140,16 @@ struct DeviceGemmMultipleDSplitKBPreShuffle : public BaseOperator
                         ck::index_t KBatch,
                         AElementwiseOperation a_element_op,
                         BElementwiseOperation b_element_op,
-                        CDEElementwiseOperation cde_element_op) = 0;
+                        CDEElementwiseOperation cde_element_op,
+                        ck::index_t Nr,
+                        ck::index_t Kr) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 
     virtual int GetPreShuffleParameters() = 0;
-#endif
 };
-
+#define KShufflePadded 256
+#define NShufflePadded 128
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck

include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp

@@ -138,6 +138,48 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
         LDSTypeB>;
 
     using Argument = typename GridwiseGemm::Argument;
+    struct DeviceArgument : public Argument
+    {
+        __host__ DeviceArgument(const ADataType* p_a_grid_,
+                                const BDataType* p_b_grid_,
+                                std::array<const void*, NumDTensor> p_ds_grid_,
+                                CDataType* p_c_grid_,
+                                index_t M_,
+                                index_t N_,
+                                index_t K_,
+                                index_t StrideA_,
+                                index_t StrideB_,
+                                std::array<index_t, NumDTensor> StrideDs_,
+                                index_t StrideC_,
+                                index_t k_batch_,
+                                AElementwiseOperation a_element_op_,
+                                BElementwiseOperation b_element_op_,
+                                CElementwiseOperation c_element_op_,
+                                index_t Nr_,
+                                index_t Kr_)
+            : Argument{p_a_grid_,
+                       p_b_grid_,
+                       p_ds_grid_,
+                       p_c_grid_,
+                       M_,
+                       N_,
+                       K_,
+                       StrideA_,
+                       StrideB_,
+                       StrideDs_,
+                       StrideC_,
+                       k_batch_,
+                       a_element_op_,
+                       b_element_op_,
+                       c_element_op_},
+              Nr{Nr_},
+              Kr{Kr_}
+        {
+        }
+
+        index_t Nr;
+        index_t Kr;
+    };
 
     int GetPreShuffleParameters() override { return NPerXDL; }
 
@@ -518,19 +560,34 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
             return false;
         }
 
-        if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding ||
-                                                       GemmSpec == GemmSpecialization::NKPadding ||
-                                                       GemmSpec == GemmSpecialization::MNKPadding ||
-                                                       GemmSpec == GemmSpecialization::KPadding))
+        constexpr bool KPadding = GemmSpec == GemmSpecialization::MKPadding ||
+                                  GemmSpec == GemmSpecialization::NKPadding ||
+                                  GemmSpec == GemmSpecialization::MNKPadding ||
+                                  GemmSpec == GemmSpecialization::KPadding;
+        constexpr bool NPadding = GemmSpec == GemmSpecialization::NPadding ||
+                                  GemmSpec == GemmSpecialization::NKPadding ||
+                                  GemmSpec == GemmSpecialization::MNKPadding;
+
+        if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !KPadding)
         {
             return false;
         }
 
-        if(arg.N % NPerBlock != 0 || arg.K % KPerBlock != 0)
+        if((arg.N % NPerBlock != 0 && !NPadding) || (arg.K % KPerBlock != 0 && !KPadding))
         {
             return false;
         }
 
+        const auto karg = dynamic_cast<const DeviceArgument*>(&arg);
+        if(NPadding && (karg->Nr != GridwiseGemm::CalculateBNShufflePadded(arg.N)))
+        {
+            return false;
+        }
+
+        if(KPadding && (karg->Kr != GridwiseGemm::CalculateBKShufflePadded(arg.K)))
+        {
+            return false;
+        }
         return GridwiseGemm::CheckValidity(arg);
     }
 
@@ -554,23 +611,27 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
                              index_t KBatch,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
-                             CElementwiseOperation c_element_op)
+                             CElementwiseOperation c_element_op,
+                             index_t Nr,
+                             index_t Kr)
     {
-        return Argument{static_cast<const ADataType*>(p_a),
-                        static_cast<const BDataType*>(p_b),
-                        p_ds,
-                        static_cast<CDataType*>(p_c),
-                        M,
-                        N,
-                        K,
-                        StrideA,
-                        StrideB,
-                        StrideDs,
-                        StrideC,
-                        KBatch,
-                        a_element_op,
-                        b_element_op,
-                        c_element_op};
+        return DeviceArgument{static_cast<const ADataType*>(p_a),
+                              static_cast<const BDataType*>(p_b),
+                              p_ds,
+                              static_cast<CDataType*>(p_c),
+                              M,
+                              N,
+                              K,
+                              StrideA,
+                              StrideB,
+                              StrideDs,
+                              StrideC,
+                              KBatch,
+                              a_element_op,
+                              b_element_op,
+                              c_element_op,
+                              Nr,
+                              Kr};
     }
 
     static auto MakeInvoker() { return Invoker{}; }
@@ -590,23 +651,27 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
                                                       index_t KBatch,
                                                       AElementwiseOperation a_element_op,
                                                       BElementwiseOperation b_element_op,
-                                                      CElementwiseOperation c_element_op) override
+                                                      CElementwiseOperation c_element_op,
+                                                      index_t Nr,
+                                                      index_t Kr) override
     {
-        return std::make_unique<Argument>(static_cast<const ADataType*>(p_a),
-                                          static_cast<const BDataType*>(p_b),
-                                          p_ds,
-                                          static_cast<CDataType*>(p_c),
-                                          M,
-                                          N,
-                                          K,
-                                          StrideA,
-                                          StrideB,
-                                          StrideDs,
-                                          StrideC,
-                                          KBatch,
-                                          a_element_op,
-                                          b_element_op,
-                                          c_element_op);
+        return std::make_unique<DeviceArgument>(static_cast<const ADataType*>(p_a),
+                                                static_cast<const BDataType*>(p_b),
+                                                p_ds,
+                                                static_cast<CDataType*>(p_c),
+                                                M,
+                                                N,
+                                                K,
+                                                StrideA,
+                                                StrideB,
+                                                StrideDs,
+                                                StrideC,
+                                                KBatch,
+                                                a_element_op,
+                                                b_element_op,
+                                                c_element_op,
+                                                Nr,
+                                                Kr);
     }
 
     // polymorphic

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp

@@ -173,6 +173,14 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
     static constexpr index_t NLane   = NPerXdl;
     static constexpr index_t NWave   = NPerBlock / NPerXdl / NXdlPerWave;
 
+    using GemmSpecialization = tensor_operation::device::GemmSpecialization;
+    static constexpr bool KPadding =
+        GemmSpec == GemmSpecialization::KPadding || GemmSpec == GemmSpecialization::MKPadding ||
+        GemmSpec == GemmSpecialization::NKPadding || GemmSpec == GemmSpecialization::MNKPadding;
+    static constexpr bool NPadding = GemmSpec == GemmSpecialization::NPadding ||
+                                     GemmSpec == GemmSpecialization::NKPadding ||
+                                     GemmSpec == GemmSpecialization::MNKPadding;
+
     static constexpr auto MakeDsGridPointer()
     {
         return generate_tuple(
@@ -212,6 +220,16 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
         return math::integer_divide_ceil(K, KLane * KPack);
     }
 
+    __host__ __device__ static auto CalculateBKShufflePadded(index_t K)
+    {
+        return (K + KShufflePadded - 1) / KShufflePadded * KShufflePadded;
+    }
+
+    __host__ __device__ static auto CalculateBNShufflePadded(index_t N)
+    {
+        return (N + NShufflePadded - 1) / NShufflePadded * NShufflePadded;
+    }
+
     __host__ __device__ static auto CalculateKPadded(index_t K)
     {
         return math::integer_divide_ceil(K, KPerBlock) * KPerBlock;
@@ -281,8 +299,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
             }
         }();
 
-        using GemmSpecialization = tensor_operation::device::GemmSpecialization;
-
         if constexpr(GemmSpec == GemmSpecialization::MKPadding ||
                      GemmSpec == GemmSpecialization::MNKPadding)
         {
@@ -349,12 +365,44 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
         }
     }
 
-    __host__ __device__ static auto MakeBGridDescriptor_Preshuffled(index_t N0, index_t K0)
+    __host__ __device__ static auto
+    MakeBGridDescriptor_Preshuffled(index_t N0, index_t K0, index_t NPadded, index_t KBatch)
     {
-        constexpr index_t NkSwizzleNumber = Number<warpSize * KPack>{};
-        return make_naive_tensor_descriptor(
-            make_tuple(N0 / NWave, NWave, K0, NkSwizzleNumber),
-            make_tuple(NWave * K0 * NkSwizzleNumber, K0 * NkSwizzleNumber, NkSwizzleNumber, I1));
+        ignore = NPadded;
+        ignore = KBatch;
+        // if N padding
+        if constexpr(GemmSpec == GemmSpecialization::NPadding ||
+                     GemmSpec == GemmSpecialization::NKPadding)
+        {
+            // origin: [N0,K0,KLane,NLane,KPack]
+            constexpr index_t NkSwizzleNumber = Number<warpSize * KPack>{};
+            const auto b_grid_desc_raw        = make_naive_tensor_descriptor(
+                make_tuple(N0 / NWave, NWave, K0 / KBatch, NkSwizzleNumber),
+                make_tuple(
+                    NWave * K0 * NkSwizzleNumber, K0 * NkSwizzleNumber, NkSwizzleNumber, I1));
+#if 0
+            auto N0new = CalculateBN0Shuffled(NPadded);
+
+            return transform_tensor_descriptor(
+                b_grid_desc_raw,
+                make_tuple(make_right_pad_transform(N0 / NWave, (N0new - N0) / NWave),
+                           make_pass_through_transform(NWave),
+                           make_pass_through_transform(K0 / KBatch),
+                           make_pass_through_transform(NkSwizzleNumber)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+#else
+            return b_grid_desc_raw;
+#endif
+        }
+        else
+        {
+            constexpr index_t NkSwizzleNumber = Number<warpSize * KPack>{};
+            return make_naive_tensor_descriptor(
+                make_tuple(N0 / NWave, NWave, K0, NkSwizzleNumber),
+                make_tuple(
+                    NWave * K0 * NkSwizzleNumber, K0 * NkSwizzleNumber, NkSwizzleNumber, I1));
+        }
     }
 
     __host__ __device__ static auto MakeBGridDescriptor_BK0_N_BK1(
@@ -371,8 +419,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
             }
         }();
 
-        using GemmSpecialization = tensor_operation::device::GemmSpecialization;
-
         if constexpr(GemmSpec == GemmSpecialization::NKPadding ||
                      GemmSpec == GemmSpecialization::MNKPadding)
         {
@@ -477,7 +523,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
                                            make_tuple(Sequence<0>{}, Sequence<1>{}),
                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
 #if 0
-        using GemmSpecialization = tensor_operation::device::GemmSpecialization;
 
         if constexpr(GemmSpec == GemmSpecialization::MNPadding ||
                      GemmSpec == GemmSpecialization::MNKPadding)
@@ -568,8 +613,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
               BK0{CalculateBK0Padded(K_, KBatch_)},
               MBlock{CalculateMBlock(M_)},
               NBlock{CalculateNBlock(N_)},
-              BN0Shuffled{CalculateBN0Shuffled(N_)},
-              BK0Shuffled{CalculateBK0Shuffled(K_)}
+              BN0Shuffled{CalculateBN0Shuffled(NPadding ? CalculateBNShufflePadded(N_) : N_)},
+              BK0Shuffled{CalculateBK0Shuffled(KPadding ? CalculateBKShufflePadded(K_) : K_)}
         {
         }
 
@@ -887,6 +932,17 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
         static_assert((MPerBlock % (MPerXdl * MXdlPerWave) == 0) &&
                           (NPerBlock % (NXdlPerWave * NPerXdl)) == 0,
                       "Invalid tuning param!");
+        // for not adding k padd operator
+        if((KPadding && (CalculateBKShufflePadded(karg.K) % KPerBlock != 0)) ||
+           (karg.BK0Shuffled % karg.KBatch != 0))
+        {
+            return false;
+        }
+
+        if((karg.N % NPerXdl != 0) || (karg.BN0Shuffled % NWave != 0))
+        {
+            return false;
+        }
 
         if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
@@ -1133,8 +1189,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
         const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1(
             problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0);
 
-        const auto b_grid_desc_bpreshuffled =
-            MakeBGridDescriptor_Preshuffled(problem.BN0Shuffled, problem.BK0Shuffled);
+        const auto b_grid_desc_bpreshuffled = MakeBGridDescriptor_Preshuffled(
+            problem.BN0Shuffled, problem.BK0Shuffled, problem.NPadded, problem.KBatch);
         const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N<CLayout>(
             problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC);
 
@@ -1570,8 +1626,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
         const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1(
             problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0);
 
-        const auto b_grid_desc_bpreshuffled =
-            MakeBGridDescriptor_Preshuffled(problem.BN0Shuffled, problem.BK0Shuffled);
+        const auto b_grid_desc_bpreshuffled = MakeBGridDescriptor_Preshuffled(
+            problem.BN0Shuffled, problem.BK0Shuffled, problem.NPadded, problem.KBatch);
         const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N<CLayout>(
             problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC);