From 18f97b2dd11f8a4e31570971c4cccd96f77f3436 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 15 Sep 2024 11:11:45 -0700 Subject: [PATCH] Add extra checks for the provided dim order and strides. Summary: Also make sure the scalar tensors stay static. Differential Revision: D62681004 --- extension/tensor/tensor_impl_ptr.cpp | 9 +++- extension/tensor/tensor_impl_ptr.h | 16 ++------ extension/tensor/tensor_ptr.h | 14 ++++++- extension/tensor/tensor_ptr_maker.h | 6 +-- .../tensor/test/tensor_impl_ptr_test.cpp | 16 ++++++++ extension/tensor/test/tensor_ptr_test.cpp | 41 ++++++++++++++++--- 6 files changed, 77 insertions(+), 25 deletions(-) diff --git a/extension/tensor/tensor_impl_ptr.cpp b/extension/tensor/tensor_impl_ptr.cpp index 01f69095b23..cbc49299b9f 100644 --- a/extension/tensor/tensor_impl_ptr.cpp +++ b/extension/tensor/tensor_impl_ptr.cpp @@ -62,6 +62,13 @@ TensorImplPtr make_tensor_impl_ptr( exec_aten::TensorShapeDynamism dynamism, std::function deleter) { const auto dim = sizes.size(); + ET_CHECK_MSG( + dim_order.empty() || dim_order.size() == dim, + "dim_order size must match sizes or be empty."); + ET_CHECK_MSG( + strides.empty() || strides.size() == dim, + "strides size must match sizes or be empty."); + if (dim_order.empty()) { dim_order.resize(dim); std::iota(dim_order.begin(), dim_order.end(), 0); @@ -89,7 +96,7 @@ TensorImplPtr make_tensor_impl_ptr( data, dim_order.data(), strides.data(), - dynamism); + dim > 0 ? dynamism : exec_aten::TensorShapeDynamism::STATIC); return TensorImplPtr( tensor_impl.release(), TensorImplPtrDeleter{ diff --git a/extension/tensor/tensor_impl_ptr.h b/extension/tensor/tensor_impl_ptr.h index 5f34f929b96..83bf534d73d 100644 --- a/extension/tensor/tensor_impl_ptr.h +++ b/extension/tensor/tensor_impl_ptr.h @@ -87,7 +87,7 @@ TensorImplPtr make_tensor_impl_ptr( * @return A TensorImplPtr that manages the newly created TensorImpl. */ template -TensorImplPtr make_tensor_impl_ptr( +inline TensorImplPtr make_tensor_impl_ptr( std::vector sizes, std::vector data, std::vector dim_order = {}, @@ -123,23 +123,13 @@ TensorImplPtr make_tensor_impl_ptr( * @return A TensorImplPtr that manages the newly created TensorImpl. */ template -TensorImplPtr make_tensor_impl_ptr( +inline TensorImplPtr make_tensor_impl_ptr( std::vector data, exec_aten::TensorShapeDynamism dynamism = exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) { - constexpr exec_aten::ScalarType scalar_type = - runtime::CppTypeToScalarType::value; std::vector sizes{exec_aten::SizesType(data.size())}; - const auto raw_data_ptr = data.data(); - auto data_ptr = std::make_shared>(std::move(data)); return make_tensor_impl_ptr( - scalar_type, - std::move(sizes), - raw_data_ptr, - {0}, - {1}, - dynamism, - [data_ptr = std::move(data_ptr)](void*) {}); + std::move(sizes), std::move(data), {0}, {1}, dynamism); } /** diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 2780a3c4076..17e18742bec 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -235,7 +235,19 @@ TensorPtr make_tensor_ptr( std::initializer_list data, exec_aten::TensorShapeDynamism dynamism = exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) { - return make_tensor_ptr(std::vector(data), dynamism); + return make_tensor_ptr(make_tensor_impl_ptr(std::vector(data), dynamism)); +} + +/** + * Creates a TensorPtr that manages a Tensor with a single scalar value. + * + * @tparam T The C++ type of the scalar value. + * @param value The scalar value to be used for the Tensor. + * @return A TensorPtr that manages the newly created TensorImpl. + */ +template +TensorPtr make_tensor_ptr(T value) { + return make_tensor_ptr(make_tensor_impl_ptr({}, std::vector{value})); } /** diff --git a/extension/tensor/tensor_ptr_maker.h b/extension/tensor/tensor_ptr_maker.h index 4e65480b7fd..8146135b153 100644 --- a/extension/tensor/tensor_ptr_maker.h +++ b/extension/tensor/tensor_ptr_maker.h @@ -421,10 +421,8 @@ inline TensorPtr full( */ inline TensorPtr scalar_tensor( exec_aten::Scalar value, - exec_aten::ScalarType type = exec_aten::ScalarType::Float, - exec_aten::TensorShapeDynamism dynamism = - exec_aten::TensorShapeDynamism::DYNAMIC_BOUND) { - return full({}, value, type, dynamism); + exec_aten::ScalarType type = exec_aten::ScalarType::Float) { + return full({}, value, type); } /** diff --git a/extension/tensor/test/tensor_impl_ptr_test.cpp b/extension/tensor/test/tensor_impl_ptr_test.cpp index f7fd062c462..642e8738705 100644 --- a/extension/tensor/test/tensor_impl_ptr_test.cpp +++ b/extension/tensor/test/tensor_impl_ptr_test.cpp @@ -377,3 +377,19 @@ TEST_F(TensorImplPtrTest, TensorImplUint8BufferTooLarge) { EXPECT_EQ(tensor_impl->strides()[0], 2); EXPECT_EQ(tensor_impl->strides()[1], 1); } + +TEST_F(TensorImplPtrTest, StridesAndDimOrderMustMatchSizes) { + float data[12] = {0}; + ET_EXPECT_DEATH( + { + auto _ = make_tensor_impl_ptr( + exec_aten::ScalarType::Float, {3, 4}, data, {}, {1}); + }, + ""); + ET_EXPECT_DEATH( + { + auto _ = make_tensor_impl_ptr( + exec_aten::ScalarType::Float, {3, 4}, data, {0}, {4, 1}); + }, + ""); +} diff --git a/extension/tensor/test/tensor_ptr_test.cpp b/extension/tensor/test/tensor_ptr_test.cpp index 96bec15cefa..7fabf9ab8f9 100644 --- a/extension/tensor/test/tensor_ptr_test.cpp +++ b/extension/tensor/test/tensor_ptr_test.cpp @@ -44,6 +44,40 @@ TEST_F(TensorPtrTest, ScalarTensorOwningData) { EXPECT_EQ(tensor->const_data_ptr()[0], 3.14f); } +TEST_F(TensorPtrTest, ScalarTensorSingleValueCreation) { + auto tensor_float = make_tensor_ptr(3.14f); + EXPECT_EQ(tensor_float->dim(), 0); + EXPECT_EQ(tensor_float->numel(), 1); + EXPECT_EQ(tensor_float->sizes().size(), 0); + EXPECT_EQ(tensor_float->strides().size(), 0); + EXPECT_EQ(tensor_float->const_data_ptr()[0], 3.14f); + EXPECT_EQ(tensor_float->scalar_type(), exec_aten::ScalarType::Float); + + auto tensor_int32 = make_tensor_ptr(42); + EXPECT_EQ(tensor_int32->dim(), 0); + EXPECT_EQ(tensor_int32->numel(), 1); + EXPECT_EQ(tensor_int32->sizes().size(), 0); + EXPECT_EQ(tensor_int32->strides().size(), 0); + EXPECT_EQ(tensor_int32->const_data_ptr()[0], 42); + EXPECT_EQ(tensor_int32->scalar_type(), exec_aten::ScalarType::Int); + + auto tensor_double = make_tensor_ptr(2.718); + EXPECT_EQ(tensor_double->dim(), 0); + EXPECT_EQ(tensor_double->numel(), 1); + EXPECT_EQ(tensor_double->sizes().size(), 0); + EXPECT_EQ(tensor_double->strides().size(), 0); + EXPECT_EQ(tensor_double->const_data_ptr()[0], 2.718); + EXPECT_EQ(tensor_double->scalar_type(), exec_aten::ScalarType::Double); + + auto tensor_int64 = make_tensor_ptr(static_cast(10000000000)); + EXPECT_EQ(tensor_int64->dim(), 0); + EXPECT_EQ(tensor_int64->numel(), 1); + EXPECT_EQ(tensor_int64->sizes().size(), 0); + EXPECT_EQ(tensor_int64->strides().size(), 0); + EXPECT_EQ(tensor_int64->const_data_ptr()[0], 10000000000); + EXPECT_EQ(tensor_int64->scalar_type(), exec_aten::ScalarType::Long); +} + TEST_F(TensorPtrTest, CreateTensorWithStridesAndDimOrder) { float data[20] = {2}; auto tensor = make_tensor_ptr( @@ -299,12 +333,7 @@ TEST_F(TensorPtrTest, TensorSharingImplModifiesSharedDataVector) { TEST_F(TensorPtrTest, TensorSharingImplResizingAffectsBothVector) { std::vector data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - auto tensor1 = make_tensor_ptr( - {3, 4}, - std::move(data), - {}, - {}, - exec_aten::TensorShapeDynamism::DYNAMIC_UNBOUND); + auto tensor1 = make_tensor_ptr({3, 4}, std::move(data)); auto tensor2 = make_tensor_ptr(tensor1); EXPECT_EQ(resize_tensor_ptr(tensor1, {2, 6}), Error::Ok);