Merge branch 'gemm_bf16_sk_muozturk' of...

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

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp

@@ -4,15 +4,15 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp"
 
 namespace ck_tile {
 
 //  A Tile Window: global memory
 //  B Tile Window: global memory
 //  C Distributed tensor: register
-template <typename Problem, typename Policy = BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy>
-struct BlockGemmPipelineAGmemBGmemCRegV1
+template <typename Problem, typename Policy = GemmPipelineAGmemBGmemCRegV1DefaultPolicy>
+struct GemmPipelineAGmemBGmemCRegV1
 {
     using ADataType      = remove_cvref_t<typename Problem::ADataType>;
     using BDataType      = remove_cvref_t<typename Problem::BDataType>;
@@ -33,6 +33,10 @@ struct BlockGemmPipelineAGmemBGmemCRegV1
     static constexpr bool kPadB = Problem::kPadB;
     static constexpr bool kPadC = Problem::kPadC;
 
+    using LayoutA = remove_cvref_t<typename Problem::LayoutA>;
+    using LayoutB = remove_cvref_t<typename Problem::LayoutB>;
+    using LayoutC = remove_cvref_t<typename Problem::LayoutC>;
+
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetStaticLdsSize()
     {
         return ck_tile::integer_divide_ceil(

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp

@@ -7,9 +7,9 @@
 
 namespace ck_tile {
 
-// Default policy for BlockGemmPipelineAGmemBGmemCRegV1
+// Default policy for GemmPipelineAGmemBGmemCRegV1
 // Default policy class should not be templated, put template on member functions instead
-struct BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy
+struct GemmPipelineAGmemBGmemCRegV1DefaultPolicy
 {
 #if 0
     // 2d

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp

@@ -4,15 +4,15 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
 
 namespace ck_tile {
 
 //  A Tile Window: global memory
 //  B Tile Window: global memory
 //  C Distributed tensor: register
-template <typename Problem, typename Policy = BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy>
-struct BlockGemmPipelineAGmemBGmemCRegV2
+template <typename Problem, typename Policy = GemmPipelineAGmemBGmemCRegV2DefaultPolicy>
+struct GemmPipelineAGmemBGmemCRegV2
 {
     using ADataType      = remove_cvref_t<typename Problem::ADataType>;
     using BDataType      = remove_cvref_t<typename Problem::BDataType>;

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp

@@ -7,12 +7,11 @@
 
 namespace ck_tile {
 
-// Default policy for BlockGemmPipelineAGmemBGmemCRegV2
+// Default policy for GemmPipelineAGmemBGmemCRegV2
 // Default policy class should not be templated, put template on member functions instead
 // NOTE: policy should be binded to its corresponding operation. It's just a coincidence that
-//   BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy is the same as
-//   BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy
-using BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy =
-    BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy;
+//   GemmPipelineAGmemBGmemCRegV2DefaultPolicy is the same as
+//   GemmPipelineAGmemBGmemCRegV1DefaultPolicy
+using GemmPipelineAGmemBGmemCRegV2DefaultPolicy = GemmPipelineAGmemBGmemCRegV1DefaultPolicy;
 
 } // namespace ck_tile

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_problem.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp

@@ -13,20 +13,23 @@ template <typename ADataType_,
           typename BDataType_,
           typename CDataType_,
           typename BlockGemmShape_,
-          bool kPadA_ = false,
-          bool kPadB_ = false,
-          bool kPadC_ = false>
-struct BlockGemmPipelineProblem
+          typename TileGemmTraits_>
+struct GemmPipelineProblem
 {
     using ADataType      = remove_cvref_t<ADataType_>;
     using BDataType      = remove_cvref_t<BDataType_>;
     using CDataType      = remove_cvref_t<CDataType_>;
     using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
+    using GemmTraits     = remove_cvref_t<TileGemmTraits_>;
 
     static constexpr index_t kBlockSize = BlockGemmShape::NumWarps * get_warp_size();
-    static constexpr bool kPadA         = kPadA_;
-    static constexpr bool kPadB         = kPadB_;
-    static constexpr bool kPadC         = kPadC_;
+    static constexpr bool kPadA         = GemmTraits::kPadA;
+    static constexpr bool kPadB         = GemmTraits::kPadB;
+    static constexpr bool kPadC         = GemmTraits::kPadC;
+
+    using LayoutA = remove_cvref_t<typename GemmTraits::LayoutA>;
+    using LayoutB = remove_cvref_t<typename GemmTraits::LayoutB>;
+    using LayoutC = remove_cvref_t<typename GemmTraits::LayoutC>;
 
     static constexpr index_t AlignmentA = kPadA ? 1 : VectorLoadSize / sizeof(ADataType);
     static constexpr index_t AlignmentB = kPadB ? 1 : VectorLoadSize / sizeof(BDataType);

include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+
+namespace ck_tile {
+
+template <bool kPadA_,
+          bool kPadB_,
+          bool kPadC_,
+          typename LayoutA_,
+          typename LayoutB_,
+          typename LayoutC_>
+struct TileGemmTraits
+{
+    static constexpr bool kPadA = kPadA_;
+    static constexpr bool kPadB = kPadB_;
+    static constexpr bool kPadC = kPadC_;
+
+    using LayoutA = LayoutA_;
+    using LayoutB = LayoutB_;
+    using LayoutC = LayoutC_;
+};
+
+} // namespace ck_tile

include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_kernel.hpp

@@ -31,8 +31,14 @@ struct Layernorm2dFwd
 
     static constexpr ck_tile::index_t kMPerBlock = Problem::BlockShape::kMPerBlock;
     static constexpr ck_tile::index_t kNPerBlock = Problem::BlockShape::kNPerBlock;
+    static constexpr bool kPadM                  = Problem::kPadM;
+    static constexpr bool kPadN                  = Problem::kPadN;
 
     static constexpr ck_tile::index_t kNThreadPerWarp = Problem::BlockShape::kNThreadPerWarp;
+    static constexpr ck_tile::index_t kNPerThread     = Problem::BlockShape::kNPerThread;
+
+    static constexpr auto I0 = number<0>{};
+    static constexpr auto I1 = number<1>{};
 
     struct Kargs
     {
@@ -96,19 +102,25 @@ struct Layernorm2dFwd
                 sequence<2>>{});
     }
 
-    template <typename Dstr>
-    CK_TILE_DEVICE static constexpr auto GetNPerThread(Dstr)
+    CK_TILE_DEVICE static int GetWelfordMaxCount(int N)
     {
-        constexpr auto nDstrSpan = Dstr::get_distributed_spans().template at<1>();
-
-        using Lengths = decltype(nDstrSpan.impl_);
+        constexpr ck_tile::index_t kNThreadPerBlock = kNPerBlock / kNPerThread;
 
-        ck_tile::index_t ret = 1;
+        int thread_id_n = get_thread_id() % kNThreadPerBlock;
+        int max_count =
+            __builtin_amdgcn_readfirstlane(N < kNPerBlock ? 0 : kNPerThread * (N / kNPerBlock));
+        int n_per_block_tail_loop =
+            __builtin_amdgcn_readfirstlane(N - max_count * kNThreadPerBlock);
 
-        ck_tile::static_for<0, Lengths::size(), 1>{}(
-            [&](auto idx) { ret *= Lengths::template at(idx); });
+        if(n_per_block_tail_loop > 0)
+        {
+            int thread_max_n = (thread_id_n + 1) * kNPerThread;
+            int delta        = thread_max_n - n_per_block_tail_loop;
+            delta            = clamp(thread_max_n - n_per_block_tail_loop, 0, kNPerThread);
+            max_count += kNPerThread - delta;
+        }
 
-        return ret;
+        return max_count;
     }
 
     template <typename DistributedTensor>
@@ -129,42 +141,29 @@ struct Layernorm2dFwd
         return out_dstr_tensor;
     }
 
-    template <bool Cond = (kHasGamma && kHasBeta)>
-    CK_TILE_DEVICE std::enable_if_t<Cond> TwoPassLayernorm2dFwd(const XDataType* p_x,
-                                                                const GammaDataType* p_gamma,
-                                                                const BetaDataType* p_beta,
-                                                                YDataType* p_y,
-                                                                MeanDataType* p_mean,
-                                                                InvStdDataType* p_invStd,
-                                                                const ComputeDataType epsilon,
-                                                                ck_tile::index_t M,
-                                                                ck_tile::index_t N) const
+    template <typename XBlockWindow,
+              typename GammaBlockWindow,
+              typename BetaBlockWindow,
+              typename YBlockWindow,
+              typename MeanBlockWindow,
+              typename InvStdBlockWindow,
+              bool Cond = (kHasGamma && kHasBeta)>
+    CK_TILE_DEVICE std::enable_if_t<Cond>
+    TwoPassLayernorm2dFwd(XBlockWindow& x_block_window,
+                          GammaBlockWindow& gamma_block_window,
+                          BetaBlockWindow& beta_block_window,
+                          YBlockWindow& y_block_window,
+                          MeanBlockWindow& mean_block_window,
+                          InvStdBlockWindow& inv_std_block_window,
+                          ComputeDataType epsilon,
+                          ck_tile::index_t N) const
     {
-        constexpr auto I0 = number<0>{};
-        constexpr auto I1 = number<1>{};
-
-        const auto x_m_n = make_naive_tensor_view<address_space_enum::global>(
-            p_x, make_tuple(M, N), make_tuple(N, 1), number<32>{}, number<1>{});
-
-        const auto gamma_n = make_naive_tensor_view<address_space_enum::global>(
-            p_gamma, make_tuple(N), make_tuple(1), number<32>{}, number<1>{});
+        // TODO - Optimize tail loop to reduce move_tile_window()
+        index_t num_n_tile_iteration =
+            __builtin_amdgcn_readfirstlane(integer_divide_ceil(N, kNPerBlock));
 
-        const auto beta_n = make_naive_tensor_view<address_space_enum::global>(
-            p_beta, make_tuple(N), make_tuple(1), number<32>{}, number<1>{});
-
-        const auto iM = get_block_id() * kMPerBlock;
-
-        constexpr auto xDstr = MakeXBlockTileDistribution();
-
-        auto x_block_window = make_tile_window(
-            x_m_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, 0}, xDstr);
-
-        index_t num_n_tile_iteration = __builtin_amdgcn_readfirstlane(N / kNPerBlock);
-
-        // TODO: padding - handle max_count if N % kNPerBlock != 0
-        constexpr auto NPerThread = GetNPerThread(xDstr);
-        ThreadWelford<ComputeDataType, XDataType> thread_welford{
-            type_convert<int>(NPerThread * N / kNPerBlock)};
+        int welford_max_count = GetWelfordMaxCount(N);
+        ThreadWelford<ComputeDataType, XDataType> thread_welford{welford_max_count};
 
         using XTensorType = decltype(load_tile(x_block_window));
         auto mean_compute_block_tensor =
@@ -190,44 +189,14 @@ struct Layernorm2dFwd
         auto inv_std_compute_block_tensor = InvSqrt(var_compute_block_tensor, epsilon);
 
         if constexpr(kSaveMean)
-        {
-            const auto mean_m = make_naive_tensor_view_packed<address_space_enum::global>(
-                p_mean, make_tuple(M), number<32>{});
-
-            auto mean_block_window =
-                make_tile_window(mean_m, make_tuple(number<kMPerBlock>{}), {iM});
-
             store_tile(mean_block_window, cast_tile<MeanDataType>(mean_compute_block_tensor));
-        }
         if constexpr(kSaveInvStd)
-        {
-            const auto inv_std_m = make_naive_tensor_view_packed<address_space_enum::global>(
-                p_invStd, make_tuple(M), number<32>{});
-
-            auto inv_std_block_window =
-                make_tile_window(inv_std_m, make_tuple(number<kMPerBlock>{}), {iM});
-
-            store_tile(inv_std_block_window, cast_tile<MeanDataType>(inv_std_compute_block_tensor));
-        }
-
-        // TODO: Extract normalize pipeline
-        const auto y_m_n = make_naive_tensor_view<address_space_enum::global>(
-            p_y, make_tuple(M, N), make_tuple(N, 1), number<32>{}, number<1>{});
-
-        auto y_block_window = make_tile_window(
-            y_m_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, 0});
-
-        constexpr auto gammaDstr = MakeGammaBetaBlockTileDistribution();
-        constexpr auto betaDstr  = gammaDstr;
-
-        auto gamma_block_window =
-            make_tile_window(gamma_n, make_tuple(number<kNPerBlock>{}), {0}, gammaDstr);
-
-        auto beta_block_window = make_tile_window(
-            beta_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {0}, betaDstr);
+            store_tile(inv_std_block_window,
+                       cast_tile<InvStdDataType>(inv_std_compute_block_tensor));
 
         // reverse read x to reuse cache
-        ck_tile::index_t stride_to_right_most_window = N - kNPerBlock;
+        ck_tile::index_t stride_to_right_most_window =
+            N % kNPerBlock == 0 ? N - kNPerBlock : N - N % kNPerBlock;
 
         move_tile_window(x_block_window, {0, -kNPerBlock});
         move_tile_window(gamma_block_window, {stride_to_right_most_window});
@@ -274,17 +243,209 @@ struct Layernorm2dFwd
         }
     }
 
+    template <typename XBlockWindow,
+              typename GammaBlockWindow,
+              typename BetaBlockWindow,
+              typename YBlockWindow,
+              typename MeanBlockWindow,
+              typename InvStdBlockWindow,
+              bool Cond = (kHasGamma && kHasBeta)>
+    CK_TILE_DEVICE std::enable_if_t<Cond>
+    OnePassLayernorm2dFwd(XBlockWindow& x_block_window,
+                          GammaBlockWindow& gamma_block_window,
+                          BetaBlockWindow& beta_block_window,
+                          YBlockWindow& y_block_window,
+                          MeanBlockWindow& mean_block_window,
+                          InvStdBlockWindow& inv_std_block_window,
+                          ComputeDataType epsilon,
+                          ck_tile::index_t N) const
+    {
+        int welford_max_count = GetWelfordMaxCount(N);
+        ThreadWelford<ComputeDataType, XDataType> thread_welford{welford_max_count};
+
+        using XTensorType = decltype(load_tile(x_block_window));
+        auto mean_compute_block_tensor =
+            thread_welford.template MakeInitialMeanVarDistributedTensor<XTensorType>();
+        auto var_compute_block_tensor =
+            thread_welford.template MakeInitialMeanVarDistributedTensor<XTensorType>();
+
+        clear_tile(mean_compute_block_tensor);
+        clear_tile(var_compute_block_tensor);
+
+        const auto x_block_tensor = load_tile(x_block_window);
+        thread_welford(x_block_tensor, mean_compute_block_tensor, var_compute_block_tensor);
+        // TODO: support cross warp Welford
+        WarpMergeWelford<ComputeDataType, true>{}(
+            mean_compute_block_tensor, var_compute_block_tensor, thread_welford.cur_count_);
+
+        auto inv_std_compute_block_tensor = InvSqrt(var_compute_block_tensor, epsilon);
+
+        if constexpr(kSaveMean)
+            store_tile(mean_block_window, cast_tile<MeanDataType>(mean_compute_block_tensor));
+        if constexpr(kSaveInvStd)
+            store_tile(inv_std_block_window,
+                       cast_tile<InvStdDataType>(inv_std_compute_block_tensor));
+
+        // normalize
+        const auto gamma_block_tensor = load_tile(gamma_block_window);
+        const auto beta_block_tensor  = load_tile(beta_block_window);
+
+        constexpr auto x_spans = decltype(x_block_tensor)::get_distributed_spans();
+
+        auto y_block_tensor =
+            make_static_distributed_tensor<YDataType>(x_block_tensor.get_tile_distribution());
+
+        sweep_tile_span(x_spans[I1], [&](auto idx1) {
+            constexpr auto j_idx = make_tuple(idx1);
+            const auto gamma     = type_convert<ComputeDataType>(gamma_block_tensor[j_idx]);
+            const auto beta      = type_convert<ComputeDataType>(beta_block_tensor[j_idx]);
+
+            sweep_tile_span(x_spans[I0], [&](auto idx0) {
+                constexpr auto i_idx   = make_tuple(idx0);
+                constexpr auto i_j_idx = make_tuple(idx0, idx1);
+
+                const auto mean    = mean_compute_block_tensor[i_idx];
+                const auto inv_std = inv_std_compute_block_tensor[i_idx];
+
+                const auto x = type_convert<ComputeDataType>(x_block_tensor[i_j_idx]);
+                auto y       = (x - mean) * inv_std * gamma + beta;
+
+                y_block_tensor(i_j_idx) = type_convert<YDataType>(y);
+            });
+        });
+
+        store_tile(y_block_window, y_block_tensor);
+    }
+
     CK_TILE_DEVICE void operator()(Kargs kargs) const
     {
-        TwoPassLayernorm2dFwd(static_cast<const XDataType*>(kargs.p_x),
-                              static_cast<const GammaDataType*>(kargs.p_gamma),
-                              static_cast<const BetaDataType*>(kargs.p_beta),
-                              static_cast<YDataType*>(kargs.p_y),
-                              static_cast<MeanDataType*>(kargs.p_mean),
-                              static_cast<InvStdDataType*>(kargs.p_invStd),
-                              static_cast<const ComputeDataType>(kargs.epsilon),
-                              kargs.M,
-                              kargs.N);
+        const auto x_m_n = [&]() {
+            const auto x_dram_naive = make_naive_tensor_view<address_space_enum::global>(
+                static_cast<const XDataType*>(kargs.p_x),
+                make_tuple(kargs.M, kargs.N),
+                make_tuple(kargs.N, 1),
+                number<kNPerThread>{},
+                number<1>{});
+
+            return pad_tensor_view(x_dram_naive,
+                                   make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
+                                   sequence<kPadM, kPadN>{});
+        }();
+
+        const auto gamma_n = [&]() {
+            const auto gamma_dram_naive = make_naive_tensor_view<address_space_enum::global>(
+                static_cast<const GammaDataType*>(kargs.p_gamma),
+                make_tuple(kargs.N),
+                make_tuple(1),
+                number<kNPerThread>{},
+                number<1>{});
+
+            return pad_tensor_view(
+                gamma_dram_naive, make_tuple(number<kNPerBlock>{}), sequence<kPadN>{});
+        }();
+
+        const auto beta_n = [&]() {
+            const auto gamma_dram_naive = make_naive_tensor_view<address_space_enum::global>(
+                static_cast<const BetaDataType*>(kargs.p_beta),
+                make_tuple(kargs.N),
+                make_tuple(1),
+                number<kNPerThread>{},
+                number<1>{});
+
+            return pad_tensor_view(
+                gamma_dram_naive, make_tuple(number<kNPerBlock>{}), sequence<kPadN>{});
+        }();
+
+        const auto iM = get_block_id() * kMPerBlock;
+
+        constexpr auto xDstr = MakeXBlockTileDistribution();
+
+        auto x_block_window = make_tile_window(
+            x_m_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, 0}, xDstr);
+
+        const auto y_m_n = [&]() {
+            const auto y_dram_naive = make_naive_tensor_view<address_space_enum::global>(
+                static_cast<YDataType*>(kargs.p_y),
+                make_tuple(kargs.M, kargs.N),
+                make_tuple(kargs.N, 1),
+                number<kNPerThread>{},
+                number<1>{});
+
+            return pad_tensor_view(y_dram_naive,
+                                   make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
+                                   sequence<kPadM, kPadN>{});
+        }();
+
+        auto y_block_window = make_tile_window(
+            y_m_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, 0});
+
+        constexpr auto gammaDstr = MakeGammaBetaBlockTileDistribution();
+        constexpr auto betaDstr  = gammaDstr;
+
+        auto gamma_block_window =
+            make_tile_window(gamma_n, make_tuple(number<kNPerBlock>{}), {0}, gammaDstr);
+
+        auto beta_block_window = make_tile_window(
+            beta_n, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {0}, betaDstr);
+
+        auto mean_block_window = [&]() {
+            if constexpr(kSaveMean)
+            {
+                const auto mean_m = [&]() {
+                    const auto mean_dram_naive =
+                        make_naive_tensor_view_packed<address_space_enum::global>(
+                            static_cast<MeanDataType*>(kargs.p_mean),
+                            make_tuple(kargs.M),
+                            number<1>{});
+
+                    return pad_tensor_view(
+                        mean_dram_naive, make_tuple(number<kMPerBlock>{}), sequence<kPadM>{});
+                }();
+
+                return make_tile_window(mean_m, make_tuple(number<kMPerBlock>{}), {iM});
+            }
+            else
+                return make_null_tile_window(make_tuple(number<kMPerBlock>{}));
+        }();
+
+        auto inv_std_block_window = [&]() {
+            if constexpr(kSaveInvStd)
+            {
+                const auto inv_std_m = [&]() {
+                    const auto inv_std_dram_naive =
+                        make_naive_tensor_view_packed<address_space_enum::global>(
+                            static_cast<InvStdDataType*>(kargs.p_invStd),
+                            make_tuple(kargs.M),
+                            number<1>{});
+
+                    return pad_tensor_view(
+                        inv_std_dram_naive, make_tuple(number<kMPerBlock>{}), sequence<kPadM>{});
+                }();
+
+                return make_tile_window(inv_std_m, make_tuple(number<kMPerBlock>{}), {iM});
+            }
+            else
+                return make_null_tile_window(make_tuple(number<kMPerBlock>{}));
+        }();
+
+        if(kargs.N <= kNPerBlock)
+            OnePassLayernorm2dFwd(x_block_window,
+                                  gamma_block_window,
+                                  beta_block_window,
+                                  y_block_window,
+                                  mean_block_window,
+                                  inv_std_block_window,
+                                  static_cast<const ComputeDataType>(kargs.epsilon),
+                                  kargs.N);
+        else
+            TwoPassLayernorm2dFwd(x_block_window,
+                                  gamma_block_window,
+                                  beta_block_window,
+                                  y_block_window,
+                                  mean_block_window,
+                                  inv_std_block_window,
+                                  static_cast<const ComputeDataType>(kargs.epsilon),
+                                  kargs.N);
     }
 };
 

include/ck_tile/ops/layernorm2d/pipeline/block_layernorm2d_fwd_problem.hpp

@@ -14,17 +14,21 @@ template <typename XDataType_,
           typename YDataType_,
           typename MeanDataType_,
           typename InvStdDataType_,
-          typename BlockShape_>
+          typename BlockShape_,
+          bool kPadM_,
+          bool kPadN_>
 struct BlockLayernorm2dFwdProblem
 {
-    using XDataType       = remove_cvref_t<XDataType_>;
-    using GammaDataType   = remove_cvref_t<GammaDataType_>;
-    using BetaDataType    = remove_cvref_t<BetaDataType_>;
-    using ComputeDataType = remove_cvref_t<ComputeDataType_>;
-    using YDataType       = remove_cvref_t<YDataType_>;
-    using MeanDataType    = remove_cvref_t<MeanDataType_>;
-    using InvStdDataType  = remove_cvref_t<InvStdDataType_>;
-    using BlockShape      = remove_cvref_t<BlockShape_>;
+    using XDataType             = remove_cvref_t<XDataType_>;
+    using GammaDataType         = remove_cvref_t<GammaDataType_>;
+    using BetaDataType          = remove_cvref_t<BetaDataType_>;
+    using ComputeDataType       = remove_cvref_t<ComputeDataType_>;
+    using YDataType             = remove_cvref_t<YDataType_>;
+    using MeanDataType          = remove_cvref_t<MeanDataType_>;
+    using InvStdDataType        = remove_cvref_t<InvStdDataType_>;
+    using BlockShape            = remove_cvref_t<BlockShape_>;
+    static constexpr bool kPadM = kPadM_;
+    static constexpr bool kPadN = kPadN_;
 };
 
 } // namespace ck_tile

