From 7b7e1d19b054605121604dbc9d6223ba1f68bee5 Mon Sep 17 00:00:00 2001
From: Vladimir Paramuzov <vladimir.paramuzov@intel.com>
Date: Wed, 7 Dec 2022 13:04:47 +0400
Subject: [PATCH] [GPU] allow softmax_bf for axis=Y && X==1 case (#14443)

---
 .../intel_gpu/src/graph/primitive_inst.cpp    |  3 ++-
 .../kernels/softmax/softmax_kernel_base.cpp   | 15 ++++++++---
 .../single_layer_tests/softmax.cpp            | 25 +++++++++++++++++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/plugins/intel_gpu/src/graph/primitive_inst.cpp b/src/plugins/intel_gpu/src/graph/primitive_inst.cpp
index 6f9e4bcc556..e814dafa91d 100644
--- a/src/plugins/intel_gpu/src/graph/primitive_inst.cpp
+++ b/src/plugins/intel_gpu/src/graph/primitive_inst.cpp
@@ -491,8 +491,9 @@ event::ptr primitive_inst::execute(const std::vector<event::ptr>& events) {
         auto out_mem = output_memory_ptr();
         auto out_alloc_type = out_mem ? out_mem->get_allocation_type() : allocation_type::unknown;
         auto out_ptr = out_mem ? out_mem->buffer_ptr() : nullptr;
+        auto impl_name = _impl->get_kernel_name();
 
-        GPU_DEBUG_COUT << id() << ": execute. Memory type: "
+        GPU_DEBUG_COUT << id() << ": execute " << impl_name << ". Memory type: "
                        << out_alloc_type << ", in_usm("
                        << in_addr.str() << "), out_usm("
                        << out_ptr << ")" << std::endl;
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_base.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_base.cpp
index 976c6f6137b..cf561fcb067 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_base.cpp
+++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_base.cpp
@@ -124,7 +124,7 @@ bool SoftmaxKernelBaseBF::Validate(const Params& p, const optional_params& o) co
         case SoftmaxDim::X:
             return input.Y().v == 1 && input.Z().v == 1 && input.Feature().v == 1;
         case SoftmaxDim::Y:
-            return input.X().v == 1 && input.Z().v == 1 && input.Feature().v == 1;
+            return input.X().v == 1 && input.Z().v == 1 && (input.Feature().v == 1 || input.GetLayout() == DataLayout::bfyx);
         case SoftmaxDim::Z:
             return input.X().v == 1 && input.Y().v == 1 && input.Feature().v == 1;
         case SoftmaxDim::FEATURE:
@@ -139,9 +139,16 @@ SoftmaxKernelBase::DispatchData SoftmaxKernelBaseBF::SetDefault(const softmax_pa
 
     DispatchData dispatchData = Parent::SetDefault(params);
 
-    auto flatten_input = input.FlattenFeatureAndSpatials();
-    dispatchData.dataSetSize = flatten_input.Feature().v;
-    dispatchData.dataSetsCount = input.Batch().v;
+    if (params.dim == SoftmaxDim::Y && input.Feature().v > 1 && input.GetLayout() == DataLayout::bfyx) {
+        // Flatten BF for such case, X is expected to be 1
+        OPENVINO_ASSERT(input.X().v == 1, "[GPU] SoftmaxKernelBaseBF: input.X() is expected to be 1 while actual value is ", input.X().v);
+        dispatchData.dataSetSize = input.Y().v;
+        dispatchData.dataSetsCount = input.Batch().v * input.Feature().v;
+    } else {
+        auto flatten_input = input.FlattenFeatureAndSpatials();
+        dispatchData.dataSetSize = flatten_input.Feature().v;
+        dispatchData.dataSetsCount = input.Batch().v;
+    }
 
     return dispatchData;
 }
diff --git a/src/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/softmax.cpp b/src/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/softmax.cpp
index 4446bf02dad..eb64e176484 100644
--- a/src/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/softmax.cpp
+++ b/src/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/softmax.cpp
@@ -44,6 +44,31 @@ INSTANTIATE_TEST_SUITE_P(
         SoftMax8LayerTest::getTestCaseName
 );
 
+const std::vector<ov::Shape> inputShapes3D = {
+    {16, 64, 64},
+};
+
+const std::vector<int64_t> axis3D = {
+    -1
+};
+
+const auto params3D = testing::Combine(
+    testing::ValuesIn(netPrecisions),
+    ::testing::Values(ov::element::undefined),
+    ::testing::Values(ov::element::undefined),
+    testing::ValuesIn(ov::test::static_shapes_to_test_representation(inputShapes3D)),
+    testing::ValuesIn(axis3D),
+    testing::Values(CommonTestUtils::DEVICE_GPU),
+    testing::Values(ov::AnyMap())
+);
+
+INSTANTIATE_TEST_SUITE_P(
+        smoke_SoftMax3D,
+        SoftMax8LayerTest,
+        params3D,
+        SoftMax8LayerTest::getTestCaseName
+);
+
 const std::vector<ov::Shape> inputShapes4D = {
     {1, 100, 1, 1},
     {1, 3, 4, 3},