From 37b39dfe3911bccaea078efedaa5d34107cc0c3a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 9 Oct 2025 23:02:16 -0700 Subject: [PATCH 1/5] Allow custom sizes, dim order and strides for tensor view. (#14944) Summary: The `make_tensor_ptr(TensrPtr)` overload creates a view on an existing `Tensor`. Here we provide a way for users to customize the shape, etc. so that they can easily do squeeze/unsqueeze and other convenient operations. Differential Revision: D84259597 (cherry picked from commit 4bdd3df3ebeb92601cbcd1e7bced254ee75cca86) --- extension/tensor/tensor_ptr.h | 81 +++++++-- extension/tensor/test/tensor_ptr_test.cpp | 204 +++++++++++++++++++++- 2 files changed, 263 insertions(+), 22 deletions(-) diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 4753ec296da..1c29e2b61f0 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -323,34 +323,77 @@ inline TensorPtr make_tensor_ptr( } /** - * Creates a TensorPtr to manage a new Tensor with the same properties - * as the given Tensor, sharing the same data without owning it. + * Creates a TensorPtr to manage a new Tensor that aliases the given Tensor's + * storage, with optional metadata overrides. Shape dynamism is inherited from + * the source tensor. * - * @param tensor The Tensor whose properties are used to create a new TensorPtr. - * @return A new TensorPtr managing a Tensor with the same properties as the - * original. + * If an override is provided (non-empty), it is passed as-is. If an override is + * empty, the corresponding metadata is reused from the source tensor when it + * fits; otherwise it is left empty for the core factory to derive a valid + * configuration. If `dim_order` is empty but `strides` is provided, `dim_order` + * is left empty so the core may infer it from the provided strides. + * + * @param tensor The source tensor to alias. + * @param sizes Optional sizes override. + * @param dim_order Optional dimension order override. + * @param strides Optional strides override. + * @return A TensorPtr aliasing the same storage with requested metadata. */ -inline TensorPtr make_tensor_ptr(const executorch::aten::Tensor& tensor) { +inline TensorPtr make_tensor_ptr( + const executorch::aten::Tensor& tensor, + std::vector sizes = {}, + std::vector dim_order = {}, + std::vector strides = {}) { + if (sizes.empty()) { + sizes.assign(tensor.sizes().begin(), tensor.sizes().end()); + } + const auto same_rank = sizes.size() == static_cast(tensor.dim()); + const auto same_shape = same_rank && + std::equal(sizes.begin(), sizes.end(), tensor.sizes().begin()); + const auto element_count = + executorch::aten::compute_numel(sizes.data(), sizes.size()); + const auto parent_element_count = tensor.numel(); + ET_CHECK_MSG( + element_count <= parent_element_count, + "Requested view has %zd elements, but source tensor only has %zd.", + static_cast(element_count), + static_cast(parent_element_count)); +#ifndef USE_ATEN_LIB + if (dim_order.empty() && strides.empty() && same_rank) { + dim_order.assign(tensor.dim_order().begin(), tensor.dim_order().end()); + } +#endif // USE_ATEN_LIB + if (strides.empty() && dim_order.empty() && same_shape) { + strides.assign(tensor.strides().begin(), tensor.strides().end()); + } return make_tensor_ptr( - std::vector( - tensor.sizes().begin(), tensor.sizes().end()), + std::move(sizes), tensor.mutable_data_ptr(), -#ifndef USE_ATEN_LIB - std::vector( - tensor.dim_order().begin(), tensor.dim_order().end()), - std::vector( - tensor.strides().begin(), tensor.strides().end()), - tensor.scalar_type(), - tensor.shape_dynamism() -#else // USE_ATEN_LIB - {}, - std::vector( - tensor.strides().begin(), tensor.strides().end()), + std::move(dim_order), + std::move(strides), tensor.scalar_type() #endif // USE_ATEN_LIB ); } +/** + * Convenience overload identical to make_tensor_ptr(*tensor_ptr, ...). + * + * @param tensor_ptr The source tensor pointer to alias. + * @param sizes Optional sizes override. + * @param dim_order Optional dimension order override. + * @param strides Optional strides override. + * @return A TensorPtr aliasing the same storage with requested metadata. + */ +inline TensorPtr make_tensor_ptr( + const TensorPtr& tensor_ptr, + std::vector sizes = {}, + std::vector dim_order = {}, + std::vector strides = {}) { + return make_tensor_ptr( + *tensor_ptr, std::move(sizes), std::move(dim_order), std::move(strides)); +} + /** * Creates a TensorPtr that manages a new Tensor with the same properties * as the given Tensor, but with a copy of the data owned by the returned diff --git a/extension/tensor/test/tensor_ptr_test.cpp b/extension/tensor/test/tensor_ptr_test.cpp index 6c98db52d41..b4da2a86a58 100644 --- a/extension/tensor/test/tensor_ptr_test.cpp +++ b/extension/tensor/test/tensor_ptr_test.cpp @@ -357,6 +357,204 @@ TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorInt32) { EXPECT_EQ(new_tensor->scalar_type(), executorch::aten::ScalarType::Int); } +TEST_F(TensorPtrTest, MakeViewOverrideSizesRankIncrease) { + std::vector data = {1, 2, 3, 4, 5, 6}; + auto tensor = make_tensor_ptr({2, 3}, std::move(data)); + auto view = make_tensor_ptr(tensor, {1, 2, 3}); + + EXPECT_EQ(view->dim(), 3); + EXPECT_EQ(view->size(0), 1); + EXPECT_EQ(view->size(1), 2); + EXPECT_EQ(view->size(2), 3); + EXPECT_EQ(view->const_data_ptr(), tensor->const_data_ptr()); + EXPECT_EQ(view->strides()[0], 6); + EXPECT_EQ(view->strides()[1], 3); + EXPECT_EQ(view->strides()[2], 1); +} + +TEST_F(TensorPtrTest, MakeViewOverrideSizesSameRankRecomputesStrides) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data); + auto view = make_tensor_ptr(tensor, {4, 3}); + + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 4); + EXPECT_EQ(view->size(1), 3); + EXPECT_EQ(view->strides()[0], 3); + EXPECT_EQ(view->strides()[1], 1); +} + +TEST_F(TensorPtrTest, MakeViewOverrideDimOrderOnly) { + float data[6] = {0}; + auto tensor = make_tensor_ptr({2, 3}, data); + auto view = make_tensor_ptr(tensor, {}, {1, 0}, {}); + + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 2); + EXPECT_EQ(view->size(1), 3); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_EQ(view->strides()[1], 2); +} + +TEST_F(TensorPtrTest, MakeViewOverrideStridesOnlyInfersDimOrder) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data); + auto view = make_tensor_ptr(tensor, {}, {}, {1, 3}); + + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 3); + EXPECT_EQ(view->size(1), 4); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_EQ(view->strides()[1], 3); +} + +TEST_F(TensorPtrTest, MakeViewReuseMetadataWhenShapeSame) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data, {1, 0}, {1, 3}); + auto view = make_tensor_ptr(tensor, {3, 4}); + + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 3); + EXPECT_EQ(view->size(1), 4); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_EQ(view->strides()[1], 3); +} + +TEST_F(TensorPtrTest, MakeViewShapeChangeWithExplicitOldStridesExpectDeath) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data); + std::vector old_strides( + tensor->strides().begin(), tensor->strides().end()); + + ET_EXPECT_DEATH( + { auto _ = make_tensor_ptr(tensor, {2, 6}, {}, old_strides); }, ""); +} + +TEST_F(TensorPtrTest, MakeViewInvalidDimOrderExpectDeath) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data); + + ET_EXPECT_DEATH( + { auto _ = make_tensor_ptr(tensor, {3, 4}, {2, 1}, {1, 4}); }, ""); +} + +TEST_F(TensorPtrTest, MakeViewFromTensorPtrConvenienceOverload) { + float data[12] = {0}; + auto tensor = make_tensor_ptr({3, 4}, data); + auto view = make_tensor_ptr(tensor, {}, {1, 0}, {}); + + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 3); + EXPECT_EQ(view->size(1), 4); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_EQ(view->strides()[1], 3); +} + +TEST_F(TensorPtrTest, MakeViewRankDecreaseFlatten) { + float data[6] = {1, 2, 3, 4, 5, 6}; + auto tensor = make_tensor_ptr( + {2, 3}, + data, + {}, + {}, + executorch::aten::ScalarType::Float, + executorch::aten::TensorShapeDynamism::DYNAMIC_UNBOUND); + auto view = make_tensor_ptr(tensor, {6}); + EXPECT_EQ(view->dim(), 1); + EXPECT_EQ(view->size(0), 6); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_NE(tensor->unsafeGetTensorImpl(), view->unsafeGetTensorImpl()); + EXPECT_EQ(resize_tensor_ptr(view, {3, 2}), Error::NotSupported); + EXPECT_EQ(view->dim(), 1); + EXPECT_EQ(view->size(0), 6); +} + +TEST_F(TensorPtrTest, MakeViewFromScalarAliasAnd1D) { + float scalar_value = 7.f; + auto tensor = make_tensor_ptr({}, &scalar_value); + auto alias = make_tensor_ptr(tensor); + EXPECT_EQ(alias->dim(), 0); + EXPECT_EQ(alias->numel(), 1); + auto reshaped = make_tensor_ptr(tensor, {1}); + EXPECT_EQ(reshaped->dim(), 1); + EXPECT_EQ(reshaped->size(0), 1); + EXPECT_EQ(reshaped->strides()[0], 1); + ET_EXPECT_DEATH({ auto unused = make_tensor_ptr(tensor, {}, {0}, {}); }, ""); + ET_EXPECT_DEATH({ auto unused = make_tensor_ptr(tensor, {}, {}, {1}); }, ""); +} + +TEST_F(TensorPtrTest, MakeViewExplicitDimOrderAndStridesShapeChange) { + float data[6] = {0}; + auto tensor = make_tensor_ptr({2, 3}, data); + auto view = make_tensor_ptr(tensor, {3, 2}, {1, 0}, {1, 3}); + EXPECT_EQ(view->dim(), 2); + EXPECT_EQ(view->size(0), 3); + EXPECT_EQ(view->size(1), 2); + EXPECT_EQ(view->strides()[0], 1); + EXPECT_EQ(view->strides()[1], 3); +} + +TEST_F(TensorPtrTest, TensorUint8dataInt16Type) { + std::vector int16_values = {-1, 2, -3, 4}; + auto byte_pointer = reinterpret_cast(int16_values.data()); + std::vector byte_data( + byte_pointer, byte_pointer + int16_values.size() * sizeof(int16_t)); + auto tensor = make_tensor_ptr( + {4}, std::move(byte_data), executorch::aten::ScalarType::Short); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 4); + auto int16_data = tensor->const_data_ptr(); + EXPECT_EQ(int16_data[0], -1); + EXPECT_EQ(int16_data[1], 2); + EXPECT_EQ(int16_data[2], -3); + EXPECT_EQ(int16_data[3], 4); +} + +TEST_F(TensorPtrTest, MakeView3DDimOrderOnly) { + float data[24] = {0}; + auto tensor = make_tensor_ptr({2, 3, 4}, data); + auto view = make_tensor_ptr(tensor, {}, {2, 0, 1}, {}); + EXPECT_EQ(view->dim(), 3); + EXPECT_EQ(view->size(0), 2); + EXPECT_EQ(view->size(1), 3); + EXPECT_EQ(view->size(2), 4); + EXPECT_EQ(view->strides()[0], 3); + EXPECT_EQ(view->strides()[1], 1); + EXPECT_EQ(view->strides()[2], 6); +} + +#ifndef USE_ATEN_LIB +TEST_F(TensorPtrTest, MakeViewDynamismPropagationResizeAlias) { + float data[12] = {0}; + auto tensor = make_tensor_ptr( + {3, 4}, + data, + {}, + {}, + executorch::aten::ScalarType::Float, + executorch::aten::TensorShapeDynamism::DYNAMIC_UNBOUND); + auto alias = make_tensor_ptr(tensor); + EXPECT_EQ(resize_tensor_ptr(alias, {2, 6}), Error::Ok); + EXPECT_EQ(alias->size(0), 2); + EXPECT_EQ(alias->size(1), 6); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->size(1), 4); +} + +TEST_F(TensorPtrTest, MakeViewSameRankShapeChangeCopiesDimOrder) { + float data[24] = {0}; + auto tensor = make_tensor_ptr({2, 3, 4}, data, {2, 0, 1}, {3, 1, 6}); + auto view = make_tensor_ptr(tensor, {4, 2, 3}); + EXPECT_EQ(view->dim(), 3); + EXPECT_EQ(view->size(0), 4); + EXPECT_EQ(view->size(1), 2); + EXPECT_EQ(view->size(2), 3); + EXPECT_EQ(view->strides()[0], 2); + EXPECT_EQ(view->strides()[1], 1); + EXPECT_EQ(view->strides()[2], 8); +} +#endif + TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorInt32) { std::vector data = {1, 2, 3, 4}; auto tensor = make_tensor_ptr({2, 2}, std::move(data)); @@ -753,7 +951,7 @@ TEST_F(TensorPtrTest, TensorDeducedScalarType) { EXPECT_EQ(tensor->const_data_ptr()[3], 4.0); } -TEST_F(TensorPtrTest, TensorUint8BufferWithFloatScalarType) { +TEST_F(TensorPtrTest, TensorUint8dataWithFloatScalarType) { std::vector data( 4 * executorch::aten::elementSize(executorch::aten::ScalarType::Float)); @@ -777,14 +975,14 @@ TEST_F(TensorPtrTest, TensorUint8BufferWithFloatScalarType) { EXPECT_EQ(tensor->const_data_ptr()[3], 4.0f); } -TEST_F(TensorPtrTest, TensorUint8BufferTooSmallExpectDeath) { +TEST_F(TensorPtrTest, TensorUint8dataTooSmallExpectDeath) { std::vector data( 2 * executorch::aten::elementSize(executorch::aten::ScalarType::Float)); ET_EXPECT_DEATH( { auto tensor = make_tensor_ptr({2, 2}, std::move(data)); }, ""); } -TEST_F(TensorPtrTest, TensorUint8BufferTooLargeExpectDeath) { +TEST_F(TensorPtrTest, TensorUint8dataTooLargeExpectDeath) { std::vector data( 5 * executorch::aten::elementSize(executorch::aten::ScalarType::Float)); ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 2}, std::move(data)); }, ""); From 4638a1182400c3b8591246f57e2f240e7eb546a7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 13 Oct 2025 13:17:47 -0700 Subject: [PATCH 2/5] Simplify make_tensor_ptr call by removing sizes --- extension/tensor/tensor_ptr.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 377c672f981..f0f586ffb56 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -272,8 +272,7 @@ inline TensorPtr make_tensor_ptr( */ template inline TensorPtr make_tensor_ptr(T value) { - return make_tensor_ptr( - std::vector{}, std::vector{value}); + return make_tensor_ptr({}, std::vector{value}); } /** @@ -373,6 +372,9 @@ inline TensorPtr make_tensor_ptr( std::move(dim_order), std::move(strides), tensor.scalar_type() +#ifndef USE_ATEN_LIB + , + tensor.shape_dynamism() #endif // USE_ATEN_LIB ); } From b0b58bbfa81b08fe93d5ceb75b8b3ad73192c22d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 13 Oct 2025 13:24:57 -0700 Subject: [PATCH 3/5] Update tensor_ptr.h --- extension/tensor/tensor_ptr.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index f0f586ffb56..5bdd922e403 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -272,7 +272,9 @@ inline TensorPtr make_tensor_ptr( */ template inline TensorPtr make_tensor_ptr(T value) { - return make_tensor_ptr({}, std::vector{value}); + return make_tensor_ptr( + std::vector{}, std::vector{value}); + } } /** From 0ccf9fc9c8bb4bd0372bdaa56a6bb6fd76451a56 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 13 Oct 2025 13:25:34 -0700 Subject: [PATCH 4/5] Fix formatting in make_tensor_ptr function --- extension/tensor/tensor_ptr.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 5bdd922e403..900252109d3 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -273,8 +273,7 @@ inline TensorPtr make_tensor_ptr( template inline TensorPtr make_tensor_ptr(T value) { return make_tensor_ptr( - std::vector{}, std::vector{value}); - } + std::vector{}, std::vector{value}); } /** From 4bd73c93434fafe8ea96046a35becf0c4eb329ba Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 13 Oct 2025 16:13:14 -0700 Subject: [PATCH 5/5] Modify make_tensor_ptr to accept empty SizesType --- extension/tensor/tensor_ptr.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index a659b02684c..d8fad857cd2 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -273,7 +273,8 @@ inline TensorPtr make_tensor_ptr( */ template inline TensorPtr make_tensor_ptr(T value) { - return make_tensor_ptr({}, std::vector{value}); + return make_tensor_ptr( + std::vector{}, std::vector{value}); } /**