[GPU] Fixed shape agnostic scatter nd update kernel (#16319)

This commit is contained in:
Vladimir Paramuzov authored and GitHub committed 2023-03-17 09:57:25 +04:00
1 parent c5ccb3e954
commit bb59672639
2 files changed
+136 -22

No files matched your search

@@ -117,24 +117,26 @@ bool ScatterNDUpdateKernelRef::Validate(const Params& p, const optional_params&
return true;
}
static std::string GetInputBlockND(const scatter_nd_update_params& params, size_t num, size_t rank) {
static std::string GetInputBlockND(const scatter_nd_update_params& params, size_t num, size_t dyn_offset, size_t rank) {
const auto& input = params.inputs[num];
auto input_dims = input.LogicalDims();
std::reverse(input_dims.begin(), input_dims.end());
auto dims = input.GetDims();
std::reverse(dims.begin(), dims.end());
std::vector<size_t> block_nd(rank + 1);
block_nd[rank] = 1;
std::vector<std::string> block_nd_s(rank + 1);
block_nd_s[rank] = "1";
size_t input_offset = num * 6;
size_t input_offset = dyn_offset * 6;
for (int32_t idx = rank - 1; idx >= 0; --idx) {
block_nd[idx] = input_dims[idx] * block_nd[idx + 1];
size_t dim_offset = idx < 2 ? idx : idx + 6 - rank;
block_nd_s[idx] = "(" + toCodeString(input.GetDims()[input.GetDims().size() - idx - 1], input_offset + dim_offset) + "*" + block_nd_s[idx + 1] + ")";
size_t dim_offset = idx < 2 ? idx : (6 - dims.size()) + idx; // convert to 6d bfwzyx idx
block_nd_s[idx] = "(" + toCodeString(dims[idx], input_offset + dim_offset) + "*" + block_nd_s[idx + 1] + ")";
}
std::string result;
@@ -180,18 +182,24 @@ KernelsData ScatterNDUpdateKernelRef::GetKernelsData(const Params& params, const
size_t input0_rank = newParams.inputs[0].LogicalDims().size();
size_t input2_rank = newParams.inputs[2].LogicalDims().size();
cldnn_jit.AddConstant(MakeJitConstant("IS_SECOND_ITER", "true"));
cldnn_jit.AddConstant(MakeJitConstant("INPUT0_BLOCK_ND", GetInputBlockND(newParams, 0, input0_rank)));
cldnn_jit.AddConstant(MakeJitConstant("INPUT1_BLOCK_ND", GetInputBlockND(newParams, 1, newParams.indices_rank - 1)));
cldnn_jit.AddConstant(MakeJitConstant("INPUT2_BLOCK_ND", GetInputBlockND(newParams, 2, input2_rank)));
size_t shape_info_offset = 0;
cldnn_jit.AddConstant(MakeJitConstant("INPUT0_BLOCK_ND", GetInputBlockND(newParams, 0, shape_info_offset, input0_rank)));
if (newParams.inputs[0].is_dynamic())
shape_info_offset++;
cldnn_jit.AddConstant(MakeJitConstant("INPUT1_BLOCK_ND", GetInputBlockND(newParams, 1, shape_info_offset, newParams.indices_rank - 1)));
if (newParams.inputs[1].is_dynamic())
shape_info_offset++;
cldnn_jit.AddConstant(MakeJitConstant("INPUT2_BLOCK_ND", GetInputBlockND(newParams, 2, shape_info_offset, input2_rank)));
cldnn_jit.AddConstant(MakeJitConstant("INDICES_RANK", newParams.indices_rank));
const auto& ind_input = newParams.inputs[1];
if (ind_input.is_dynamic()) {
size_t last_idx = newParams.indices_rank - 1;
size_t dim_offset = last_idx < 2 ? last_idx : 5;
size_t input_idx = last_idx < 2 ? ind_input.GetDims().size() - last_idx : 0;
auto dims = ind_input.GetDims();
std::reverse(dims.begin(), dims.end());
auto indices_last_dim = toCodeString(ind_input.GetDims()[input_idx], 6 + dim_offset);
size_t last_idx = newParams.indices_rank - 1;
size_t dim_offset = last_idx < 2 ? last_idx : last_idx + 6 - newParams.indices_rank;
auto indices_last_dim = toCodeString(dims[last_idx], dim_offset + (newParams.inputs[0].is_dynamic() ? 6 : 0));
cldnn_jit.AddConstant(MakeJitConstant("INDICES_LAST_DIM", indices_last_dim));
} else {
cldnn_jit.AddConstant(MakeJitConstant("INDICES_LAST_DIM", dispatchData.indicesLastDim));
@@ -22,6 +22,18 @@
using namespace cldnn;
using namespace ::tests;
namespace {
template<typename T>
T generate_random_val(int min, int max, int k = 8) {
static std::default_random_engine generator(random_seed);
// 1/k is the resolution of the floating point numbers
std::uniform_int_distribution<int> distribution(k * min, k * max);
T val = (T)distribution(generator);
val /= k;
return val;
}
}
struct scatter_nd_update_basic_test_params
{
@@ -51,17 +63,6 @@ struct scatter_nd_update_random_test : testing::TestWithParam<scatter_nd_update_
return cldnn::format::bfwzyx;
}
template<typename T>
T generate_random_val(int min, int max, int k = 8) {
static std::default_random_engine generator(random_seed);
// 1/k is the resolution of the floating point numbers
std::uniform_int_distribution<int> distribution(k * min, k * max);
T val = (T)distribution(generator);
val /= k;
return val;
}
template <typename T>
std::vector<T> generate_unique_indices(const scatter_nd_update_basic_test_params& p) {
std::set<std::vector<T>> unique_indices;
@@ -4460,6 +4461,111 @@ TEST(scatter_nd_update_gpu, dynamic) {
}
}
TEST(scatter_nd_update_gpu, dynamic_5d) {
auto& engine = get_test_engine();
auto input1_layout = layout{{ 8, -1, -1, 384}, data_types::f32, format::bfyx };
auto input2_layout = layout{{-1, -1, -1, -1, -1}, data_types::i32, format::bfzyx };
auto input3_layout = layout{{-1, -1, -1, 384}, data_types::f32, format::bfyx };
topology topology;
topology.add(input_layout("data", input1_layout));
topology.add(input_layout("indices", input2_layout));
topology.add(input_layout("updates", input3_layout));
topology.add(scatter_nd_update("scatter_nd_update", input_info("data"), input_info("indices"), input_info("updates"), 5));
ExecutionConfig config;
config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
network network(engine, topology, config);
auto get_expected_res = [](const std::vector<float>& input,
const std::vector<int32_t>& indices,
const std::vector<float>& updates,
ov::Shape input_shape,
ov::Shape indices_shape,
ov::Shape updates_shape) -> std::vector<float> {
size_t count = std::accumulate(input_shape.begin(), input_shape.end(), 1, std::multiplies<size_t>());
auto outputs_ref = std::vector<float>(count);
ngraph::runtime::reference::scatterNdUpdate<float, int32_t>(input.data(),
indices.data(),
updates.data(),
outputs_ref.data(),
input_shape,
indices_shape,
updates_shape);
return outputs_ref;
};
auto generate_unique_indices = [](ov::Shape data_shape, ov::Shape indices_shape) -> std::vector<int32_t>{
std::set<std::vector<int32_t>> unique_indices;
std::vector<int32_t> result;
size_t last_indices_dim = indices_shape.at(indices_shape.size() - 1);
size_t count = std::accumulate(indices_shape.begin(), indices_shape.end(), 1, std::multiplies<size_t>()) / last_indices_dim;
while (unique_indices.size() != count) {
std::vector<int32_t> indices;
for (size_t i = 0; i < last_indices_dim; i++) {
indices.push_back(static_cast<int32_t>(generate_random_val<int>(0, data_shape[i] - 1)));
}
unique_indices.insert(indices);
}
std::for_each(unique_indices.begin(),
unique_indices.end(),
[&](const std::vector<int32_t>& indices) {
result.insert(result.end(), indices.begin(), indices.end());
});
return result;
};
std::vector<std::vector<ov::Shape>> test_shapes = {
{ { 8, 3, 1, 384 }, { 1, 3, 1, 384, 4 }, { 1, 3, 1, 384 } },
{ { 8, 3, 2, 384 }, { 1, 3, 1, 384, 4 }, { 1, 3, 1, 384 } },
};
for (auto& shapes : test_shapes) {
ov::Shape in1_shape = shapes[0];
ov::Shape in2_shape = shapes[1];
ov::Shape in3_shape = shapes[2];
auto input1 = engine.allocate_memory({ in1_shape, data_types::f32, format::bfyx }); // Dictionary
auto input2 = engine.allocate_memory({ in2_shape, data_types::i32, format::bfzyx }); // Indexes
auto input3 = engine.allocate_memory({ in3_shape, data_types::f32, format::bfyx }); // Updates
std::vector<float> input_data = generate_random_1d<float>(input1->count(), 1, 100);
std::vector<int32_t> indices = generate_unique_indices(in1_shape, in2_shape);
std::vector<float> updates = generate_random_1d<float>(input3->count(), 100, 200);
auto expected_res = get_expected_res(input_data, indices, updates, in1_shape, in2_shape, in3_shape);
set_values<float>(input1, input_data);
set_values<int32_t>(input2, indices);
set_values<float>(input3, updates);
network.set_input_data("data", input1);
network.set_input_data("indices", input2);
network.set_input_data("updates", input3);
auto inst = network.get_primitive("scatter_nd_update");
auto impl = inst->get_impl();
ASSERT_TRUE(impl != nullptr);
ASSERT_TRUE(impl->is_dynamic());
auto outputs = network.execute();
auto output = outputs.at("scatter_nd_update").get_memory();
ASSERT_EQ(output->get_layout().get_partial_shape(), input1->get_layout().get_partial_shape());
cldnn::mem_lock<float> output_ptr(output, get_test_stream());
for (size_t i = 0; i < expected_res.size(); ++i) {
ASSERT_EQ(expected_res[i], output_ptr[i]) << " i = " << i;
}
}
}
#ifdef RUN_ALL_MODEL_CACHING_TESTS
TEST_P(scatter_nd_update_random_test, random_cached)
{