From 7cdd7248d1fd1d4f523310d34eb2090ef0236f7e Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Thu, 25 Apr 2024 16:13:43 +0000
Subject: [PATCH 1/6] use output buffer as intermediate buffer instead of
 my_chunk to avoid buffer overfill when dynamic shape

---
 .../cl_kernels/softmax_gpu_bf.cl              | 172 ++++++++++++++----
 .../single_layer_tests/dynamic/softmax.cpp    |  27 ++-
 2 files changed, 162 insertions(+), 37 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
index b70fe7a5173..aeb6d8834c7 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
@@ -77,6 +77,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     const uint actual_leftovers = leftovers - aligned_offset;
     const uint leftover_idx = data_set_offset + aligned_offset + workers_per_data_set * items_num + in_data_set_idx;
 
+#if IS_DYNAMIC
+    // use output buffer as intermediate variable instead of my_chunk when (item_num+2) is bigger than STACK_SIZE
+    // this happens when data_set_size > 16384 (engineInfo.maxWorkGroupSize=512) or data_set_size > 32768 (engineInfo.maxWorkGroupSize=1024)
+    const bool use_output_buffer = (items_num + 2) > STACK_SIZE;
+#else
+    const bool use_output_buffer = false;
+#endif
+
     INPUT0_TYPE my_chunk[STACK_SIZE];
     INPUT0_TYPE my_sum = UNIT_VAL_ZERO;
 
@@ -93,14 +101,20 @@ KERNEL (softmax_gpu_continuous_bfyx)(
             BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
             {
-                my_chunk[input_idx+j] = vec_tmp[j];
+                if (use_output_buffer)
+                    output[aligned_data_offset + get_sub_group_local_id() + (input_idx + j) * get_sub_group_size()] = vec_tmp[j];
+                else
+                    my_chunk[input_idx+j] = vec_tmp[j];
             }
         }
 
         for (; input_idx < items_num; input_idx++)
         {
             BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
-            my_chunk[input_idx] = vec_tmp;
+            if (use_output_buffer)
+                output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()] = vec_tmp;
+            else
+                my_chunk[input_idx] = vec_tmp;
         }
     }
 #else
@@ -124,19 +138,29 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     for (; input_idx < items_num; input_idx++)
     {
-        my_chunk[input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+        if (use_output_buffer) {
+            output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()]
+            = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+        } else
+            my_chunk[input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
     }
 
     if (in_data_set_idx < aligned_offset)
     {
         INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
-        my_chunk[input_idx++] = tmp;
+        if (use_output_buffer)
+            output[data_set_offset + in_data_set_idx] = tmp;
+        else
+            my_chunk[input_idx++] = tmp;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
         INPUT0_TYPE tmp = input[leftover_idx];
-        my_chunk[input_idx++] = tmp;
+        if (use_output_buffer)
+            output[leftover_idx] = tmp;
+        else
+            my_chunk[input_idx++] = tmp;
     }
 
     INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
@@ -146,6 +170,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             my_maximum = max(my_maximum, my_chunk[j]);
         }
+        if (use_output_buffer) {
+            if (in_data_set_idx < aligned_offset) {
+                my_maximum = max(my_maximum, output[data_set_offset + in_data_set_idx]);
+            }
+            if (in_data_set_idx < actual_leftovers) {
+                my_maximum = max(my_maximum, output[leftover_idx]);
+            }
+        }
     }
 
     my_maximum = sub_group_reduce_max(my_maximum);
@@ -170,11 +202,32 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     barrier(CLK_LOCAL_MEM_FENCE);
 
     const uint num_iters = input_idx;
-    for (uint j=0; j<num_iters; ++j)
-    {
-        INPUT0_TYPE tmp = native_exp(my_chunk[j] - my_maximum);
-        my_sum += tmp;
-        my_chunk[j] = tmp;
+
+    if (use_output_buffer) {
+        for (uint j=0; j<num_iters; ++j) {
+            INPUT0_TYPE tmp = native_exp(output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] - my_maximum);
+            my_sum += tmp;
+            output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] = tmp;
+        }
+
+        if (in_data_set_idx < aligned_offset) {
+            INPUT0_TYPE tmp = native_exp(output[data_set_offset + in_data_set_idx] - my_maximum);
+            my_sum += tmp;
+            output[data_set_offset + in_data_set_idx] = tmp;
+        }
+
+        if (in_data_set_idx < actual_leftovers) {
+            INPUT0_TYPE tmp = native_exp(output[leftover_idx] - my_maximum);
+            my_sum += tmp;
+            output[leftover_idx] = tmp;
+        }
+    } else {
+        for (uint j=0; j<num_iters; ++j)
+        {
+            INPUT0_TYPE tmp = native_exp(my_chunk[j] - my_maximum);
+            my_sum += tmp;
+            my_chunk[j] = tmp;
+        }
     }
 
     my_sum = sub_group_reduce_add(my_sum);
@@ -204,20 +257,32 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
         {
             BLOCK_TYPE_OPT vec_tmp;
-            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
-            {
-                ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
-                FUSED_OPS_MAIN;
-                vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
+            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
+                if (use_output_buffer) {
+                    ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum;
+                    FUSED_OPS_MAIN;
+                    vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
+                } else {
+                    ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
+                    FUSED_OPS_MAIN;
+                    vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
+                }
             }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
 
         for (; output_idx<items_num; output_idx++)
         {
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
-            FUSED_OPS_MAIN;
-            vec_tmp = FUSED_OPS_RESULT_MAIN;
+            BLOCK_TYPE vec_tmp;
+            if (use_output_buffer) {
+                ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
+                FUSED_OPS_MAIN;
+                vec_tmp = FUSED_OPS_RESULT_MAIN;
+            } else {
+                ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+                FUSED_OPS_MAIN;
+                vec_tmp = FUSED_OPS_RESULT_MAIN;
+            }
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -245,23 +310,41 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
-        FUSED_OPS_MAIN;
-        output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
+        if (use_output_buffer) {
+            ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
+            FUSED_OPS_MAIN;
+            output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
+        } else {
+            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+            FUSED_OPS_MAIN;
+            output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
+        }
     }
 
     if (in_data_set_idx < aligned_offset)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
-        FUSED_OPS_LEFTOVERS;
-        output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        if (use_output_buffer) {
+            ACTIVATION_TYPE dequantized = output[data_set_offset + in_data_set_idx] / my_sum;
+            FUSED_OPS_LEFTOVERS;
+            output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        } else {
+            ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+            FUSED_OPS_LEFTOVERS;
+            output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        }
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
-        FUSED_OPS_LEFTOVERS;
-        output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        if (use_output_buffer) {
+            ACTIVATION_TYPE dequantized = output[leftover_idx] / my_sum;
+            FUSED_OPS_LEFTOVERS;
+            output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        } else {
+            ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+            FUSED_OPS_LEFTOVERS;
+            output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+        }
     }
 #else
 #if IS_DYNAMIC
