Merge branch 'develop' into ck_host_lib

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

@@ -44,6 +44,48 @@ struct DeviceBatchedGemm : public BaseOperator
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };
 
+template <typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename ADataType,
+          typename BDataType,
+          typename BScaleType,
+          typename CDataType,
+          index_t ScaleBlockN,
+          index_t ScaleBlockK,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation>
+struct DeviceBatchedGemmV2BScale : public BaseOperator
+{
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_a,
+                        const void* p_b,
+                        void* p_c,
+                        ck::index_t M,
+                        ck::index_t N,
+                        ck::index_t K,
+                        ck::index_t StrideA,
+                        ck::index_t StrideB,
+                        ck::index_t StrideC,
+                        ck::index_t StrideScaleB,
+                        ck::index_t BatchStrideA,
+                        ck::index_t BatchStrideB,
+                        ck::index_t BatchStrideC,
+                        ck::index_t BatchStrideScaleB,
+                        const void* p_b_scale,
+                        ck::index_t Batch,
+                        ck::index_t KBatch,
+                        AElementwiseOperation a_element_op,
+                        BElementwiseOperation b_element_op,
+                        CElementwiseOperation c_element_op) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+
+    virtual bool GetPermuteB()         = 0;
+    virtual ck::index_t GetKPerBlock() = 0;
+};
+
 template <typename ALayout,
           typename BLayout,
           typename CLayout,

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

@@ -89,7 +89,8 @@ struct DeviceBatchedGemmV2MultiD : public BaseOperator
                         index_t BatchStrideE,
                         AElementwiseOperation a_element_op,
                         BElementwiseOperation b_element_op,
-                        CDEElementwiseOperation cde_element_op) = 0;
+                        CDEElementwiseOperation cde_element_op,
+                        index_t KBatch) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };

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

@@ -36,6 +36,10 @@ struct DeviceGemmV2 : public BaseOperator
                         CElementwiseOperation c_element_op) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+
+    virtual bool GetPermuteA()         = 0;
+    virtual bool GetPermuteB()         = 0;
+    virtual ck::index_t GetKPerBlock() = 0;
 };
 
 template <typename ALayout,
@@ -73,6 +77,43 @@ struct DeviceGemmV2R1 : public BaseOperator
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };
 
+template <typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename ADataType,
+          typename BDataType,
+          typename BScaleType,
+          typename CDataType,
+          index_t ScaleBlockN,
+          index_t ScaleBlockK,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation>
+struct DeviceGemmV2BScale : public BaseOperator
+{
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_a,
+                        const void* p_b,
+                        void* p_c,
+                        ck::index_t M,
+                        ck::index_t N,
+                        ck::index_t K,
+                        ck::index_t StrideA,
+                        ck::index_t StrideB,
+                        ck::index_t StrideC,
+                        ck::index_t StrideScaleB,
+                        const void* p_b_scale,
+                        ck::index_t KSplit,
+                        AElementwiseOperation a_element_op,
+                        BElementwiseOperation b_element_op,
+                        CElementwiseOperation c_element_op) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+
+    virtual bool GetPermuteB()         = 0;
+    virtual ck::index_t GetKPerBlock() = 0;
+};
+
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck

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

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2023-2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2023-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
+#ifndef CK_CODE_GEN_RTC
 #include <array>
+#endif
 
 #include "ck/tensor_operation/gpu/device/device_base.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_utils.hpp"
@@ -13,8 +15,13 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 
+#ifdef CK_CODE_GEN_RTC
+template <typename T>
+using is_tuple = decltype(ck::declval<T&>().IsTuple());
+#else
 template <typename T>
 using is_tuple = decltype(std::declval<T&>().IsTuple());
