From 17cf39a96499e7cd805a40b97939c2cf68d38c1f Mon Sep 17 00:00:00 2001
From: Sunghoon <35605090+hanbitmyths@users.noreply.github.com>
Date: Thu, 28 Oct 2021 10:01:46 -0700
Subject: [PATCH] Clean up unnecessary codes in softmax and hardmax kernel
 (#9580)

* add p50 in test

* remove unnecessary codes from softmax

* remove unnecessary codes from hardmax

Co-authored-by: Yulong Wang <yulongw@microsoft.com>
---
 onnxruntime/core/providers/cpu/math/hardmax.cc  | 6 +-----
 onnxruntime/core/providers/cpu/math/softmax.cc  | 6 +-----
 onnxruntime/core/providers/cuda/math/softmax.cc | 6 +-----
 onnxruntime/core/providers/rocm/math/softmax.cc | 6 +-----
 4 files changed, 4 insertions(+), 20 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/math/hardmax.cc b/onnxruntime/core/providers/cpu/math/hardmax.cc
index 3c2f6461ff..c62994e22e 100644
--- a/onnxruntime/core/providers/cpu/math/hardmax.cc
+++ b/onnxruntime/core/providers/cpu/math/hardmax.cc
@@ -114,12 +114,8 @@ Status Hardmax<float>::Compute(OpKernelContext* ctx) const {
   }
 
   if (is_transpose_required) {
-    std::vector<size_t> reverse_permutation(rank);
-    for (size_t i = 0, end = permutation.size(); i < end; ++i) {
-      reverse_permutation[permutation[i]] = i;
-    }
     // Perform the transpose to get the axes back to the original ordering
-    ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(reverse_permutation, intermediate_output, *Y));
+    ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(permutation, intermediate_output, *Y));
   }
 
   return Status::OK();
diff --git a/onnxruntime/core/providers/cpu/math/softmax.cc b/onnxruntime/core/providers/cpu/math/softmax.cc
index 4c8f977b2d..20ac514aa5 100644
--- a/onnxruntime/core/providers/cpu/math/softmax.cc
+++ b/onnxruntime/core/providers/cpu/math/softmax.cc
@@ -180,12 +180,8 @@ Status Softmax<T>::ComputeImplOpset13(const Tensor& input, Tensor& output, size_
                                     log_softmax_, thread_pool));
 
   if (is_transpose_required) {
-    std::vector<size_t> reverse_permutation(rank);
-    for (size_t i = 0, end = permutation.size(); i < end; ++i) {
-      reverse_permutation[permutation[i]] = i;
-    }
     // Perform the transpose to get the axes back to the original ordering
-    ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(reverse_permutation, intermediate_output, output));
+    ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(permutation, intermediate_output, output));
   }
 
   return Status::OK();
diff --git a/onnxruntime/core/providers/cuda/math/softmax.cc b/onnxruntime/core/providers/cuda/math/softmax.cc
index a9f2f61f63..1ff5230acd 100644
--- a/onnxruntime/core/providers/cuda/math/softmax.cc
+++ b/onnxruntime/core/providers/cuda/math/softmax.cc
@@ -223,15 +223,11 @@ SPECIALIZED_SOFTMAX_HELPER_IMPL_BFloat16(true)
     return status;
 
   if (is_transpose_required) {
-    std::vector<size_t> reverse_permutation(rank);
-    for (size_t i = 0, end = permutation.size(); i < end; ++i) {
-      reverse_permutation[permutation[i]] = i;
-    }
     // Perform the transpose to get the axes back to the original ordering
     ORT_RETURN_IF_ERROR(Transpose::DoTranspose(cuda_ep_->GetDeviceProp(),
                                                Stream(),
                                                CublasHandle(),
-                                               reverse_permutation, *intermediate_output, *Y));
+                                               permutation, *intermediate_output, *Y));
   }
 
   return Status::OK();
diff --git a/onnxruntime/core/providers/rocm/math/softmax.cc b/onnxruntime/core/providers/rocm/math/softmax.cc
index 8805992e0a..6ab85d9d4a 100644
--- a/onnxruntime/core/providers/rocm/math/softmax.cc
+++ b/onnxruntime/core/providers/rocm/math/softmax.cc
@@ -199,15 +199,11 @@ Status Softmax<T>::ComputeInternal(OpKernelContext* ctx) const {
     return status;
 
   if (is_transpose_required) {
-    std::vector<size_t> reverse_permutation(rank);
-    for (size_t i = 0, end = permutation.size(); i < end; ++i) {
-      reverse_permutation[permutation[i]] = i;
-    }
     // Perform the transpose to get the axes back to the original ordering
     ORT_RETURN_IF_ERROR(Transpose::DoTranspose(rocm_ep_->GetDeviceProp(),
                                                Stream(),
                                                RocblasHandle(),
-                                               reverse_permutation, *intermediate_output, *Y));
+                                               permutation, *intermediate_output, *Y));
   }
 
   return Status::OK();