@@ -271,15 +354,22 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
         {
             BLOCK_TYPE_OPT vec_tmp;
-            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
-                vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
+            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
+                if (use_output_buffer)
+                    vec_tmp[j] = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+                else
+                    vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
+            }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
 
         for (; output_idx < items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+            if (use_output_buffer)
+                vec_tmp = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+            else
+                vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -301,14 +391,26 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+        if (use_output_buffer) {
+            output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()]
+            = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+        } else
+            output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
     }
 
-    if (in_data_set_idx < aligned_offset)
-        output[data_set_offset + in_data_set_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+    if (in_data_set_idx < aligned_offset) {
+        if (use_output_buffer)
+            output[data_set_offset + in_data_set_idx] = ACTIVATION(output[data_set_offset + in_data_set_idx] / my_sum, ACTIVATION_PARAMS);
+        else
+            output[data_set_offset + in_data_set_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+    }
 
-    if (in_data_set_idx < actual_leftovers)
-        output[leftover_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+    if (in_data_set_idx < actual_leftovers) {
+        if (use_output_buffer)
+            output[leftover_idx] = ACTIVATION(output[leftover_idx] / my_sum, ACTIVATION_PARAMS);
+        else
+            output[leftover_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+    }
 #endif
 }
 #ifdef CALC_POWER
diff --git a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
index 01b6aad0afe..be08309ca37 100644
--- a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
+++ b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
@@ -56,7 +56,7 @@ protected:
         for (auto&& shape : inputDynamicShapes)
             params.push_back(std::make_shared<ov::op::v0::Parameter>(model_type, shape));
 
-        const auto softMax = std::make_shared<ov::op::v1::Softmax>(params.at(0), axis);
+        const auto softMax = std::make_shared<ov::op::v8::Softmax>(params.at(0), axis);
         auto makeFunction = [](ov::ParameterVector &params, const std::shared_ptr<ov::Node> &lastNode) {
             ov::ResultVector results;
 
@@ -74,7 +74,7 @@ TEST_P(SoftMaxLayerGPUTest, Inference) {
 }
 
 const std::vector<ov::element::Type> netPrecisions = {
-       ov::element::f32, ov::element::f16
+    ov::element::f32, ov::element::f16
 };
 
 const std::vector<int64_t> axis2D = {0, 1};
@@ -107,6 +107,15 @@ const std::vector<ov::test::InputShape> inputShapes4D = {
     },
 };
 
+const std::vector<int64_t> axis3D = {-1};
+
+const std::vector<ov::test::InputShape> inputShapes3D = {
+    {
+        {-1, -1, -1},
+        {{16, 64, 64}, {4, 6, 27200}}
+    },
+};
+
 const std::vector<int64_t> axis5D = {0, 1, 2, 4};
 
 const std::vector<ov::test::InputShape> inputShapes5D = {
@@ -134,6 +143,20 @@ INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest4D,
                                             testing::ValuesIn(axis4D)),
                          SoftMaxLayerGPUTest::getTestCaseName);
 
+INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest3D,
+        SoftMaxLayerGPUTest,
+                         ::testing::Combine(testing::ValuesIn(netPrecisions),
+                                            testing::ValuesIn(inputShapes3D),
+                                            testing::ValuesIn(axis3D)),
+                         SoftMaxLayerGPUTest::getTestCaseName);
+
+INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest3D_small,
+        SoftMaxLayerGPUTest,
+                         ::testing::Combine(testing::ValuesIn(netPrecisions),
+                                            testing::ValuesIn(inputShapes3Dsmall),
+                                            testing::ValuesIn(axis3D)),
+                         SoftMaxLayerGPUTest::getTestCaseName);
+
 INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest5D,
         SoftMaxLayerGPUTest,
                          ::testing::Combine(testing::ValuesIn(netPrecisions),

From dd4fd236e76abf8d11efc4d8bb1e09865c8f0900 Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Thu, 25 Apr 2024 16:18:58 +0000
Subject: [PATCH 2/6] fix single layer test bug

---
 .../single_layer_tests/dynamic/softmax.cpp    | 31 +++++++------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
index be08309ca37..533b1d0be85 100644
--- a/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
+++ b/src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/softmax.cpp
@@ -90,6 +90,15 @@ const std::vector<ov::test::InputShape> inputShapes2D = {
     }
 };
 
+const std::vector<int64_t> axis3D = {-1};
+
+const std::vector<ov::test::InputShape> inputShapes3D = {
+    {
+        {-1, -1, -1},
+        {{16, 64, 64}, {4, 6, 27200}}
+    },
+};
+
 const std::vector<int64_t> axis4D = {1, 2, 3};
 
 const std::vector<ov::test::InputShape> inputShapes4D = {
@@ -107,15 +116,6 @@ const std::vector<ov::test::InputShape> inputShapes4D = {
     },
 };
 
-const std::vector<int64_t> axis3D = {-1};
-
-const std::vector<ov::test::InputShape> inputShapes3D = {
-    {
-        {-1, -1, -1},
-        {{16, 64, 64}, {4, 6, 27200}}
-    },
-};
-
 const std::vector<int64_t> axis5D = {0, 1, 2, 4};
 
 const std::vector<ov::test::InputShape> inputShapes5D = {
@@ -136,13 +136,6 @@ INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest2D,
                                             testing::ValuesIn(axis2D)),
                          SoftMaxLayerGPUTest::getTestCaseName);
 
-INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest4D,
-        SoftMaxLayerGPUTest,
-                         ::testing::Combine(testing::ValuesIn(netPrecisions),
-                                            testing::ValuesIn(inputShapes4D),
-                                            testing::ValuesIn(axis4D)),
-                         SoftMaxLayerGPUTest::getTestCaseName);
-
 INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest3D,
         SoftMaxLayerGPUTest,
                          ::testing::Combine(testing::ValuesIn(netPrecisions),
@@ -150,11 +143,11 @@ INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest3D,
                                             testing::ValuesIn(axis3D)),
                          SoftMaxLayerGPUTest::getTestCaseName);
 
-INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest3D_small,
+INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest4D,
         SoftMaxLayerGPUTest,
                          ::testing::Combine(testing::ValuesIn(netPrecisions),
-                                            testing::ValuesIn(inputShapes3Dsmall),
-                                            testing::ValuesIn(axis3D)),
+                                            testing::ValuesIn(inputShapes4D),
+                                            testing::ValuesIn(axis4D)),
                          SoftMaxLayerGPUTest::getTestCaseName);
 
 INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest5D,

