Merge branch 'codegen_hiprtc' of github.com:ROCm/composable_kernel into codegen_hiprtc

example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp

@@ -41,6 +41,7 @@ float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_conf
                                    t.prec_sq,
                                    t.prec_kw,
                                    t.block_m,
+                                   t.activation,
                                    t.gate_only,
                                    t.fused_quant};
     auto a1 = fused_moegemm_args{

example/ck_tile/15_fused_moe/instances/fused_moegemm_api.cpp

@@ -17,15 +17,67 @@ float fused_moegemm(fused_moegemm_traits t, fused_moegemm_args a, const ck_tile:
     // clang-format off
     float r = -1;
     if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" &&
-       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1)
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1 && t.activation == 0)
     {
-        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>;
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 1;
+        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 0 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 0;
+        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "fp16" && t.prec_w == "fp16" && t.prec_o == "fp16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 1;
+        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "fp16" && t.prec_w == "fp16" && t.prec_o == "fp16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 0 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 0;
+        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 1;
+        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 0 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 0;
+        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "fp16" && t.prec_w == "fp16" && t.prec_o == "fp16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 1;
+        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
         r = fused_moegemm_<t_>(s, a);
     }
     else if(t.prec_i == "fp16" && t.prec_w == "fp16" && t.prec_o == "fp16" && t.prec_st == "fp32" &&
-       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1)
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 0 && t.activation == 1)
     {
-        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>;
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 0;
+        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, act_, go_, 0>;
         r = fused_moegemm_<t_>(s, a);
     }
     // clang-format on

example/ck_tile/15_fused_moe/instances/fused_moegemm_api_internal.hpp

@@ -21,8 +21,18 @@ float fused_moegemm_(const ck_tile::stream_config& s, fused_moegemm_args a)
                                                typename Ts_::BlockTile_1,
                                                typename Ts_::WarpPerBlock_0,
                                                typename Ts_::WarpTile_0>;
-    using f_problem =
-        ck_tile::FusedMoeGemmPipelineProblem<typename Ts_::ADataType,
+
+    constexpr auto get_activation_ = []() {
+        if constexpr(Ts_::Activation == 0)
+        {
+            return ck_tile::element_wise::FastGeluAsm{};
+        }
+        else
+            return ck_tile::element_wise::Silu{};
+    };
+    using f_act_ = ck_tile::remove_cvref_t<decltype(get_activation_())>;
+
+    using f_problem = ck_tile::FusedMoeGemmPipelineProblem<typename Ts_::ADataType,
                                                            typename Ts_::GDataType,
                                                            typename Ts_::DDataType,
                                                            typename Ts_::AccDataType,
@@ -33,7 +43,7 @@ float fused_moegemm_(const ck_tile::stream_config& s, fused_moegemm_args a)
                                                            typename Ts_::YSmoothScaleDataType,
                                                            typename Ts_::TopkWeightDataType,
                                                            typename Ts_::IndexDataType,
-                                             ck_tile::element_wise::FastGeluAsm, // TODO: hardcoded
+                                                           f_act_, // TODO: hardcoded
                                                            f_shape,
                                                            f_traits>;
 

example/ck_tile/15_fused_moe/instances/fused_moegemm_api_traits.hpp

@@ -16,6 +16,7 @@ template <typename I,
           typename BlockTIle_, // seq<b_token, b_interm, b_hidden, b_down>
           typename WarpPerBlock_,
           typename WarpTile_,               // seq<*,*,*>, used to select mfma
+          ck_tile::index_t Activation_ = 0, // 0: Gelu 1: Silu
           ck_tile::index_t GateOnly_   = 0,
           ck_tile::index_t FusedQuant_ = 0>
 struct fmoe_ // traits, ugly name, only used for internal
@@ -44,10 +45,11 @@ struct fmoe_ // traits, ugly name, only used for internal
     using WarpPerBlock_0 = ck_tile::remove_cvref_t<WarpPerBlock_>;
     using WarpTile_0     = ck_tile::remove_cvref_t<WarpTile_>;
 
