Layernorm4d (#1022)

* Rename folder

* Add layernorm 4d fwd example

* Rename original layernorm example

* Add layernorm 4d f16  test

* Add layernorm4d_fwd client example

* Support layernorm4D in ckProfiler

* Rename groupnorm to groupnorm fwd in example

* Rename layernorm and group fwd in test

* Rename normalization to normalization_fwd (instances)

* Add fwd to DeviceNormalization

* Rename external api header

* Rename folder, because we can also add bwd in this folder

* Add fwd in layernorm and groupnorm (profiler

* Fix compile error

---------
Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>

library/src/tensor_operation_instance/gpu/normalization/device_layernorm2d_f16_instance.cpp → library/src/tensor_operation_instance/gpu/normalization_fwd/device_layernorm2d_fwd_f16_instance.cpp

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
-#include "normalization_instance_common.hpp"
+#include "normalization_fwd_instance_common.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -10,8 +10,8 @@ namespace instance {
 
 using Pass = ck::tensor_operation::element_wise::PassThrough;
 
-void add_device_normalization_rank_2_1_f16_instances(
-    std::vector<std::unique_ptr<DeviceNormalization<F16, F16, F16, F16, F32, Pass, 2, 1>>>&
+void add_device_normalization_fwd_rank_2_1_f16_instances(
+    std::vector<std::unique_ptr<DeviceNormalizationFwd<F16, F16, F16, F16, F32, Pass, 2, 1>>>&
         instances)
 {
     add_device_operation_instances(instances,

library/src/tensor_operation_instance/gpu/normalization/device_layernorm2d_f32_instance.cpp → library/src/tensor_operation_instance/gpu/normalization_fwd/device_layernorm2d_fwd_f32_instance.cpp

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
-#include "normalization_instance_common.hpp"
+#include "normalization_fwd_instance_common.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -10,8 +10,8 @@ namespace instance {
 
 using Pass = ck::tensor_operation::element_wise::PassThrough;
 
-void add_device_normalization_rank_2_1_f32_instances(
-    std::vector<std::unique_ptr<DeviceNormalization<F32, F32, F32, F32, F32, Pass, 2, 1>>>&
+void add_device_normalization_fwd_rank_2_1_f32_instances(
+    std::vector<std::unique_ptr<DeviceNormalizationFwd<F32, F32, F32, F32, F32, Pass, 2, 1>>>&
         instances)
 {
     add_device_operation_instances(instances,

library/src/tensor_operation_instance/gpu/normalization/device_layernorm4d_f16_instance.cpp → library/src/tensor_operation_instance/gpu/normalization_fwd/device_layernorm4d_fwd_f16_instance.cpp

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
-#include "normalization_instance_common.hpp"
+#include "normalization_fwd_instance_common.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -10,8 +10,8 @@ namespace instance {
 
 using Pass = ck::tensor_operation::element_wise::PassThrough;
 
-void add_device_normalization_rank_4_3_f16_instances(
-    std::vector<std::unique_ptr<DeviceNormalization<F16, F16, F16, F16, F32, Pass, 4, 3>>>&
+void add_device_normalization_fwd_rank_4_3_f16_instances(
+    std::vector<std::unique_ptr<DeviceNormalizationFwd<F16, F16, F16, F16, F32, Pass, 4, 3>>>&
         instances)
 {
     add_device_operation_instances(instances,

library/src/tensor_operation_instance/gpu/normalization/device_layernorm4d_f32_instance.cpp → library/src/tensor_operation_instance/gpu/normalization_fwd/device_layernorm4d_fwd_f32_instance.cpp

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
-#include "normalization_instance_common.hpp"
+#include "normalization_fwd_instance_common.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -10,8 +10,8 @@ namespace instance {
 
 using Pass = ck::tensor_operation::element_wise::PassThrough;
 
-void add_device_normalization_rank_4_3_f32_instances(
-    std::vector<std::unique_ptr<DeviceNormalization<F32, F32, F32, F32, F32, Pass, 4, 3>>>&
+void add_device_normalization_fwd_rank_4_3_f32_instances(
+    std::vector<std::unique_ptr<DeviceNormalizationFwd<F32, F32, F32, F32, F32, Pass, 4, 3>>>&
         instances)
 {
     add_device_operation_instances(instances,

profiler/include/profiler/profile_groupnorm_impl.hpp → profiler/include/profiler/profile_groupnorm_fwd_impl.hpp

@@ -7,7 +7,7 @@
 
 #include "ck/ck.hpp"
 
-#include "ck/library/tensor_operation_instance/gpu/normalization.hpp"
+#include "ck/library/tensor_operation_instance/gpu/normalization_fwd.hpp"
 
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
@@ -88,14 +88,14 @@ bool profile_groupnorm_impl(int do_verification,
     beta_dev.ToDevice(beta.mData.data());
 
     // add device normalization instances
-    using DeviceOp = ck::tensor_operation::device::DeviceNormalization<XDataType,
-                                                                       GammaDataType,
-                                                                       BetaDataType,
-                                                                       YDataType,
-                                                                       SaveMeanInvStdDataType,
-                                                                       PassThrough,
-                                                                       5,
-                                                                       3>;
+    using DeviceOp = ck::tensor_operation::device::DeviceNormalizationFwd<XDataType,
+                                                                          GammaDataType,
+                                                                          BetaDataType,
+                                                                          YDataType,
+                                                                          SaveMeanInvStdDataType,
+                                                                          PassThrough,
+                                                                          5,
+                                                                          3>;
 
     // get device op instances
     const auto instance_ptrs =

profiler/include/profiler/profile_layernorm_impl.hpp → profiler/include/profiler/profile_layernorm_fwd_impl.hpp

@@ -6,7 +6,7 @@
 #include <iomanip>
 
 #include "ck/ck.hpp"
-#include "ck/library/tensor_operation_instance/gpu/normalization.hpp"
+#include "ck/library/tensor_operation_instance/gpu/normalization_fwd.hpp"
 #include "ck/library/utility/check_err.hpp"
 #include "ck/library/utility/device_memory.hpp"
 #include "ck/library/utility/host_tensor.hpp"
@@ -94,14 +94,14 @@ bool profile_layernorm_impl(int do_verification,
     constexpr int NumReduceDim = Rank - 1;
 
     // add device normalization instances
-    using DeviceOp = ck::tensor_operation::device::DeviceNormalization<XDataType,
-                                                                       GammaDataType,
-                                                                       BetaDataType,
-                                                                       YDataType,
-                                                                       SaveMeanInvStdDataType,
-                                                                       PassThrough,
-                                                                       Rank,
-                                                                       NumReduceDim>;
+    using DeviceOp = ck::tensor_operation::device::DeviceNormalizationFwd<XDataType,
+                                                                          GammaDataType,
+                                                                          BetaDataType,
+                                                                          YDataType,
+                                                                          SaveMeanInvStdDataType,
+                                                                          PassThrough,
+                                                                          Rank,
+                                                                          NumReduceDim>;
 
     // get device op instances
     const auto instance_ptrs =

profiler/src/CMakeLists.txt

@@ -16,8 +16,8 @@ set(PROFILER_SOURCES
     profile_grouped_conv_fwd.cpp
     profile_grouped_conv_bwd_weight.cpp
     profile_reduce.cpp
-    profile_groupnorm.cpp
-    profile_layernorm.cpp
+    profile_groupnorm_fwd.cpp
+    profile_layernorm_fwd.cpp
     profile_max_pool3d_fwd.cpp
     profile_avg_pool3d_bwd.cpp
     profile_max_pool3d_bwd.cpp
@@ -77,7 +77,7 @@ target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_grouped_conv2d_bwd_w
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_grouped_conv3d_bwd_weight_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_conv2d_fwd_bias_relu_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_conv2d_fwd_bias_relu_add_instance)
-target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_normalization_instance)
+target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_normalization_fwd_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_softmax_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_reduce_instance)
 target_link_libraries(${PROFILER_EXECUTABLE} PRIVATE device_batchnorm_instance)

profiler/src/profile_groupnorm.cpp → profiler/src/profile_groupnorm_fwd.cpp

@@ -6,7 +6,7 @@
 #include <unordered_map>
 
 #include "profiler/data_type_enum.hpp"
-#include "profiler/profile_groupnorm_impl.hpp"
+#include "profiler/profile_groupnorm_fwd_impl.hpp"
 #include "profiler_operation_registry.hpp"
 
 using ck::index_t;

profiler/src/profile_layernorm.cpp → profiler/src/profile_layernorm_fwd.cpp

@@ -6,7 +6,7 @@
 #include <unordered_map>
 
 #include "profiler/data_type_enum.hpp"
-#include "profiler/profile_layernorm_impl.hpp"
+#include "profiler/profile_layernorm_fwd_impl.hpp"
 #include "profiler_operation_registry.hpp"
 
 using ck::index_t;
@@ -76,19 +76,46 @@ int profile_layernorm(int argc, char* argv[])
     arg_parser(argc, argv);
     const std::vector<index_t> length = arg_parser.long_opts["length"];
 
-    using F16          = ck::half_t;
-    using F32          = float;
-    constexpr int rank = 2;
+    using F16 = ck::half_t;
+    using F32 = float;
 
-    if(data_type == ck::DataTypeEnum::Half)
+    if(length.size() == 2)
     {
-        ck::profiler::profile_layernorm_impl<F16, F16, F16, F32, F16, F32, false, rank>(
-            do_verification, init_method, do_log, time_kernel, length);
+        constexpr int rank = 2;
+
+        if(data_type == ck::DataTypeEnum::Half)
+        {
+            ck::profiler::profile_layernorm_impl<F16, F16, F16, F32, F16, F32, false, rank>(
+                do_verification, init_method, do_log, time_kernel, length);
+        }
+        else if(data_type == ck::DataTypeEnum::Float)
+        {
+            ck::profiler::profile_layernorm_impl<F32, F32, F32, F32, F32, F32, false, rank>(
+                do_verification, init_method, do_log, time_kernel, length);
+        }
+        else
+        {
+            throw std::runtime_error("not implemented yet");
+        }
     }
-    else if(data_type == ck::DataTypeEnum::Float)
+    else if(length.size() == 4)
     {
-        ck::profiler::profile_layernorm_impl<F32, F32, F32, F32, F32, F32, false, rank>(
-            do_verification, init_method, do_log, time_kernel, length);
+        constexpr int rank = 4;
+
+        if(data_type == ck::DataTypeEnum::Half)
+        {
+            ck::profiler::profile_layernorm_impl<F16, F16, F16, F32, F16, F32, false, rank>(
+                do_verification, init_method, do_log, time_kernel, length);
+        }
+        else if(data_type == ck::DataTypeEnum::Float)
+        {
+            ck::profiler::profile_layernorm_impl<F32, F32, F32, F32, F32, F32, false, rank>(
+                do_verification, init_method, do_log, time_kernel, length);
+        }
+        else
+        {
+            throw std::runtime_error("not implemented yet");
+        }
     }
     else
     {

test/CMakeLists.txt

@@ -139,7 +139,7 @@ add_subdirectory(grouped_convnd_fwd)
 add_subdirectory(grouped_convnd_bwd_weight)
 add_subdirectory(block_to_ctile_map)
 add_subdirectory(softmax)
-add_subdirectory(normalization)
+add_subdirectory(normalization_fwd)
 add_subdirectory(data_type)
 add_subdirectory(elementwise_normalization)
 add_subdirectory(batchnorm)

test/normalization/CMakeLists.txt

-add_custom_target(test_normalization)
-add_gtest_executable(test_layernorm2d_fp32 test_layernorm2d_fp32.cpp)
-if(result EQUAL 0)
-  target_link_libraries(test_layernorm2d_fp32 PRIVATE utility device_normalization_instance)
-  add_dependencies(test_normalization test_layernorm2d_fp32)
-endif()
-add_gtest_executable(test_groupnorm_fp32 test_groupnorm_fp32.cpp)
-if(result EQUAL 0)
-  target_link_libraries(test_groupnorm_fp32 PRIVATE utility device_normalization_instance)
-  add_dependencies(test_normalization test_groupnorm_fp32)
-endif()
-add_gtest_executable(test_layernorm2d_fp16 test_layernorm2d_fp16.cpp)
-if(result EQUAL 0)
-  target_link_libraries(test_layernorm2d_fp16 PRIVATE utility device_normalization_instance)
-  add_dependencies(test_normalization test_layernorm2d_fp16)
-endif()
-add_gtest_executable(test_groupnorm_fp16 test_groupnorm_fp16.cpp)
-if(result EQUAL 0)
-  target_link_libraries(test_groupnorm_fp16 PRIVATE utility device_normalization_instance)
-  add_dependencies(test_normalization test_groupnorm_fp16)
-endif()

test/normalization_fwd/CMakeLists.txt

+add_custom_target(test_normalization_fwd)
+add_gtest_executable(test_layernorm2d_fwd_fp32 test_layernorm2d_fwd_fp32.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_layernorm2d_fwd_fp32 PRIVATE utility device_normalization_fwd_instance)
+  add_dependencies(test_normalization_fwd test_layernorm2d_fwd_fp32)
+endif()
+
+add_gtest_executable(test_groupnorm_fwd_fp32 test_groupnorm_fwd_fp32.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_groupnorm_fwd_fp32 PRIVATE utility device_normalization_fwd_instance)
+  add_dependencies(test_normalization_fwd test_groupnorm_fwd_fp32)
+endif()
+
+add_gtest_executable(test_layernorm2d_fwd_fp16 test_layernorm2d_fwd_fp16.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_layernorm2d_fwd_fp16 PRIVATE utility device_normalization_fwd_instance)
+  add_dependencies(test_normalization_fwd test_layernorm2d_fwd_fp16)
+endif()
+
+add_gtest_executable(test_layernorm4d_fwd_fp16 test_layernorm4d_fwd_fp16.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_layernorm4d_fwd_fp16 PRIVATE utility device_normalization_fwd_instance)
+  add_dependencies(test_normalization_fwd test_layernorm4d_fwd_fp16)
+endif()
+
+add_gtest_executable(test_groupnorm_fwd_fp16 test_groupnorm_fwd_fp16.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_groupnorm_fwd_fp16 PRIVATE utility device_normalization_fwd_instance)
+  add_dependencies(test_normalization_fwd test_groupnorm_fwd_fp16)
+endif()

test/normalization/test_groupnorm_fp16.cpp → test/normalization_fwd/test_groupnorm_fwd_fp16.cpp

@@ -2,7 +2,7 @@
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "gtest/gtest.h"
-#include "profiler/profile_groupnorm_impl.hpp"
+#include "profiler/profile_groupnorm_fwd_impl.hpp"
 
 using F16 = ck::half_t;
 using F32 = float;

test/normalization/test_groupnorm_fp32.cpp → test/normalization_fwd/test_groupnorm_fwd_fp32.cpp

@@ -2,7 +2,7 @@
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "gtest/gtest.h"
-#include "profiler/profile_groupnorm_impl.hpp"
+#include "profiler/profile_groupnorm_fwd_impl.hpp"
 
 using F16 = ck::half_t;
 using F32 = float;

test/normalization/test_layernorm2d_fp16.cpp → test/normalization_fwd/test_layernorm2d_fwd_fp16.cpp

@@ -2,7 +2,7 @@
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "gtest/gtest.h"
-#include "profiler/profile_layernorm_impl.hpp"
+#include "profiler/profile_layernorm_fwd_impl.hpp"
 
 using F16 = ck::half_t;
 using F32 = float;

test/normalization/test_layernorm2d_fp32.cpp → test/normalization_fwd/test_layernorm2d_fwd_fp32.cpp

@@ -2,7 +2,7 @@
 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "gtest/gtest.h"
-#include "profiler/profile_layernorm_impl.hpp"
+#include "profiler/profile_layernorm_fwd_impl.hpp"
 
 using F16 = ck::half_t;
 using F32 = float;

test/normalization_fwd/test_layernorm4d_fwd_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gtest/gtest.h"
+#include "profiler/profile_layernorm_fwd_impl.hpp"
+
+using F16 = ck::half_t;
+using F32 = float;
+using ck::index_t;
+
+template <typename Tuple>
+class TestLayernorm4d : public ::testing::Test
+{
+    protected:
+    using XDataType              = std::tuple_element_t<0, Tuple>;
+    using GammaDataType          = std::tuple_element_t<1, Tuple>;
+    using BetaDataType           = std::tuple_element_t<2, Tuple>;
+    using ComputeDataType        = std::tuple_element_t<3, Tuple>;
+    using YDataType              = std::tuple_element_t<4, Tuple>;
+    using SaveMeanInvStdDataType = std::tuple_element_t<5, Tuple>;
+
+    void Run()
+    {
+        // [N, D], reduce D
+        std::vector<std::vector<ck::index_t>> lengths = {
+            {1, 1, 1, 1}, {7, 7, 7, 7}, {256, 16, 16, 8}};
+
+        for(auto length : lengths)
+        {
+            bool success = ck::profiler::profile_layernorm_impl<XDataType,
+                                                                GammaDataType,
+                                                                BetaDataType,
+                                                                ComputeDataType,
+                                                                YDataType,
+                                                                SaveMeanInvStdDataType,
+                                                                true,
+                                                                4>(true, 2, false, false, length);
+            EXPECT_TRUE(success);
+        }
+    }
+};
+
+using KernelTypes = ::testing::Types<
+    // XDataType, GammaDataType, BetaDataType, ComputeDataType, YDataType>
+    std::tuple<F16, F16, F16, F32, F16, F32>>;
+
+TYPED_TEST_SUITE(TestLayernorm4d, KernelTypes);
+TYPED_TEST(TestLayernorm4d, Test_FP16) { this->Run(); }