From 41dc86d3e2a29401e33889aacd247b457403b661 Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Thu, 25 Apr 2024 17:38:54 +0000
Subject: [PATCH 3/6] fix max calcuation bug

---
 .../kernel_selector/cl_kernels/softmax_gpu_bf.cl   | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
index aeb6d8834c7..8748b16b764 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
@@ -166,17 +166,23 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
     {
         const uint num_iters = input_idx;
-        for (uint j=0; j<num_iters; ++j)
-        {
-            my_maximum = max(my_maximum, my_chunk[j]);
-        }
+
         if (use_output_buffer) {
+            for (uint j=0; j<num_iters; ++j)
+            {
+                my_maximum = max(my_maximum, output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()]);
+            }
             if (in_data_set_idx < aligned_offset) {
                 my_maximum = max(my_maximum, output[data_set_offset + in_data_set_idx]);
             }
             if (in_data_set_idx < actual_leftovers) {
                 my_maximum = max(my_maximum, output[leftover_idx]);
             }
+        } else {
+            for (uint j=0; j<num_iters; ++j)
+            {
+                my_maximum = max(my_maximum, my_chunk[j]);
+            }
         }
     }
 

From 816ea7eb3ad4f9f11f0f1f9521356f4fda0a1d2f Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Mon, 29 Apr 2024 14:05:08 +0000
Subject: [PATCH 4/6] reduce if condition

---
 .../cl_kernels/softmax_gpu_bf.cl              | 364 ++++++++++++------
 1 file changed, 236 insertions(+), 128 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
index 8748b16b764..d1efd90e9ad 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
@@ -92,6 +92,12 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     // Read inputs and Get maximum value from data set
     uint input_idx=0;
+
+
+#if IS_DYNAMIC
+    // Case for my_chunk[] when (items_num + 2) <= STACK_SIZE
+    if (!use_output_buffer) {
+#endif
 #if IS_DYNAMIC
     if (workers_per_data_set > SUB_GROUP_SIZE)
     {
@@ -101,20 +107,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
             BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
             {
-                if (use_output_buffer)
-                    output[aligned_data_offset + get_sub_group_local_id() + (input_idx + j) * get_sub_group_size()] = vec_tmp[j];
-                else
-                    my_chunk[input_idx+j] = vec_tmp[j];
+                my_chunk[input_idx+j] = vec_tmp[j];
             }
         }
 
         for (; input_idx < items_num; input_idx++)
         {
             BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
-            if (use_output_buffer)
-                output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()] = vec_tmp;
-            else
-                my_chunk[input_idx] = vec_tmp;
+            my_chunk[input_idx] = vec_tmp;
         }
     }
 #else
@@ -138,51 +138,28 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     for (; input_idx < items_num; input_idx++)
     {
-        if (use_output_buffer) {
-            output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()]
-            = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
-        } else
-            my_chunk[input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+        my_chunk[input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
     }
 
     if (in_data_set_idx < aligned_offset)
     {
         INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
-        if (use_output_buffer)
-            output[data_set_offset + in_data_set_idx] = tmp;
-        else
-            my_chunk[input_idx++] = tmp;
+        my_chunk[input_idx++] = tmp;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
         INPUT0_TYPE tmp = input[leftover_idx];
-        if (use_output_buffer)
-            output[leftover_idx] = tmp;
-        else
-            my_chunk[input_idx++] = tmp;
+        my_chunk[input_idx++] = tmp;
     }
 
     INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
     {
         const uint num_iters = input_idx;
 
-        if (use_output_buffer) {
-            for (uint j=0; j<num_iters; ++j)
-            {
-                my_maximum = max(my_maximum, output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()]);
-            }
-            if (in_data_set_idx < aligned_offset) {
-                my_maximum = max(my_maximum, output[data_set_offset + in_data_set_idx]);
-            }
-            if (in_data_set_idx < actual_leftovers) {
-                my_maximum = max(my_maximum, output[leftover_idx]);
-            }
-        } else {
-            for (uint j=0; j<num_iters; ++j)
-            {
-                my_maximum = max(my_maximum, my_chunk[j]);
-            }
+        for (uint j=0; j<num_iters; ++j)
+        {
+            my_maximum = max(my_maximum, my_chunk[j]);
         }
     }
 
@@ -209,31 +186,11 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     const uint num_iters = input_idx;
 
-    if (use_output_buffer) {
-        for (uint j=0; j<num_iters; ++j) {
-            INPUT0_TYPE tmp = native_exp(output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] - my_maximum);
-            my_sum += tmp;
-            output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] = tmp;
-        }
-
-        if (in_data_set_idx < aligned_offset) {
-            INPUT0_TYPE tmp = native_exp(output[data_set_offset + in_data_set_idx] - my_maximum);
-            my_sum += tmp;
-            output[data_set_offset + in_data_set_idx] = tmp;
-        }
-
-        if (in_data_set_idx < actual_leftovers) {
-            INPUT0_TYPE tmp = native_exp(output[leftover_idx] - my_maximum);
-            my_sum += tmp;
-            output[leftover_idx] = tmp;
-        }
-    } else {
-        for (uint j=0; j<num_iters; ++j)
-        {
-            INPUT0_TYPE tmp = native_exp(my_chunk[j] - my_maximum);
-            my_sum += tmp;
-            my_chunk[j] = tmp;
-        }
+    for (uint j=0; j<num_iters; ++j)
+    {
+        INPUT0_TYPE tmp = native_exp(my_chunk[j] - my_maximum);
+        my_sum += tmp;
+        my_chunk[j] = tmp;
     }
 
     my_sum = sub_group_reduce_add(my_sum);
@@ -264,15 +221,9 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
-                if (use_output_buffer) {
-                    ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum;
-                    FUSED_OPS_MAIN;
-                    vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
-                } else {
-                    ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
-                    FUSED_OPS_MAIN;
-                    vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
-                }
+                ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
+                FUSED_OPS_MAIN;
+                vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
             }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
