Merge branch 'amd-develop' into amd-master

example/20_grouped_conv_bwd_weight/grouped_conv_bwd_weight_xdl_fp16_comp_bf8_fp8.cpp

@@ -83,4 +83,23 @@ using HostConvBwdWeightInstance = ck::tensor_operation::host::ReferenceConvBwdWe
 
 #include "run_grouped_conv_bwd_weight_example.inc"
 
-int main(int argc, char* argv[]) { return !run_grouped_conv_bwd_weight_example(argc, argv); }
+int main(int argc, char* argv[])
+{
+    ExecutionConfig config;
+    ck::utils::conv::ConvParam conv_param = DefaultConvParam;
+
+    if(!parse_cmd_args(argc, argv, config, conv_param))
+    {
+        return 1;
+    }
+
+    switch(conv_param.num_dim_spatial_)
+    {
+    case 1: return !run_grouped_conv_bwd_weight<1>(config, conv_param);
+    case 2: return !run_grouped_conv_bwd_weight<2>(config, conv_param);
+    case 3: return !run_grouped_conv_bwd_weight<3>(config, conv_param);
+    default: break;
+    }
+
+    return 1;
+}

example/20_grouped_conv_bwd_weight/run_grouped_conv_bwd_weight_example.inc

@@ -5,7 +5,7 @@ template <ck::index_t NDimSpatial>
 bool run_grouped_conv_bwd_weight(const ExecutionConfig& config,
                                  const ck::utils::conv::ConvParam& conv_param)
 {
-    // Dl op doesn't support split_k > 1
+    // Dl and WMMA ops don't support split_k > 1
     constexpr ck::index_t split_k = 1;
 
     const auto in_g_n_c_wis_desc =
@@ -143,23 +143,3 @@ bool run_grouped_conv_bwd_weight(const ExecutionConfig& config,
 
     return true;
 }
-
-bool run_grouped_conv_bwd_weight_example(int argc, char* argv[])
-{
-    ExecutionConfig config;
-    ck::utils::conv::ConvParam conv_param = DefaultConvParam;
-
-    if(!parse_cmd_args(argc, argv, config, conv_param))
-    {
-        return false;
-    }
-
-    switch(conv_param.num_dim_spatial_)
-    {
-    case 1: return run_grouped_conv_bwd_weight<1>(config, conv_param);
-    case 2: return run_grouped_conv_bwd_weight<2>(config, conv_param);
-    case 3: return run_grouped_conv_bwd_weight<3>(config, conv_param);
-    }
-
-    return false;
-}

example/21_gemm_layernorm/gemm_bias_relu_add_layernorm_xdl_welford_fp16.cpp

@@ -114,12 +114,15 @@ void host_gemm_layernorm(Tensor<HDataType>& h_m_n,
                                                                               BetaDataType,
                                                                               HDataType,
                                                                               AccDataType,
+                                                                              AccDataType,
                                                                               HElementOp,
                                                                               2,
                                                                               1>;
 
     Tensor<EMeanVarDataType> e_m_n(HostTensorDescriptor{M, N});
     Tensor<AccDataType> c_m_n(HostTensorDescriptor{M, N});
+    Tensor<AccDataType> save_mean({M});
+    Tensor<AccDataType> save_inv_std({M});
 
     auto ref_gemm         = ReferenceGemm{};
     auto ref_gemm_invoker = ref_gemm.MakeInvoker();
@@ -145,7 +148,7 @@ void host_gemm_layernorm(Tensor<HDataType>& h_m_n,
     auto ref_layernorm_invoker = ref_layernorm.MakeInvoker();
 
     auto ref_layernorm_argument = ref_layernorm.MakeArgument(
-        e_m_n, gamma_n, beta_n, h_m_n, h_element_op, {M, N}, {1}, epsilon);
+        e_m_n, gamma_n, beta_n, h_m_n, save_mean, save_inv_std, h_element_op, {M, N}, {1}, epsilon);
     ref_layernorm_invoker.Run(ref_layernorm_argument);
 }
 

example/22_cgemm/CMakeLists.txt

 add_custom_target(example_cgemm_xdl)
 
 add_example_executable(example_cgemm_xdl_bf16 cgemm_xdl_bf16.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_cgemm_xdl example_cgemm_xdl_bf16)
-endif()
+add_example_dependencies(example_cgemm_xdl example_cgemm_xdl_bf16)
+
 add_example_executable(example_cgemm_xdl_fp16 cgemm_xdl_fp16.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_cgemm_xdl example_cgemm_xdl_fp16)
-endif()
+add_example_dependencies(example_cgemm_xdl example_cgemm_xdl_fp16)
+
 add_example_executable(example_cgemm_xdl_fp32 cgemm_xdl_fp32.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_cgemm_xdl example_cgemm_xdl_fp32)
-endif()
+add_example_dependencies(example_cgemm_xdl example_cgemm_xdl_fp32)
+
 add_example_executable(example_cgemm_xdl_int8 cgemm_xdl_int8.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_cgemm_xdl example_cgemm_xdl_int8)
-endif()
+add_example_dependencies(example_cgemm_xdl example_cgemm_xdl_int8)
+
 if(USE_BITINT_EXTENSION_INT4)
     add_example_executable(example_cgemm_xdl_int4 cgemm_xdl_int4.cpp)
-  add_dependencies(example_cgemm_xdl example_cgemm_xdl_int4)
+    add_example_dependencies(example_cgemm_xdl example_cgemm_xdl_int4)
 endif()

example/24_batched_gemm/CMakeLists.txt

 add_custom_target(example_batched_gemm_xdl)
+
 add_example_executable(example_batched_gemm_xdl_fp32 batched_gemm_xdl_fp32.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_fp32)
-endif()
+add_example_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_fp32)
+
 add_example_executable(example_batched_gemm_xdl_fp16 batched_gemm_xdl_fp16.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_fp16)