+#endif
 
 /**
  * \brief Grouped Convolution Forward
@@ -72,12 +79,18 @@ struct DeviceGroupedConvFwdMultipleABD : public BaseOperator
     static constexpr index_t NumDTensor = DsDataType::Size();
 
     static_assert(NumDTensor == DsLayout::Size(), "wrong! Inconsistent NumDTensor");
-
+#ifdef CK_CODE_GEN_RTC
+    using APointers = ck::conditional_t<isMultiA, ck::Array<const void*, NumATensor>&, const void*>;
+    using BPointers = ck::conditional_t<isMultiB, ck::Array<const void*, NumBTensor>&, const void*>;
+#else
     // If DataType is tuple, user has to pass std::array with pointers.
     using APointers =
-        std::conditional_t<isMultiA, std::array<const void*, NumATensor>&, const void*>;
+        ck::conditional_t<isMultiA, std::array<const void*, NumATensor>&, const void*>;
     using BPointers =
-        std::conditional_t<isMultiB, std::array<const void*, NumBTensor>&, const void*>;
+        ck::conditional_t<isMultiB, std::array<const void*, NumBTensor>&, const void*>;
+#endif
+
+#ifndef CK_CODE_GEN_RTC
 
     /**
      * \brief Make argument pointer for grouped conv fwd.
@@ -150,6 +163,7 @@ struct DeviceGroupedConvFwdMultipleABD : public BaseOperator
                         const CDEElementwiseOperation& cde_element_op) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+#endif
 };
 
 } // namespace device

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

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
+#include <array>
 #include <iostream>
+#include <sstream>
+#include <stdexcept>
 #include <vector>
 
 #include "device_base.hpp"
+#include "ck/utility/ignore.hpp"
 
 namespace ck {
 namespace tensor_operation {
 namespace device {
 
+///
+/// @brief      Structure representing single GEMM problem arguments.
+///
+///             The pointer to the vector of those structures is passed to the GroupedGEMM entry
+///             point kernel.
+///
+/// @tparam     NumDTensor  The number of D input tensors.
+///
+template <index_t NumDTensor = 0>
+struct GroupedGemmKernelArgument
+{
+    __host__ __device__ GroupedGemmKernelArgument(const void* p_a_grid_,
+                                                  const void* p_b_grid_,
+                                                  std::array<const void*, NumDTensor> p_ds_grid_,
+                                                  void* p_e_grid_,
+                                                  index_t M_,
+                                                  index_t N_,
+                                                  index_t K_,
+                                                  index_t StrideA_,
+                                                  index_t StrideB_,
+                                                  std::array<index_t, NumDTensor> StrideDs_,
+                                                  index_t StrideE_)
+        : p_a_grid{p_a_grid_},
+          p_b_grid{p_b_grid_},
+          p_ds_grid{p_ds_grid_},
+          p_e_grid{p_e_grid_},
+          M{M_},
+          N{N_},
+          K{K_},
+          StrideA{StrideA_},
+          StrideB{StrideB_},
+          StrideDs{StrideDs_},
+          StrideE{StrideE_}
+    {
+    }
+
+    const void* p_a_grid;
+    const void* p_b_grid;
+    std::array<const void*, NumDTensor> p_ds_grid;
+    void* p_e_grid;
+    index_t M;
+    index_t N;
+    index_t K;
+    index_t StrideA;
+    index_t StrideB;
+    std::array<index_t, NumDTensor> StrideDs;
+    index_t StrideE;
+
+    void Print() const
+    {
+        std::stringstream str;
+        for(auto sd : StrideDs)
+            str << sd << ",";
+
+        std::cout << "arg {"
+                  << "M:" << M << ", "
+                  << "N:" << N << ", "
+                  << "K:" << K << ", "
+                  << "SA:" << StrideA << ", "
+                  << "SB:" << StrideB << ", "
+                  << "SE:" << StrideE << ", "
+                  << "SDs: {" << str.str() << "}"
+                  << "}" << std::endl;
+    }
+};
+
 struct GemmDesc
 {
     ck::index_t M_, N_, K_;
@@ -48,6 +118,66 @@ struct DeviceGroupedGemm : public BaseOperator
                         CElementwiseOperation c_element_op) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+
+    //---------------------------------------------------------------------------------------------
+    /// @brief      Sets the device kernel arguments pointer and may copy data to device.
+    ///
+    /// TODO: Add which kernels are using this (TileLoop * FixedNK ??)
+    ///
+    /// @param      p_arg               The pointer to the Argument we're going to update.
+    /// @param[in]  p_dev_kernel_args   The pointer to the device memory which will contain kernel
+    ///                                 arguments.
+    /// @param[in]  p_host_kernel_args  The pointer to the host memory which contains kernel
+    ///                                 arguments that should be copied to device memory.
+    ///
+    virtual void SetDeviceKernelArgs(BaseArgument* p_arg,
+                                     void* p_dev_kernel_args,
+                                     const void* p_host_kernel_args) const
+    {
+        ignore = p_arg;
+        ignore = p_dev_kernel_args;
+        ignore = p_host_kernel_args;
+
+        std::ostringstream err;
+        err << "This function is not implemented by the kernel: " << this->GetTypeString()
+            << __FILE__ << ":" << __LINE__ << ", in function: " << __func__;
+        throw std::runtime_error(err.str());
+    }
+
+    //----------------------------------------------------------------------------------------------
+    /// @brief      Sets the device kernel arguments pointer and may copy data to device.
+    ///
+    /// @param      p_arg              The pointer to the Argument we're going to update.
+    /// @param[in]  p_dev_kernel_args  The pointer to the device memory which contains kernel
+    ///                                arguments.
+    ///
+    virtual void SetDeviceKernelArgs(BaseArgument* p_arg, void* p_dev_kernel_args) const
+    {
+        ignore = p_arg;
+        ignore = p_dev_kernel_args;
+
+        std::ostringstream err;
+        err << "This function is not implemented by the kernel: " << this->GetTypeString()
+            << __FILE__ << ":" << __LINE__ << ", in function: " << __func__;
+        throw std::runtime_error(err.str());
+    }
+
+    //----------------------------------------------------------------------------------------------
+    /// @brief      Gets the device kernel argument size.
+    ///
+    /// @param[in]  p_arg  The pointer to the Device op Argument.
+    ///
+    /// @return     The device kernel argument size.
+    ///
+    virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const
+    {
+        ignore = p_arg;
+
+        std::ostringstream err;
+        err << "This function is not implemented by the kernel: " << this->GetTypeString()
+            << __FILE__ << ":" << __LINE__ << ", in function: " << __func__;
+        throw std::runtime_error(err.str());
+    }
 };
 
 } // namespace device

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

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
-#include <iostream>
-#include <array>
-
-#include "device_grouped_gemm.hpp"
+#include "device_grouped_gemm_splitk.hpp"
 
 namespace ck {
 namespace tensor_operation {
 namespace device {
 
-template <index_t NumDTensor = 0>
-struct GroupedGemmKernelArgument
-{
-    const void* p_a_grid;
-    const void* p_b_grid;
-    std::array<const void*, NumDTensor> p_ds_grid;
-    void* p_e_grid;
-
-    index_t M;
-    index_t N;
-    index_t K;
-
-    index_t StrideA;
-    index_t StrideB;
-    std::array<index_t, NumDTensor> StrideDs;
-    index_t StrideE;
-};
-
 template <typename ALayout,
           typename BLayout,
           typename DsLayout,
@@ -41,21 +20,18 @@ template <typename ALayout,
           typename AElementwiseOperation,
           typename BElementwiseOperation,
           typename CElementwiseOperation>
-struct DeviceGroupedGemmFixedNK : DeviceGroupedGemm<ALayout,
-                                                    BLayout,
-                                                    DsLayout,
-                                                    ELayout,
-                                                    ADataType,
-                                                    BDataType,
-                                                    DsDataType,
-                                                    EDataType,
-                                                    AElementwiseOperation,
-                                                    BElementwiseOperation,
-                                                    CElementwiseOperation>
+struct DeviceGroupedGemmFixedNK : DeviceGroupedGemmSplitK<ALayout,
+                                                          BLayout,
+                                                          DsLayout,
+                                                          ELayout,
+                                                          ADataType,
+                                                          BDataType,
+                                                          DsDataType,
+                                                          EDataType,
+                                                          AElementwiseOperation,
+                                                          BElementwiseOperation,
+                                                          CElementwiseOperation>
 {
-    virtual void SetDeviceKernelArgs(BaseArgument* p_arg, const void* kernel_args) const = 0;
-    virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const               = 0;
-    virtual void SetKBatch(BaseArgument* p_arg, index_t k_batch) const                   = 0;
 };
 
 } // namespace device

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

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2023-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include <array>
-#include <iostream>
-#include <vector>
-#include <sstream>
-
-#include "device_grouped_gemm.hpp"
-
-namespace ck {
-namespace tensor_operation {
-namespace device {
-
-///
-/// @brief      Structure representing single GEMM problem arguments.
-///
-///             The pointer to the vector of those structures is passed to the GroupedGEMM entry
-///             point kernel.
-///
-/// @tparam     NumDTensor  The number of D input tensors.
-///
-template <index_t NumDTensor = 0>
-struct GroupedGemmMultipleDKernelArguments
-{
-    __host__ __device__
-    GroupedGemmMultipleDKernelArguments(const void* p_a_grid_,
-                                        const void* p_b_grid_,
-                                        std::array<const void*, NumDTensor> p_ds_grid_,
-                                        void* p_e_grid_,
-                                        index_t M_,
-                                        index_t N_,
-                                        index_t K_,
-                                        index_t StrideA_,
-                                        index_t StrideB_,
-                                        std::array<index_t, NumDTensor> StrideDs_,
-                                        index_t StrideE_)
-        : p_a_grid{p_a_grid_},
-          p_b_grid{p_b_grid_},
-          p_ds_grid{p_ds_grid_},
-          p_e_grid{p_e_grid_},
-          M{M_},
-          N{N_},
-          K{K_},
-          StrideA{StrideA_},
-          StrideB{StrideB_},
-          StrideDs{StrideDs_},
-          StrideE{StrideE_}
-    {
-    }
-
-    const void* p_a_grid;
-    const void* p_b_grid;
-    std::array<const void*, NumDTensor> p_ds_grid;
-    void* p_e_grid;
-    index_t M;
-    index_t N;
-    index_t K;
-    index_t StrideA;
-    index_t StrideB;
-    std::array<index_t, NumDTensor> StrideDs;
-    index_t StrideE;
-
-    void Print() const
-    {
-        std::stringstream str;
-        for(auto sd : StrideDs)
-            str << sd << ",";
-
-        std::cout << "arg {"
-                  << "M:" << M << ", "
-                  << "N:" << N << ", "
-                  << "K:" << K << ", "
-                  << "SA:" << StrideA << ", "
-                  << "SB:" << StrideB << ", "
-                  << "SE:" << StrideE << ", "
-                  << "SDs: {" << str.str() << "}"
-                  << "}" << std::endl;
-    }
-};
-
-template <typename ALayout,
-          typename BLayout,
-          typename DsLayout,
-          typename ELayout,
-          typename ADataType,
-          typename BDataType,
-          typename DsDataType,
-          typename EDataType,
-          typename AElementwiseOperation,
-          typename BElementwiseOperation,
-          typename CDEElementwiseOperation>
-struct DeviceGroupedGemmMultipleDSplitK : public DeviceGroupedGemm<ALayout,
-                                                                   BLayout,
-                                                                   DsLayout,
-                                                                   ELayout,
-                                                                   ADataType,
-                                                                   BDataType,
-                                                                   DsDataType,
-                                                                   EDataType,
-                                                                   AElementwiseOperation,
-                                                                   BElementwiseOperation,
-                                                                   CDEElementwiseOperation>
-{
-    //----------------------------------------------------------------------------------------------
-    /// @brief      Sets the k batch size.
-    ///
-    /// @param      p_arg   Pointer to the Argument we're going to change.
-    /// @param[in]  kbatch  The kbatch value.
-    ///
-    virtual void SetKBatchSize(BaseArgument* p_arg, index_t kbatch) const = 0;
-
-    //----------------------------------------------------------------------------------------------
-    /// @brief      Sets the device kernel arguments pointer.
-    ///
-    /// @param      p_arg              The pointer to the Argument we're going to update.
-    /// @param[in]  p_dev_kernel_args  The pointer to the device memory which contains kernel
-    ///                                arguments.
-    ///
-    virtual void SetDeviceKernelArgs(BaseArgument* p_arg, void* p_dev_kernel_args) const = 0;
-
-    //----------------------------------------------------------------------------------------------
-    /// @brief      Gets the device kernel argument size.
-    ///
-    /// @param[in]  p_arg  The pointer to the Device op Argument.
-    ///
-    /// @return     The device kernel argument size.
-    ///
-    virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const = 0;
-};
-
-} // namespace device
-} // namespace tensor_operation
-} // namespace ck

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 #pragma once
-#include <iostream>
-#include <vector>
 
 #include "device_grouped_gemm.hpp"
 
@@ -31,7 +31,23 @@ struct DeviceGroupedGemmSplitK : public DeviceGroupedGemm<ALayout,
                                                           BElementwiseOperation,
                                                           CElementwiseOperation>
 {
+    //----------------------------------------------------------------------------------------------
+    /// @brief      Sets the k batch size.
+    ///
+    /// @param      p_arg   Pointer to the Argument we're going to change.
+    /// @param[in]  kbatch  The kbatch value.
+    ///
     virtual void SetKBatchSize(BaseArgument* p_arg, index_t kbatch) const = 0;
+    //----------------------------------------------------------------------------------------------
+    /// @brief      Sets the k batch size.
+    ///
+    /// @param      p_arg   Pointer to the Argument we're going to change.
+    /// @param[in]  kbatch  The kbatch value.
+    ///
+    virtual void SetKBatch(BaseArgument* p_arg, index_t kbatch) const
+    {
+        this->SetKBatchSize(p_arg, kbatch);
+    };
 };
 
 } // namespace device

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

@@ -3,83 +3,20 @@
 
 #pragma once
 
-#include <array>
-#include <iostream>
-#include <vector>
-#include <sstream>
-
 #include "device_grouped_gemm.hpp"
 
 namespace ck {
 namespace tensor_operation {
 namespace device {
 
+/// @brief Grouped GEMM kernel using output Tile Looping algorithm
 ///
-/// @brief      Structure representing single GEMM problem arguments.
-///
-///             The pointer to the vector of those structures is passed to the GroupedGEMM entry
-///             point kernel.
-///
-/// @tparam     NumDTensor  The number of D input tensors.
+/// @par This kernel does not require any knowledge about input data sizes (GEMM M/N/K)
+///       It requires only the number of groups to launch. Other information like
+///       data pointers and GEMM sizes, packed into gemm kernel args may be all dynamic
+///       (known only at kernel run-time).
 ///
-template <index_t NumDTensor = 0>
-struct GroupedGemmTileLoopKernelArguments
-{
-    __host__ __device__
-    GroupedGemmTileLoopKernelArguments(const void* p_a_grid_,
-                                       const void* p_b_grid_,
-                                       std::array<const void*, NumDTensor> p_ds_grid_,
-                                       void* p_e_grid_,
-                                       index_t M_,
-                                       index_t N_,
-                                       index_t K_,
-                                       index_t StrideA_,
-                                       index_t StrideB_,
-                                       std::array<index_t, NumDTensor> StrideDs_,
-                                       index_t StrideE_)
-        : p_a_grid{p_a_grid_},
-          p_b_grid{p_b_grid_},
-          p_ds_grid{p_ds_grid_},
-          p_e_grid{p_e_grid_},
-          M{M_},
-          N{N_},
-          K{K_},
-          StrideA{StrideA_},
-          StrideB{StrideB_},
-          StrideDs{StrideDs_},
-          StrideE{StrideE_}
-    {
-    }
-
-    const void* p_a_grid;
-    const void* p_b_grid;
-    std::array<const void*, NumDTensor> p_ds_grid;
-    void* p_e_grid;
-    index_t M;
-    index_t N;
-    index_t K;
-    index_t StrideA;
-    index_t StrideB;
-    std::array<index_t, NumDTensor> StrideDs;
-    index_t StrideE;
-
-    void Print() const
-    {
-        std::stringstream str;
-        for(auto sd : StrideDs)
-            str << sd << ",";
-
-        std::cout << "arg {"
-                  << "M:" << M << ", "
-                  << "N:" << N << ", "
-                  << "K:" << K << ", "
-                  << "SA:" << StrideA << ", "
-                  << "SB:" << StrideB << ", "
-                  << "SE:" << StrideE << ", "
-                  << "SDs: {" << str.str() << "}"
-                  << "}" << std::endl;
-    }
-};
+/// @note This kernel does not support SplitK.
 
 template <typename ALayout,
           typename BLayout,
@@ -104,23 +41,6 @@ struct DeviceGroupedGemmTileLoop : public DeviceGroupedGemm<ALayout,
                                                             BElementwiseOperation,
                                                             CDEElementwiseOperation>
 {
-    //----------------------------------------------------------------------------------------------
-    /// @brief      Sets the device kernel arguments pointer.
-    ///
-    /// @param      p_arg              The pointer to the Argument we're going to update.
-    /// @param[in]  p_dev_kernel_args  The pointer to the device memory which contains kernel
-    ///                                arguments.
-    ///
-    virtual void SetDeviceKernelArgs(BaseArgument* p_arg, void* p_dev_kernel_args) const = 0;
-
-    //----------------------------------------------------------------------------------------------
-    /// @brief      Gets the device kernel argument size.
-    ///
-    /// @param[in]  p_arg  The pointer to the Device op Argument.
-    ///
-    /// @return     The device kernel argument size.
-    ///
-    virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const = 0;
 };
 
 } // namespace device

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

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -29,6 +29,7 @@ enum struct GemmSpecialization
     MNKOPadding,
 };
 
+#ifndef CK_CODE_GEN_RTC
 inline std::string getGemmSpecializationString(const GemmSpecialization& s)
 {
     switch(s)
@@ -52,6 +53,7 @@ inline std::string getGemmSpecializationString(const GemmSpecialization& s)
     default: return "Unrecognized specialization!";
     }
 }
+#endif
 
 } // namespace device
 } // namespace tensor_operation

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

@@ -3,11 +3,17 @@
 
 #pragma once
 
+#ifndef CK_CODE_GEN_RTC
 #include <functional>
 #include <iostream>
 #include <iterator>
 #include <numeric>
 #include <sstream>
+#include <stdio.h>
+
+#include "ck/host_utility/device_prop.hpp"
+#include "ck/host_utility/kernel_launch.hpp"
+#endif
 
 #include "ck/utility/common_header.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
@@ -15,15 +21,12 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
 #include "ck/tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp"
-#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_multiple_d_xdl_cshuffle.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_multiple_abd_xdl_cshuffle.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_utils.hpp"
-#include "ck/host_utility/device_prop.hpp"
-#include "ck/host_utility/kernel_launch.hpp"
-#include "ck/host_utility/io.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -91,8 +94,7 @@ __device__ void device_grouped_conv_fwd_multiple_abd_xdl_cshuffle(
     const Block2ETileMap block_2_ctile_map,
     const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     // offset base pointer for each work-group
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
@@ -259,8 +261,13 @@ __global__ void
 
 } // namespace
 
+#ifdef CK_CODE_GEN_RTC
+template <typename T>
+using is_tuple = decltype(ck::declval<T&>().IsTuple());
+#else
 template <typename T>
 using is_tuple = decltype(std::declval<T&>().IsTuple());
+#endif
 
 //
 // @brief      Device Convolution operation.
@@ -429,8 +436,8 @@ struct CodegenDeviceGroupedConvFwdMultipleABD_Xdl_CShuffle
 
     // If we are using multiAB and one of the template datatype parameters is not a tuple, convert
     // it to it
-    using GemmADataType = std::conditional_t<!isMultiA && isMultiB, Tuple<ADataType>, ADataType>;
-    using GemmBDataType = std::conditional_t<!isMultiB && isMultiA, Tuple<BDataType>, BDataType>;
+    using GemmADataType = ck::conditional_t<!isMultiA && isMultiB, Tuple<ADataType>, ADataType>;
+    using GemmBDataType = ck::conditional_t<!isMultiB && isMultiA, Tuple<BDataType>, BDataType>;
 
 #define GridwiseGemmTemplateParameters                                                          \
     GemmADataType, GemmBDataType, ComputeDataType, AccDataType, CShuffleDataType, DsDataType,   \
@@ -449,15 +456,13 @@ struct CodegenDeviceGroupedConvFwdMultipleABD_Xdl_CShuffle
         CDEBlockTransferScalarPerVector_NPerBlock, LoopSched
     // Use appropriate gridwise gemm
     using GridwiseGemm =
-        std::conditional_t<isMultiA || isMultiB,
-                           GridwiseGemmMultipleABD_xdl_cshuffle<GridwiseGemmTemplateParameters>,
-                           GridwiseGemmMultipleD_xdl_cshuffle<GridwiseGemmTemplateParameters>>;
+        ck::conditional_t<isMultiA || isMultiB,
+                          GridwiseGemmMultipleABD_xdl_cshuffle<GridwiseGemmTemplateParameters>,
+                          GridwiseGemmMultipleD_xdl_cshuffle<GridwiseGemmTemplateParameters>>;
 
     // If ADataTypes or BDataTypes is tuple, user has to pass ck::Array with pointers.
-    using APointers =
-        std::conditional_t<isMultiA, ck::Array<const void*, NumATensor>&, const void*>;
-    using BPointers =
-        std::conditional_t<isMultiB, ck::Array<const void*, NumBTensor>&, const void*>;
+    using APointers = ck::conditional_t<isMultiA, ck::Array<const void*, NumATensor>&, const void*>;
+    using BPointers = ck::conditional_t<isMultiB, ck::Array<const void*, NumBTensor>&, const void*>;
     // Use Tuple for the both cases for GridPointer to initialize it in Argument constructor (not
     // in initializer list what is required for single const pointer).
     using AGridPointer = remove_cvref_t<
@@ -812,7 +817,6 @@ struct CodegenDeviceGroupedConvFwdMultipleABD_Xdl_CShuffle
 
         static_for<0, NumDTensor, 1>{}([&](auto i) {
             using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
-
             // FIXME: layout
             if constexpr(is_same_v<DLayout, ctc::G_NW_K> || is_same_v<DLayout, ctc::G_NHW_K> ||
                          is_same_v<DLayout, ctc::G_NDHW_K> || is_same_v<DLayout, ctc::GNWK> ||
@@ -965,18 +969,18 @@ struct CodegenDeviceGroupedConvFwdMultipleABD_Xdl_CShuffle
         const BElementwiseOperation& b_element_op,
         const CDEElementwiseOperation& cde_element_op)
     {
-        std::array<index_t, NDimSpatial + 3> a_g_n_c_wis_lengths_i32;
-        std::array<index_t, NDimSpatial + 3> a_g_n_c_wis_strides_i32;
-        std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_lengths_i32;
-        std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_strides_i32;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_lengths_i32;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_strides_i32;
-        std::array<index_t, NDimSpatial + 3> e_g_n_k_wos_lengths_i32;
-        std::array<index_t, NDimSpatial + 3> e_g_n_k_wos_strides_i32;
-        std::array<index_t, NDimSpatial> conv_filter_strides_i32;
-        std::array<index_t, NDimSpatial> conv_filter_dilations_i32;
-        std::array<index_t, NDimSpatial> input_left_pads_i32;
-        std::array<index_t, NDimSpatial> input_right_pads_i32;
+        ck::Array<index_t, NDimSpatial + 3> a_g_n_c_wis_lengths_i32;
+        ck::Array<index_t, NDimSpatial + 3> a_g_n_c_wis_strides_i32;
+        ck::Array<index_t, NDimSpatial + 3> b_g_k_c_xs_lengths_i32;
+        ck::Array<index_t, NDimSpatial + 3> b_g_k_c_xs_strides_i32;
+        ck::Array<ck::Array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_lengths_i32;
+        ck::Array<ck::Array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_strides_i32;
+        ck::Array<index_t, NDimSpatial + 3> e_g_n_k_wos_lengths_i32;
+        ck::Array<index_t, NDimSpatial + 3> e_g_n_k_wos_strides_i32;
+        ck::Array<index_t, NDimSpatial> conv_filter_strides_i32;
+        ck::Array<index_t, NDimSpatial> conv_filter_dilations_i32;
+        ck::Array<index_t, NDimSpatial> input_left_pads_i32;
+        ck::Array<index_t, NDimSpatial> input_right_pads_i32;
 
         array_convert(a_g_n_c_wis_lengths_i32, a_g_n_c_wis_lengths);
         array_convert(a_g_n_c_wis_strides_i32, a_g_n_c_wis_strides);

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

@@ -56,8 +56,7 @@ __global__ void
             const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch,
             const Block2ETileMap block_2_etile_map)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     const index_t num_blocks_per_batch =

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

@@ -74,8 +74,7 @@ __global__ void
                                           const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch,
                                           const Block2ETileMap block_2_etile_map)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
     const index_t g_idx = __builtin_amdgcn_readfirstlane(get_block_1d_id() / num_blocks_per_batch);

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

@@ -60,8 +60,7 @@ __global__ void
             const index_t batch_count,
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
@@ -108,7 +107,7 @@ __global__ void
     ignore = block_2_ctile_map;
     ignore = batch_count;
     ignore = compute_base_ptr_of_batch;
-#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+#endif // end of if (defined(__gfx9__))
 }
 
 // Computes C = A * B0 * B1

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

@@ -83,8 +83,7 @@ __global__ void
                                 const Block2ETileMap block_2_etile_map)
 {
 
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
     const index_t g_idx = __builtin_amdgcn_readfirstlane(get_block_1d_id() / num_blocks_per_batch);

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

@@ -68,8 +68,7 @@ __global__ void
             const index_t batch_count,
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);

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

@@ -41,12 +41,15 @@ __global__ void
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     const index_t g_idx = blockIdx.z % karg.Batch;
+    const index_t k_idx = blockIdx.z / karg.Batch;
 
     const auto a_batch_offset  = karg.compute_ptr_offset_of_batch.GetAPtrOffset(g_idx);
     const auto b_batch_offset  = karg.compute_ptr_offset_of_batch.GetBPtrOffset(g_idx);
     const auto ds_batch_offset = karg.compute_ptr_offset_of_batch.GetDsPtrOffset(g_idx);
     const auto c_batch_offset  = karg.compute_ptr_offset_of_batch.GetCPtrOffset(g_idx);
 
+    auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, k_idx);
+
     // populate pointer, desc for Ds
     static_for<0, GridwiseGemm::NumDTensor, 1>{}([&](auto i) {
         // D pointer
@@ -54,8 +57,8 @@ __global__ void
     });
 
     GridwiseGemm::template Run<HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(
-        karg.p_a_grid + a_batch_offset,
-        karg.p_b_grid + b_batch_offset,
+        karg.p_a_grid + a_batch_offset + splitk_batch_offset.a_k_split_offset,
+        karg.p_b_grid + b_batch_offset + splitk_batch_offset.b_k_split_offset,
         karg.p_ds_grid,
         karg.p_c_grid + c_batch_offset,
         p_shared,
@@ -87,12 +90,15 @@ __global__ void
     __shared__ char p_shared_1[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     const index_t g_idx = blockIdx.z % karg.Batch;
+    const index_t k_idx = blockIdx.z / karg.Batch;
 
     const auto a_batch_offset  = karg.compute_ptr_offset_of_batch.GetAPtrOffset(g_idx);
     const auto b_batch_offset  = karg.compute_ptr_offset_of_batch.GetBPtrOffset(g_idx);
     const auto ds_batch_offset = karg.compute_ptr_offset_of_batch.GetDsPtrOffset(g_idx);
     const auto c_batch_offset  = karg.compute_ptr_offset_of_batch.GetCPtrOffset(g_idx);
 
+    auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, k_idx);
+
     // populate pointer, desc for Ds
     static_for<0, GridwiseGemm::NumDTensor, 1>{}([&](auto i) {
         // D pointer
@@ -100,8 +106,8 @@ __global__ void
     });
 
     GridwiseGemm::template Run_2Lds<HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(
-        karg.p_a_grid + a_batch_offset,
-        karg.p_b_grid + b_batch_offset,
+        karg.p_a_grid + a_batch_offset + splitk_batch_offset.a_k_split_offset,
+        karg.p_b_grid + b_batch_offset + splitk_batch_offset.b_k_split_offset,
         karg.p_ds_grid,
         karg.p_c_grid + c_batch_offset,
         p_shared_0,
@@ -303,7 +309,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                  index_t Batch_,
                  AElementwiseOperation a_element_op_,
                  BElementwiseOperation b_element_op_,
-                 CElementwiseOperation c_element_op_)
+                 CElementwiseOperation c_element_op_,
+                 index_t KBatch_)
             : GridwiseGemm::Argument{p_a_grid_,
                                      p_b_grid_,
                                      p_ds_grid_,
@@ -315,7 +322,7 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                                      StrideB_,
                                      StrideDs_,
                                      StrideE_,
-                                     1,
+                                     KBatch_,
                                      a_element_op_,
                                      b_element_op_,
                                      c_element_op_},
@@ -336,13 +343,14 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                 arg.Print();
             }
 
-            if(!GridwiseGemm::CheckValidity(arg) || arg.KBatch > 1)
+            if(!GridwiseGemm::CheckValidity(arg))
             {
                 throw std::runtime_error("wrong! GridwiseGemm has invalid setting");
             }
 
             index_t gdx, gdy, gdz;
-            std::tie(gdx, gdy, gdz) = GridwiseGemm::CalculateGridSize(arg.M, arg.N, arg.Batch);
+            std::tie(gdx, gdy, gdz) =
+                GridwiseGemm::CalculateGridSize(arg.M, arg.N, arg.Batch * arg.KBatch);
 
             float ave_time = 0;
 
@@ -387,10 +395,11 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                         rotating_mem.Next();
                         // clear c mem
                         if(arg_.KBatch > 1)
-                            hipGetErrorString(hipMemsetAsync(arg_.p_c_grid,
-                                                             0,
-                                                             arg_.M * arg_.N * sizeof(CDataType),
-                                                             stream_config.stream_id_));
+                            hipGetErrorString(
+                                hipMemsetAsync(arg_.p_c_grid,
+                                               0,
+                                               arg.Batch * arg_.M * arg_.N * sizeof(CDataType),
+                                               stream_config.stream_id_));
                     };
 
                     ave_time = ck::utility::launch_and_time_kernel_with_preprocess<false>(
@@ -889,7 +898,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                              index_t BatchStrideE,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
-                             CElementwiseOperation c_element_op)
+                             CElementwiseOperation c_element_op,
+                             index_t KBatch = 1)
     {
         return Argument{static_cast<const ADataType*>(p_a),
                         static_cast<const BDataType*>(p_b),
@@ -909,7 +919,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                         Batch,
                         a_element_op,
                         b_element_op,
-                        c_element_op};
+                        c_element_op,
+                        KBatch};
     }
 
     static auto MakeInvoker() { return Invoker{}; }
@@ -934,7 +945,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                         index_t BatchStrideE,
                         AElementwiseOperation a_element_op,
                         BElementwiseOperation b_element_op,
-                        CElementwiseOperation c_element_op) override
+                        CElementwiseOperation c_element_op,
+                        index_t KBatch = 1) override
     {
         return std::make_unique<Argument>(static_cast<const ADataType*>(p_a),
                                           static_cast<const BDataType*>(p_b),
@@ -954,7 +966,8 @@ struct DeviceBatchedGemmMultiD_Xdl_CShuffle_V3
                                           Batch,
                                           a_element_op,
                                           b_element_op,
-                                          c_element_op);
+                                          c_element_op,
+                                          KBatch);
     }
 
     // polymorphic

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

@@ -59,8 +59,7 @@ __global__ void
             const ComputeBasePrtOfBatch compute_base_ptr_of_batch_,
             const Block2CTileMap block_2_ctile_map)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
     const index_t g_idx = __builtin_amdgcn_readfirstlane(get_block_1d_id() / num_blocks_per_batch);

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

@@ -67,8 +67,7 @@ __global__ void
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch,
             const C0MatrixMask c0_matrix_mask)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
@@ -127,7 +126,7 @@ __global__ void
     ignore = batch_count;
     ignore = compute_base_ptr_of_batch;
     ignore = c0_matrix_mask;
-#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+#endif // end of if (defined(__gfx9__))
 }
 
 // Computes C = A * B0 * B1

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

@@ -62,8 +62,7 @@ __global__ void
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch,
             const C0MatrixMask c0_matrix_mask)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx94__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
@@ -112,7 +111,7 @@ __global__ void
     ignore = batch_count;
     ignore = compute_base_ptr_of_batch;
     ignore = c0_matrix_mask;
-#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+#endif // end of if (defined(__gfx9__))
 }
 
 // Computes C = A * B0 * B1