@@ -280,15 +231,9 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx<items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            if (use_output_buffer) {
-                ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
-                FUSED_OPS_MAIN;
-                vec_tmp = FUSED_OPS_RESULT_MAIN;
-            } else {
-                ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
-                FUSED_OPS_MAIN;
-                vec_tmp = FUSED_OPS_RESULT_MAIN;
-            }
+            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+            FUSED_OPS_MAIN;
+            vec_tmp = FUSED_OPS_RESULT_MAIN;
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -316,41 +261,23 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        if (use_output_buffer) {
-            ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
-            FUSED_OPS_MAIN;
-            output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
-        } else {
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
-            FUSED_OPS_MAIN;
-            output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
-        }
+        ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+        FUSED_OPS_MAIN;
+        output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
     }
 
     if (in_data_set_idx < aligned_offset)
     {
-        if (use_output_buffer) {
-            ACTIVATION_TYPE dequantized = output[data_set_offset + in_data_set_idx] / my_sum;
-            FUSED_OPS_LEFTOVERS;
-            output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-        } else {
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
-            FUSED_OPS_LEFTOVERS;
-            output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-        }
+        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+        FUSED_OPS_LEFTOVERS;
+        output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
-        if (use_output_buffer) {
-            ACTIVATION_TYPE dequantized = output[leftover_idx] / my_sum;
-            FUSED_OPS_LEFTOVERS;
-            output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-        } else {
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
-            FUSED_OPS_LEFTOVERS;
-            output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-        }
+        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+        FUSED_OPS_LEFTOVERS;
+        output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
 #else
 #if IS_DYNAMIC
@@ -361,10 +288,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
-                if (use_output_buffer)
-                    vec_tmp[j] = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-                else
-                    vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
+                vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
             }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
@@ -372,10 +296,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx < items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            if (use_output_buffer)
-                vec_tmp = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-            else
-                vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+            vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -397,25 +318,212 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        if (use_output_buffer) {
-            output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()]
-            = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-        } else
-            output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < aligned_offset) {
-        if (use_output_buffer)
-            output[data_set_offset + in_data_set_idx] = ACTIVATION(output[data_set_offset + in_data_set_idx] / my_sum, ACTIVATION_PARAMS);
-        else
-            output[data_set_offset + in_data_set_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+        output[data_set_offset + in_data_set_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < actual_leftovers) {
-        if (use_output_buffer)
-            output[leftover_idx] = ACTIVATION(output[leftover_idx] / my_sum, ACTIVATION_PARAMS);
-        else
-            output[leftover_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+        output[leftover_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+    }
+#endif
+#if IS_DYNAMIC
+    } else { // Case for output[] directly when (items_num + 2) > STACK_SIZE
+    if (workers_per_data_set > SUB_GROUP_SIZE)
+    {
+        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
+        for (; input_idx < num_iters; input_idx += OPT_BLOCK_SIZE)
+        {
+            BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
+            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
+            {
+                output[aligned_data_offset + get_sub_group_local_id() + (input_idx + j) * get_sub_group_size()] = vec_tmp[j];
+            }
+        }
+
+        for (; input_idx < items_num; input_idx++)
+        {
+            BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
+            output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()] = vec_tmp;
+        }
+    }
+
+    for (; input_idx < items_num; input_idx++)
+    {
+        output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()]
+        = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+    }
+
+    if (in_data_set_idx < aligned_offset)
+    {
+        INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
+        output[data_set_offset + in_data_set_idx] = tmp;
+    }
+
+    if (in_data_set_idx < actual_leftovers)
+    {
+        INPUT0_TYPE tmp = input[leftover_idx];
+        output[leftover_idx] = tmp;
+    }
+
+    INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
+    {
+        const uint num_iters = input_idx;
+
+        for (uint j=0; j<num_iters; ++j)
+        {
+            my_maximum = max(my_maximum, output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()]);
+        }
+        if (in_data_set_idx < aligned_offset) {
+            my_maximum = max(my_maximum, output[data_set_offset + in_data_set_idx]);
+        }
+        if (in_data_set_idx < actual_leftovers) {
+            my_maximum = max(my_maximum, output[leftover_idx]);
+        }
+    }
+
+    my_maximum = sub_group_reduce_max(my_maximum);
+
+    if (get_sub_group_local_id() == 0)
+        lg_storage[get_sub_group_id()] = my_maximum;
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (in_data_set_idx == 0)
+    {
+        for (uint j=1; j<get_num_sub_groups(); ++j)
+            my_maximum = max(my_maximum, lg_storage[j]);
+
+        lg_storage[0] = my_maximum;
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    //my_maximum from this point is in fact global maximum
+    my_maximum = lg_storage[0];
+
+    // Get exp(x-max) and sum of exp(x-max)
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const uint num_iters = input_idx;
+
+    for (uint j=0; j<num_iters; ++j) {
+        INPUT0_TYPE tmp = native_exp(output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] - my_maximum);
+        my_sum += tmp;
+        output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] = tmp;
+    }
+
+    if (in_data_set_idx < aligned_offset) {
+        INPUT0_TYPE tmp = native_exp(output[data_set_offset + in_data_set_idx] - my_maximum);
+        my_sum += tmp;
+        output[data_set_offset + in_data_set_idx] = tmp;
+    }
+
+    if (in_data_set_idx < actual_leftovers) {
+        INPUT0_TYPE tmp = native_exp(output[leftover_idx] - my_maximum);
+        my_sum += tmp;
+        output[leftover_idx] = tmp;
+    }
+
+    my_sum = sub_group_reduce_add(my_sum);
+
+    if (get_sub_group_local_id() == 0)
+        lg_storage[get_sub_group_id()] = my_sum;
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (in_data_set_idx == 0)
+    {
+        for (uint j=1; j<get_num_sub_groups(); ++j)
+            my_sum += lg_storage[j];
+
+        lg_storage[0] = my_sum;
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    my_sum = lg_storage[0];
+
+    // Write outputs
+    uint output_idx = 0;
+#if HAS_FUSED_OPS
+    if (workers_per_data_set > SUB_GROUP_SIZE)
+    {
+        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
+        for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
+        {
+            BLOCK_TYPE_OPT vec_tmp;
+            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
+                ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum;
+                FUSED_OPS_MAIN;
+                vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
+            }
+            BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
+        }
+
+        for (; output_idx<items_num; output_idx++)
+        {
+            BLOCK_TYPE vec_tmp;
+            ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
+            FUSED_OPS_MAIN;
+            vec_tmp = FUSED_OPS_RESULT_MAIN;
+            BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
+        }
+    }
+
+    for (; output_idx < items_num; output_idx++)
+    {
+        ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
+        FUSED_OPS_MAIN;
+        output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
+    }
+
+    if (in_data_set_idx < aligned_offset)
+    {
+        ACTIVATION_TYPE dequantized = output[data_set_offset + in_data_set_idx] / my_sum;
+        FUSED_OPS_LEFTOVERS;
+        output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+    }
+
+    if (in_data_set_idx < actual_leftovers)
+    {
+        ACTIVATION_TYPE dequantized = output[leftover_idx] / my_sum;
+        FUSED_OPS_LEFTOVERS;
+        output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
+    }
+#else
+    if (workers_per_data_set > SUB_GROUP_SIZE)
+    {
+        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
+        for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
+        {
+            BLOCK_TYPE_OPT vec_tmp;
+            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
+                vec_tmp[j] = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+            }
+            BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
+        }
+
+        for (; output_idx < items_num; output_idx++)
+        {
+            BLOCK_TYPE vec_tmp;
+            vec_tmp = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+            BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
+        }
+    }
+
+    for (; output_idx < items_num; output_idx++)
+    {
+        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()]
+        = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
+    }
+
+    if (in_data_set_idx < aligned_offset) {
+        output[data_set_offset + in_data_set_idx] = ACTIVATION(output[data_set_offset + in_data_set_idx] / my_sum, ACTIVATION_PARAMS);
+    }
+
+    if (in_data_set_idx < actual_leftovers) {
+        output[leftover_idx] = ACTIVATION(output[leftover_idx] / my_sum, ACTIVATION_PARAMS);
+    }
+#endif
     }
 #endif
 }