-endif()
+add_example_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_fp16)
+
 add_example_executable(example_batched_gemm_xdl_bf16 batched_gemm_xdl_bf16.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_bf16)
-endif()
+add_example_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_bf16)
+
 add_example_executable(example_batched_gemm_xdl_int8 batched_gemm_xdl_int8.cpp)
-if(result EQUAL 0)
-  add_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_int8)
-endif()
+add_example_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_int8)
+
 if(USE_BITINT_EXTENSION_INT4)
     add_example_executable(example_batched_gemm_xdl_int4 batched_gemm_xdl_int4.cpp)
-  if(result EQUAL 0)
-    add_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_int4)
-  endif()
+    add_example_dependencies(example_batched_gemm_xdl example_batched_gemm_xdl_int4)
 endif()

example/27_layernorm/layernorm_fp16.cpp

@@ -7,9 +7,12 @@ using XDataType       = ck::half_t;
 using GammaDataType          = ck::half_t;
 using BetaDataType           = ck::half_t;
 using YDataType              = ck::half_t;
+using SaveMeanInvStdDataType = float;
 using ComputeDataType        = float;
 using PassThrough            = ck::tensor_operation::element_wise::PassThrough;
 
+#define SAVE_MEAN_INV_STD
+
 constexpr int Rank         = 2;
 constexpr int NumReduceDim = 1;
 
@@ -19,6 +22,7 @@ using DeviceInstance =
                                                           BetaDataType,
                                                           ComputeDataType,
                                                           YDataType,
+                                                          SaveMeanInvStdDataType,
                                                           PassThrough,
                                                           Rank,
                                                           NumReduceDim,
@@ -33,7 +37,8 @@ using DeviceInstance =
                                                           8,   // GammaScalarPerVector
                                                           1,   // BetaVecDim (0=M, 1=K)
                                                           8,   // BetaScalarPerVector
-                                                          8>;  // OutScalarPerVector
+                                                          8,   // YScalarPerVector
+                                                          1>;  // SaveMeanInvStdScalarPerVector
 #include "run_layernorm_example.inc"
 
 int main() { return run_groupnorm_example<DeviceInstance>(); }

example/27_layernorm/layernorm_splitk_fp16.cpp

@@ -7,9 +7,12 @@ using XDataType       = ck::half_t;
 using GammaDataType          = ck::half_t;
 using BetaDataType           = ck::half_t;
 using YDataType              = ck::half_t;
+using SaveMeanInvStdDataType = float;
 using ComputeDataType        = float;
 using PassThrough            = ck::tensor_operation::element_wise::PassThrough;
 
+#define SAVE_MEAN_INV_STD
+
 constexpr int Rank         = 2;
 constexpr int NumReduceDim = 1;
 
