Merge branch 'develop' into layernorm/instance_support

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_bf16_n512_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 8,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  2,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  4,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  8,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_bf16_n64_n128_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_bf16_n768_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  3,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  6,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1, 12,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n1024_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+#if 0
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 8,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  8,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1, 16,  4,  64, 1,  true , false, false>>(const S&, A);
+
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  1, 256, 4,  true , false, false>>(const S&, A);
+#endif
+
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 1, 2,  128, 8,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 2,  128, 4,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 2,  128, 2,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 1,  true,  false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n1536_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 4,   64, 8,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 2,  128, 4,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  256, 2,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 6, 1,  256, 1,  true,  false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n2048_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 1, 1,  256, 8,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 4,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 2,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 8, 1,  256, 1,  true,  false, false>>(const S&, A);
+
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n256_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n3072_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  128, 8,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1,  256, 4,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 6, 1,  256, 2,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 3, 1, 1024, 1,  true,  false, false>>(const S&, A);
+
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n4096_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true,  false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true,  false, false>>(const S&, A);
+
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n4096_tp_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true,  false, true>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true,  false, true>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true,  false, true>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true,  false, true>>(const S&, A);
+
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n512_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 8,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  4,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  8,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n64_n128_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n768_instance.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "layernorm2d_fwd_instance_common.hpp"
+
+// clang-format off
+//                                                       rm  rn  tm  tn  vn  pd     mv     2p
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  3,  4,  64, 4,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  6,  4,  64, 2,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1, 12,  4,  64, 1,  true , false, false>>(const S&, A);
+// clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_instance_common.hpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <ck_tile/core.hpp>
+#include "layernorm2d_fwd.hpp"
+#include <iostream>
+
+#pragma once
+
+using S = ck_tile::stream_config;
+using A = layernorm2d_fwd_args;
+
+template <typename DataType_,
+          ck_tile::index_t Repeat_M_,         // each thread repeat along M
+          ck_tile::index_t Repeat_N_,         // each thread repeat along N
+          ck_tile::index_t ThreadPerBlock_M_, // num threads along M
+          ck_tile::index_t ThreadPerBlock_N_, // num threads along N
+          ck_tile::index_t Vector_N_,         // vector size along N
+          bool kPadN_,
+          bool kSaveMeanInvStd_,
+          bool kTwoPass_>
+using trait_ = layernorm2d_fwd_traits_<DataType_,
+                                       Repeat_M_,
+                                       Repeat_N_,
+                                       ThreadPerBlock_M_,
+                                       ThreadPerBlock_N_,
+                                       Vector_N_,
+                                       kPadN_,
+                                       kSaveMeanInvStd_,
+                                       kTwoPass_>;
+
+template <typename Traits_>
+float layernorm2d_fwd_(const S& s, A a)
+{
+    using DataType = typename Traits_::DataType;
+
+    using PipelineProblem = ck_tile::Layernorm2dFwdPipelineProblem<
+        typename LayerNormTypeConfig<DataType>::XDataType,
+        typename LayerNormTypeConfig<DataType>::GammaDataType,
+        typename LayerNormTypeConfig<DataType>::BetaDataType,
+        typename LayerNormTypeConfig<DataType>::ComputeDataType,
+        typename LayerNormTypeConfig<DataType>::YDataType,
+        typename LayerNormTypeConfig<DataType>::MeanDataType,
+        typename LayerNormTypeConfig<DataType>::InvStdDataType,
+        typename Traits_::Shape,
+        Traits_::kPadN,
+        Traits_::kSaveMeanInvStd,
+        Traits_::kTwoPass>;
+
+    using OnePassPipeline = ck_tile::Layernorm2dFwdPipelineOnePass<PipelineProblem>;
+    using TwoPassPipeline = ck_tile::Layernorm2dFwdPipelineTwoPass<PipelineProblem>;
+    using Pipeline        = std::conditional_t<Traits_::kTwoPass, TwoPassPipeline, OnePassPipeline>;
+
+    using Kernel = ck_tile::Layernorm2dFwd<Pipeline>;
+
+    const dim3 grids                       = Kernel::GridSize(a);
+    constexpr dim3 blocks                  = Kernel::BlockSize();
+    constexpr ck_tile::index_t kBlockPerCu = 1;
+
+    auto kargs = Kernel::MakeKargs(a);
+    if(s.log_level_ > 0)
+        std::cout << ", " << Kernel::GetName() << std::flush;
+
+    return ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+}

