Skip to content

Commit

Permalink
Fixed Tensor::set_shape to rely on capacity
Browse files Browse the repository at this point in the history
  • Loading branch information
@ilya-lavrenov
ilya-lavrenov committed Dec 1, 2023
1 parent abfbdd1 commit 6c1092b
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 26 deletions.
49 changes: 36 additions & 13 deletions src/core/tests/ov_tensor_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -369,25 +369,26 @@ TEST_F(OVTensorTest, saveDimsAndSizeAfterMoveStringTensor) {
ASSERT_THROW(t.data<std::string>(), ov::Exception);
}

// SetShape
// set_shape
TEST_F(OVTensorTest, canSetShape) {
const ov::Shape origShape({1, 2, 3});
ov::Tensor t{ov::element::f32, {1, 2, 3}};
const ov::Shape newShape({4, 5, 6});
ov::Tensor t{ov::element::f32, origShape};
const ov::Shape newShape({4, 5, 6}), newShape2({4, 5, 6, 7});

const void* orig_data = t.data();
ASSERT_EQ(t.get_shape(), origShape);
ASSERT_NO_THROW(t.set_shape({4, 5, 6}));
ASSERT_NO_THROW(t.set_shape(newShape));
ASSERT_EQ(newShape, t.get_shape());
ASSERT_EQ(byteStrides(ov::row_major_strides(newShape), t.get_element_type()), t.get_strides());
ASSERT_NE(orig_data, t.data());

// check that setShape for copy changes original Tensor
// check that set_shape for copy changes original Tensor
{
ov::Tensor t2 = t;
ASSERT_NO_THROW(t2.set_shape(newShape));
ASSERT_EQ(newShape, t.get_shape());
ASSERT_NO_THROW(t2.set_shape(newShape2));
ASSERT_EQ(newShape2, t.get_shape());
ASSERT_EQ(t2.get_shape(), t.get_shape());
ASSERT_EQ(t2.data(), t.data());
orig_data = t.data();
}

Expand All @@ -402,30 +403,30 @@ TEST_F(OVTensorTest, canSetShape) {
TEST_F(OVTensorTest, canSetShapeStringTensor) {
const ov::Shape origShape({1, 2, 3});
ov::Tensor t{ov::element::string, {1, 2, 3}};
const ov::Shape newShape({4, 5, 6});
const ov::Shape newShape({4, 5, 6}), newShape2({4, 5, 6, 7});

const void* orig_data = t.data();
ASSERT_EQ(t.get_shape(), origShape);
ASSERT_NO_THROW(t.set_shape(newShape));
ASSERT_EQ(newShape, t.get_shape());
ASSERT_EQ(byteStrides(ov::row_major_strides(newShape), t.get_element_type()), t.get_strides());
ASSERT_NE(orig_data, t.data());
const void* new_data = t.data();

// check that setShape for copy changes original Tensor
{
ov::Tensor t2 = t;
ASSERT_NO_THROW(t2.set_shape(origShape));
ASSERT_EQ(origShape, t2.get_shape());
ASSERT_EQ(origShape, t.get_shape());
ASSERT_NO_THROW(t2.set_shape(newShape2));
ASSERT_EQ(newShape2, t2.get_shape());
ASSERT_EQ(t2.get_shape(), t.get_shape());
ASSERT_EQ(t2.data(), t.data());
orig_data = t.data();
}

// set_shape for smaller memory - does not perform reallocation
{
ASSERT_NO_THROW(t.set_shape(origShape));
ASSERT_EQ(origShape, t.get_shape());
ASSERT_EQ(new_data, t.data());
ASSERT_EQ(orig_data, t.data());
}
}

Expand Down Expand Up @@ -497,6 +498,28 @@ TEST_F(OVTensorTest, canSetShapeOfOriginalSizeAfterDecreasingOnPreallocatedMemor
ASSERT_NO_THROW(t.set_shape(originalShape));
}

TEST_F(OVTensorTest, canSetShapeOfOriginalSizeAfterDecreasing) {
const ov::Shape shape({4, 5, 6}), small_shape({1, 2, 3});
ov::Tensor t{ov::element::f32, shape};
void* data = t.data();

ASSERT_NO_THROW(t.set_shape(small_shape));
EXPECT_EQ(data, t.data());
ASSERT_NO_THROW(t.set_shape(shape));
EXPECT_EQ(data, t.data());
}

TEST_F(OVTensorTest, canSetShapeOfOriginalSizeAfterDecreasingStringTensor) {
const ov::Shape shape({4, 5, 6}), small_shape({1, 2, 3});
ov::Tensor t{ov::element::string, shape};
void* data = t.data();

ASSERT_NO_THROW(t.set_shape(small_shape));
EXPECT_EQ(data, t.data());
ASSERT_NO_THROW(t.set_shape(shape));
EXPECT_EQ(data, t.data());
}

TEST_F(OVTensorTest, canChangeShapeOnStridedTensor) {
float data[64 * 4];
ov::Tensor t{ov::element::f32, {4, 2, 2}, data, {64, 16, 4}};
Expand Down
22 changes: 9 additions & 13 deletions src/inference/src/dev/make_tensor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,43 +182,39 @@ class AllocatedTensor : public ViewTensor {
OPENVINO_ASSERT(allocator, "Allocator was not initialized");
auto num_elements = shape_size(shape);
auto data = const_cast<Allocator&>(allocator).allocate(element_type.size() * num_elements);
init(data, element_type, shape);
initialize_elements(data, element_type, shape);
return data;
}()},
m_allocator{allocator} {}

~AllocatedTensor() {
auto num_elements = get_size();
destroy(0, num_elements);
destroy_elements(0, num_elements);
m_allocator.deallocate(m_ptr, get_byte_size());
}

void set_shape(ov::Shape new_shape) override {
if (m_shape == new_shape)
return;

auto old_num_elements = get_size();
auto old_byte_size = get_byte_size();
auto current_capacity = shape_size(m_capacity), current_num_elements = get_size(), current_byte_size = get_byte_size();
m_shape = std::move(new_shape);
auto new_num_elements = get_size();

if (get_byte_size() > old_byte_size) {
if (new_num_elements > current_capacity) {
// allocate buffer and initialize objects from scratch
destroy(0, old_num_elements);
m_allocator.deallocate(m_ptr, old_byte_size);
destroy_elements(0, current_num_elements);
m_allocator.deallocate(m_ptr, current_byte_size);
m_ptr = m_allocator.allocate(get_byte_size());
init(m_ptr, m_element_type, m_shape);
} else {
// destroy only not needed objects
destroy(new_num_elements, old_num_elements);
initialize_elements(m_ptr, m_element_type, m_shape);
}

m_strides.clear();
update_strides();
}

private:
void destroy(size_t begin_ind, size_t end_ind) {
void destroy_elements(size_t begin_ind, size_t end_ind) {
// it removes elements from tail
if (get_element_type() == element::Type_t::string) {
auto strings = static_cast<std::string*>(m_ptr);
Expand All @@ -229,7 +225,7 @@ class AllocatedTensor : public ViewTensor {
}
}

void init(void* data, const element::Type& element_type, const Shape& shape) {
void initialize_elements(void* data, const element::Type& element_type, const Shape& shape) {
if (element_type == element::Type_t::string) {
auto num_elements = shape_size(shape);
auto string_ptr = static_cast<std::string*>(data);
Expand Down

0 comments on commit 6c1092b

Please sign in to comment.