@@ -19,6 +22,7 @@ using DeviceInstance =
                                                                 BetaDataType,
                                                                 ComputeDataType,
                                                                 YDataType,
+                                                                SaveMeanInvStdDataType,
                                                                 PassThrough,
                                                                 Rank,
                                                                 NumReduceDim,
@@ -33,7 +37,8 @@ using DeviceInstance =
                                                                 8,   // GammaScalarPerVector
                                                                 1,   // BetaVecDim (0=M, 1=K)
                                                                 8,   // BetaScalarPerVector
-                                                                8>;  // YScalarPerVector
+                                                                8,   // YScalarPerVector
+                                                                1>; // SaveMeanInvStdScalarPerVector
 
 #include "run_layernorm_example.inc"
 

example/27_layernorm/run_layernorm_example.inc

@@ -10,22 +10,13 @@ int run_groupnorm_example()
 
     ck::index_t M      = 1024;
     ck::index_t N      = 1024;
-    ck::index_t Stride = N;
 
-    auto f_host_tensor_descriptor1d = [](std::size_t len, std::size_t stride) {
-        return HostTensorDescriptor({len}, {stride});
-    };
-
-    auto f_host_tensor_descriptor2d = [](std::size_t row, std::size_t col, std::size_t stride) {
-        using namespace ck::literals;
-
-        return HostTensorDescriptor({row, col}, {stride, 1_uz});
-    };
-
-    Tensor<XDataType> x(f_host_tensor_descriptor2d(M, N, Stride));
-    Tensor<GammaDataType> gamma(f_host_tensor_descriptor1d(N, 1));
-    Tensor<BetaDataType> beta(f_host_tensor_descriptor1d(N, 1));
-    Tensor<YDataType> y(f_host_tensor_descriptor2d(M, N, Stride));
+    Tensor<XDataType> x({M, N});
+    Tensor<GammaDataType> gamma({N});
+    Tensor<BetaDataType> beta({N});
+    Tensor<YDataType> y({M, N});
+    Tensor<SaveMeanInvStdDataType> save_mean({M});
+    Tensor<SaveMeanInvStdDataType> save_inv_std({M});
 
     x.GenerateTensorValue(GeneratorTensor_3<XDataType>{0.0, 1.0});
     gamma.GenerateTensorValue(GeneratorTensor_3<GammaDataType>{0.0, 1.0});
@@ -35,6 +26,11 @@ int run_groupnorm_example()
     DeviceMem gamma_dev(sizeof(GammaDataType) * gamma.mDesc.GetElementSpaceSize());
     DeviceMem beta_dev(sizeof(BetaDataType) * beta.mDesc.GetElementSpaceSize());
     DeviceMem y_dev(sizeof(YDataType) * y.mDesc.GetElementSpaceSize());
+#ifdef SAVE_MEAN_INV_STD
+    DeviceMem save_mean_dev(sizeof(SaveMeanInvStdDataType) * save_mean.mDesc.GetElementSpaceSize());
+    DeviceMem save_inv_std_dev(sizeof(SaveMeanInvStdDataType) *
+                               save_inv_std.mDesc.GetElementSpaceSize());
+#endif
 
     x_dev.ToDevice(x.mData.data());
     gamma_dev.ToDevice(gamma.mData.data());
@@ -47,14 +43,23 @@ int run_groupnorm_example()
         {0, 1},
         {0, 1},
         std::vector<ck::index_t>{y.mDesc.GetStrides().begin(), y.mDesc.GetStrides().end()},
+        std::vector<ck::index_t>{save_mean.mDesc.GetStrides().begin(),
+                                    save_mean.mDesc.GetStrides().end()},
+        std::vector<ck::index_t>{save_mean.mDesc.GetStrides().begin(),
+                                    save_mean.mDesc.GetStrides().end()},
         {1},
         1e-4,
         x_dev.GetDeviceBuffer(),
         gamma_dev.GetDeviceBuffer(),
         beta_dev.GetDeviceBuffer(),
         y_dev.GetDeviceBuffer(),
+#ifdef SAVE_MEAN_INV_STD
+        save_mean_dev.GetDeviceBuffer(),
+        save_inv_std_dev.GetDeviceBuffer(),
+#else
         nullptr,
         nullptr,
+#endif
         PassThrough{});
 
     if(!device_instance.IsSupportedArgument(argument_ptr.get()))