example/ck_tile/02_layernorm2d/layernorm2d_fwd.cpp

@@ -2,161 +2,120 @@
 #include "layernorm2d_fwd.hpp"
 #include <cstring>
 
-// Host API implementation
-float layernorm2d_fwd(layernorm2d_fwd_traits t,
-                      layernorm2d_fwd_args a,
-                      const ck_tile::stream_config& s)
+// different threshold for different dtype
+template <typename DataType>
+auto get_elimit()
 {
-    if(t.data_type.compare("fp16") == 0)
-    {
-        using XDataType     = ck_tile::half_t;
-        using YDataType     = ck_tile::half_t;
-        using GammaDataType = ck_tile::half_t;
-        using BetaDataType  = ck_tile::half_t;
-#ifdef SAVE_MEAN_INV_STD
-        using MeanDataType   = ck_tile::half_t;
-        using InvStdDataType = ck_tile::half_t;
-#else
-        using MeanDataType   = ck_tile::null_type;
-        using InvStdDataType = ck_tile::null_type;
-#endif
-        using ComputeDataType = float;
-
-        using thread_tile = ck_tile::sequence<4, 4>;
-        using warp_tile   = ck_tile::sequence<8, 128>;
-        using block_tile  = ck_tile::sequence<32, 128>;
-
-        using Shape = ck_tile::TileLayernorm2dShape<thread_tile, warp_tile, block_tile>;
-
-        using PipelineProblem = ck_tile::BlockLayernorm2dFwdProblem<XDataType,
-                                                                    GammaDataType,
-                                                                    BetaDataType,
-                                                                    ComputeDataType,
-                                                                    YDataType,
-                                                                    MeanDataType,
-                                                                    InvStdDataType,
-                                                                    Shape,
-                                                                    true,
-                                                                    true>;
-
-        using Kernel = ck_tile::Layernorm2dFwd<PipelineProblem>;
-
-        auto kargs = Kernel::MakeKargs(
-            a.p_x, a.p_gamma, a.p_beta, a.p_y, a.p_mean, a.p_invStd, a.epsilon, a.M, a.N);
-
-        const dim3 grids      = Kernel::GridSize(a.M);
-        constexpr dim3 blocks = Kernel::BlockSize();
-
-        constexpr ck_tile::index_t kBlockPerCu = Shape::kMWarpPerBlock * Shape::kNWarpPerBlock;
-
-        float ave_time = ck_tile::launch_kernel(
-            s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-
-        return ave_time;
-    }
+    double rtol = 1e-2;
+    double atol = 1e-2;
+    return ck_tile::make_tuple(rtol, atol);
+}
 
-    return 0;
+template <>
+auto get_elimit<ck_tile::bf16_t>()
+{
+    double rtol = 1e-2;
+    double atol = 1e-2;
+    return ck_tile::make_tuple(rtol, atol);
 }
 
 auto create_args(int argc, char* argv[])
 {
     ck_tile::ArgParser arg_parser;
     arg_parser.insert("m", "3328", "m dimension")
-        .insert("n", "4096", "m dimension")
+        .insert("n", "4096", "n dimension")
+        .insert("stride", "-1", "stride per row, if -1 then equal to n")
         .insert("e", "1e-5", "epsilon")
+        .insert("save_mv", "0", "save mean/variance(invstd) or not. set to 1 in training case")
         .insert("v", "1", "cpu validation or not")
-        .insert("prec", "fp16", "precision");
+        .insert("kname", "1", "print kernel name or not")
+        .insert("prec", "fp16", "precision")
+        .insert("warmup", "5", "cold iter")
+        .insert("repeat", "20", "hot iter");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
 }
 
-int main(int argc, char* argv[])
+template <typename DataType, bool SaveMeanVar>
+bool run(const ck_tile::ArgParser& arg_parser)
 {
-
-    auto [result, arg_parser] = create_args(argc, argv);
-    if(!result)
-        return -1;
-
+    ck_tile::index_t m      = arg_parser.get_int("m");
+    ck_tile::index_t n      = arg_parser.get_int("n");
+    ck_tile::index_t stride = arg_parser.get_int("stride");
+    if(stride < 0)
+        stride = n;
     float epsilon         = arg_parser.get_float("e");
-    ck_tile::index_t M    = arg_parser.get_int("m");
-    ck_tile::index_t N    = arg_parser.get_int("n");
     std::string data_type = arg_parser.get_str("prec");
+    int kname             = arg_parser.get_int("kname");
     int do_validation     = arg_parser.get_int("v");
+    int warmup            = arg_parser.get_int("warmup");
+    int repeat            = arg_parser.get_int("repeat");
 
-    using XDataType     = ck_tile::half_t;
-    using YDataType     = ck_tile::half_t;
-    using GammaDataType = ck_tile::half_t;
-    using BetaDataType  = ck_tile::half_t;
-#ifdef SAVE_MEAN_INV_STD
-    using MeanDataType   = ck_tile::half_t;
-    using InvStdDataType = ck_tile::half_t;
-#else
-    using MeanDataType = ck_tile::null_type;
-    using InvStdDataType = ck_tile::null_type;
-#endif
-    using ComputeDataType = float;
+    assert(stride >= n);
 
-    // host verify
-    ck_tile::HostTensor<XDataType> x_host({M, N});
-    ck_tile::HostTensor<GammaDataType> gamma_host({N});
-    ck_tile::HostTensor<BetaDataType> beta_host({N});
+    using TypeConfig = LayerNormTypeConfig<DataType>;
+
+    using XDataType     = typename TypeConfig::XDataType;
+    using YDataType     = typename TypeConfig::YDataType;
+    using GammaDataType = typename TypeConfig::GammaDataType;
+    using BetaDataType  = typename TypeConfig::BetaDataType;
+
+    using MeanDataType =
+        std::conditional_t<SaveMeanVar, typename TypeConfig::MeanDataType, ck_tile::null_type>;
+    using InvStdDataType =
+        std::conditional_t<SaveMeanVar, typename TypeConfig::InvStdDataType, ck_tile::null_type>;
 
-    ck_tile::HostTensor<YDataType> y_host_ref({M, N});
-    ck_tile::HostTensor<YDataType> y_host_dev({M, N});
+    using ComputeDataType = typename TypeConfig::ComputeDataType;
 
-    ck_tile::HostTensor<MeanDataType> mean_host_ref({M});
-    ck_tile::HostTensor<InvStdDataType> invStd_host_ref({M});
+    // host verify
+    ck_tile::HostTensor<XDataType> x_host({m, n}, {stride, 1});
+    ck_tile::HostTensor<GammaDataType> gamma_host({n});
+    ck_tile::HostTensor<BetaDataType> beta_host({n});
+
+    ck_tile::HostTensor<YDataType> y_host_ref({m, n}, {stride, 1});
+    ck_tile::HostTensor<YDataType> y_host_dev({m, n}, {stride, 1});
 
-#ifdef SAVE_MEAN_INV_STD
-    ck_tile::HostTensor<MeanDataType> mean_host_dev({M});
-    ck_tile::HostTensor<InvStdDataType> invStd_host_dev({M});
-#endif
+    ck_tile::HostTensor<MeanDataType> mean_host_ref({m});
+    ck_tile::HostTensor<InvStdDataType> invStd_host_ref({m});
 
-    ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host);
-    ck_tile::FillUniformDistribution<GammaDataType>{-5.f, 5.f}(gamma_host);
-    ck_tile::FillUniformDistribution<BetaDataType>{-5.f, 5.f}(beta_host);
+    ck_tile::FillUniformDistribution<XDataType>{-.5f, .5f}(x_host);
+    ck_tile::FillUniformDistribution<GammaDataType>{-.5f, .5f}(gamma_host);
+    ck_tile::FillUniformDistribution<BetaDataType>{-.5f, .5f}(beta_host);
 
     ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem gamma_buf(gamma_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem beta_buf(beta_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem y_buf(y_host_dev.get_element_space_size_in_bytes());
 
-#ifdef SAVE_MEAN_INV_STD
-    ck_tile::DeviceMem mean_buf(mean_host_dev.get_element_space_size_in_bytes());
-    ck_tile::DeviceMem invStd_buf(invStd_host_dev.get_element_space_size_in_bytes());
-#endif
-
     x_buf.ToDevice(x_host.data());
     gamma_buf.ToDevice(gamma_host.data());
     beta_buf.ToDevice(beta_host.data());
 
-    layernorm2d_fwd_traits traits{data_type};
+    std::cout << "[" << data_type << "]"
+              << " m:" << m << ", n:" << n << ", stride:" << stride << std::flush;
+
+    layernorm2d_fwd_traits traits{data_type, SaveMeanVar};
 
     layernorm2d_fwd_args args{x_buf.GetDeviceBuffer(),
                               gamma_buf.GetDeviceBuffer(),
                               beta_buf.GetDeviceBuffer(),
                               y_buf.GetDeviceBuffer(),
-#ifdef SAVE_MEAN_INV_STD
-                              mean_buf.GetDeviceBuffer(),
-                              invStd_buf.GetDeviceBuffer(),
-#else
                               nullptr,
                               nullptr,
-#endif
                               epsilon,
-                              M,
-                              N};
+                              m,
+                              n,
+                              stride};
 
-    float ave_time = layernorm2d_fwd(traits, args, ck_tile::stream_config{nullptr, true});
+    float ave_time = layernorm2d_fwd(
+        traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});
 
-    std::size_t num_byte = sizeof(XDataType) * M * N + sizeof(GammaDataType) * N +
-                           sizeof(BetaDataType) * N + sizeof(YDataType) * M * N;
+    std::size_t num_byte = sizeof(XDataType) * m * n + sizeof(GammaDataType) * n +
+                           sizeof(BetaDataType) * n + sizeof(YDataType) * m * n;
 
     float gb_per_sec = num_byte / 1.E6 / ave_time;
-    std::cout << "[" << data_type << "]"
-              << " m:" << M << ", n:" << N << ", " << ave_time << " ms, " << gb_per_sec << " GB/s"
-              << std::flush;
+    std::cout << ", " << ave_time * 1.E3 << " us, " << gb_per_sec << " GB/s" << std::flush;
 
     bool pass = true;
 
@@ -174,20 +133,59 @@ int main(int argc, char* argv[])
 
         y_buf.FromDevice(y_host_dev.data());
 
-        pass = ck_tile::check_err(y_host_dev, y_host_ref);
+        auto [rtol, atol] = get_elimit<DataType>();
+        if(stride == n)
+        {
+            pass = ck_tile::check_err(
+                y_host_dev, y_host_ref, std::string("OUT Error: Incorrect results!"), rtol, atol);
+        }
+        else
+        {
+            for(int i_r = 0; i_r < m; i_r++)
+            {
+                std::vector<YDataType> y_host_dev_row(y_host_dev.begin() + i_r * stride,
+                                                      y_host_dev.begin() + i_r * stride + n);
+                std::vector<YDataType> y_host_ref_row(y_host_ref.begin() + i_r * stride,
+                                                      y_host_ref.begin() + i_r * stride + n);
+                pass &= ck_tile::check_err(y_host_dev_row,
+                                           y_host_ref_row,
+                                           std::string("OUT[") + std::to_string(i_r) +
+                                               std::string("] Error: Incorrect results!"),
+                                           rtol,
+                                           atol);
+            }
+        }
+
+        std::cout << ", valid:" << (pass ? "y" : "n") << std::flush << std::endl;
+    }
 