-    using BlockTile_1    = ck_tile::sequence<BT_, BD_, BI_ / (GateOnly_ ? 1 : 2)>;
+    using BlockTile_1    = ck_tile::sequence<BT_, BD_, BI_>;
     using WarpPerBlock_1 = ck_tile::remove_cvref_t<WarpPerBlock_>;
     using WarpTile_1     = ck_tile::remove_cvref_t<WarpTile_>;
 
+    static constexpr ck_tile::index_t Activation = Activation_; // 0: Gelu 1: Silu
     static constexpr ck_tile::index_t GateOnly   = GateOnly_;
     static constexpr ck_tile::index_t FusedQuant = FusedQuant_;
 };

example/ck_tile/15_fused_moe/instances/fused_moegemm_bf16_m32.cpp

@@ -8,7 +8,18 @@
 
 // clang-format off
 template float fused_moegemm_<
-    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>
+    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 0, 0, 0>
 >(const ck_tile::stream_config& s, fused_moegemm_args a);
 
+template float fused_moegemm_<
+    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 0, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template float fused_moegemm_<
+    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template float fused_moegemm_<
+    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
 // clang-format on

example/ck_tile/15_fused_moe/instances/fused_moegemm_fp16_m32.cpp

@@ -8,7 +8,19 @@
 
 // clang-format off
 template float fused_moegemm_<
-    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>
+    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 0, 0, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template float fused_moegemm_<
+    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 0, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template float fused_moegemm_<
+    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template float fused_moegemm_<
+    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 1, 0>
 >(const ck_tile::stream_config& s, fused_moegemm_args a);
 
 // clang-format on

example/ck_tile/15_fused_moe/main.cpp

@@ -108,12 +108,14 @@ auto create_args(int argc, char* argv[])
         .insert(
             "gate_only", "1", "w0(gate/up) style, 0:gate+up will double interm size, 1:only gate")
         .insert("api", "0", "benchmark api set: 0:fused-moe(moe-gemm+moe-sorting), 1:moe-gemm")
+        .insert("act", "0", "activation after first gemm. 0:gelu, 1:silu")
         .insert("balance",
                 "0",
                 "if set to 1, will try balance the expert in topk-ids(convenient for testing)")
         .insert("init",
-                "2",
-                "init method. 0:random stepped float(fast). 1: random uniform, 2:rand normalized"
+                "1",
+                "init method. 0:random stepped float(fast). 1: random uniform[-0.5, 0.5], 2:rand "
+                "normalized[0, 1]"
                 "normalized(slow)")
         .insert("seed", "11939", "seed used to do random")
         .insert("warmup", "5", "cold iter")
@@ -135,6 +137,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
     ck_tile::index_t intermediate_size = arg_parser.get_int("i");
     ck_tile::index_t stride            = arg_parser.get_int("stride");
     ck_tile::index_t block_m           = arg_parser.get_int("bm");
+    ck_tile::index_t activation        = arg_parser.get_int("act");
     if(stride < 0)
         stride = hidden_size;
     std::string prec_i  = arg_parser.get_str("prec_i");
@@ -194,11 +197,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
             return std::string(", st:") + std::to_string(stride);
     }();
 
-    std::cout << "[" << api_str << "|" << prec_str << "]"
+    std::cout
+        << "[" << api_str << "|" << prec_str << "]"
         << " t:" << tokens << ", e:" << experts << ", k:" << topk << stride_str
         << ", hidden:" << hidden_size << ", interm:" << intermediate_size << ", tp:" << tp
-              << ", shrd_interm:" << shared_intermediate_size_0 << "|" << shared_intermediate_size_1
-              << ", go:" << gate_only << ", q:" << fused_quant << std::flush;
+        << ", act:"
+        << activation
+        // << ", shrd_interm:" << shared_intermediate_size_0 << "|" << shared_intermediate_size_1
+        << (gate_only ? ", g1u0" : ", g1u1") << ", q:" << fused_quant << std::flush;
 
     using TypeConfig           = FusedMoeGemmTypeConfig<I, W, O, ST, SW, SQ, KW>;
     using ADataType            = typename TypeConfig::ADataType;
