1 files changed, 14 insertions, 76 deletions

diff --git a/src/cpu/operators/CpuSoftmax.cpp b/src/cpu/operators/CpuSoftmax.cpp
index ae14381ad9..fecee7d765 100644
--- a/src/cpu/operators/CpuSoftmax.cpp
+++ b/src/cpu/operators/CpuSoftmax.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, 2023 Arm Limited.
+ * Copyright (c) 2021, 2023-2024 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -41,15 +41,7 @@ namespace arm_compute
 {
 namespace cpu
 {
-CpuSoftmaxGeneric::CpuSoftmaxGeneric()
-    : _permute_input(),
-      _permute_output(),
-      _softmax_kernel(),
-      _tmp(),
-      _input_permuted(),
-      _output_permuted(),
-      _needs_permute(false),
-      _aux_mem(InternalTensorIdx::COUNT)
+CpuSoftmaxGeneric::CpuSoftmaxGeneric() : _softmax_kernel(), _tmp(), _aux_mem(InternalTensorIdx::COUNT)
 {
 }
 
@@ -63,17 +55,9 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa
     const unsigned int actual_axis =
         static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(src->num_dimensions())));
 
-    _needs_permute = actual_axis > 0;
+    _axis = actual_axis;
 
-    if (_needs_permute)
-    {
-        _permute_input.configure(src, &_input_permuted,
-                                 softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis));
-    }
-
-    // We want to deal with a 2D input. Either it is the permuted version of the original input (4D case)
-    // or it is the original input case (2D case)
-    const ITensorInfo *tmp_input = (_needs_permute ? &_input_permuted : src);
+    const ITensorInfo *tmp_input = src;
 
     TensorInfo tensor_info_tmp;
     if (is_data_type_quantized_asymmetric(src->data_type()))
@@ -88,20 +72,10 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa
 
     // Configure kernels
     auto sm = std::make_unique<kernels::CpuSoftmaxKernel>();
-    if (_needs_permute)
-    {
-        // The normalization kernel stores the result in a permuted output tensor
-        sm->configure(tmp_input, &_output_permuted, beta, is_log, &_tmp);
 
-        // Re-permute the permuted output into the requested (4D) output
-        _permute_output.configure(&_output_permuted, dst,
-                                  softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis));
-    }
-    else
-    {
-        // Softmax 2D case
-        sm->configure(tmp_input, dst, beta, is_log, &_tmp);
-    }
+    // Softmax 2D case
+    sm->configure(tmp_input, dst, beta, is_log, actual_axis, &_tmp);
+
     _softmax_kernel = std::move(sm);
 
     if (_tmp.total_size() > 0)
@@ -109,11 +83,6 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa
         _aux_mem[InternalTensorIdx::TMP] =
             MemoryInfo(offset_int_vec(InternalTensorIdx::TMP), MemoryLifetime::Temporary, _tmp.total_size());
     }
-
-    _aux_mem[InternalTensorIdx::PERMUTED_SRC] = MemoryInfo(offset_int_vec(InternalTensorIdx::PERMUTED_SRC),
-                                                           MemoryLifetime::Temporary, _input_permuted.total_size());
-    _aux_mem[InternalTensorIdx::PERMUTED_DST] = MemoryInfo(offset_int_vec(InternalTensorIdx::PERMUTED_DST),
-                                                           MemoryLifetime::Temporary, _output_permuted.total_size());
 }
 
 Status
@@ -133,25 +102,11 @@ CpuSoftmaxGeneric::validate(const ITensorInfo *src, const ITensorInfo *dst, floa
     {
         tensor_info_tmp = src->clone()->set_data_type(DataType::F32).set_is_resizable(true);
     }
-
     const unsigned int actual_axis =
         static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(src->num_dimensions())));
 
-    const bool needs_permute = actual_axis > 0;
-
-    if (needs_permute)
-    {
-        const PermutationVector permutation_vector =
-            softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis);
-        const TensorShape permuted_shape =
-            misc::shape_calculator::compute_permutation_output_shape(*src, permutation_vector);
-        TensorInfo input_permuted(src->clone()->set_tensor_shape(permuted_shape));
-        ARM_COMPUTE_RETURN_ON_ERROR(CpuPermute::validate(src, &input_permuted, permutation_vector));
-        TensorInfo output_permuted(dst->clone()->set_tensor_shape(permuted_shape));
-        ARM_COMPUTE_RETURN_ON_ERROR(CpuPermute::validate(&output_permuted, dst, permutation_vector));
-    }
-
-    ARM_COMPUTE_RETURN_ON_ERROR(kernels::CpuSoftmaxKernel::validate(src, dst, beta, is_log, &tensor_info_tmp));
+    ARM_COMPUTE_RETURN_ON_ERROR(
+        kernels::CpuSoftmaxKernel::validate(src, dst, beta, actual_axis, is_log, &tensor_info_tmp));
 
     return Status{};
 }
@@ -165,34 +120,17 @@ void CpuSoftmaxGeneric::run(ITensorPack &tensors)
 
     CpuAuxTensorHandler tmp(offset_int_vec(InternalTensorIdx::TMP), _tmp, tensors, true);
 
-    CpuAuxTensorHandler input_permuted(offset_int_vec(InternalTensorIdx::PERMUTED_SRC), _input_permuted, tensors, true);
-    CpuAuxTensorHandler output_permuted(offset_int_vec(InternalTensorIdx::PERMUTED_DST), _output_permuted, tensors,
-                                        true);
-
     ITensorPack softmax_pack;
 
-    if (_needs_permute)
-    {
-        ITensorPack permute_in_pack = {{TensorType::ACL_SRC, src}, {TensorType::ACL_DST, input_permuted.get()}};
-        _permute_input.run(permute_in_pack);
+    softmax_pack = {{TensorType::ACL_SRC_0, src}, {TensorType::ACL_DST_0, dst}, {TensorType::ACL_DST_1, tmp.get()}};
 
-        softmax_pack = {{TensorType::ACL_SRC_0, input_permuted.get()},
-                        {TensorType::ACL_DST_0, output_permuted.get()},
-                        {TensorType::ACL_DST_1, tmp.get()}};
-    }
-    else
+    if (_axis == 0)
     {
-        softmax_pack = {{TensorType::ACL_SRC_0, src}, {TensorType::ACL_DST_0, dst}, {TensorType::ACL_DST_1, tmp.get()}};
+        NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimY, _softmax_kernel->window(), softmax_pack);
     }
-
-    NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimY, _softmax_kernel->window(), softmax_pack);
-
-    if (_needs_permute)
+    else
     {
-        ITensorPack permute_out_pack;
-        permute_out_pack.add_tensor(TensorType::ACL_SRC, output_permuted.get());
-        permute_out_pack.add_tensor(TensorType::ACL_DST, dst);
-        _permute_output.run(permute_out_pack);
+        NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimX, _softmax_kernel->window(), softmax_pack);
     }
 }