-#ifdef SAVE_MEAN_INV_STD
-        mean_buf.FromDevice(mean_host_dev.data());
-        pass &= ck_tile::check_err(mean_host_dev, mean_host_ref);
+    return pass;
+}
 
-        invStd_buf.FromDevice(invStd_host_dev.data());
-        pass &= ck_tile::check_err(invStd_host_dev, invStd_host_ref);
-#endif
+int main(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
 
-        std::cout << ", valid:" << (pass ? "y" : "n") << std::flush;
+    const std::string data_type = arg_parser.get_str("prec");
+    int save_mv                 = arg_parser.get_int("save_mv");
+    if(data_type == "fp16" && save_mv)
+    {
+        return run<ck_tile::half_t, true>(arg_parser) ? 0 : -2;
+    }
+    else if(data_type == "fp16" && !save_mv)
+    {
+        return run<ck_tile::half_t, false>(arg_parser) ? 0 : -2;
+    }
+    else if(data_type == "bf16" && save_mv)
+    {
+        return run<ck_tile::bf16_t, true>(arg_parser) ? 0 : -2;
+    }
+    else if(data_type == "bf16" && !save_mv)
+    {
+        return run<ck_tile::bf16_t, true>(arg_parser) ? 0 : -2;
     }
 
-    std::cout << std::endl << std::flush;
-
-    return !pass;
+    return -3;
 }