@@ -370,6 +376,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                 prec_sq,
                                 prec_kw,
                                 block_m,
+                                activation,
                                 gate_only,
                                 fused_quant};
 
@@ -389,7 +396,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                             num_sorted_tiles_buf.GetDeviceBuffer(),
                             block_m,
                             hidden_size,
-                            shared_intermediate_size_0,
+                            intermediate_size / tp,
                             tokens,
                             experts,
                             topk,
@@ -408,6 +415,28 @@ bool run(const ck_tile::ArgParser& arg_parser)
                   << cal_tbps(ave_time) << " TB/s" << std::flush;
         bool pass = true;
 
+#define CPU_FUSED_MOE(act_type_)                                                 \
+    ck_tile::reference_fused_moe<AccDataType, act_type_>(a_host,                 \
+                                                         g_host,                 \
+                                                         d_host,                 \
+                                                         sa_host,                \
+                                                         sg_host,                \
+                                                         sd_host,                \
+                                                         sy_host,                \
+                                                         o_host,                 \
+                                                         sorted_token_ids_host,  \
+                                                         sorted_weight_host,     \
+                                                         sorted_expert_ids_host, \
+                                                         num_sorted_tiles_host,  \
+                                                         topk_ids_host,          \
+                                                         block_m,                \
+                                                         tokens,                 \
+                                                         experts,                \
+                                                         hidden_size,            \
+                                                         intermediate_size / tp, \
+                                                         topk,                   \
+                                                         gate_only)
+
         if(do_validation)
         {
             ck_tile::reference_moe_sorting<TopkWeightDataType, IndexDataType>(
@@ -419,28 +448,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
                 num_sorted_tiles_host.mData[0],
                 experts,
                 block_m);
-
-            ck_tile::reference_fused_moe<AccDataType, ck_tile::element_wise::Gelu>(
-                a_host,
-                g_host,
-                d_host,
-                sa_host,
-                sg_host,
-                sd_host,
-                sy_host,
-                o_host,
-                sorted_token_ids_host,
-                sorted_weight_host,
-                sorted_expert_ids_host,
-                num_sorted_tiles_host,
-                topk_ids_host,
-                block_m,
-                tokens,
-                experts,
-                hidden_size,
-                shared_intermediate_size_0,
-                topk,
-                gate_only);
+            if(activation == 0)
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Gelu);
+            }
+            else
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Silu);
+            }
 
             auto o_dev = o_buf.ToHost<ODataType>();
             // o_dev.savetxt("gpu-out.txt", "float");
@@ -491,6 +506,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                     prec_sq,
                                     prec_kw,
                                     block_m,
+                                    activation,
                                     gate_only,
                                     fused_quant};
 
@@ -507,7 +523,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                 sorted_expert_ids_buf.GetDeviceBuffer(),
                                 num_sorted_tiles_buf.GetDeviceBuffer(),
                                 hidden_size,
-                                shared_intermediate_size_0,
+                                intermediate_size / tp,
                                 tokens,
                                 experts,
                                 topk,
@@ -529,27 +545,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
         if(do_validation)
         {
-            ck_tile::reference_fused_moe<AccDataType, ck_tile::element_wise::Gelu>(
-                a_host,
-                g_host,
-                d_host,
-                sa_host,
-                sg_host,
-                sd_host,
-                sy_host,
-                o_host,
-                sorted_token_ids_host,
-                sorted_weight_host,
-                sorted_expert_ids_host,
-                num_sorted_tiles_host,
-                topk_ids_host,
-                block_m,
-                tokens,
-                experts,
-                hidden_size,
-                shared_intermediate_size_0,
-                topk,
-                gate_only);
+            if(activation == 0)
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Gelu);
+            }
+            else
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Silu);
+            }
 
             auto o_dev = o_buf.ToHost<ODataType>();
             // o_dev.savetxt("gpu-out.txt", "float");

