[GPU] Adjust in/out layouts in codition prim to handle prealloc (#20740)
This commit is contained in:
@@ -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...>
|
||||
|
||||
Reference in New Issue
Block a user