example/ck_tile/02_layernorm2d/layernorm2d_fwd.hpp

@@ -8,23 +8,114 @@
 #include "ck_tile/ops/layernorm2d.hpp"
 #include <string>
 
-struct layernorm2d_fwd_traits
+template <typename DataType>
+struct LayerNormTypeConfig;
+
+template <>
+struct LayerNormTypeConfig<ck_tile::half_t>
 {
-    std::string data_type;
+    using XDataType       = ck_tile::half_t;
+    using YDataType       = ck_tile::half_t;
+    using GammaDataType   = ck_tile::half_t;
+    using BetaDataType    = ck_tile::half_t;
+    using MeanDataType    = ck_tile::half_t;
+    using InvStdDataType  = ck_tile::half_t;
+    using ComputeDataType = float;
+};
+
+template <>
+struct LayerNormTypeConfig<ck_tile::bf16_t>
+{
+    using XDataType       = ck_tile::bf16_t;
+    using YDataType       = ck_tile::bf16_t;
+    using GammaDataType   = ck_tile::bf16_t;
+    using BetaDataType    = ck_tile::bf16_t;
+    using MeanDataType    = ck_tile::bf16_t;
+    using InvStdDataType  = ck_tile::bf16_t;
+    using ComputeDataType = float;
+};
+
+// runtime args
+struct layernorm2d_fwd_args : public ck_tile::Layernorm2dFwdHostArgs
+{
+};
+
+// this is used to pattern-match internl kernel implementation, not to instantiate kernel
+template <typename DataType_,
+          ck_tile::index_t Repeat_M_,         // each thread repeat along M
+          ck_tile::index_t Repeat_N_,         // each thread repeat along N
+          ck_tile::index_t ThreadPerBlock_M_, // num threads along M
+          ck_tile::index_t ThreadPerBlock_N_, // num threads along N
+          ck_tile::index_t Vector_N_,         // vector size along N
+          bool kPadN_,
+          bool kSaveMeanInvStd_,
+          bool kTwoPass_>
+struct layernorm2d_fwd_traits_
+{
+    using DataType = ck_tile::remove_cvref_t<DataType_>;
+
+    static constexpr bool is_warp_per_row = ThreadPerBlock_N_ <= warpSize;
+    static_assert((ThreadPerBlock_M_ * ThreadPerBlock_N_) % warpSize == 0);
+    static constexpr ck_tile::index_t total_warps =
+        (ThreadPerBlock_M_ * ThreadPerBlock_N_) / warpSize;
+
+    // num of warps along m
+    static constexpr ck_tile::index_t BlockWarps_M = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(warpSize % ThreadPerBlock_N_ == 0);
+            return total_warps * (warpSize / ThreadPerBlock_N_);
+        }
+        else
+        {
+            // static_assert(warpSize % ThreadPerBlock_M_ == 0);
+            return total_warps / (ThreadPerBlock_N_ / warpSize);
+        }
+    }();
+
+    // num of warps along n
+    static constexpr ck_tile::index_t BlockWarps_N = []() {
+        if constexpr(is_warp_per_row)
+        {
+            static_assert(warpSize % ThreadPerBlock_N_ == 0);
+            return 1;
+        }
+        else
+        {
+            static_assert(ThreadPerBlock_N_ % warpSize == 0);
+            return ThreadPerBlock_N_ / warpSize;
+        }
+    }();
+
+    static constexpr ck_tile::index_t Repeat_M = Repeat_M_;
+    static constexpr ck_tile::index_t Repeat_N = Repeat_N_;
+
+    static constexpr ck_tile::index_t Block_M = Repeat_M_ * ThreadPerBlock_M_;
+    static constexpr ck_tile::index_t Block_N = Repeat_N_ * ThreadPerBlock_N_ * Vector_N_;
+
+    static constexpr ck_tile::index_t Warp_M = ThreadPerBlock_M_ / BlockWarps_M;
+    static constexpr ck_tile::index_t Warp_N = ThreadPerBlock_N_ / BlockWarps_N * Vector_N_;
+
+    using BlockTile  = ck_tile::sequence<Block_M, Block_N>;
+    using BlockWarps = ck_tile::sequence<BlockWarps_M, BlockWarps_N>;
+    using WarpTile   = ck_tile::sequence<Warp_M, Warp_N>;
+    using Vector     = ck_tile::sequence<1, Vector_N_>;
+
+    using Shape = ck_tile::Layernorm2dShape<BlockTile, BlockWarps, WarpTile, Vector>;
+
+    static constexpr bool kPadN           = kPadN_;
+    static constexpr bool kSaveMeanInvStd = kSaveMeanInvStd_;
+    static constexpr bool kTwoPass        = kTwoPass_;
 };
 