example/ck_tile/16_batched_gemm/batched_gemm.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <hip/hip_runtime.h>
 
@@ -51,7 +51,7 @@ float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stre
                                ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
                                ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
 
-    using TilePartitioner = ck_tile::GemmTilePartitioner<CodegenGemmShape>;
+    using TilePartitioner = ck_tile::GemmTile2DPartitioner<CodegenGemmShape>;
 
     using GemmEpilogue = std::conditional_t<
         CShuffleEpilogue,
@@ -63,8 +63,8 @@ float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stre
                                                                    kOutputRank,
                                                                    1,
                                                                    0,
-                                                                   TilePartitioner::kM,
-                                                                   TilePartitioner::kN>>,
+                                                                   TilePartitioner::MPerBlock,
+                                                                   TilePartitioner::NPerBlock>>,
         ck_tile::Default2DEpilogue<
             ck_tile::Default2DEpilogueProblem<AccDataType, CDataType, kPadM, kPadN>>>;
 
@@ -72,9 +72,7 @@ float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stre
         ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout>;
     using CodegenPipelineProblem = ck_tile::
         GemmPipelineProblem<ADataType, BDataType, AccDataType, CodegenGemmShape, CodegenGemmTraits>;
-    using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
-    using CodegenGemmPipeline =
-        ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem, CodegenGemmPolicy>;
+    using CodegenGemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
     // ToDo: Will add the codegen part to test different pipeline policies in GEMM.
     // Now we only use the BlockGemmASmemBSmemCRegV1DefaultPolicy.
     using Kernel = ck_tile::BatchedGemmKernel<TilePartitioner, CodegenGemmPipeline, GemmEpilogue>;

example/ck_tile/16_batched_gemm/batched_gemm.hpp

@@ -39,7 +39,7 @@ auto create_args(int argc, char* argv[])
         .insert("stride_b", "0", "Tensor B stride")
         .insert("stride_c", "0", "Tensor C stride")
         .insert("a_layout", "R", "A tensor data layout - Row by default")
-        .insert("b_layout", "R", "B tensor data layout - Row by default")
+        .insert("b_layout", "C", "B tensor data layout - Row by default")
         .insert("c_layout", "R", "C tensor data layout - Row by default")
         .insert("batch_stride_a", "32768", "Batch A stride")
         .insert("batch_stride_b", "16384", "Batch B stride")

example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
+auto calculate_rtol_atol(const ck_tile::index_t K,
+                         const ck_tile::index_t kbatch,
+                         const float max_accumulated_value)
+{
+    using ComputeType =
+        std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
+    // Calculate thresholds
+    const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType, AccDataType>(
+        ck_tile::integer_divide_ceil(K, kbatch));
+    const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType, AccDataType>(
+        max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
+    // Calculate error due to split_k accumulation
+    const auto rtol_split_k =
+        ck_tile::get_relative_threshold<CDataType, CDataType, CDataType>(kbatch);
+    const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(
+        max_accumulated_value, kbatch);
+    // Use higher threshold
+    return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
+}
+
 template <typename ALayout, typename BLayout, typename CLayout>
 float invoke_batched_gemm(ck_tile::DeviceMem& a_m_k_dev_buf,
                           ck_tile::DeviceMem& b_k_n_dev_buf,
@@ -179,8 +199,18 @@ int run_batched_gemm_example_with_layouts(int argc,
 
         ck_tile::reference_batched_gemm<ADataType, BDataType, AccDataType, CDataType>(
             a_m_k, b_n_k, c_m_n_host_ref);
-
-        pass = ck_tile::check_err(c_m_n_dev_result, c_m_n_host_ref);
+        const float max_accumulated_value =
+            *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end());
+        const auto rtol_atol = calculate_rtol_atol(K, kbatch, max_accumulated_value);
+        pass                 = ck_tile::check_err(c_m_n_dev_result,
+                                  c_m_n_host_ref,
+                                  "Error: Incorrect results!",
+                                  rtol_atol.at(ck_tile::number<0>{}),
+                                  rtol_atol.at(ck_tile::number<1>{}));
+
+        std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{})
+                  << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{})
+                  << std::endl;
 
         std::cout << "The CPU veification result is:" << (pass ? "correct" : "fail") << std::endl;
     }