From 60c8c4d0508901f22b91f652dfa358f29609c4f4 Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Fri, 17 May 2024 17:01:28 +0000
Subject: [PATCH 5/6] reduce code duplication

---
 .../cl_kernels/softmax_gpu_bf.cl              | 308 +++++-------------
 1 file changed, 79 insertions(+), 229 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
index d1efd90e9ad..eff78073125 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
@@ -35,6 +35,30 @@
 
 #endif
 
+inline uint FUNC(index_calc)(uint idx, bool use_output_buffer, uint aligned_data_offset) {
+#if IS_DYNAMIC
+    if (use_output_buffer) {
+        return (aligned_data_offset + get_sub_group_local_id() + idx * get_sub_group_size());
+    } else {
+        return idx;
+    }
+#else
+    return idx;
+#endif
+}
+
+inline uint FUNC(index_calc_leftover)(uint idx, bool use_output_buffer, uint leftover_idx) {
+#if IS_DYNAMIC
+    if (use_output_buffer) {
+        return leftover_idx;
+    } else {
+        return idx;
+    }
+#else
+    return idx;
+#endif
+}
+
 REQD_SUB_GROUP_SIZE(SUB_GROUP_SIZE)
 KERNEL (softmax_gpu_continuous_bfyx)(
     OPTIONAL_SHAPE_INFO_ARG
@@ -93,11 +117,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     // Read inputs and Get maximum value from data set
     uint input_idx=0;
 
+    OUTPUT_TYPE* tmp_out_ptr;
 
-#if IS_DYNAMIC
     // Case for my_chunk[] when (items_num + 2) <= STACK_SIZE
-    if (!use_output_buffer) {
-#endif
+    if (use_output_buffer)
+        tmp_out_ptr = output;
+    else
+        tmp_out_ptr = &my_chunk[0];
+
 #if IS_DYNAMIC
     if (workers_per_data_set > SUB_GROUP_SIZE)
     {
@@ -107,14 +134,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
             BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
             {
-                my_chunk[input_idx+j] = vec_tmp[j];
+                tmp_out_ptr[FUNC_CALL(index_calc)(input_idx+j, use_output_buffer, aligned_data_offset)] = vec_tmp[j];
             }
         }
 
         for (; input_idx < items_num; input_idx++)
         {
             BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
-            my_chunk[input_idx] = vec_tmp;
+            tmp_out_ptr[FUNC_CALL(index_calc)(input_idx, use_output_buffer, aligned_data_offset)] = vec_tmp;
         }
     }
 #else
@@ -124,12 +151,12 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
 #if SUBGROUP_BLOCK_SIZE == 1
-            my_chunk[input_idx] = vec_tmp;
+            tmp_out_ptr[input_idx] = vec_tmp;
 #else
             unroll_for (int j = 0; j < SUBGROUP_BLOCK_SIZE; j++)
             {
                 INPUT0_TYPE tmp = vec_tmp[j];
-                my_chunk[input_idx+j] = tmp;
+                tmp_out_ptr[input_idx+j] = tmp;
             }
 #endif
         }
@@ -138,28 +165,35 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     for (; input_idx < items_num; input_idx++)
     {
-        my_chunk[input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+        tmp_out_ptr[FUNC_CALL(index_calc)(input_idx, use_output_buffer, aligned_data_offset)] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
     }
+    const uint num_iter_even = input_idx;
 
     if (in_data_set_idx < aligned_offset)
     {
         INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
-        my_chunk[input_idx++] = tmp;
+        tmp_out_ptr[FUNC_CALL(index_calc_leftover)(input_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] = tmp;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
         INPUT0_TYPE tmp = input[leftover_idx];
-        my_chunk[input_idx++] = tmp;
+        tmp_out_ptr[FUNC_CALL(index_calc_leftover)(input_idx++, use_output_buffer, leftover_idx)] = tmp;
     }
 
     INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
     {
-        const uint num_iters = input_idx;
+        uint num_iters = num_iter_even;
 
         for (uint j=0; j<num_iters; ++j)
         {
-            my_maximum = max(my_maximum, my_chunk[j]);
+            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)]);
+        }
+        if (in_data_set_idx < aligned_offset) {
+            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, data_set_offset + in_data_set_idx)]);
+        }
+        if (in_data_set_idx < actual_leftovers) {
+            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, leftover_idx)]);
         }
     }
 
@@ -184,13 +218,26 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     // Get exp(x-max) and sum of exp(x-max)
     barrier(CLK_LOCAL_MEM_FENCE);
 
-    const uint num_iters = input_idx;
-
-    for (uint j=0; j<num_iters; ++j)
     {
-        INPUT0_TYPE tmp = native_exp(my_chunk[j] - my_maximum);
-        my_sum += tmp;
-        my_chunk[j] = tmp;
+        uint num_iters = num_iter_even;
+
+        for (uint j=0; j<num_iters; ++j)
+        {
+            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)] - my_maximum);
+            my_sum += tmp;
+            tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)] = tmp;
+        }
+        if (in_data_set_idx < aligned_offset) {
+            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters, use_output_buffer, data_set_offset + in_data_set_idx)] - my_maximum);
+            my_sum += tmp;
+            tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, data_set_offset + in_data_set_idx)] = tmp;
+        }
+
+        if (in_data_set_idx < actual_leftovers) {
+            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters, use_output_buffer, leftover_idx)] - my_maximum);
+            my_sum += tmp;
+            tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, leftover_idx)] = tmp;
+        }
     }
 
     my_sum = sub_group_reduce_add(my_sum);