library/include/ck/library/reference_tensor_operation/gpu/reference_gemm.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+
+#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+
+namespace ck {
+
+template <typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename ADataType,
+          typename BDataType,
+          typename CDataType,
+          typename AccDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CDEElementwiseOperation,
+          typename ComputeTypeA,
+          typename ComputeTypeB>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        naive_gemm_kernel(const ADataType* __restrict__ p_a_grid,
+                          const BDataType* __restrict__ p_b_grid,
+                          CDataType* __restrict__ p_c_grid,
+                          index_t m,
+                          index_t n,
+                          index_t k,
+                          const AElementwiseOperation a_element_op,
+                          const BElementwiseOperation b_element_op,
+                          const CDEElementwiseOperation c_element_op)
+{
+    using RowMajor = ck::tensor_layout::gemm::RowMajor;
+
+    const int row_idx = blockIdx.x * blockDim.x + threadIdx.x;
+    const int col_idx = blockIdx.y * blockDim.y + threadIdx.y;
+
+    if(row_idx < m && col_idx < n)
+    {
+
+        AccDataType v_acc = static_cast<AccDataType>(0.0);
+        ComputeTypeA v_a  = static_cast<ComputeTypeA>(0.0);
+        ComputeTypeB v_b  = static_cast<ComputeTypeB>(0.0);
+        CDataType v_c     = static_cast<CDataType>(0.0);
+
+        for(int k_idx = 0; k_idx < k; ++k_idx)
+        {
+            // check input matrices layout
+            int element_idx_a = 0;
+            int element_idx_b = 0;
+            if constexpr(std::is_same_v<ALayout, RowMajor>)
+            {
+                element_idx_a = row_idx * k + k_idx;
+            }
+            else
+            {
+                element_idx_a = row_idx + m * k_idx;
+            }
+            if constexpr(std::is_same_v<BLayout, RowMajor>)
+            {
+                element_idx_b = k_idx * n + col_idx;
+            }
+            else
+            {
+                element_idx_b = k_idx + k * col_idx;
+            }
+            // apply a_element_op
+            a_element_op(v_a, p_a_grid[element_idx_a]);
+            // apply b_element_op
+            b_element_op(v_b, p_b_grid[element_idx_b]);
+            // multiply and accumulate
+            v_acc += static_cast<AccDataType>(v_a) * static_cast<AccDataType>(v_b);
+        }
+        // apply c_element_op
+        c_element_op(v_c, v_acc);
+        // check output matrix layout
+        int element_idx_c = 0;
+        if constexpr(std::is_same_v<CLayout, RowMajor>)
+        {
+            element_idx_c = row_idx * n + col_idx;
+        }
+        else
+        {
+            element_idx_c = row_idx + m * col_idx;
+        }
+        // prepare output
+        p_c_grid[element_idx_c] = v_c;
+    }
+}
+
+} // namespace ck
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename ALayout,
+          typename BLayout,
+          typename CLayout,
+          typename ADataType,
+          typename BDataType,
+          typename CDataType,
+          typename AccDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation,
+          typename ComputeTypeA = CDataType,
+          typename ComputeTypeB = ComputeTypeA>
+struct ReferenceGemm : public device::BaseOperator
+{
+    // Argument
+    struct Argument : public device::BaseArgument
+    {
+        Argument(const void* p_a_grid,
+                 const void* p_b_grid,
+                 void* p_c_grid,
+                 index_t m,
+                 index_t n,
+                 index_t k,
+                 AElementwiseOperation a_element_op,
+                 BElementwiseOperation b_element_op,
+                 CElementwiseOperation c_element_op)
+            : p_a_grid_{static_cast<const ADataType*>(p_a_grid)},
+              p_b_grid_{static_cast<const BDataType*>(p_b_grid)},
+              p_c_grid_{static_cast<CDataType*>(p_c_grid)},
+              m_{m},
+              n_{n},
+              k_{k},
+              a_element_op_{a_element_op},
+              b_element_op_{b_element_op},
+              c_element_op_{c_element_op}
+        {
+        }
+
+        const ADataType* p_a_grid_;
+        const BDataType* p_b_grid_;
+        CDataType* p_c_grid_;
+
+        index_t m_;
+        index_t n_;
+        index_t k_;
+
+        AElementwiseOperation a_element_op_;
+        BElementwiseOperation b_element_op_;
+        CElementwiseOperation c_element_op_;
+    };
+
+    // Invoker
+    struct Invoker : public device::BaseInvoker
+    {
+        using Argument = ReferenceGemm::Argument;
+
+        float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
+        {
+            int block_size = 16;
+            dim3 block_dim(block_size, block_size, 1);
+            dim3 grid_dim(
+                (arg.m_ + block_size - 1) / block_size, (arg.n_ + block_size - 1) / block_size, 1);
+
+            auto launch_kernel = [&]() {
+                const auto kernel = naive_gemm_kernel<ALayout,
+                                                      BLayout,
+                                                      CLayout,
+                                                      ADataType,
+                                                      BDataType,
+                                                      CDataType,
+                                                      AccDataType,
+                                                      AElementwiseOperation,
+                                                      BElementwiseOperation,
+                                                      CElementwiseOperation,
+                                                      ComputeTypeA,
+                                                      ComputeTypeB>;
+
+                return launch_and_time_kernel(stream_config,
+                                              kernel,
+                                              grid_dim,
+                                              block_dim,
+                                              0,
+                                              arg.p_a_grid_,
+                                              arg.p_b_grid_,
+                                              arg.p_c_grid_,
+                                              arg.m_,
+                                              arg.n_,
+                                              arg.k_,
+                                              arg.a_element_op_,
+                                              arg.b_element_op_,
+                                              arg.c_element_op_);
+            };
+
+            return launch_kernel();
+        }
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& stream_config = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
+        }
+    };
+
+    bool IsSupportedArgument(const device::BaseArgument*) override { return true; }
+
+    static auto MakeArgument(const void* p_a_grid,
+                             const void* p_b_grid,
+                             void* p_c_grid,
+                             index_t m,
+                             index_t n,
+                             index_t k,
+                             AElementwiseOperation a_element_op,
+                             BElementwiseOperation b_element_op,
+                             CElementwiseOperation c_element_op)
+    {
+        return Argument{
+            p_a_grid, p_b_grid, p_c_grid, m, n, k, a_element_op, b_element_op, c_element_op};
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    virtual std::unique_ptr<device::BaseInvoker> MakeInvokerPointer()
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "Device Reference Gemm"
+            << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/CMakeLists.txt

@@ -37,11 +37,7 @@ function(add_instance_library INSTANCE_NAME)
         endforeach()
     endif()
 
-    if(INSTANCES_ONLY)
-        set(INST_TARGETS ${DEFAULT_GPU_TARGETS})
-    else()
-        set(INST_TARGETS ${GPU_TARGETS})
-    endif()
+    set(INST_TARGETS ${SUPPORTED_GPU_TARGETS})
 
     # Do not build DL instances if DL_KERNELS macro is not set
     foreach(source IN LISTS ARGN)
@@ -75,11 +71,7 @@ function(add_instance_library INSTANCE_NAME)
     if(ARGN)
         set(INST_OBJ)
         foreach(source IN LISTS ARGN)
-            if(INSTANCES_ONLY)
-                set(INST_TARGETS ${DEFAULT_GPU_TARGETS})
-            else()
-                set(INST_TARGETS ${GPU_TARGETS})
-            endif()
+            set(INST_TARGETS ${SUPPORTED_GPU_TARGETS})
             if(source MATCHES "_xdl")
                 list(REMOVE_ITEM INST_TARGETS gfx900 gfx906 gfx1030 gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201)
             elseif(ARGN MATCHES "_wmma")
@@ -191,12 +183,7 @@ FOREACH(subdir_path ${dir_list})
             set(add_inst 1)
         endif()
 
-        if(INSTANCES_ONLY)
-            set(INST_TARGETS ${DEFAULT_GPU_TARGETS})
-        else()
-            set(INST_TARGETS ${GPU_TARGETS})
-        endif()
-
+        set(INST_TARGETS ${SUPPORTED_GPU_TARGETS})
 
         if(("${cmake_instance}" MATCHES "quantization") AND (DEFINED DTYPES) AND (NOT DTYPES MATCHES "int8"))
             message("quantization instances will not be built!")

profiler/src/CMakeLists.txt

@@ -24,7 +24,7 @@ set(PROFILER_SOURCES
     profile_permute_scale.cpp
 )
 
-if(GPU_TARGETS MATCHES "gfx9")
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx9")
   if(DTYPES MATCHES "fp32" OR DTYPES MATCHES "fp64" OR NOT DEFINED DTYPES)
     list(APPEND PROFILER_SOURCES profile_contraction_bilinear.cpp)
     list(APPEND PROFILER_SOURCES profile_contraction_scale.cpp)
@@ -49,7 +49,7 @@ if(GPU_TARGETS MATCHES "gfx9")
     list(APPEND PROFILER_SOURCES profile_grouped_gemm_multiply_tile_loop.cpp)
   endif()
   list(APPEND PROFILER_SOURCES profile_gemm_multiply_add.cpp)
-  if(GPU_TARGETS MATCHES "gfx94")
+  if(SUPPORTED_GPU_TARGETS MATCHES "gfx94")
     list(APPEND PROFILER_SOURCES profile_gemm_multiply_multiply.cpp)
     list(APPEND PROFILER_SOURCES profile_gemm_ab_scale.cpp)
   endif()
@@ -69,7 +69,7 @@ if(GPU_TARGETS MATCHES "gfx9")
 
 endif()
 
-if(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12" OR GPU_TARGETS MATCHES "gfx9")
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12" OR SUPPORTED_GPU_TARGETS MATCHES "gfx9")
   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
     list(APPEND PROFILER_SOURCES profile_gemm_bilinear.cpp)
   endif()
@@ -111,7 +111,7 @@ target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_column_to_image_inst
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_transpose_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_permute_scale_instance)
 
-if(GPU_TARGETS MATCHES "gfx9")
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx9")
   if(DTYPES MATCHES "fp32" OR DTYPES MATCHES "fp64" OR NOT DEFINED DTYPES)
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_contraction_bilinear_instance)
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_contraction_scale_instance)
@@ -135,7 +135,7 @@ if(GPU_TARGETS MATCHES "gfx9")
   target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_batched_gemm_instance)
   target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_batched_gemm_reduce_instance)
   target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_multiply_add_instance)