@@ -240,7 +270,18 @@ int run_batched_gemm_example_with_layouts(int argc,
         ck_tile::hip_check_error(hipFree(d_C));
 
         c_m_n_gpu_buf_ref.FromDevice(c_m_n_gpu_ref.data());
-        pass = ck_tile::check_err(c_m_n_dev_result, c_m_n_gpu_ref);
+        const float max_accumulated_value =
+            *std::max_element(c_m_n_gpu_ref.mData.begin(), c_m_n_gpu_ref.mData.end());
+        const auto rtol_atol = calculate_rtol_atol(K, kbatch, max_accumulated_value);
+        pass                 = ck_tile::check_err(c_m_n_dev_result,
+                                  c_m_n_gpu_ref,
+                                  "Error: Incorrect results!",
+                                  rtol_atol.at(ck_tile::number<0>{}),
+                                  rtol_atol.at(ck_tile::number<1>{}));
+
+        std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{})
+                  << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{})
+                  << std::endl;
 
         std::cout << "The GPU verification result is: " << (pass ? "correct" : "fail") << std::endl;
     }
@@ -260,11 +301,11 @@ int run_batched_gemm_example(int argc, char* argv[])
     std::string a_layout = arg_parser.get_str("a_layout");
     std::string b_layout = arg_parser.get_str("b_layout");
 
-    if(a_layout == "R" && b_layout == "R")
-    {
-        return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
-    }
-    else if(a_layout == "R" && b_layout == "C")
+    // if(a_layout == "R" && b_layout == "R")
+    // {
+    //     return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
+    // }
+    if(a_layout == "R" && b_layout == "C")
     {
         return run_batched_gemm_example_with_layouts(argc, argv, Row{}, Col{}, Row{});
     }

example/ck_tile/17_grouped_gemm/grouped_gemm.cpp

@@ -15,7 +15,6 @@
 #include "ck_tile/ops/gemm.hpp"
 #include "ck_tile/host.hpp"
 #include "grouped_gemm.hpp"
-#include "utils.hpp"
 
 namespace {
 
@@ -89,12 +88,9 @@ using CodegenPipelineProblem =
                                  CodegenGemmShape,
                                  CodegenGemmTraits<ALayout, BLayout, CLayout>>;
 
-using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
-
 template <typename ALayout, typename BLayout, typename CLayout>
 using CodegenGemmPipeline =
-    ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>,
-                                          CodegenGemmPolicy>;
+    ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>>;
 
 template <typename ALayout, typename BLayout, typename CLayout>
 using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner,
@@ -102,7 +98,7 @@ using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner,
                                           GemmEpilogue<CLayout>>;
 }; // namespace
 
-std::size_t GetWorkspaceSize(const std::vector<grouped_gemm_kargs>& gemm_descs)
+std::size_t get_workspace_size(const std::vector<grouped_gemm_kargs>& gemm_descs)
 {
     return ::Kernel<std::nullptr_t, std::nullptr_t, std::nullptr_t>::GetWorkSpaceSize(gemm_descs);
 }

example/ck_tile/17_grouped_gemm/grouped_gemm.hpp

@@ -41,7 +41,7 @@ auto create_args(int argc, char* argv[])
         .insert("stride_Bs", "", "Tensor B strides - it is empty by default.")
         .insert("stride_Cs", "", "Tensor C strides - it is empty by default.")
         .insert("a_layout", "R", "A tensor data layout - Row by default.")
