From 28f470c26c2c2f47ee884230914ec582fbccf332 Mon Sep 17 00:00:00 2001
From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com>
Date: Tue, 25 Oct 2022 13:00:24 +0800
Subject: [PATCH] [ROCm] Use SkipLayerNorm original implementation in kernel
 explorer (#13382)

### Description
<!-- Describe your changes. -->

Wrap SkipLayerNormoriginal implementation as a function.
Use it as part of SkipLayerNormTunableOp.
Use it in Kernel explorer to compare the gap between TunableOp and
Original implementation.

the profile output like below:
`float16 8 512 768 <class
'_kernel_explorer.SkipLayerNorm_half_Original'> 23.48 us 804.04 GB/s

float16 8 512 768 <class '_kernel_explorer.SkipLayerNorm_half_Tunable'>
20.41 us 925.00 GB/s
...`

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Co-authored-by: peixuanzuo <peixuanzuo@linmif39a000004.zvflicr54joexhdgnhvmxrxygg.phxx.internal.cloudapp.net>
---
 .../rocm/bert/skip_layer_norm_impl.cu         | 66 ++-------------
 .../rocm/bert/skip_layer_norm_tunable_op.h    | 82 ++++++++++++++++++-
 .../kernels/skip_layer_norm.cc                | 49 ++++++++---
 3 files changed, 121 insertions(+), 76 deletions(-)

diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu
index a4c94236d1..f324caf919 100644
--- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu
+++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_impl.cu
@@ -46,71 +46,15 @@ Status LaunchSkipLayerNormKernel(
   // this must be true because element_count is the total size of the tensor
   assert(element_count % ld == 0);
 
+  SkipLayerNormParams<T> op_params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count);
+  static SkipLayerNormTunableOp<T> op;
+
+  // If disable tuning, the default implementation is SkipLayerNormStaticSelection.
   if (tuning) {
-    static SkipLayerNormTunableOp<T> op;
     op.EnableTuning();
-
-    SkipLayerNormParams<T> op_params(stream, output, input, skip, gamma, beta, bias, epsilon, ld, element_count);
-    return op(&op_params);
   }
 
-  bool hasBias = (bias == nullptr) ? false : true;
-  if (0 == (ld % 4)) {
-    const int grid_size = element_count / ld;
-    if (ld <= 32) {
-      constexpr int block_size = 32;
-      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 64) {
-      constexpr int block_size = 64 / 2;
-      SkipLayerNormKernelSmall<T, block_size, 2><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 128) {
-      constexpr int block_size = 128 / 4;
-      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 384) {
-      constexpr int block_size = 384 / 4;
-      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 768) {
-      constexpr int block_size = 768 / 4;
-      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 1024) {
-      constexpr int block_size = 1024 / 4;
-      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else {
-      constexpr int block_size = 256;
-      SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output);
-    }
-  } else {
-    const int grid_size = element_count / ld;
-    if (ld <= 32) {
-      constexpr int block_size = 32;
-      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 64) {
-      constexpr int block_size = 64;
-      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld <= 128) {
-      constexpr int block_size = 128;
-      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else if (ld == 384) {
-      constexpr int block_size = 384;
-      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output, hasBias);
-    } else {
-      constexpr int block_size = 256;
-      SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, stream>>>(
-          ld, input, skip, beta, gamma, bias, maybe2half<T>(epsilon), output);
-    }
-  }
-  return HIP_CALL(hipPeekAtLastError());
+  return op(&op_params);
 }
 
 template Status LaunchSkipLayerNormKernel<float>(hipStream_t stream, float* output, const float* input,
diff --git a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h
index 115ddaf1b4..299ce6f804 100644
--- a/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h
+++ b/onnxruntime/contrib_ops/rocm/bert/skip_layer_norm_tunable_op.h
@@ -60,8 +60,8 @@ template <typename T, int ThreadsPerBlock, int VecSize>
 Status SkipLayerNormRegularOp(const SkipLayerNormParams<T>* params) {
   TUNABLE_OP_RETURN_UNSUPPOTED_ARGUMENT_IF(
       !((params->ld > 0 && params->ld % VecSize == 0 &&
-       (params->ld >= ThreadsPerBlock * VecSize ||
-       (params->ld < 64 && params->ld > (ThreadsPerBlock - GPU_WARP_SIZE) * VecSize)))));
+         (params->ld >= ThreadsPerBlock * VecSize ||
+          (params->ld < GPU_WARP_SIZE && params->ld > (ThreadsPerBlock - GPU_WARP_SIZE) * VecSize)))));
   SkipLayerNormKernelVec<T, ThreadsPerBlock, VecSize><<<dim3(CeilDiv(params->element_count, params->ld)),
                                                         dim3(ThreadsPerBlock),
                                                         0, params->stream>>>(
@@ -71,6 +71,79 @@ Status SkipLayerNormRegularOp(const SkipLayerNormParams<T>* params) {
   return HIP_CALL(hipGetLastError());
 }
 
+template <typename T>
+Status SkipLayerNormStaticSelection(const SkipLayerNormParams<T>* params) {
+  bool hasBias = (params->bias == nullptr) ? false : true;
+  if (0 == (params->ld % 4)) {
+    const int grid_size = params->element_count / params->ld;
+    if (params->ld <= 32) {
+      constexpr int block_size = 32;
+      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 64) {
+      constexpr int block_size = 64 / 2;
+      SkipLayerNormKernelSmall<T, block_size, 2><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 128) {
+      constexpr int block_size = 128 / 4;
+      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 384) {
+      constexpr int block_size = 384 / 4;
+      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 768) {
+      constexpr int block_size = 768 / 4;
+      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 1024) {
+      constexpr int block_size = 1024 / 4;
+      SkipLayerNormKernelSmall<T, block_size, 4><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else {
+      constexpr int block_size = 256;
+      SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output);
+    }
+  } else {
+    const int grid_size = params->element_count / params->ld;
+    if (params->ld <= 32) {
+      constexpr int block_size = 32;
+      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 64) {
+      constexpr int block_size = 64;
+      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld <= 128) {
+      constexpr int block_size = 128;
+      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else if (params->ld == 384) {
+      constexpr int block_size = 384;
+      SkipLayerNormKernelSmall<T, block_size, 1><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output, hasBias);
+    } else {
+      constexpr int block_size = 256;
+      SkipLayerNormKernel<T, block_size><<<grid_size, block_size, 0, params->stream>>>(
+          params->ld, params->input, params->skip, params->beta, params->gamma, params->bias,
+          maybe2half<T>(params->epsilon), params->output);
+    }
+  }
+  return HIP_CALL(hipPeekAtLastError());
+}
+
 #define ADD_OP_FOR_ALL_VEC_SIZE(name, threads_per_block)  \
   this->ops_.emplace_back(name<T, threads_per_block, 1>); \
   this->ops_.emplace_back(name<T, threads_per_block, 2>); \
@@ -90,11 +163,12 @@ template <typename T>
 class SkipLayerNormTunableOp : public onnxruntime::rocm::tunable::TunableOp<SkipLayerNormParams<T>> {
  public:
   SkipLayerNormTunableOp() {
+    this->ops_.emplace_back(SkipLayerNormStaticSelection<T>);
     ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormSmallOp)
     ADD_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormRegularOp)
 
-    // NOTE: the 30-th kernel is SkipLayerNormRegularOp ThreadsPerBlock=64 VecSize=1
-    this->SetDefaultId(30);
+    // NOTE: the 1st kernel is SkipLayerNorm Original implementation.
+    this->SetDefaultId(0);
   }
 };
 
diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm.cc b/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm.cc
index 8a12669e85..2f20e09932 100644
--- a/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm.cc
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/skip_layer_norm.cc
@@ -62,6 +62,30 @@ class SkipLayerNormRegular : public IKernelExplorer {
   ParamsT params_{};
 };
 
+template <typename T>
+class SkipLayerNormStaticSelection : public IKernelExplorer {
+ public:
+  SkipLayerNormStaticSelection(DeviceArray& output, DeviceArray& input, DeviceArray& skip,
+                               DeviceArray& gamma, DeviceArray& beta, DeviceArray& bias,
+                               float epsilon, int hidden_size, int element_count)
+      : params_(this->Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
+                static_cast<T*>(skip.ptr()), static_cast<T*>(gamma.ptr()), static_cast<T*>(beta.ptr()),
+                static_cast<T*>(bias.ptr()), epsilon, hidden_size, element_count) {}
+
+  void Run() override {
+    ORT_THROW_IF_ERROR((contrib::rocm::SkipLayerNormStaticSelection<T>(&params_)));
+  }
+
+  bool IsSupported() {
+    Status status = contrib::rocm::SkipLayerNormStaticSelection<T>(&params_);
+    return status.IsOK();
+  }
+
+ private:
+  using ParamsT = contrib::rocm::SkipLayerNormParams<T>;
+  ParamsT params_{};
+};
+
 template <typename T>
 class SkipLayerNormTunable : public IKernelExplorer {
  public:
@@ -113,15 +137,15 @@ class SkipLayerNormTunable : public IKernelExplorer {
   REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 320)                        \
   REGISTER_OP_FOR_ALL_VEC_SIZE(name, type, 384)
 
-#define REGISTER_TUNABLE_OP(type)                                              \
-  py::class_<SkipLayerNormTunable<type>>(m, "SkipLayerNorm_" #type "_Tunable") \
-      .def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&,    \
-                    DeviceArray&, DeviceArray&,                                \
-                    float, int, int>())                                        \
-      .def("SetRepeats", &SkipLayerNormTunable<type>::SetRepeats)              \
-      .def("Profile", &SkipLayerNormTunable<type>::Profile)                    \
-      .def("Run", &SkipLayerNormTunable<type>::Run)                            \
-      .def("IsSupported", &SkipLayerNormTunable<type>::IsSupported);
+#define REGISTER_OP_TYPED(name, type)                                       \
+  py::class_<name<type>>(m, #name "_" #type)                                \
+      .def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, DeviceArray&, \
+                    DeviceArray&, DeviceArray&,                             \
+                    float, int, int>())                                     \
+      .def("SetRepeats", &name<type>::SetRepeats)                           \
+      .def("Profile", &name<type>::Profile)                                 \
+      .def("Run", &name<type>::Run)                                         \
+      .def("IsSupported", &name<type>::IsSupported);
 
 void InitSkipLayerNorm(py::module m) {
   REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormSmall, half);
@@ -129,8 +153,11 @@ void InitSkipLayerNorm(py::module m) {
   REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormRegular, half);
   REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK_ALL_VEC_SIZE(SkipLayerNormRegular, float);
 
-  REGISTER_TUNABLE_OP(half);
-  REGISTER_TUNABLE_OP(float);
+  REGISTER_OP_TYPED(SkipLayerNormTunable, half);
+  REGISTER_OP_TYPED(SkipLayerNormTunable, float);
+
+  REGISTER_OP_TYPED(SkipLayerNormStaticSelection, half);
+  REGISTER_OP_TYPED(SkipLayerNormStaticSelection, float);
 }
 
 }  // namespace onnxruntime