@@ -221,7 +268,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
-                ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
+                ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx + j, use_output_buffer, aligned_data_offset)] / my_sum;
                 FUSED_OPS_MAIN;
                 vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
             }
@@ -231,7 +278,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx<items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+            ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum;
             FUSED_OPS_MAIN;
             vec_tmp = FUSED_OPS_RESULT_MAIN;
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
@@ -244,13 +291,13 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE vec_tmp;
 #if SUBGROUP_BLOCK_SIZE == 1
-            ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+            ACTIVATION_TYPE dequantized = tmp_out_ptr[output_idx] / my_sum;
             FUSED_OPS_MAIN;
             vec_tmp = FUSED_OPS_RESULT_MAIN;
 #else
             for (int j = 0; j < SUBGROUP_BLOCK_SIZE; j++)
             {
-                ACTIVATION_TYPE dequantized = my_chunk[output_idx + j] / my_sum;
+                ACTIVATION_TYPE dequantized = tmp_out_ptr[output_idx + j] / my_sum;
                 FUSED_OPS_MAIN;
                 vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
             }
@@ -261,21 +308,21 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum;
         FUSED_OPS_MAIN;
         output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
     }
 
     if (in_data_set_idx < aligned_offset)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] / my_sum;
         FUSED_OPS_LEFTOVERS;
         output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
-        ACTIVATION_TYPE dequantized = my_chunk[output_idx++] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, leftover_idx)] / my_sum;
         FUSED_OPS_LEFTOVERS;
         output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
@@ -288,7 +335,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
-                vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
+                vec_tmp[j] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx + j, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
             }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
@@ -296,7 +343,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx < items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+            vec_tmp = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -307,10 +354,10 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE vec_tmp;
 #if SUBGROUP_BLOCK_SIZE == 1
-            vec_tmp = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+            vec_tmp = ACTIVATION(tmp_out_ptr[output_idx] / my_sum, ACTIVATION_PARAMS);
 #else
             for (int j = 0; j < SUBGROUP_BLOCK_SIZE; j++)
-                vec_tmp[j] = ACTIVATION(my_chunk[output_idx + j] / my_sum, ACTIVATION_PARAMS);
+                vec_tmp[j] = ACTIVATION(tmp_out_ptr[output_idx + j] / my_sum, ACTIVATION_PARAMS);
 #endif
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
@@ -318,212 +365,15 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(my_chunk[output_idx] / my_sum, ACTIVATION_PARAMS);
+        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < aligned_offset) {
-        output[data_set_offset + in_data_set_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
+        output[data_set_offset + in_data_set_idx] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < actual_leftovers) {
-        output[leftover_idx] = ACTIVATION(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
-    }
-#endif
-#if IS_DYNAMIC
-    } else { // Case for output[] directly when (items_num + 2) > STACK_SIZE
-    if (workers_per_data_set > SUB_GROUP_SIZE)
-    {
-        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
-        for (; input_idx < num_iters; input_idx += OPT_BLOCK_SIZE)
-        {
-            BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
-            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
-            {
-                output[aligned_data_offset + get_sub_group_local_id() + (input_idx + j) * get_sub_group_size()] = vec_tmp[j];
-            }
-        }
-
-        for (; input_idx < items_num; input_idx++)
-        {
-            BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
-            output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()] = vec_tmp;
-        }
-    }
-
-    for (; input_idx < items_num; input_idx++)
-    {
-        output[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()]
-        = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
-    }
-
-    if (in_data_set_idx < aligned_offset)
-    {
-        INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
-        output[data_set_offset + in_data_set_idx] = tmp;
-    }
-
-    if (in_data_set_idx < actual_leftovers)
-    {
-        INPUT0_TYPE tmp = input[leftover_idx];
-        output[leftover_idx] = tmp;
-    }
-
-    INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
-    {
-        const uint num_iters = input_idx;
-
-        for (uint j=0; j<num_iters; ++j)
-        {
-            my_maximum = max(my_maximum, output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()]);
-        }
-        if (in_data_set_idx < aligned_offset) {
-            my_maximum = max(my_maximum, output[data_set_offset + in_data_set_idx]);
-        }
-        if (in_data_set_idx < actual_leftovers) {
-            my_maximum = max(my_maximum, output[leftover_idx]);
-        }
-    }
-
-    my_maximum = sub_group_reduce_max(my_maximum);
-
-    if (get_sub_group_local_id() == 0)
-        lg_storage[get_sub_group_id()] = my_maximum;
-
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (in_data_set_idx == 0)
-    {
-        for (uint j=1; j<get_num_sub_groups(); ++j)
-            my_maximum = max(my_maximum, lg_storage[j]);
-
-        lg_storage[0] = my_maximum;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    //my_maximum from this point is in fact global maximum
-    my_maximum = lg_storage[0];
-
-    // Get exp(x-max) and sum of exp(x-max)
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    const uint num_iters = input_idx;
-
-    for (uint j=0; j<num_iters; ++j) {
-        INPUT0_TYPE tmp = native_exp(output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] - my_maximum);
-        my_sum += tmp;
-        output[aligned_data_offset + get_sub_group_local_id() + j * get_sub_group_size()] = tmp;
-    }
-
-    if (in_data_set_idx < aligned_offset) {
-        INPUT0_TYPE tmp = native_exp(output[data_set_offset + in_data_set_idx] - my_maximum);
-        my_sum += tmp;
-        output[data_set_offset + in_data_set_idx] = tmp;
-    }
-
-    if (in_data_set_idx < actual_leftovers) {
-        INPUT0_TYPE tmp = native_exp(output[leftover_idx] - my_maximum);
-        my_sum += tmp;
-        output[leftover_idx] = tmp;
-    }
-
-    my_sum = sub_group_reduce_add(my_sum);
-
-    if (get_sub_group_local_id() == 0)
-        lg_storage[get_sub_group_id()] = my_sum;
-
-    barrier(CLK_LOCAL_MEM_FENCE);
-    if (in_data_set_idx == 0)
-    {
-        for (uint j=1; j<get_num_sub_groups(); ++j)
-            my_sum += lg_storage[j];
-
-        lg_storage[0] = my_sum;
-    }
-    barrier(CLK_LOCAL_MEM_FENCE);
-
-    my_sum = lg_storage[0];
-
-    // Write outputs
-    uint output_idx = 0;
-#if HAS_FUSED_OPS
-    if (workers_per_data_set > SUB_GROUP_SIZE)
-    {
-        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
-        for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
-        {
-            BLOCK_TYPE_OPT vec_tmp;
-            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
-                ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum;
-                FUSED_OPS_MAIN;
-                vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
-            }
-            BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
-        }
-
-        for (; output_idx<items_num; output_idx++)
-        {
-            BLOCK_TYPE vec_tmp;
-            ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
-            FUSED_OPS_MAIN;
-            vec_tmp = FUSED_OPS_RESULT_MAIN;
-            BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
-        }
-    }
-
-    for (; output_idx < items_num; output_idx++)
-    {
-        ACTIVATION_TYPE dequantized = output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum;
-        FUSED_OPS_MAIN;
-        output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
-    }
-
-    if (in_data_set_idx < aligned_offset)
-    {
-        ACTIVATION_TYPE dequantized = output[data_set_offset + in_data_set_idx] / my_sum;
-        FUSED_OPS_LEFTOVERS;
-        output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-    }
-
-    if (in_data_set_idx < actual_leftovers)
-    {
-        ACTIVATION_TYPE dequantized = output[leftover_idx] / my_sum;
-        FUSED_OPS_LEFTOVERS;
-        output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
-    }
-#else
-    if (workers_per_data_set > SUB_GROUP_SIZE)
-    {
-        const uint num_iters = items_num - (items_num % OPT_BLOCK_SIZE);
-        for (; output_idx < num_iters; output_idx += OPT_BLOCK_SIZE)
-        {
-            BLOCK_TYPE_OPT vec_tmp;
-            unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
-                vec_tmp[j] = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + (output_idx + j) * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-            }
-            BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
-        }
-
-        for (; output_idx < items_num; output_idx++)
-        {
-            BLOCK_TYPE vec_tmp;
-            vec_tmp = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-            BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
-        }
-    }
-
-    for (; output_idx < items_num; output_idx++)
-    {
-        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()]
-        = ACTIVATION(output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] / my_sum, ACTIVATION_PARAMS);
-    }
-
-    if (in_data_set_idx < aligned_offset) {
-        output[data_set_offset + in_data_set_idx] = ACTIVATION(output[data_set_offset + in_data_set_idx] / my_sum, ACTIVATION_PARAMS);
-    }
-
-    if (in_data_set_idx < actual_leftovers) {
-        output[leftover_idx] = ACTIVATION(output[leftover_idx] / my_sum, ACTIVATION_PARAMS);
-    }
-#endif
+        output[leftover_idx] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, leftover_idx)] / my_sum, ACTIVATION_PARAMS);
     }
 #endif
 }