-  if(GPU_TARGETS MATCHES "gfx94")
+  if(SUPPORTED_GPU_TARGETS MATCHES "gfx94")
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_multiply_multiply_instance)
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_ab_scale_instance)
   endif()
@@ -159,7 +159,7 @@ if(GPU_TARGETS MATCHES "gfx9")
   target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_grouped_conv3d_fwd_convinvscale_instance)
 endif()
 
-if(GPU_TARGETS MATCHES "gfx9" OR GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12")
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx9" OR SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12")
   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
     target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_gemm_bilinear_instance)
   endif()

script/cmake-ck-dev.sh

@@ -7,15 +7,16 @@ MY_PROJECT_SOURCE=$1
 
 if [ $# -ge 2 ] ; then
     GPU_TARGETS=$2
-    REST_ARGS=${@:3}
+    shift 2
+    REST_ARGS=$@
 else
     GPU_TARGETS="gfx908;gfx90a;gfx940"
     REST_ARGS=
 fi
 
 cmake                                                                                             \
--D CMAKE_PREFIX_PATH=/opt/rocm                                                                    \
--D CMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc                                                         \
+-D CMAKE_PREFIX_PATH=/opt/rocm-6.2.1/                                                                 \
+-D CMAKE_CXX_COMPILER=/opt/rocm-6.2.1/bin/hipcc                                                        \
 -D CMAKE_CXX_FLAGS="-Xclang -mllvm -Xclang -enable-post-misched=0 -std=c++17 -O3 -ftemplate-backtrace-limit=0  -fPIE  -Wno-gnu-line-marker"     \
 -D CMAKE_BUILD_TYPE=Release                                                                       \
 -D BUILD_DEV=ON                                                                                   \
@@ -24,3 +25,7 @@ cmake
 -D USE_BITINT_EXTENSION_INT4=OFF                                                                  \
 $REST_ARGS                                                                                        \
 ${MY_PROJECT_SOURCE}
+
+
+# -D CMAKE_PREFIX_PATH=/opt/rocm                                                                    \
+# -D CMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc     
\ No newline at end of file

script/cmake-ck-release.sh

@@ -7,7 +7,8 @@ MY_PROJECT_SOURCE=$1
 
 if [ $# -ge 2 ] ; then
     GPU_TARGETS=$2
-    REST_ARGS=${@:3}
+    shift 2
+    REST_ARGS=$@
 else
     GPU_TARGETS="gfx908;gfx90a;gfx940"
     REST_ARGS=

test/CMakeLists.txt

@@ -41,11 +41,7 @@ function(add_test_executable TEST_NAME)
         endforeach()
     endif()
 
-    if(INSTANCES_ONLY)
-        set(TEST_TARGETS ${DEFAULT_GPU_TARGETS})
-    else()
-        set(TEST_TARGETS ${GPU_TARGETS})
-    endif()
+    set(TEST_TARGETS ${SUPPORTED_GPU_TARGETS})
 
     foreach(source IN LISTS ARGN)
         if(NOT DEFINED DL_KERNELS AND source MATCHES "_dl")
@@ -122,11 +118,7 @@ function(add_gtest_executable TEST_NAME)
         endforeach()
     endif()
 
-    if(INSTANCES_ONLY)
-        set(TEST_TARGETS ${DEFAULT_GPU_TARGETS})
-    else()
-        set(TEST_TARGETS ${GPU_TARGETS})
-    endif()
+    set(TEST_TARGETS ${SUPPORTED_GPU_TARGETS})
 
     foreach(source IN LISTS ARGN)
         if(NOT DEFINED DL_KERNELS AND source MATCHES "_dl")
@@ -211,10 +203,10 @@ add_subdirectory(conv_tensor_rearrange)
 add_subdirectory(transpose)
 add_subdirectory(permute_scale)
 add_subdirectory(wrapper)
-if(GPU_TARGETS MATCHES "gfx11")
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx11")
     add_subdirectory(wmma_op)
 endif()
-if(GPU_TARGETS MATCHES "gfx942" AND CK_HIP_VERSION_MAJOR GREATER_EQUAL 6 AND CK_HIP_VERSION_MINOR GREATER_EQUAL 2) # smfmac needs ROCm6.2
+if(SUPPORTED_GPU_TARGETS MATCHES "gfx942" AND CK_HIP_VERSION_MAJOR GREATER_EQUAL 6 AND CK_HIP_VERSION_MINOR GREATER_EQUAL 2) # smfmac needs ROCm6.2
     add_subdirectory(smfmac_op)
 endif()
 add_subdirectory(position_embedding)