[GPU] Adjust in/out layouts in codition prim to handle prealloc (#20740)

This commit is contained in:
Vladimir Paramuzov
2023-10-27 16:54:57 +00:00
committed by GitHub
parent 79f90839b9
commit cde757d66a
3 changed files with 118 additions and 2 deletions
@@ -258,6 +258,14 @@ void condition_inst::postprocess_output_memory(network::ptr executed_net, cldnn:
auto out_mem_idx = out_mem_map.first;
auto inner_out_id = out_mem_map.second;
auto mem_ptr = executed_net->get_output(inner_out_id).get_memory();
if (mem_ptr) {
auto layout = _impl_params->get_output_layout(out_mem_idx);
GPU_DEBUG_LOG << "Reshape output from " << mem_ptr->get_layout().to_short_string()
<< " to " << layout.to_short_string() << std::endl;
// Preallocation logic may allocate more memory than actually produced on current iteration, so we need to adjust output buffers layout
mem_ptr = get_network().get_engine().reinterpret_buffer(*mem_ptr, layout);
}
_outputs[out_mem_idx] = mem_ptr;
GPU_DEBUG_LOG << "Inner net - Outputs[" << out_mem_idx << "]" << mem_ptr->get_layout().to_short_string() << std::endl;
}
@@ -48,6 +48,14 @@ struct condition_impl : typed_primitive_impl<condition> {
if (iter != branch.input_map.end()) {
const primitive_id& input_internal_id = iter->second;
auto mem_ptr = instance.input_memory_ptr(mem_idx);
if (mem_ptr) {
auto dep = instance.dependencies()[mem_idx];
auto layout = dep.first->get_impl_params()->get_output_layout(dep.second);
GPU_DEBUG_LOG << "Reshape input from " << mem_ptr->get_layout().to_short_string()
<< " to " << layout.to_short_string() << std::endl;
// Preallocation logic may allocate more memory than actually produced on current iteration, so we need to adjust input buffers layout
mem_ptr = instance.get_network().get_engine().reinterpret_buffer(*mem_ptr, layout);
}
executed_net->set_input_data(input_internal_id, mem_ptr);
GPU_DEBUG_LOG << "Inner net - Inputs[" << mem_idx << "]" << mem_ptr->get_layout().to_short_string() << std::endl;
}
@@ -2,6 +2,9 @@
// SPDX-License-Identifier: Apache-2.0
//
#include "intel_gpu/primitives/permute.hpp"
#include "intel_gpu/runtime/internal_properties.hpp"
#include "random_generator.hpp"
#include "test_utils.h"
#include <intel_gpu/primitives/input_layout.hpp>
@@ -19,8 +22,9 @@ using namespace ::tests;
namespace {
template <class T>
bool is_output_equal(const cldnn::memory::ptr mem, const std::vector<T>& ref)
{
bool is_output_equal(const cldnn::memory::ptr mem, const std::vector<T>& ref) {
if (mem->count() != ref.size())
return false;
cldnn::mem_lock<T> ptr(mem, get_test_stream());
for (size_t i = 0; i < mem->get_layout().count(); i++) {
if (!are_equal(ptr[i], ref[i])) return false;
@@ -239,6 +243,102 @@ TEST(condition_gpu, basic_range_equal_comp) {
ASSERT_TRUE(is_output_equal(out_data_false, pooling_when_false_data));
}
TEST(condition_gpu, dynamic_shapes) {
auto& engine = get_test_engine();
ExecutionConfig config = get_test_default_config(engine);
config.set_property(ov::intel_gpu::optimize_data(true));
config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
const int64_t d1 = 2;
const int64_t d2 = 4;
layout input_lay = {{-1, d1, -1, d2}, data_types::f32, format::bfyx};
auto predicate = engine.allocate_memory({{ 1 }, data_types::u8, format::bfyx });
const primitive_id condition_id = "condition";
const primitive_id condition_id_true = condition_id + "_when_true";
const primitive_id condition_id_false = condition_id + "_when_false";
const primitive_id branch_input_id = "branch_input";
const primitive_id model_input = "input";
const primitive_id predicate_input = "predicate";
const primitive_id tranpose = "transpose";
cldnn::topology topology;
topology.add(input_layout(model_input, input_lay));
topology.add(input_layout(predicate_input, predicate->get_layout()));
topology.add(permute(tranpose, model_input, {1, 0, 2, 3}));
const float shift = 4.f;
auto generate_simple_branch = [&](bool branch_true_false, const primitive_id& input_id, const data_types dt) {
auto mem = engine.allocate_memory(layout{{d1, 1, 1, d2}, dt, format::bfyx});
{
cldnn::mem_lock<float> l(mem, get_test_stream());
for (size_t i = 0; i < mem->count(); i++) {
l.data()[i] = shift;
}
}
primitive_id const_id = "const_input";
eltwise_mode mode = branch_true_false ? eltwise_mode::sum : eltwise_mode::sub;
auto id = branch_true_false ? condition_id_true : condition_id_false;
cldnn::topology branch_topology(input_layout(input_id, { {d1, -1, -1, d2}, dt, format::bfyx }),
data(const_id, mem),
eltwise(id, {input_id, const_id}, mode)
);
condition::branch branch;
branch.inner_program = program::build_program(engine, branch_topology, config, false, false, true);
branch.input_map.insert({tranpose, branch_input_id});
branch.output_map.insert({0, id});
return branch;
};
condition::branch branch_true = generate_simple_branch(true, branch_input_id, data_types::f32);
condition::branch branch_false = generate_simple_branch(false, branch_input_id, data_types::f32);
topology.add(condition(condition_id, { input_info(predicate_input), tranpose }, branch_true, branch_false));
tests::random_generator rg(GET_SUITE_NAME);
std::vector<uint8_t> predicate_data_true = { 1 };
std::vector<uint8_t> predicate_data_false = { 0 };
network net(engine, topology, config);
auto check_output = [](const cldnn::memory::ptr mem, const std::vector<float>& ref, ov::Shape expected_shape) {
ASSERT_EQ(mem->get_layout().get_shape(), expected_shape);
ASSERT_EQ(mem->count(), ref.size());
cldnn::mem_lock<float> ptr(mem, get_test_stream());
for (size_t i = 0; i < mem->get_layout().count(); i++) {
ASSERT_EQ(ptr[i], ref[i]) << "i = " << i;
}
};
for (size_t i = 0; i < 10; i++) {
layout l = {{1, d1, 1 + static_cast<int64_t>(i), d2}, data_types::f32, format::bfyx};
std::vector<float> input_data = rg.generate_random_1d<float>(l.count(), -10, 10);
auto mem = engine.allocate_memory(l);
std::vector<float> expected_result_when_true = input_data;
std::vector<float> expected_result_when_false = input_data;
for (size_t i = 0; i < input_data.size(); i++) {
expected_result_when_true[i] += shift;
expected_result_when_false[i] -= shift;
}
set_values(mem, input_data);
set_values(predicate, predicate_data_true);
net.set_input_data(model_input, mem);
net.set_input_data(predicate_input, predicate);
auto outputs = net.execute();
check_output(outputs.at(condition_id).get_memory(), expected_result_when_true, {d1, 1, 1+i, d2});
set_values(predicate, predicate_data_false);
net.set_input_data(model_input, mem);
net.set_input_data(predicate_input, predicate);
outputs = net.execute();
check_output(outputs.at(condition_id).get_memory(), expected_result_when_false, {d1, 1, 1+i, d2});
}
}
TEST(condition_gpu, basic_stacked_ifs) {
/*
<prims...>