From ced40dae307f71d459f9ea8760833fa5a14ae7c2 Mon Sep 17 00:00:00 2001
From: wilson-seok <wilson.seok@intel.com>
Date: Tue, 21 May 2024 16:13:57 +0000
Subject: [PATCH 6/6] simplify index calculation

---
 .../cl_kernels/softmax_gpu_bf.cl              | 101 +++++++-----------
 1 file changed, 41 insertions(+), 60 deletions(-)

diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
index eff78073125..d01f23249be 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/softmax_gpu_bf.cl
@@ -35,30 +35,6 @@
 
 #endif
 
-inline uint FUNC(index_calc)(uint idx, bool use_output_buffer, uint aligned_data_offset) {
-#if IS_DYNAMIC
-    if (use_output_buffer) {
-        return (aligned_data_offset + get_sub_group_local_id() + idx * get_sub_group_size());
-    } else {
-        return idx;
-    }
-#else
-    return idx;
-#endif
-}
-
-inline uint FUNC(index_calc_leftover)(uint idx, bool use_output_buffer, uint leftover_idx) {
-#if IS_DYNAMIC
-    if (use_output_buffer) {
-        return leftover_idx;
-    } else {
-        return idx;
-    }
-#else
-    return idx;
-#endif
-}
-
 REQD_SUB_GROUP_SIZE(SUB_GROUP_SIZE)
 KERNEL (softmax_gpu_continuous_bfyx)(
     OPTIONAL_SHAPE_INFO_ARG
@@ -125,6 +101,9 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     else
         tmp_out_ptr = &my_chunk[0];
 
+    uint data_offset = (use_output_buffer) ? aligned_data_offset + get_sub_group_local_id() : 0;
+    uint idx_stride = (use_output_buffer) ? get_sub_group_size() : 1;
+
 #if IS_DYNAMIC
     if (workers_per_data_set > SUB_GROUP_SIZE)
     {
@@ -134,14 +113,14 @@ KERNEL (softmax_gpu_continuous_bfyx)(
             BLOCK_TYPE_OPT vec_tmp = BLOCK_READ_OPT(input, aligned_data_offset + input_idx * get_sub_group_size());
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++)
             {
-                tmp_out_ptr[FUNC_CALL(index_calc)(input_idx+j, use_output_buffer, aligned_data_offset)] = vec_tmp[j];
+                tmp_out_ptr[data_offset + idx_stride * (input_idx + j)] = vec_tmp[j];
             }
         }
 
         for (; input_idx < items_num; input_idx++)
         {
             BLOCK_TYPE vec_tmp = BLOCK_READ(input, aligned_data_offset + input_idx * get_sub_group_size());
-            tmp_out_ptr[FUNC_CALL(index_calc)(input_idx, use_output_buffer, aligned_data_offset)] = vec_tmp;
+            tmp_out_ptr[data_offset + idx_stride * input_idx] = vec_tmp;
         }
     }
 #else
@@ -165,35 +144,36 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 
     for (; input_idx < items_num; input_idx++)
     {
-        tmp_out_ptr[FUNC_CALL(index_calc)(input_idx, use_output_buffer, aligned_data_offset)] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
+        tmp_out_ptr[data_offset + idx_stride * input_idx] = input[aligned_data_offset + get_sub_group_local_id() + input_idx * get_sub_group_size()];
     }
-    const uint num_iter_even = input_idx;
 
     if (in_data_set_idx < aligned_offset)
     {
         INPUT0_TYPE tmp = input[data_set_offset + in_data_set_idx];
-        tmp_out_ptr[FUNC_CALL(index_calc_leftover)(input_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] = tmp;
+        tmp_out_ptr[use_output_buffer ? data_set_offset + in_data_set_idx : input_idx++] = tmp;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
         INPUT0_TYPE tmp = input[leftover_idx];
-        tmp_out_ptr[FUNC_CALL(index_calc_leftover)(input_idx++, use_output_buffer, leftover_idx)] = tmp;
+        tmp_out_ptr[use_output_buffer ? leftover_idx : input_idx++] = tmp;
     }
 
     INPUT0_TYPE my_maximum = -UNIT_VAL_MAX;
     {
-        uint num_iters = num_iter_even;
+        const uint num_iters = input_idx;
 
         for (uint j=0; j<num_iters; ++j)
         {
-            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)]);
+            my_maximum = max(my_maximum, tmp_out_ptr[data_offset + idx_stride * j]);
         }