-struct layernorm2d_fwd_args
+template <typename Traits_>
+float layernorm2d_fwd_(const ck_tile::stream_config& s, layernorm2d_fwd_args a);
+
+// This is the public API, will be generated by script
+struct layernorm2d_fwd_traits
 {
-    const void* p_x;
-    const void* p_gamma;
-    const void* p_beta;
-    void* p_y;
-    void* p_mean;
-    void* p_invStd;
-    float epsilon;
-    ck_tile::index_t M;
-    ck_tile::index_t N;
+    std::string data_type;
+    bool save_mean_var;
 };
 
-// host API
 float layernorm2d_fwd(layernorm2d_fwd_traits, layernorm2d_fwd_args, const ck_tile::stream_config&);

example/ck_tile/02_layernorm2d/script/perf_test.sh

+
+# run from top of ck folder
+EXE=build/bin/tile_example_layernorm2d_fwd
+
+$EXE -m=1 -n=1 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=80 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=128 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=144 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=168 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=184 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=256 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=288 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=344 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=376 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=448 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=512 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=924 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=1024 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=1078 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=1996 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+$EXE -m=700 -n=4080 -e=1e-12 -v=1 -prec=bf16 -repeat=1000
+
+$EXE -m=700 -n=80 -e=1e-12 -v=1  -prec=fp16 -repeat=1000
+$EXE -m=700 -n=128 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=144 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=168 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=184 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=256 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=288 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=344 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=376 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=448 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=512 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=924 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=1024 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=1078 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=1996 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
+$EXE -m=700 -n=4080 -e=1e-12 -v=1 -prec=fp16 -repeat=1000
\ No newline at end of file

