From 4afdc19958aa2814deb2af995c5e56b34f8989b9 Mon Sep 17 00:00:00 2001
From: Aswin John Mathews <81309834+amathews-amd@users.noreply.github.com>
Date: Thu, 13 May 2021 17:54:06 -0500
Subject: [PATCH] ROCm optimized layernorm for MI100 (#7682)

* layernorm optimizations

* Changed HIP flag from HIP_VERSION to __HIP_PLATFORM_HCC__
---
 .../contrib_ops/cuda/layer_norm_impl.cu       |  6 ++-
 .../training_ops/cuda/nn/layer_norm_impl.cu   | 39 ++++++++++++++++++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
index ce4da0720c..740aae4d45 100644
--- a/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/layer_norm_impl.cu
@@ -364,7 +364,11 @@ void HostApplyLayerNorm(
   const int warp_size = prop.warpSize;
   ORT_ENFORCE(warp_size == GPU_WARP_SIZE);
 
-  const dim3 threads(warp_size, 4, 1);
+  dim3 threads(warp_size, 4, 1);
+#ifdef __HIP_PLATFORM_HCC__
+  // Optimization for ROCm MI100
+  threads.y = 1;
+#endif
   const dim3 blocks(1, std::min<unsigned int>(n1, maxGridY), 1);
   int nshared =
       threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;
diff --git a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
index 66a06a9634..574e57ea62 100644
--- a/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
+++ b/orttraining/orttraining/training_ops/cuda/nn/layer_norm_impl.cu
@@ -306,6 +306,7 @@ __global__ void cuComputeGradInput(
     const int numx = blockDim.x * blockDim.y;
     const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
     if (use_gamma) {
+#ifndef __HIP_PLATFORM_HCC__
       int l = 4 * thrx;
       for (; l + 3 < n2; l += 4 * numx) {
         for (int k = 0; k < 4; ++k) {
@@ -331,7 +332,24 @@ __global__ void cuComputeGradInput(
           sum_loss2 += c_loss * (c_output - U(beta[l]));
         }
       }
+#else
+      // Optimization for ROCm MI100
+      for( int l = 0; l < n2 ; l += numx) {
+        int idx = l + thrx;
+        T gamma_idx = (idx<n2)?gamma[ idx ]:T(0); 
+        const U c_loss = static_cast<U>( (idx<n2)?k_dout[ idx ]:T(0) );
+        sum_loss1 += c_loss * U( gamma_idx );
+        if (use_mean) {
+          const U c_h = static_cast<U>( (idx<n2)?k_input[ idx ]:T(0) );
+          sum_loss2 += c_loss * U(gamma_idx) * (c_h - c_mean) * c_invvar;
+        } else {
+          const U c_output = static_cast<U>( (idx<n2)?k_output[idx]:T(0) );
+          sum_loss2 += c_loss * (c_output - U( (idx<n2)?beta[idx]:T(0) ));
+        }
+      }
+#endif
     } else {
+#ifndef __HIP_PLATFORM_HCC__
       int l = 4 * thrx;
       for (; l + 3 < n2; l += 4 * numx) {
         for (int k = 0; k < 4; ++k) {
@@ -357,6 +375,21 @@ __global__ void cuComputeGradInput(
           sum_loss2 += c_loss * c_output;
         }
       }
+#else
+      // Optimization for ROCm MI100
+      for( int l = 0; l < n2 ; l += numx) {
+          int idx = l + thrx;
+          const U c_loss = static_cast<U>((idx<n2)?k_dout[idx]:T(0));
+          sum_loss1 += c_loss;
+          if (use_mean) {
+            const U c_h = static_cast<U>((idx<n2)?k_input[idx]:T(0));
+            sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
+          } else {
+            const U c_output = static_cast<U>((idx<n2)?k_output[idx]:T(0));
+            sum_loss2 += c_loss * c_output;
+          }
+      }
+#endif
     }
     // intra-warp reductions
     for (int mask = blockDim.x / 2; mask > 0; mask /= 2) {
@@ -504,7 +537,11 @@ void HostLayerNormGradient(
   // compute grad_input
   const uint64_t maxGridY = prop.maxGridSize[1];
   const dim3 blocks1(1, std::min<unsigned int>(static_cast<unsigned int>(n1), static_cast<unsigned int>(maxGridY)), 1);
-  const dim3 threads1(warp_size, 4, 1);
+  dim3 threads1(warp_size, 4, 1);
+#ifdef __HIP_PLATFORM_HCC__
+  // Optimization for ROCm MI100
+  threads1.y = 2;
+#endif
   int nshared =
       threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;
   if (mean == nullptr && !simplified) {