-        if (in_data_set_idx < aligned_offset) {
-            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, data_set_offset + in_data_set_idx)]);
-        }
-        if (in_data_set_idx < actual_leftovers) {
-            my_maximum = max(my_maximum, tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, leftover_idx)]);
+        if (use_output_buffer) {
+            if (in_data_set_idx < aligned_offset) {
+                my_maximum = max(my_maximum, tmp_out_ptr[data_set_offset + in_data_set_idx]);
+            }
+            if (in_data_set_idx < actual_leftovers) {
+                my_maximum = max(my_maximum, tmp_out_ptr[leftover_idx]);
+            }
         }
     }
 
@@ -219,24 +199,25 @@ KERNEL (softmax_gpu_continuous_bfyx)(
     barrier(CLK_LOCAL_MEM_FENCE);
 
     {
-        uint num_iters = num_iter_even;
+        const uint num_iters = input_idx;
 
         for (uint j=0; j<num_iters; ++j)
         {
-            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)] - my_maximum);
+            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[data_offset + idx_stride * j] - my_maximum);
             my_sum += tmp;
-            tmp_out_ptr[FUNC_CALL(index_calc)(j, use_output_buffer, aligned_data_offset)] = tmp;
+            tmp_out_ptr[data_offset + idx_stride * j] = tmp;
         }
-        if (in_data_set_idx < aligned_offset) {
-            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters, use_output_buffer, data_set_offset + in_data_set_idx)] - my_maximum);
-            my_sum += tmp;
-            tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, data_set_offset + in_data_set_idx)] = tmp;
-        }
-
-        if (in_data_set_idx < actual_leftovers) {
-            INPUT0_TYPE tmp = native_exp(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters, use_output_buffer, leftover_idx)] - my_maximum);
-            my_sum += tmp;
-            tmp_out_ptr[FUNC_CALL(index_calc_leftover)(num_iters++, use_output_buffer, leftover_idx)] = tmp;
+        if (use_output_buffer) {
+            if (in_data_set_idx < aligned_offset) {
+                INPUT0_TYPE tmp = native_exp(tmp_out_ptr[data_set_offset + in_data_set_idx] - my_maximum);
+                my_sum += tmp;
+                tmp_out_ptr[data_set_offset + in_data_set_idx] = tmp;
+            }
+            if (in_data_set_idx < actual_leftovers) {
+                INPUT0_TYPE tmp = native_exp(tmp_out_ptr[leftover_idx] - my_maximum);
+                my_sum += tmp;
+                tmp_out_ptr[leftover_idx] = tmp;
+            }
         }
     }
 
@@ -268,7 +249,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
-                ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx + j, use_output_buffer, aligned_data_offset)] / my_sum;
+                ACTIVATION_TYPE dequantized = tmp_out_ptr[data_offset + idx_stride * (output_idx + j)] / my_sum;
                 FUSED_OPS_MAIN;
                 vec_tmp[j] = FUSED_OPS_RESULT_MAIN;
             }
@@ -278,7 +259,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx<items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum;
+            ACTIVATION_TYPE dequantized = tmp_out_ptr[data_offset + idx_stride * output_idx] / my_sum;
             FUSED_OPS_MAIN;
             vec_tmp = FUSED_OPS_RESULT_MAIN;
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
@@ -308,21 +289,21 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[data_offset + idx_stride * output_idx] / my_sum;
         FUSED_OPS_MAIN;
         output[aligned_data_offset + get_sub_group_local_id() + i * get_sub_group_size()] = FUSED_OPS_RESULT_MAIN;
     }
 
     if (in_data_set_idx < aligned_offset)
     {
-        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[use_output_buffer ? data_set_offset + in_data_set_idx : output_idx++] / my_sum;
         FUSED_OPS_LEFTOVERS;
         output[data_set_offset + in_data_set_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
 
     if (in_data_set_idx < actual_leftovers)
     {
-        ACTIVATION_TYPE dequantized = tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, leftover_idx)] / my_sum;
+        ACTIVATION_TYPE dequantized = tmp_out_ptr[use_output_buffer ? leftover_idx : output_idx++] / my_sum;
         FUSED_OPS_LEFTOVERS;
         output[leftover_idx] = FUSED_OPS_RESULT_LEFTOVERS;
     }
@@ -335,7 +316,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         {
             BLOCK_TYPE_OPT vec_tmp;
             unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++){
-                vec_tmp[j] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx + j, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
+                vec_tmp[j] = ACTIVATION(tmp_out_ptr[data_offset + idx_stride * (output_idx + j)] / my_sum, ACTIVATION_PARAMS);
             }
             BLOCK_WRITE_OPT(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
@@ -343,7 +324,7 @@ KERNEL (softmax_gpu_continuous_bfyx)(
         for (; output_idx < items_num; output_idx++)
         {
             BLOCK_TYPE vec_tmp;
-            vec_tmp = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
+            vec_tmp = ACTIVATION(tmp_out_ptr[data_offset + idx_stride * output_idx] / my_sum, ACTIVATION_PARAMS);
             BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
         }
     }
@@ -365,15 +346,15 @@ KERNEL (softmax_gpu_continuous_bfyx)(
 #endif
     for (; output_idx < items_num; output_idx++)
     {
-        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc)(output_idx, use_output_buffer, aligned_data_offset)] / my_sum, ACTIVATION_PARAMS);
+        output[aligned_data_offset + get_sub_group_local_id() + output_idx * get_sub_group_size()] = ACTIVATION(tmp_out_ptr[data_offset + idx_stride * output_idx] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < aligned_offset) {
-        output[data_set_offset + in_data_set_idx] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, data_set_offset + in_data_set_idx)] / my_sum, ACTIVATION_PARAMS);
+        output[data_set_offset + in_data_set_idx] = ACTIVATION(tmp_out_ptr[use_output_buffer ? data_set_offset + in_data_set_idx : output_idx++] / my_sum, ACTIVATION_PARAMS);
     }
 
     if (in_data_set_idx < actual_leftovers) {
-        output[leftover_idx] = ACTIVATION(tmp_out_ptr[FUNC_CALL(index_calc_leftover)(output_idx++, use_output_buffer, leftover_idx)] / my_sum, ACTIVATION_PARAMS);
+        output[leftover_idx] = ACTIVATION(tmp_out_ptr[use_output_buffer? leftover_idx : output_idx++] / my_sum, ACTIVATION_PARAMS);
     }
 #endif
 }