@@ -72,24 +77,45 @@ int run_groupnorm_example()
 
     bool pass = true;
     {
-        Tensor<YDataType> host_y(f_host_tensor_descriptor2d(M, N, Stride));
-        using ReferenceInstance = ck::tensor_operation::host::ReferenceLayernorm<XDataType,
+        Tensor<YDataType> host_y({M, N});
+        Tensor<SaveMeanInvStdDataType> host_save_mean({M});
+        Tensor<SaveMeanInvStdDataType> host_save_inv_std({M});
+
+        using ReferenceInstance =
+            ck::tensor_operation::host::ReferenceLayernorm<XDataType,
                                                            GammaDataType,
                                                            BetaDataType,
                                                            YDataType,
+                                                           SaveMeanInvStdDataType,
                                                            ComputeDataType,
                                                            PassThrough,
                                                            Rank,
                                                            NumReduceDim>;
 
         ReferenceInstance ref;
-        auto ref_argument =
-            ref.MakeArgument(x, gamma, beta, host_y, PassThrough{}, {M, N}, {1}, 1e-4);
+        auto ref_argument = ref.MakeArgument(x,
+                                             gamma,
+                                             beta,
+                                             host_y,
+                                             host_save_mean,
+                                             host_save_inv_std,
+                                             PassThrough{},
+                                             {M, N},
+                                             {1},
+                                             1e-4);
         auto ref_invoker  = ref.MakeInvoker();
         ref_invoker.Run(ref_argument);
 
         y_dev.FromDevice(y.mData.data());
-        pass &= ck::utils::check_err(y, host_y, "Error: Incorrect results", 1e-3, 1e-3);
+        pass &= ck::utils::check_err(y, host_y, "Error: Incorrect results (y)", 1e-3, 1e-3);
+#ifdef SAVE_MEAN_INV_STD
+        save_mean_dev.FromDevice(save_mean.mData.data());
+        save_inv_std_dev.FromDevice(save_inv_std.mData.data());
+        pass &= ck::utils::check_err(
+            save_mean, host_save_mean, "Error: Incorrect results (mean)", 1e-3, 1e-3);
+        pass &= ck::utils::check_err(
+            save_inv_std, host_save_inv_std, "Error: Incorrect results (inv_std)", 1e-3, 1e-3);
+#endif
     }
 
     return (pass ? 0 : 1);

example/30_grouped_conv_fwd_multiple_d/CMakeLists.txt

@@ -5,31 +5,25 @@ set(target 0)
 foreach(gpu IN LISTS GPU_TARGETS)
     if(gpu IN_LIST gpu_list1 AND target EQUAL 0)
         add_custom_target(example_grouped_conv_fwd_multiple_d)
+
         add_example_executable(example_grouped_conv_fwd_bias_relu_add_xdl_fp16 grouped_conv_fwd_bias_relu_add_xdl_fp16.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_fp16)
-   endif()
+        add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_fp16)
+
         add_example_executable(example_grouped_conv_fwd_xdl_fp16 grouped_conv_fwd_xdl_fp16.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_xdl_fp16)
-   endif()
+        add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_xdl_fp16)
+
         add_example_executable(example_grouped_conv_fwd_bias_relu_add_xdl_fp32 grouped_conv_fwd_bias_relu_add_xdl_fp32.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_fp32)
-   endif()
+        add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_fp32)
+
         add_example_executable(example_grouped_conv_fwd_bias_relu_add_xdl_bf16 grouped_conv_fwd_bias_relu_add_xdl_bf16.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_bf16)
-   endif()
+        add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_bf16)
+
         add_example_executable(example_grouped_conv_fwd_bias_relu_add_xdl_int8 grouped_conv_fwd_bias_relu_add_xdl_int8.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_int8)
-   endif()
+        add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_int8)
+
         if(USE_BITINT_EXTENSION_INT4)
             add_example_executable(example_grouped_conv_fwd_bias_relu_add_xdl_int4 grouped_conv_fwd_bias_relu_add_xdl_int4.cpp)
-      if(result EQUAL 0)
-         add_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_int4)
-      endif()
+            add_example_dependencies(example_grouped_conv_fwd_multiple_d example_grouped_conv_fwd_bias_relu_add_xdl_int4)
         endif() # USE_BITINT_EXTENSION_INT4
 
         set(target 1)

example/30_grouped_conv_fwd_multiple_d/README.md