example/ck_tile/02_layernorm2d/script/smoke_test.sh

+#!/bin/sh
+# call from top of CK folder
+EXE=./build/bin/tile_example_layernorm2d_fwd
+
+for pr_i in "fp16" "bf16" ; do
+$EXE -prec=$pr_i -m=99  -n=13
+$EXE -prec=$pr_i -m=17  -n=16
+$EXE -prec=$pr_i -m=1   -n=100
+$EXE -prec=$pr_i -m=4   -n=128
+$EXE -prec=$pr_i -m=80  -n=127
+$EXE -prec=$pr_i -m=22  -n=255 -stride=256
+$EXE -prec=$pr_i -m=7   -n=599
+$EXE -prec=$pr_i -m=19  -n=512
+$EXE -prec=$pr_i -m=33  -n=313 -stride=1000
+$EXE -prec=$pr_i -m=11  -n=510
+$EXE -prec=$pr_i -m=171 -n=676 -stride=818
+$EXE -prec=$pr_i -m=91  -n=636
+$EXE -prec=$pr_i -m=12  -n=768 -stride=800
+$EXE -prec=$pr_i -m=100 -n=766 -stride=812
+$EXE -prec=$pr_i -m=31  -n=1024
+$EXE -prec=$pr_i -m=64  -n=1000 -stride=1004
+$EXE -prec=$pr_i -m=8   -n=1501
+$EXE -prec=$pr_i -m=3   -n=1826
+$EXE -prec=$pr_i -m=5   -n=2040
+$EXE -prec=$pr_i -m=7   -n=2734
+$EXE -prec=$pr_i -m=1   -n=3182
+$EXE -prec=$pr_i -m=9   -n=4096
+$EXE -prec=$pr_i -m=3   -n=8192
+$EXE -prec=$pr_i -m=1   -n=10547
+$EXE -prec=$pr_i -m=3   -n=17134
+done

