diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 27e2e3451ce..f0f586ffb56 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -323,26 +323,54 @@ 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) { - return make_tensor_ptr( - std::vector( - tensor.sizes().begin(), tensor.sizes().end()), - tensor.mutable_data_ptr(), +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 - std::vector( - tensor.dim_order().begin(), tensor.dim_order().end()), -#else // 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 - std::vector( - tensor.strides().begin(), tensor.strides().end()), + if (strides.empty() && dim_order.empty() && same_shape) { + strides.assign(tensor.strides().begin(), tensor.strides().end()); + } + return make_tensor_ptr( + std::move(sizes), + tensor.mutable_data_ptr(), + std::move(dim_order), + std::move(strides), tensor.scalar_type() #ifndef USE_ATEN_LIB , @@ -352,21 +380,21 @@ inline TensorPtr make_tensor_ptr(const executorch::aten::Tensor& tensor) { } /** - * Creates a TensorPtr to manage a new Tensor with the same properties - * as the Tensor referenced by the given TensorPtr, sharing the same data - * without owning it. + * Convenience overload identical to make_tensor_ptr(*tensor_ptr, ...). * - * This is a convenience overload equivalent to make_tensor_ptr(*tensor_ptr). - * It does not extend the lifetime of the underlying buffer; if the original - * owner releases the storage, all views aliasing it become dangling. - * - * @param tensor_ptr The TensorPtr whose underlying Tensor is used to initialize - * the returned view. - * @return A new TensorPtr managing a Tensor with the same properties as the - * original. + * @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) { - return make_tensor_ptr(*tensor_ptr); +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)); } /** diff --git a/extension/tensor/test/tensor_ptr_test.cpp b/extension/tensor/test/tensor_ptr_test.cpp index 04356875867..9156a0c4b10 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)); @@ -803,7 +1001,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)); @@ -827,14 +1025,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)); }, "");