ckProfiler for layernorm (#330)

* Refine parameter

* Add base class for layernorm

* Add layernorm instance

* Add layernorm to ckProfiler

* Remove redundant

* Add verification

* Fix compile error due to merge

[ROCm/composable_kernel commit: fdfd7eb597]

example/27_layernorm/layernorm_blockwise.cpp

@@ -46,7 +46,7 @@ using DeviceInstance = ck::tensor_operation::device::DeviceLayernorm<XDataType,
                                                                      8,   // SrcScalarPerVector
                                                                      8,   // GammaScalarPerVector
                                                                      8,   // BetaScalarPerVector
-                                                                     1>;  // OutScalarPerVector
+                                                                     8>;  // OutScalarPerVector
 
 int main()
 {

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

@@ -7,7 +7,7 @@
 #include <sstream>
 
 #include "ck/utility/reduction_operator.hpp"
-#include "ck/tensor_operation/gpu/device/device_base.hpp"
+#include "ck/tensor_operation/gpu/device/device_normalization.hpp"
 #include "ck/tensor_operation/gpu/device/device_reduce.hpp"
 #include "ck/tensor_operation/gpu/device/device_reduce_multiblock.hpp"
 #include "ck/tensor_operation/gpu/device/device_reduce_common.hpp"
@@ -39,7 +39,14 @@ template <typename XDataType,
           index_t GammaSrcVectorSize,
           index_t BetaSrcVectorSize,
           index_t YDstVectorSize>
-struct DeviceLayernorm : public BaseOperator
+struct DeviceLayernorm : public DeviceNormalization2<XDataType,
+                                                     GammaDataType,
+                                                     BetaDataType,
+                                                     AccDataType,
+                                                     YDataType,
+                                                     AccElementwiseOperation,
+                                                     Rank,
+                                                     NumReduceDim>
 {
     static_assert(
         (KThreadSliceSize % GammaSrcVectorSize == 0),
@@ -297,17 +304,18 @@ struct DeviceLayernorm : public BaseOperator
         return true;
     };
 
-    std::unique_ptr<BaseArgument> MakeArgumentPointer(const std::vector<index_t> lengths,
-                                                      const std::vector<index_t> xStrides,
-                                                      const std::vector<index_t> gammaStrides,
-                                                      const std::vector<index_t> betaStrides,
-                                                      const std::vector<index_t> reduceDims,
-                                                      AccDataType epsilon,
-                                                      const void* p_x,
-                                                      const void* p_gamma,
-                                                      const void* p_beta,
-                                                      void* p_y,
-                                                      AccElementwiseOperation acc_elementwise_op)
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<index_t> lengths,
+                        const std::vector<index_t> xStrides,
+                        const std::vector<index_t> gammaStrides,
+                        const std::vector<index_t> betaStrides,
+                        const std::vector<index_t> reduceDims,
+                        AccDataType epsilon,
+                        const void* p_x,
+                        const void* p_gamma,
+                        const void* p_beta,
+                        void* p_y,
+                        AccElementwiseOperation acc_elementwise_op) override
     {
         return std::make_unique<Argument>(lengths,
                                           xStrides,
@@ -322,7 +330,10 @@ struct DeviceLayernorm : public BaseOperator
                                           static_cast<YDataType*>(p_y));
     };
 
-    std::unique_ptr<BaseInvoker> MakeInvokerPointer() { return std::make_unique<Invoker>(); };
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
 
     std::string GetTypeString() const override
     {
@@ -332,7 +343,6 @@ struct DeviceLayernorm : public BaseOperator
         str << "DeviceLayernorm<" << BlockSize << ",";
         str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
         str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
-        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
         str << "XYSrcVectorDim_" << XYSrcVectorDim  << ",";
         str << "VectorSize_X" << XSrcVectorSize << "_Gamma" << GammaSrcVectorSize << "_Beta" << BetaSrcVectorSize << "_Y" << YDstVectorSize << ">";
         // clang-format on

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

@@ -38,6 +38,49 @@ struct DeviceNormalization : public BaseOperator
 
 using DeviceNormalizationPtr = std::unique_ptr<DeviceNormalization>;
 
+template <typename XDataType,
+          typename GammaDataType,
+          typename BetaDataType,
+          typename AccDataType,
+          typename YDataType,
+          typename AccElementwiseOperation,
+          index_t Rank,
+          index_t NumReduceDim>
+struct DeviceNormalization2 : public BaseOperator
+{
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<index_t> lengths,
+                        const std::vector<index_t> xStrides,
+                        const std::vector<index_t> gammaStrides,
+                        const std::vector<index_t> betaStrides,
+                        const std::vector<index_t> reduceDims,
+                        AccDataType epsilon,
+                        const void* p_x,
+                        const void* p_gamma,
+                        const void* p_beta,
+                        void* p_y,
+                        AccElementwiseOperation acc_elementwise_op) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
+template <typename XDataType,
+          typename GammaDataType,
+          typename BetaDataType,
+          typename AccDataType,
+          typename YDataType,
+          typename AccElementwiseOperation,
+          index_t Rank,
+          index_t NumReduceDim>
+using DeviceNormalization2Ptr = std::unique_ptr<DeviceNormalization2<XDataType,
+                                                                     GammaDataType,
+                                                                     BetaDataType,
+                                                                     AccDataType,
+                                                                     YDataType,
+                                                                     AccElementwiseOperation,
+                                                                     Rank,
+                                                                     NumReduceDim>>;
+
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck

library/src/tensor_operation_instance/gpu/normalization/CMakeLists.txt

@@ -1,5 +1,7 @@
 # device_normalization_instance
 set(DEVICE_NORMALIZATION_INSTANCE_SOURCE
+    device_layernorm_f16_instance.cpp
+    device_layernorm_f32_instance.cpp
     device_softmax_f32_f32_instance.cpp
     device_softmax_f16_f16_instance.cpp
 )

library/src/tensor_operation_instance/gpu/normalization/device_layernorm_f16_instance.cpp

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/device_layernorm.hpp"
+#include "ck/utility/data_type.hpp"
+
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Pass = ck::tensor_operation::element_wise::PassThrough;
+
+template <index_t Rank, index_t Reduce>
+using device_layernorm_f16_instances = std::tuple<
+    // clang-format off
+        // XDataType, GammaDataType, BetaDataType, AccDataType, YDataType, Rank, NumReduceDim, BlockSize, MThreadClusterSize, KThreadClusterSize, MThreadSliceSize, KThreadSliceSize, XYSrcVectorDim, XSrcVectorSize, GammaSrcVectorSize, BetaSrcVectorSize, YDstVectorSize>
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 1, 1, 1, 1>, // fallback kernel
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 2, 2, 2, 2>, // fallback kernel
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 4, 4, 4, 4>, // fallback kernel
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 4, 64, 1, 8, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 2, 128, 1, 8, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 2, 128, 1, 16, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 2, 128, 1, 32, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 1, 256, 1, 8, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 1, 256, 1, 16, 1, 8, 8, 8, 8>,
+        DeviceLayernorm<F16, F16, F16, F32, F16, Pass, Rank, Reduce, 256, 1, 256, 1, 32, 1, 8, 8, 8, 8>
+    // clang-format on
+    >;
+
+void add_device_layernorm_f16_rank2_instances(
+    std::vector<DeviceNormalization2Ptr<F16, F16, F16, F32, F16, Pass, 2, 1>>& instances)
+{
+    add_device_operation_instances(instances, device_layernorm_f16_instances<2, 1>{});
+}
+
+void add_device_layernorm_f16_rank4_instances(
+    std::vector<DeviceNormalization2Ptr<F16, F16, F16, F32, F16, Pass, 4, 3>>& instances)
+{
+    add_device_operation_instances(instances, device_layernorm_f16_instances<4, 3>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/normalization/device_layernorm_f32_instance.cpp

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/device_layernorm.hpp"
+#include "ck/utility/data_type.hpp"
+
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using F32 = float;
+
+using Pass = ck::tensor_operation::element_wise::PassThrough;
+
+template <index_t Rank, index_t Reduce>
+using device_layernorm_f32_instances = std::tuple<
+    // clang-format off
+        // XDataType, GammaDataType, BetaDataType, AccDataType, YDataType, Rank, NumReduceDim, BlockSize, MThreadClusterSize, KThreadClusterSize, MThreadSliceSize, KThreadSliceSize, XYSrcVectorDim, XSrcVectorSize, GammaSrcVectorSize, BetaSrcVectorSize, YDstVectorSize>
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 1, 1, 1, 1>, // fallback kernel
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 2, 2, 2, 2>, // fallback kernel
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 8, 32, 1, 8, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 4, 64, 1, 8, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 2, 128, 1, 8, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 2, 128, 1, 16, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 2, 128, 1, 32, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 1, 256, 1, 8, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 1, 256, 1, 16, 1, 4, 4, 4, 4>,
+        DeviceLayernorm<F32, F32, F32, F32, F32, Pass, Rank, Reduce, 256, 1, 256, 1, 32, 1, 4, 4, 4, 4>
+    // clang-format on
+    >;
+
+void add_device_layernorm_f32_rank2_instances(
+    std::vector<DeviceNormalization2Ptr<F32, F32, F32, F32, F32, Pass, 2, 1>>& instances)
+{
+    add_device_operation_instances(instances, device_layernorm_f32_instances<2, 1>{});
+}
+
+void add_device_layernorm_f32_rank4_instances(
+    std::vector<DeviceNormalization2Ptr<F32, F32, F32, F32, F32, Pass, 4, 3>>& instances)
+{
+    add_device_operation_instances(instances, device_layernorm_f32_instances<4, 3>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

profiler/CMakeLists.txt

@@ -21,6 +21,7 @@ set(PROFILER_SOURCE
     src/profile_conv_bwd_weight.cpp
     src/profile_grouped_conv_fwd.cpp
     src/profile_reduce.cpp
+    src/profile_layernorm.cpp
     src/profile_normalization.cpp
 )
 

profiler/include/profile_layernorm_impl.hpp

@@ -0,0 +1,238 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iomanip>
+
+#include "ck/ck.hpp"
+#include "profiler/include/data_type_enum.hpp"
+#include "ck/tensor_operation/gpu/device/device_layernorm.hpp"
+
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_layernorm.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using F16         = ck::half_t;
+using F32         = float;
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+void add_device_layernorm_f16_rank2_instances(
+    std::vector<DeviceNormalization2Ptr<F16, F16, F16, F32, F16, PassThrough, 2, 1>>&);
+
+void add_device_layernorm_f32_rank2_instances(
+    std::vector<DeviceNormalization2Ptr<F32, F32, F32, F32, F32, PassThrough, 2, 1>>&);
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+
+namespace ck {
+namespace profiler {
+
+template <typename XDataType,
+          typename GammaDataType,
+          typename BetaDataType,
+          typename AccDataType,
+          typename YDataType,
+          index_t Rank>
+void profile_layernorm_impl(int do_verification,
+                            int init_method,
+                            bool do_log,
+                            bool time_kernel,
+                            std::vector<index_t> length,
+                            std::vector<index_t> strideXY,
+                            std::vector<index_t> strideGamma,
+                            std::vector<index_t> strideBeta)
+{
+    using F16         = ck::half_t;
+    using F32         = float;
+    using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+    if(length.size() < 2)
+        return;
+
+    // Assume normalize dimension except for first dimension
+    std::vector<index_t> reduce_length{length.begin() + 1, length.end()};
+    std::vector<index_t> reduce_dim;
+    for(int i = 1; i < Rank; ++i)
+        reduce_dim.push_back(i);
+
+    Tensor<XDataType> x(length);
+    Tensor<GammaDataType> gamma(reduce_length, strideGamma);
+    Tensor<BetaDataType> beta(reduce_length, strideBeta);
+    Tensor<YDataType> y(length, strideXY);
+    Tensor<YDataType> host_y(length, strideXY);
+
+    switch(init_method)
+    {
+    // case 0: break;
+    case 0:
+        x.GenerateTensorValue(GeneratorTensor_1<XDataType>{});
+        gamma.GenerateTensorValue(GeneratorTensor_1<GammaDataType>{});
+        beta.GenerateTensorValue(GeneratorTensor_1<BetaDataType>{});
+        y.GenerateTensorValue(GeneratorTensor_1<YDataType>{});
+        break;
+    case 1:
+        x.GenerateTensorValue(GeneratorTensor_2<XDataType>{-5, 5});
+        gamma.GenerateTensorValue(GeneratorTensor_2<GammaDataType>{-5, 5});
+        beta.GenerateTensorValue(GeneratorTensor_2<BetaDataType>{-5, 5});
+        y.GenerateTensorValue(GeneratorTensor_2<YDataType>{-5, 5});
+        break;
+    default:
+        x.GenerateTensorValue(GeneratorTensor_3<XDataType>{0, 1});
+        gamma.GenerateTensorValue(GeneratorTensor_3<GammaDataType>{-0.5, 0.5});
+        beta.GenerateTensorValue(GeneratorTensor_3<BetaDataType>{-0.5, 0.5});
+        y.GenerateTensorValue(GeneratorTensor_3<YDataType>{-0.5, 0.5});
+    }
+
+    DeviceMem x_dev(sizeof(XDataType) * x.mDesc.GetElementSpaceSize());
+    DeviceMem gamma_dev(sizeof(GammaDataType) * gamma.mDesc.GetElementSpaceSize());
+    DeviceMem beta_dev(sizeof(BetaDataType) * beta.mDesc.GetElementSpaceSize());
+    DeviceMem y_dev(sizeof(YDataType) * y.mDesc.GetElementSpaceSize());
+
+    x_dev.ToDevice(x.mData.data());
+    gamma_dev.ToDevice(gamma.mData.data());
+    beta_dev.ToDevice(beta.mData.data());
+
+    // add device normalization instances
+    constexpr int NumReduceDim = Rank - 1;
+    std::vector<tensor_operation::device::DeviceNormalization2Ptr<XDataType,
+                                                                  GammaDataType,
+                                                                  BetaDataType,
+                                                                  AccDataType,
+                                                                  YDataType,
+                                                                  PassThrough,
+                                                                  Rank,
+                                                                  NumReduceDim>>
+        instances;
+
+    if constexpr(is_same<XDataType, F16>::value && is_same<GammaDataType, F16>::value &&
+                 is_same<BetaDataType, F16>::value && is_same<YDataType, F16>::value &&
+                 is_same<AccDataType, F32>::value)
+    {
+        if(length.size() == 2)
+            tensor_operation::device::instance::add_device_layernorm_f16_rank2_instances(instances);
+    }
+    else if constexpr(is_same<XDataType, F32>::value && is_same<GammaDataType, F32>::value &&
+                      is_same<BetaDataType, F32>::value && is_same<YDataType, F32>::value &&
+                      is_same<AccDataType, F32>::value)
+    {
+        if(length.size() == 2)
+            tensor_operation::device::instance::add_device_layernorm_f32_rank2_instances(instances);
+    }
+
+    if(instances.size() <= 0)
+    {
+        throw std::runtime_error("wrong! no device normalization instance found");
+    }
+
+    std::string best_instance_name;
+    float best_avg_time   = std::numeric_limits<float>::max();
+    float best_gb_per_sec = 0;
+
+    if(do_verification)
+    {
+        using ReferenceInstance = ck::tensor_operation::host::ReferenceLayernorm<XDataType,
+                                                                                 GammaDataType,
+                                                                                 BetaDataType,
+                                                                                 YDataType,
+                                                                                 AccDataType,
+                                                                                 PassThrough,
+                                                                                 Rank,
+                                                                                 NumReduceDim>;
+
+        ReferenceInstance ref;
+        auto ref_argument =
+            ref.MakeArgument(x, gamma, beta, host_y, PassThrough{}, length, reduce_dim, 1e-4);
+        auto ref_invoker = ref.MakeInvoker();
+        ref_invoker.Run(ref_argument);
+    }
+
+    for(auto& inst_ptr : instances)
+    {
+        auto argument_ptr = inst_ptr->MakeArgumentPointer(length,
+                                                          strideXY,
+                                                          strideGamma,
+                                                          strideBeta,
+                                                          reduce_dim,
+                                                          1e-4,
+                                                          x_dev.GetDeviceBuffer(),
+                                                          gamma_dev.GetDeviceBuffer(),
+                                                          beta_dev.GetDeviceBuffer(),
+                                                          y_dev.GetDeviceBuffer(),
+                                                          PassThrough{});
+
+        if(!inst_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            std::cout << inst_ptr->GetTypeString() << " skipped due to unsupported argument: ";
+            LogRange(std::cout << "input lengths = [", length, "], ") << std::endl;
+
+            return;
+        }
+
+        auto invoker_ptr = inst_ptr->MakeInvokerPointer();
+
+        float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
+
+        std::size_t num_bytes = x.mDesc.GetElementSize() * sizeof(XDataType) +
+                                gamma.mDesc.GetElementSize() * sizeof(GammaDataType) +
+                                beta.mDesc.GetElementSize() * sizeof(BetaDataType) +
+                                y.mDesc.GetElementSize() * sizeof(YDataType);
+
+        float gb_per_sec = num_bytes / 1.E6 / avg_time;
+
+        std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << gb_per_sec << " GB/s, "
+                  << inst_ptr->GetTypeString() << std::endl;
+
+        if(avg_time < best_avg_time)
+        {
+            best_instance_name = inst_ptr->GetTypeString();
+            best_avg_time      = avg_time;
+            best_gb_per_sec    = gb_per_sec;
+        }
+
+        if(do_verification)
+        {
+            y_dev.FromDevice(y.mData.data());
+
+            bool pass = ck::utils::check_err(
+                y.mData, host_y.mData, "Error: Incorrect results d1", 1e-3, 1e-3);
+
+            if(do_log)
+            {
+                LogRangeAsType<float>(std::cout << "x  : ", x.mData, ",") << std::endl;
+                LogRangeAsType<float>(std::cout << "host_y  : ", host_y.mData, ",") << std::endl;
+                LogRangeAsType<float>(std::cout << "y  : ", y.mData, ",") << std::endl;
+            }
+
+            if(!pass)
+            {
+                std::cout << inst_ptr->GetTypeString() << " failed verification: ";
+                LogRange(std::cout << "lengths = [", length, ", ") << "]." << std::endl;
+                return;
+            }
+            else
+            {
+                std::cout << "pass" << std::endl;
+            }
+        }
+    }
+
+    LogRange(std::cout << "length = ", length, ",") << ", ";
+    LogRange(std::cout << "stride = ", strideXY, ",") << ", ";
+    LogRange(std::cout << "reduce dims ", reduce_dim, ",") << std::endl;
+    std::cout << "best perf = " << best_avg_time << " ms, " << best_gb_per_sec << " GB/s, "
+              << best_instance_name << std::endl;
+}
+
+} // namespace profiler
+} // namespace ck

profiler/include/profile_normalization_impl.hpp

@@ -36,7 +36,6 @@ namespace profiler {
 
 enum struct NormType
 {
-    LAYERNORM,
     BATCHNORM,
     SOFTMAX,
 };

profiler/src/profile_layernorm.cpp

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <vector>
+#include <unordered_map>
+
+#include "profiler/include/profile_layernorm_impl.hpp"
+
+using ck::index_t;
+
+struct LayernormArgParser
+{
+    std::unordered_map<std::string, std::vector<int>> long_opts = {
+        {"length", {}}, {"strideXY", {}}, {"strideGamma", {}}, {"strideBeta", {}}};
+
+    bool parse_opt(int argc, char* argv[], const std::string& key, int i)
+    {
+        if(std::string("--") + key == argv[i])
+        {
+            int pos = i;
+            while(++i < argc && argv[i][0] != '-') {}
+            int end = i;
+            for(int j = pos + 1; j < end; j++)
+            {
+                long_opts[key].push_back(std::stoi(argv[j]));
+            }
+            return true;
+        }
+        return false;
+    }
+
+    void operator()(int argc, char* argv[])
+    {
+        for(auto& kv : long_opts)
+        {
+            for(int i = 1; i < argc; i++)
+            {
+                if(parse_opt(argc, argv, kv.first, i))
+                    break;
+            }
+        }
+    }
+};
+
+void print_help_layernorm()
+{
+    std::cout << "arg1: data type (0: fp16; 1: fp32)\n"
+              << "arg2: verification (0: no; 1: yes)\n"
+              << "arg3: initialization (0: no init; 1: integer value; 2: decimal value)\n"
+              << "arg4: print tensor value (0: no; 1: yes)\n"
+              << "arg5: time kernel (0=n0, 1=yes)\n"
+              << "--length: tensor extents (e.g, --length 1024 1024) \n"
+              << "--strideXY: tensor strides (e.g, --strideXY 1024 1)\n"
+              << "--strideGamma: tensor strides (e.g, --strideGamma 1)\n"
+              << "--strideBeta: tensor strides (e.g, --strideBeta 1)\n"
+              << std::endl;
+}
+
+int profile_layernorm(int argc, char* argv[])
+{
+    if(argc <= 2)
+    {
+        print_help_layernorm();
+        return 0;
+    }
+
+    LayernormArgParser arg_parser;
+
+    // short unnamed options
+    const ck::DataTypeEnum data_type = static_cast<ck::DataTypeEnum>(std::stoi(argv[2]));
+    const bool do_verification       = std::stoi(argv[3]);
+    const int init_method            = std::stoi(argv[4]);
+    const bool do_log                = std::stoi(argv[5]);
+    const bool time_kernel           = std::stoi(argv[6]);
+
+    // parse the long options
+    arg_parser(argc, argv);
+    const std::vector<index_t> length      = arg_parser.long_opts["length"];
+    const std::vector<index_t> strideXY    = arg_parser.long_opts["strideXY"];
+    const std::vector<index_t> strideGamma = arg_parser.long_opts["strideGamma"];
+    const std::vector<index_t> strideBeta  = arg_parser.long_opts["strideBeta"];
+
+    using F16          = ck::half_t;
+    using F32          = float;
+    constexpr int rank = 2;
+
+    if(data_type == ck::DataTypeEnum::Half)
+    {
+        ck::profiler::profile_layernorm_impl<F16, F16, F16, F32, F16, rank>(do_verification,
+                                                                            init_method,
+                                                                            do_log,
+                                                                            time_kernel,
+                                                                            length,
+                                                                            strideXY,
+                                                                            strideGamma,
+                                                                            strideBeta);
+    }
+    else if(data_type == ck::DataTypeEnum::Float)
+    {
+        ck::profiler::profile_layernorm_impl<F32, F32, F32, F32, F32, rank>(do_verification,
+                                                                            init_method,
+                                                                            do_log,
+                                                                            time_kernel,
+                                                                            length,
+                                                                            strideXY,
+                                                                            strideGamma,
+                                                                            strideBeta);
+    }
+    else
+    {
+        throw std::runtime_error("not implemented yet");
+    }
+
+    return 0;
+}
+
+// hijack main() for quick debugging
+// int main(int argc, char* argv[])
+// {
+//     profile_layernorm(argc, argv);
+//     return 0;
+// }

profiler/src/profile_normalization.cpp

@@ -13,8 +13,7 @@ using ck::profiler::NormType;
 
 struct ArgParser
 {
-    std::unordered_map<std::string, NormType> norm_dict = {{"layernorm", NormType::LAYERNORM},
-                                                           {"batchnorm", NormType::BATCHNORM},
+    std::unordered_map<std::string, NormType> norm_dict = {{"batchnorm", NormType::BATCHNORM},
                                                            {"softmax", NormType::SOFTMAX}};
 
     std::unordered_map<std::string, std::vector<int>> long_opts = {

profiler/src/profiler.cpp

@@ -19,6 +19,7 @@ int profile_conv_bwd_data(int, char*[]);
 int profile_conv_bwd_weight(int, char*[]);
 int profile_grouped_conv_fwd(int, char*[]);
 int profile_normalization(int, char*[]);
+int profile_layernorm(int, char*[]);
 int profile_reduce(int, char*[]);
 
 static void print_helper_message()
@@ -115,11 +116,14 @@ int main(int argc, char* argv[])
     {
         return profile_reduce(argc, argv);
     }
-    else if(strcmp(argv[1], "batchnorm") == 0 || strcmp(argv[1], "layernorm") == 0 ||
-            strcmp(argv[1], "softmax") == 0)
+    else if(strcmp(argv[1], "batchnorm") == 0 || strcmp(argv[1], "softmax") == 0)
     {
         return profile_normalization(argc, argv);
     }
+    else if(strcmp(argv[1], "layernorm") == 0)
+    {
+        return profile_layernorm(argc, argv);
+    }
     else
     {
         print_helper_message();