example/ck_tile/05_reduce/CMakeLists.txt

+set(EXAMPLE_REDUCE "tile_example_reduce")
+# not using add_example_executable() to add this target, since we don't want this to have
+# to be included in "make all/install/check"
+message("adding example ${EXAMPLE_REDUCE}")
+
+add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL reduce.cpp)
+target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
+
+# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
+list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
+
+target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS})
+
+# TODO: we have to turn off this global prop, otherwise the progress bar generated
+# by cmake will print too many files, execvp: /bin/sh: Argument list too long
+# however, this property may affect global
+# TODO: consider codegen a makefile by us
+set_property(GLOBAL PROPERTY RULE_MESSAGES OFF)
\ No newline at end of file

example/ck_tile/05_reduce/reduce.cpp

+#include "ck_tile/host.hpp"
+#include "reduce.hpp"
+#include <cstring>
+
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("m", "3328", "m dimension")
+        .insert("n", "4096", "n dimension")
+        .insert("v", "1", "cpu validation or not")
+        .insert("prec", "fp16", "precision")
+        .insert("warmup", "5", "cold iter")
+        .insert("repeat", "20", "hot iter");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+template <typename DataType>
+bool run(const ck_tile::ArgParser& arg_parser)
+{
+    using ADataType   = DataType;
+    using AccDataType = float;
+    using BDataType   = DataType;
+
+    ck_tile::index_t m = arg_parser.get_int("m");
+    ck_tile::index_t n = arg_parser.get_int("n");
+    int do_validation  = arg_parser.get_int("v");
+    int warmup         = arg_parser.get_int("warmup");
+    int repeat         = arg_parser.get_int("repeat");
+
+    ck_tile::HostTensor<ADataType> a_host({m, n});
+    ck_tile::HostTensor<BDataType> b_host_ref({m});
+    ck_tile::HostTensor<BDataType> b_host_dev({m});
+
+    ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_host);
+
+    ck_tile::DeviceMem a_buf(a_host.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem b_buf(b_host_dev.get_element_space_size_in_bytes());
+
+    a_buf.ToDevice(a_host.data());
+
+    using BlockWarps = ck_tile::sequence<4, 1>;
+    using BlockTile  = ck_tile::sequence<128, 128>;
+    using WarpTile   = ck_tile::sequence<32, 128>;
+    using ThreadTile = ck_tile::sequence<8, 8>;
+
+    constexpr ck_tile::index_t kBlockSize  = 256;
+    constexpr ck_tile::index_t kBlockPerCu = 1;
+    ck_tile::index_t kGridSize             = (m / BlockTile::at(ck_tile::number<0>{}));
+    std::cout << "grid size " << kGridSize << std::endl;
+
+    using Kernel = ck_tile::Reduce<ADataType,
+                                   AccDataType,
+                                   BDataType,
+                                   kBlockSize,
+                                   BlockWarps,
+                                   BlockTile,
+                                   WarpTile,
+                                   ThreadTile>;
+
+    float ave_time = launch_kernel(ck_tile::stream_config{nullptr, true, 0, warmup, repeat},
+                                   ck_tile::make_kernel<kBlockSize, kBlockPerCu>(
+                                       Kernel{},
+                                       kGridSize,
+                                       kBlockSize,
+                                       0,
+                                       static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
+                                       static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
+                                       m,
+                                       n));
+
+    std::size_t num_btype = sizeof(ADataType) * m * n + sizeof(BDataType) * m;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl;
+
+    bool pass = true;
+
+    if(do_validation)
+    {
+        // reference
+        ck_tile::reference_reduce<ADataType, AccDataType, BDataType>(a_host, b_host_ref);
+        b_buf.FromDevice(b_host_dev.mData.data());
+        pass = ck_tile::check_err(b_host_dev, b_host_ref);
+
+        std::cout << "valid:" << (pass ? "y" : "n") << std::flush << std::endl;
+    }
+
+    return pass;
+}
+
+int main(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    const std::string data_type = arg_parser.get_str("prec");
+
+    if(data_type == "fp16")
+    {
+        return run<ck_tile::half_t>(arg_parser) ? 0 : -2;
+    }
+    if(data_type == "bf16")
+    {
+        return run<ck_tile::bf16_t>(arg_parser) ? 0 : -2;
+    }
+}