@@ -4,7 +4,7 @@ arg1: verification (0=no, 1=yes)
 arg2: initialization (0=no init, 1=integer value, 2=decimal value)
 arg3: time kernel (0=no, 1=yes)
 Following arguments (depending on number of spatial dims):
- Number of spatial dimensions (1=Conv1d, 2=Conv2d, 3=Conv3d)
+ Number of spatial dimensions (1=Conv1D, 2=Conv2D, 3=Conv3D)
  G, N, K, C,
  <filter spatial dimensions>, (ie Y, X for 2D)
  <input image spatial dimensions>, (ie Hi, Wi for 2D)

example/32_batched_gemm_scale_softmax_gemm/CMakeLists.txt

 add_custom_target(example_gemm_scale_softmax_gemm)
 
 add_example_executable(example_batched_gemm_scale_softmax_gemm_xdl_fp16 batched_gemm_scale_softmax_gemm_xdl_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_xdl_fp16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_xdl_fp16)
+
 add_example_executable(example_batched_gemm_scale_softmax_gemm_permute_xdl_fp16 batched_gemm_scale_softmax_gemm_permute_xdl_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_permute_xdl_fp16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_permute_xdl_fp16)
+
 add_example_executable(example_grouped_gemm_scale_softmax_gemm_permute_xdl_fp16 grouped_gemm_scale_softmax_gemm_permute_xdl_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_grouped_gemm_scale_softmax_gemm_permute_xdl_fp16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_grouped_gemm_scale_softmax_gemm_permute_xdl_fp16)
+
 add_example_executable(example_batched_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16 batched_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16)
+
 add_example_executable(example_grouped_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16 grouped_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_grouped_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_grouped_gemm_lower_triangle_scale_softmax_gemm_permute_xdl_fp16)
+
 add_example_executable(example_batched_gemm_scale_softmax_gemm_xdl_bf16 batched_gemm_scale_softmax_gemm_xdl_bf16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_xdl_bf16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_xdl_bf16)
+
 add_example_executable(example_batched_gemm_scale_softmax_gemm_permute_xdl_bf16 batched_gemm_scale_softmax_gemm_permute_xdl_bf16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_permute_xdl_bf16)
-endif()
+add_example_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_permute_xdl_bf16)
 

example/35_splitK_gemm/CMakeLists.txt

@@ -5,27 +5,22 @@ foreach(gpu IN LISTS GPU_TARGETS)
    add_custom_target(example_splitK_gemm_xdl)
 
    add_example_executable(example_splitK_gemm_xdl_fp32 splitK_gemm_xdl_fp32.cpp)
-  if(result EQUAL 0)
-    add_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_fp32)
-   endif()
+   add_example_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_fp32)
+
    add_example_executable(example_splitK_gemm_xdl_fp16 splitK_gemm_xdl_fp16.cpp)
-  if(result EQUAL 0)
-    add_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_fp16)
-   endif()
+   add_example_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_fp16)
+
    add_example_executable(example_splitK_gemm_xdl_bf16 splitK_gemm_xdl_bf16.cpp)
-  if(result EQUAL 0)
-    add_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_bf16)
-   endif()
+   add_example_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_bf16)
+
    add_example_executable(example_splitK_gemm_xdl_int8 splitK_gemm_xdl_int8.cpp)
-  if(result EQUAL 0)
-    add_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_int8)
-   endif()
+   add_example_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_int8)
+
    if(USE_BITINT_EXTENSION_INT4)
        add_example_executable(example_splitK_gemm_xdl_int4 splitK_gemm_xdl_int4.cpp)
-      if(result EQUAL 0)
-        add_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_int4)
-      endif()
+       add_example_dependencies(example_splitK_gemm_xdl example_splitK_gemm_xdl_int4)
    endif()
+
    set(target 1)
  endif()
 endforeach()

example/38_grouped_conv_bwd_data_multiple_d/CMakeLists.txt

@@ -4,14 +4,13 @@ set(target 0)
 foreach(gpu IN LISTS GPU_TARGETS)
     if(gpu IN_LIST gpu_list_xdl AND target EQUAL 0)
         add_custom_target(example_grouped_conv_bwd_data)
+
         add_example_executable(example_grouped_conv_bwd_data_xdl_fp16 grouped_conv_bwd_data_xdl_fp16.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_xdl_fp16)