-        .insert("b_layout", "R", "B tensor data layout - Row by default.")
+        .insert("b_layout", "C", "B tensor data layout - Row by default.")
         .insert("c_layout", "R", "C tensor data layout - Row by default.")
         .insert("validate", "1", "0. No validation, 1. Validation on CPU.")
         .insert("warmup", "10", "number of iterations before benchmark the kernel.")
@@ -52,8 +52,8 @@ auto create_args(int argc, char* argv[])
     return std::make_tuple(result, arg_parser);
 }
 
-std::size_t GetWorkspaceSize(const std::vector<grouped_gemm_kargs>& gemm_descs);
+std::size_t get_workspace_size(const std::vector<grouped_gemm_kargs>& gemm_descs);
 
-float grouped_gemm_calc(const std::vector<grouped_gemm_kargs>& gemm_descs,
+float grouped_gemm(const std::vector<grouped_gemm_kargs>& gemm_descs,
                    const ck_tile::stream_config& s,
                    void* p_workspace_);

example/ck_tile/17_grouped_gemm/utils.hpp

-// SPDX-License-Identifier: MIT
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-template <typename TLayout>
-constexpr auto
-f_host_tensor_descriptor(std::size_t row, std::size_t col, std::size_t stride, TLayout layout)
-{
-    using namespace ck_tile::literals;
-
-    if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
-    {
-        return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz});
-    }
-    else
-    {
-        return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride});
-    }
-}
-template <typename TLayout>
-constexpr auto
-f_get_default_stride(std::size_t row, std::size_t col, std::size_t stride, TLayout layout)
-{
-    if(stride == 0)
-    {
-        if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
-        {
-            return col;
-        }
-        else
-        {
-            return row;
-        }
-    }
-    else
-        return stride;
-}

include/ck/ck.hpp

@@ -17,7 +17,9 @@ CK_DECLARE_ENV_VAR_BOOL(CK_LOGGING)
 #endif
 // to do: add various levels of logging with CK_LOG_LEVEL
 
+#ifndef CK_TIME_KERNEL
 #define CK_TIME_KERNEL 1
+#endif
 
 // constant address space for kernel parameter
 // https://llvm.org/docs/AMDGPUUsage.html#address-spaces
@@ -155,6 +157,9 @@ CK_DECLARE_ENV_VAR_BOOL(CK_LOGGING)
 // LDS direct loads using inline assembly
 #define CK_USE_AMD_LDS_DIRECT_LOAD_INLINE_ASM 0
 
+// set rounding to nearest even as default for bf16 conversions
+#define CK_USE_RNE_BF16_CONVERSION 1
+
 // set rounding to nearest even as default for f8 conversions
 #define CK_USE_SR_F8_CONVERSION 0
 

include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.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
 
@@ -122,19 +122,6 @@ __global__ void
     static_for<0, NumDTensor, 1>{}(
         [&](auto i) { p_ds_grid_grp(i) = p_ds_grid[i] + ds_group_offset[i]; });
 
-    if constexpr(is_same_v<AElementwiseOperation, element_wise::DynamicUnaryOp>)
-    {
-        a_element_op.InitUnaryOpPtrOnDevice();
-    }
-    if constexpr(is_same_v<BElementwiseOperation, element_wise::DynamicUnaryOp>)
-    {
-        b_element_op.InitUnaryOpPtrOnDevice();
-    }
-    if constexpr(is_same_v<CDEElementwiseOperation, element_wise::DynamicUnaryOp>)
-    {
-        cde_element_op.InitUnaryOpPtrOnDevice();
-    }
-
     if constexpr(isMultiA || isMultiB)
     {
         AsPointer p_as_grid_grp;

include/ck/utility/data_type.hpp

@@ -31,8 +31,6 @@ struct pk_i4_t
     type data;
     __host__ __device__ constexpr pk_i4_t() : data{type{}} {}
     __host__ __device__ constexpr pk_i4_t(type init) : data{init} {}
-
-    __host__ __device__ constexpr operator float() const { return static_cast<int8_t>(data); }
 };
 
 inline constexpr auto next_pow2(uint32_t x)