From 05a24b17769cb634f16343d551e25f542f0855b1 Mon Sep 17 00:00:00 2001
From: Steve Yoo <steve.yoo@intel.com>
Date: Fri, 1 Sep 2023 08:39:36 +0900
Subject: [PATCH] [GPU] Try to use softmax_ref when types are mismatched
 (#19209)

* Remove support key for UINT8 and INT8
---
 .../kernels/softmax/softmax_kernel_bf.cpp     |  2 --
 .../softmax_kernel_items_class_optimized.cpp  | 32 ++++++++++++++++++-
 .../unit/fusions/softmax_fusion_test.cpp      |  4 +++
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_bf.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_bf.cpp
index 30171493dae..d5304e4c784 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_bf.cpp
+++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_bf.cpp
@@ -14,8 +14,6 @@ ParamsKey SoftmaxKernel_bf::GetSupportedKey() const {
     k.EnableInputDataType(Datatype::F32);
     k.EnableOutputDataType(Datatype::F16);
     k.EnableOutputDataType(Datatype::F32);
-    k.EnableOutputDataType(Datatype::UINT8);
-    k.EnableOutputDataType(Datatype::INT8);
     k.EnableInputLayout(DataLayout::bfyx);
     k.EnableInputLayout(DataLayout::bf);
     k.EnableOutputLayout(DataLayout::bfyx);
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_items_class_optimized.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_items_class_optimized.cpp
index a8b31a2c72b..71d8c809e43 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_items_class_optimized.cpp
+++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/softmax/softmax_kernel_items_class_optimized.cpp
@@ -35,7 +35,37 @@ inline static size_t GetItemClassCount(const DataTensor& input, SoftmaxDim dim)
     return item_class_count;
 }
 
-ParamsKey SoftmaxKerneItemsClassOptimized::GetSupportedKey() const { return GetDefaultSupportedKey(); }
+ParamsKey SoftmaxKerneItemsClassOptimized::GetSupportedKey() const {
+    ParamsKey k;
+    k.EnableInputDataType(Datatype::F16);
+    k.EnableInputDataType(Datatype::F32);
+    k.EnableOutputDataType(Datatype::F16);
+    k.EnableOutputDataType(Datatype::F32);
+    k.EnableInputLayout(DataLayout::byxf);
+    k.EnableInputLayout(DataLayout::bfyx);
+    k.EnableInputLayout(DataLayout::yxfb);
+    k.EnableInputLayout(DataLayout::bf);
+    k.EnableInputLayout(DataLayout::fb);
+    k.EnableInputLayout(DataLayout::bfzyx);
+    k.EnableInputLayout(DataLayout::f);
+    k.EnableOutputLayout(DataLayout::f);
+    k.EnableOutputLayout(DataLayout::bfyx);
+    k.EnableOutputLayout(DataLayout::byxf);
+    k.EnableOutputLayout(DataLayout::yxfb);
+    k.EnableOutputLayout(DataLayout::bf);
+    k.EnableOutputLayout(DataLayout::fb);
+    k.EnableOutputLayout(DataLayout::bfzyx);
+    k.EnableSoftmaxDim(SoftmaxDim::X);
+    k.EnableSoftmaxDim(SoftmaxDim::Y);
+    k.EnableSoftmaxDim(SoftmaxDim::Z);
+    k.EnableSoftmaxDim(SoftmaxDim::FEATURE);
+    k.EnableSoftmaxDim(SoftmaxDim::BATCH);
+    k.EnableDifferentTypes();
+    k.EnableTensorOffset();
+    k.EnableTensorPitches();
+    k.EnableBatching();
+    return k;
+}
 
 DeviceFeaturesKey SoftmaxKerneItemsClassOptimized::get_required_device_features_key(const Params& params, const optional_params& options) const {
     DeviceFeaturesKey k;
diff --git a/src/plugins/intel_gpu/tests/unit/fusions/softmax_fusion_test.cpp b/src/plugins/intel_gpu/tests/unit/fusions/softmax_fusion_test.cpp
index 342841af66a..8eee305eed3 100644
--- a/src/plugins/intel_gpu/tests/unit/fusions/softmax_fusion_test.cpp
+++ b/src/plugins/intel_gpu/tests/unit/fusions/softmax_fusion_test.cpp
@@ -67,9 +67,11 @@ public:
 /* ----------------------------------------------------------------------------------------------------- */
 #define CASE_SOFTMAX_FP32_1 {1, 15, 4, 5}, data_types::f32, format::bfyx, 1, data_types::f32, format::bfyx
 #define CASE_SOFTMAX_FP32_2 {1, 15, 4, 5}, data_types::f32, format::bfyx, 3, data_types::f32, format::bfyx
+#define CASE_SOFTMAX_FP32_3 {1, 15, 1, 1}, data_types::f32, format::bfyx, 1, data_types::f32, format::bfyx
 
 #define CASE_SOFTMAX_FP16_1 {1, 15, 4, 5}, data_types::f16, format::bfyx, 1, data_types::f16, format::bfyx
 #define CASE_SOFTMAX_FP16_2 {1, 15, 4, 5}, data_types::f16, format::bfyx, 3, data_types::f16, format::bfyx
+#define CASE_SOFTMAX_FP16_3 {1, 15, 1, 1}, data_types::f16, format::bfyx, 1, data_types::f16, format::bfyx
 
 class softmax_quantize : public SoftmaxPrimitiveFusingTest {};
 TEST_P(softmax_quantize, basic) {
@@ -93,9 +95,11 @@ INSTANTIATE_TEST_SUITE_P(fusings_gpu, softmax_quantize,
     ::testing::ValuesIn(std::vector<softmax_test_params>{
                         softmax_test_params{ CASE_SOFTMAX_FP32_1, 2, 3 },
                         softmax_test_params{ CASE_SOFTMAX_FP32_2, 3, 3 },
+                        softmax_test_params{ CASE_SOFTMAX_FP32_3, 2, 3 },
 
                         softmax_test_params{ CASE_SOFTMAX_FP16_1, 2, 3 },
                         softmax_test_params{ CASE_SOFTMAX_FP16_2, 3, 3 },
+                        softmax_test_params{ CASE_SOFTMAX_FP16_3, 2, 3 },
 }));
 
 class softmax_quantize_fusing_through : public SoftmaxPrimitiveFusingTest {};