-   endif()
+        add_example_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_xdl_fp16)
+
         add_example_executable(example_grouped_conv_bwd_data_bias_relu_xdl_fp16 grouped_conv_bwd_data_bias_relu_xdl_fp16.cpp)
-   if(result EQUAL 0)
-     add_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_bias_relu_xdl_fp16)
-   endif()
+        add_example_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_bias_relu_xdl_fp16)
+
         set(target 1)
     endif()
 endforeach()
@@ -19,10 +18,10 @@ endforeach()
 foreach(gpu IN LISTS GPU_TARGETS)
     if(gpu IN_LIST gpu_list_wmma AND target EQUAL 0)
         add_custom_target(example_grouped_conv_bwd_data)
+
         add_example_executable(example_grouped_conv_bwd_data_wmma_fp16 grouped_conv_bwd_data_wmma_fp16.cpp)
-   if(result EQUAL 0)
-      add_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_wmma_fp16)
-   endif()
+        add_example_dependencies(example_grouped_conv_bwd_data example_grouped_conv_bwd_data_wmma_fp16)
+
         set(target 1)
     endif()
 endforeach()

example/39_permute/CMakeLists.txt

 add_custom_target(example_permute)
 
 add_example_executable(example_permute_1xHxW_fp16 permute_1xHxW_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_permute example_permute_1xHxW_fp16)
-endif()
+add_example_dependencies(example_permute example_permute_1xHxW_fp16)
+
 add_example_executable(example_permute_NxHxW_fp16 permute_NxHxW_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_permute example_permute_NxHxW_fp16)
-endif()
+add_example_dependencies(example_permute example_permute_NxHxW_fp16)
+
 add_example_executable(example_permute_HxWx4_fp16 permute_HxWx4_fp16.cpp)
-if(result EQUAL 0)
-    add_dependencies(example_permute example_permute_HxWx4_fp16)
-endif()
+add_example_dependencies(example_permute example_permute_HxWx4_fp16)

example/42_groupnorm/groupnorm_sigmoid_mul_fp16.cpp

@@ -10,8 +10,11 @@ using XDataType       = ck::half_t;
 using GammaDataType          = ck::half_t;
 using BetaDataType           = ck::half_t;
 using YDataType              = ck::half_t;
+using SaveMeanInvStdDataType = float;
 using ComputeDataType        = float;
 
