This commit is contained in:
Wilson Seok 2024-06-12 03:26:57 +02:00 committed by GitHub
commit 824b92487d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 97 additions and 34 deletions

View File

@ -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;
@ -84,6 +92,18 @@ KERNEL (softmax_gpu_continuous_bfyx)(
// Read inputs and Get maximum value from data set
uint input_idx=0;
OUTPUT_TYPE* tmp_out_ptr;
// Case for my_chunk[] when (items_num + 2) <= STACK_SIZE
if (use_output_buffer)
tmp_out_ptr = output;
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)
{
@ -93,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++)
{
my_chunk[input_idx+j] = 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());
my_chunk[input_idx] = vec_tmp;
tmp_out_ptr[data_offset + idx_stride * input_idx] = vec_tmp;
}
}
#else
@ -110,12 +130,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
}
@ -124,27 +144,36 @@ 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[data_offset + idx_stride * 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;
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];
my_chunk[input_idx++] = tmp;
tmp_out_ptr[use_output_buffer ? leftover_idx : input_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, my_chunk[j]);
my_maximum = max(my_maximum, tmp_out_ptr[data_offset + idx_stride * j]);
}
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]);
}
}
}
@ -169,12 +198,27 @@ 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;
const uint num_iters = input_idx;
for (uint j=0; j<num_iters; ++j)
{
INPUT0_TYPE tmp = native_exp(tmp_out_ptr[data_offset + idx_stride * j] - my_maximum);
my_sum += tmp;
tmp_out_ptr[data_offset + idx_stride * j] = 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;
}
}
}
my_sum = sub_group_reduce_add(my_sum);
@ -204,9 +248,8 @@ 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;
unroll_for (int j = 0; j < OPT_BLOCK_SIZE; j++) {
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;
}
@ -215,7 +258,8 @@ KERNEL (softmax_gpu_continuous_bfyx)(
for (; output_idx<items_num; output_idx++)
{
ACTIVATION_TYPE dequantized = my_chunk[output_idx] / my_sum;
BLOCK_TYPE vec_tmp;
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);
@ -228,13 +272,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;
}
@ -245,21 +289,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[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 = my_chunk[output_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 = my_chunk[output_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;
}
@ -271,15 +315,16 @@ 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++){
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);
}
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[data_offset + idx_stride * output_idx] / my_sum, ACTIVATION_PARAMS);
BLOCK_WRITE(output, aligned_data_offset + output_idx * get_sub_group_size(), vec_tmp);
}
}
@ -290,10 +335,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);
}
@ -301,14 +346,16 @@ 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[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(my_chunk[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[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(my_chunk[output_idx++] / my_sum, ACTIVATION_PARAMS);
if (in_data_set_idx < actual_leftovers) {
output[leftover_idx] = ACTIVATION(tmp_out_ptr[use_output_buffer? leftover_idx : output_idx++] / my_sum, ACTIVATION_PARAMS);
}
#endif
}
#ifdef CALC_POWER

View File

@ -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};
@ -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 = {
@ -127,6 +136,13 @@ INSTANTIATE_TEST_SUITE_P(softMaxGPUDynamicTest2D,
testing::ValuesIn(axis2D)),
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(softMaxGPUDynamicTest4D,
SoftMaxLayerGPUTest,
::testing::Combine(testing::ValuesIn(netPrecisions),