+#define SAVE_MEAN_INV_STD
+
 struct YElementOp
 {
     template <typename Y, typename X>
@@ -39,6 +42,7 @@ using DeviceInstance =
                                                           BetaDataType,
                                                           ComputeDataType,
                                                           YDataType,
+                                                          SaveMeanInvStdDataType,
                                                           YElementOp,
                                                           Rank,
                                                           NumReduceDim,
@@ -53,7 +57,8 @@ using DeviceInstance =
                                                           2,    // GammaScalarPerVector
                                                           1,    // BetaVecDim (0=M, 1=K)
                                                           2,    // BetaScalarPerVector
-                                                          2>;   // OutScalarPerVector
+                                                          2,    // YScalarPerVector
+                                                          1>;   // SaveMeanInvStdScalarPerVector
 
 #include "run_groupnorm_example.inc"
 

example/42_groupnorm/groupnorm_splitk_fp16.cpp

@@ -10,15 +10,19 @@ using XDataType       = ck::half_t;
 using GammaDataType          = ck::half_t;
 using BetaDataType           = ck::half_t;
 using YDataType              = ck::half_t;
+using SaveMeanInvStdDataType = float;
 using ComputeDataType        = float;
 using YElementOp             = ck::tensor_operation::element_wise::Swish;
 
+#define SAVE_MEAN_INV_STD
+
 using DeviceInstance =
     ck::tensor_operation::device::DeviceNormalizationSplitKImpl<XDataType,
                                                                 GammaDataType,
                                                                 BetaDataType,
                                                                 ComputeDataType,
                                                                 YDataType,
+                                                                SaveMeanInvStdDataType,
                                                                 YElementOp,
                                                                 Rank,
                                                                 NumReduceDim,
@@ -33,7 +37,8 @@ using DeviceInstance =
                                                                 2,   // GammaScalarPerVector
                                                                 1,   // BetaVecDim (0=M, 1=K)
                                                                 2,   // BetaScalarPerVector
-                                                                2>;  // OutScalarPerVector
+                                                                2,   // YScalarPerVector
+                                                                1>; // SaveMeanInvStdScalarPerVector
 
 #include "run_groupnorm_example.inc"
 

example/42_groupnorm/groupnorm_swish_fp16.cpp

@@ -10,15 +10,19 @@ using XDataType       = ck::half_t;
 using GammaDataType          = ck::half_t;
 using BetaDataType           = ck::half_t;
 using YDataType              = ck::half_t;
+using SaveMeanInvStdDataType = float;
 using ComputeDataType        = float;
 using YElementOp             = ck::tensor_operation::element_wise::Swish;
 
+#define SAVE_MEAN_INV_STD
+
 using DeviceInstance =
     ck::tensor_operation::device::DeviceNormalizationImpl<XDataType,
                                                           GammaDataType,
                                                           BetaDataType,
                                                           ComputeDataType,
                                                           YDataType,
+                                                          SaveMeanInvStdDataType,
                                                           YElementOp,
                                                           Rank,
                                                           NumReduceDim,
@@ -33,7 +37,8 @@ using DeviceInstance =
                                                           2,    // GammaScalarPerVector
                                                           1,    // BetaVecDim (0=M, 1=K)
                                                           2,    // BetaScalarPerVector
-                                                          2>;   // OutScalarPerVector
+                                                          2,    // YScalarPerVector
+                                                          1>;   // SaveMeanInvStdScalarPerVector
 
 #include "run_groupnorm_example.inc"
 

example/42_groupnorm/run_groupnorm_example.inc

@@ -34,6 +34,8 @@ int run_groupnorm_example(int argc, char* argv[])
     Tensor<YDataType> y({N, H, W, G, C});
     Tensor<GammaDataType> gamma({G, C});
     Tensor<BetaDataType> beta({G, C});
+    Tensor<SaveMeanInvStdDataType> save_mean({N, G});
+    Tensor<SaveMeanInvStdDataType> save_inv_std({N, G});
 
     ck::utils::FillUniformDistribution<XDataType>{0.f, 1.f}(x);
     ck::utils::FillUniformDistribution<GammaDataType>{0.f, 1.f}(gamma);
@@ -43,6 +45,11 @@ int run_groupnorm_example(int argc, char* argv[])
     DeviceMem gamma_dev(sizeof(GammaDataType) * gamma.mDesc.GetElementSpaceSize());
     DeviceMem beta_dev(sizeof(BetaDataType) * beta.mDesc.GetElementSpaceSize());
     DeviceMem y_dev(sizeof(YDataType) * y.mDesc.GetElementSpaceSize());
+#ifdef SAVE_MEAN_INV_STD
+    DeviceMem save_mean_dev(sizeof(SaveMeanInvStdDataType) * save_mean.mDesc.GetElementSpaceSize());
+    DeviceMem save_inv_std_dev(sizeof(SaveMeanInvStdDataType) *
+                               save_inv_std.mDesc.GetElementSpaceSize());
+#endif
 
     x_dev.ToDevice(x.mData.data());
     gamma_dev.ToDevice(gamma.mData.data());
@@ -57,14 +64,23 @@ int run_groupnorm_example(int argc, char* argv[])
         {0, 0, 0, C, 1},
         {0, 0, 0, C, 1},
         std::vector<ck::index_t>{y.mDesc.GetStrides().begin(), y.mDesc.GetStrides().end()},
+        std::vector<ck::index_t>{save_mean.mDesc.GetStrides().begin(),
+                                    save_mean.mDesc.GetStrides().end()},
+        std::vector<ck::index_t>{save_mean.mDesc.GetStrides().begin(),
+                                    save_mean.mDesc.GetStrides().end()},
         {1, 2, 4}, // reduction dimension: [H, W, C]
         1e-6,
         x_dev.GetDeviceBuffer(),
         gamma_dev.GetDeviceBuffer(),
         beta_dev.GetDeviceBuffer(),
         y_dev.GetDeviceBuffer(),
+#ifdef SAVE_MEAN_INV_STD
+        save_mean_dev.GetDeviceBuffer(),
+        save_inv_std_dev.GetDeviceBuffer(),
+#else
         nullptr,
         nullptr,
+#endif
         y_element_op);
 
     if(!device_instance.IsSupportedArgument(argument_ptr.get()))
@@ -92,21 +108,40 @@ int run_groupnorm_example(int argc, char* argv[])
     bool pass = true;
     {
         Tensor<YDataType> host_y({N, H, W, G, C});
-        using ReferenceInstance = ck::tensor_operation::host::ReferenceGroupnorm<XDataType,
+        Tensor<SaveMeanInvStdDataType> host_save_mean(HostTensorDescriptor{N, G});
+        Tensor<SaveMeanInvStdDataType> host_save_inv_std(HostTensorDescriptor{N, G});
+        using ReferenceInstance =
+            ck::tensor_operation::host::ReferenceGroupnorm<XDataType,
                                                            GammaDataType,
                                                            BetaDataType,
                                                            YDataType,
+                                                           SaveMeanInvStdDataType,
                                                            ComputeDataType,
                                                            YElementOp>;
 
         ReferenceInstance ref;
-        auto ref_argument =
-            ref.MakeArgument(x, gamma, beta, host_y, y_element_op, {N, H, W, G, C}, 1e-6);
+        auto ref_argument = ref.MakeArgument(x,
+                                             gamma,
+                                             beta,
+                                             host_y,
+                                             host_save_mean,
+                                             host_save_inv_std,
+                                             y_element_op,
+                                             {N, H, W, G, C},
+                                             1e-6);
         auto ref_invoker  = ref.MakeInvoker();
         ref_invoker.Run(ref_argument);
 
         y_dev.FromDevice(y.mData.data());
         pass &= ck::utils::check_err(y, host_y, "Error: Incorrect results", 1e-3, 1e-3);
+#ifdef SAVE_MEAN_INV_STD
+        save_mean_dev.FromDevice(save_mean.mData.data());
+        save_inv_std_dev.FromDevice(save_inv_std.mData.data());
+        pass &= ck::utils::check_err(
+            save_mean, host_save_mean, "Error: Incorrect results (mean)", 1e-3, 1e-3);
+        pass &= ck::utils::check_err(
+            save_inv_std, host_save_inv_std, "Error: Incorrect results (inv_std)", 1e-3, 1e-3);
+#endif
     }
 
     return (pass ? 0 : 1);

example/45_elementwise_normalization/elementwise_layernorm_blockwise.cpp

@@ -167,19 +167,30 @@ int main()
                            XElementwiseOperation>(x, a, b, mn, XElementwiseOperation{});
 
         Tensor<YDataType> host_y(f_host_tensor_descriptor2d(M, N, Stride));
+        Tensor<AccDataType> host_save_mean({M});
+        Tensor<AccDataType> host_save_inv_std({M});
         using ReferenceInstance =
             ck::tensor_operation::host::ReferenceLayernorm<XDataType,
                                                            GammaDataType,
                                                            BetaDataType,
                                                            YDataType,
                                                            AccDataType,
+                                                           AccDataType,
                                                            YElementwiseOperation,
                                                            Rank,
                                                            NumReduceDim>;
 
         ReferenceInstance ref;
-        auto ref_argument =
-            ref.MakeArgument(x, gamma, beta, host_y, YElementwiseOperation{}, {M, N}, {1}, 1e-4);
+        auto ref_argument = ref.MakeArgument(x,
+                                             gamma,
+                                             beta,
+                                             host_y,
+                                             host_save_mean,
+                                             host_save_inv_std,
+                                             YElementwiseOperation{},
+                                             {M, N},
+                                             {1},
+                                             1e-4);
         auto ref_invoker  = ref.MakeInvoker();
         ref_invoker.Run(ref_argument);
 

example/52_im2col_col2im/CMakeLists.txt

@@ -3,10 +3,13 @@ set(target 0)
 foreach(gpu IN LISTS GPU_TARGETS)
     if(gpu IN_LIST gpu_list AND target EQUAL 0)
         add_custom_target(example_im2col_col2im)
+
         add_example_executable(example_image_to_column_f32 image_to_column_f32.cpp)
-   add_dependencies(example_im2col_col2im example_image_to_column_f32)
+        add_example_dependencies(example_im2col_col2im example_image_to_column_f32)
+
         add_example_executable(example_column_to_image_f32 column_to_image_f32.cpp)
-   add_dependencies(example_im2col_col2im example_column_to_image_f32)
+        add_example_dependencies(example_im2col_col2im example_column_to_image_f32)
+
         set(target 1)
     endif()
 endforeach()