diff --git a/docs/source/extension-tensor.md b/docs/source/extension-tensor.md index 131eca00c7d..260692404de 100644 --- a/docs/source/extension-tensor.md +++ b/docs/source/extension-tensor.md @@ -120,7 +120,12 @@ auto tensor = make_tensor_ptr( ScalarType::Int); // int scalar type ``` -The `data` vector must be large enough to accommodate all the elements according to the provided sizes and scalar type. +The `data` vector must hold exactly as many bytes as the provided sizes and scalar type describe. Its bytes are raw memory, not `Byte` elements: like a `void*` passed to `from_blob`, the scalar type defaults to `ScalarType::Float` when you omit it. To make a `Byte` tensor from a `std::vector`, pass `ScalarType::Byte`: + +```cpp +std::vector pixels(224 * 224 * 3); +auto image = make_tensor_ptr({224, 224, 3}, std::move(pixels), ScalarType::Byte); +``` #### Non-Owning Data from Raw Pointer @@ -155,6 +160,25 @@ auto tensor = make_tensor_ptr( The `TensorPtr` will call the custom deleter when it is destroyed, i.e., when the smart pointer is reset and no more references to the underlying `Tensor` exist. +#### Checking a Layout Without Aborting + +`make_tensor_ptr` aborts the process if the dim order or strides are invalid for the sizes, or if the data does not match the sizes. When they come from outside the program, for example from a file or over the network, check them first: + +```cpp +if (validate_tensor_layout(sizes, dim_order, strides) != Error::Ok) { + return Error::InvalidArgument; // reject the input instead of aborting +} +// After the check above, this product cannot overflow. +const auto numel = std::accumulate( + sizes.begin(), sizes.end(), size_t(1), std::multiplies()); +if (data.size() != numel) { + return Error::InvalidArgument; +} +auto tensor = make_tensor_ptr(sizes, std::move(data), dim_order, strides); +``` + +`validate_tensor_layout` does not look at the data, so check its size yourself, as above. For a `std::vector` of raw bytes, the expected size is `numel * elementSize(type)` bytes. Compute it with an overflow check, because this product can overflow `size_t` even when `numel` fits. + #### Sharing Existing Tensor Since `TensorPtr` is a `std::shared_ptr`, you can easily create a `TensorPtr` that shares an existing `Tensor`. Any changes made to the shared data are reflected across all instances sharing the same data. diff --git a/extension/tensor/tensor_ptr.cpp b/extension/tensor/tensor_ptr.cpp index 9a73d50b446..2164a98b770 100644 --- a/extension/tensor/tensor_ptr.cpp +++ b/extension/tensor/tensor_ptr.cpp @@ -65,23 +65,31 @@ struct Storage final { } }; #endif // USE_ATEN_LIB -} // namespace -TensorPtr make_tensor_ptr( - std::vector sizes, - void* data, - std::vector dim_order, - std::vector strides, - executorch::aten::ScalarType type, - executorch::aten::Device device, - executorch::aten::TensorShapeDynamism dynamism, - std::function deleter) { +// Fills an empty dim_order and computes the strides. Shared by make_tensor_ptr +// and validate_tensor_layout so the check rejects every layout the factory +// aborts on. A stride that overflows wraps around unless +// reject_stride_overflow is set. +runtime::Error resolve_layout( + const std::vector& sizes, + std::vector& dim_order, + const std::vector& strides, + bool reject_stride_overflow, + std::vector& computed_strides) { const auto dim = sizes.size(); - ET_CHECK_MSG( + ET_CHECK_OR_RETURN_ERROR( + dim <= runtime::kTensorDimensionLimit, + InvalidArgument, + "Tensor has %zu dims, more than the limit of %zu", + dim, + runtime::kTensorDimensionLimit); + ET_CHECK_OR_RETURN_ERROR( dim_order.empty() || dim_order.size() == dim, + InvalidArgument, "dim_order size must match sizes or be empty."); - ET_CHECK_MSG( + ET_CHECK_OR_RETURN_ERROR( strides.empty() || strides.size() == dim, + InvalidArgument, "strides size must match sizes or be empty."); if (dim_order.empty()) { @@ -93,16 +101,32 @@ TensorPtr make_tensor_ptr( }); } } - std::vector computed_strides(dim); - - auto error = runtime::dim_order_to_stride( - sizes.data(), dim_order.data(), dim, computed_strides.data()); - ET_CHECK_MSG(error == runtime::Error::Ok, "Failed to compute strides."); + ET_CHECK_OR_RETURN_ERROR( + runtime::validate_dim_order(dim_order.data(), dim), + InvalidArgument, + "Invalid dim order: values must be a permutation of [0, %zu)", + dim); + // Same strides as runtime::dim_order_to_stride, but an overflow is detected. + computed_strides.assign(dim, 1); + for (size_t i = dim; i > 1; --i) { + const auto inner = dim_order[i - 1]; + const auto outer = dim_order[i - 2]; + const executorch::aten::StridesType size = + sizes[inner] == 0 ? 1 : sizes[inner]; + const bool overflow = c10::mul_overflows( + computed_strides[inner], size, &computed_strides[outer]); + ET_CHECK_OR_RETURN_ERROR( + !(overflow && reject_stride_overflow), + InvalidArgument, + "Stride of dim %zu overflows", + static_cast(outer)); + } if (!strides.empty()) { for (size_t i = 0; i < dim; i++) { - ET_CHECK_MSG( + ET_CHECK_OR_RETURN_ERROR( strides[i] == computed_strides[i] || sizes[i] == 1, + InvalidArgument, "invalid strides for dim %zu: %" ET_PRI_SIZES_AND_STRIDES "!= %" ET_PRI_SIZES_AND_STRIDES " while its size is %" ET_PRI_SIZES_AND_STRIDES " != 1", @@ -112,10 +136,33 @@ TensorPtr make_tensor_ptr( sizes[i]); } } + return runtime::Error::Ok; +} +} // namespace +TensorPtr make_tensor_ptr( + std::vector sizes, + void* data, + std::vector dim_order, + std::vector strides, + executorch::aten::ScalarType type, + executorch::aten::Device device, + executorch::aten::TensorShapeDynamism dynamism, + std::function deleter) { + std::vector computed_strides; + // Accept strides that overflow: the runtime builds such tensors, and + // aliasing or cloning one calls this. validate_tensor_layout rejects them. + const auto error = resolve_layout( + sizes, + dim_order, + strides, + /*reject_stride_overflow=*/false, + computed_strides); + ET_CHECK_MSG(error == runtime::Error::Ok, "Invalid tensor layout."); strides = std::move(computed_strides); #ifndef USE_ATEN_LIB + const auto dim = sizes.size(); executorch::aten::TensorImpl tensor_impl( type, dim, @@ -296,6 +343,22 @@ runtime::Error resize_tensor_ptr( sizes.data(), sizes.size())); } +runtime::Error validate_tensor_layout( + const std::vector& sizes, + const std::vector& dim_order, + const std::vector& strides) { + ET_CHECK_OK_OR_RETURN_ERROR( + executorch::aten::safe_numel(sizes.data(), sizes.size()).error()); + auto mutable_dim_order = dim_order; + std::vector computed_strides; + return resolve_layout( + sizes, + mutable_dim_order, + strides, + /*reject_stride_overflow=*/true, + computed_strides); +} + // ---- Device tensor helper ---- // // This helper relies on the ExecuTorch DeviceAllocator and the portable tensor diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 5786a22c631..5723d3ec548 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -38,6 +38,12 @@ using TensorPtr = std::shared_ptr; * the runtime environment and pass it in. To copy CPU data to a device, use * `clone_tensor_ptr_to` instead. * + * Aborts if `dim_order` or `strides` are invalid for `sizes`. Dense strides + * that do not fit in `StridesType` are not an error here: they wrap around, as + * in tensors the runtime builds. To reject a layout that comes from outside + * the program instead, including such strides, check it with + * `validate_tensor_layout` first. + * * @param sizes A vector specifying the size of each dimension. * @param data A pointer to the data buffer (CPU or device, see device). * @param dim_order A vector specifying the order of dimensions. @@ -105,7 +111,9 @@ inline TensorPtr make_tensor_ptr( * Overload for a typed data pointer. The scalar type is deduced from the * pointer's element type, the same way the vector overload deduces it from the * vector's element type, so a buffer of one type does not become a Tensor of - * another. Pass `type` explicitly to reinterpret the buffer as a different + * another. The exception is `uint8_t`: a `uint8_t*` gives a Byte tensor, but + * `make_tensor_ptr(sizes, std::vector)` takes raw bytes and defaults + * to Float. Pass `type` explicitly to reinterpret the buffer as a different * type. * * @tparam T The C++ type of the tensor elements, deduced from the pointer. @@ -413,6 +421,10 @@ inline TensorPtr make_tensor_ptr(T value) { * memory's lifetime is tied to the TensorImpl. The result is always a CPU * tensor. * + * The bytes are not Byte elements: like a `void*` passed to `from_blob`, they + * are read as `type`, which defaults to Float. To get a Byte tensor of the + * same values, pass `ScalarType::Byte`. + * * @param sizes A vector specifying the size of each dimension. * @param data A vector containing the raw memory for the tensor's data. * @param dim_order A vector specifying the order of dimensions. @@ -618,6 +630,38 @@ runtime::Error resize_tensor_ptr( TensorPtr& tensor, const std::vector& sizes); +/** + * Checks a tensor layout without creating a tensor. It returns an error for + * every layout that the `make_tensor_ptr` overloads taking sizes and data + * abort on. It also rejects negative sizes, sizes whose product overflows, and + * dense strides that do not fit in `StridesType` (`int32_t`, or `int64_t` in + * ATen mode). + * + * A layout has at most `kTensorDimensionLimit` (16) dimensions. The dim + * order and strides must each be empty or have one entry per dimension. A + * non-empty dim order must be a permutation of the dimensions. An empty dim + * order follows the strides from largest to smallest, or is the default order + * when the strides are empty too. If two strides are equal, pass the dim order + * to set their order. Non-empty strides must be the dense strides for that dim + * order, except in dimensions of size 1. + * + * It takes no data and no scalar type, so it cannot check the data or its + * size in bytes. The vector overloads of `make_tensor_ptr` also need exactly + * as many elements, or bytes for raw bytes, as the sizes and scalar type + * describe. + * + * @param sizes A vector specifying the size of each dimension. + * @param dim_order A vector specifying the order of dimensions, or empty. + * @param strides A vector specifying the strides of each dimension, or empty. + * @return Error::Ok if the layout is valid, or Error::InvalidArgument + * otherwise. + */ +ET_NODISCARD +runtime::Error validate_tensor_layout( + const std::vector& sizes, + const std::vector& dim_order, + const std::vector& strides); + /** * Clones a TensorPtr's data onto the given target device, allocating and * copying as needed. diff --git a/extension/tensor/test/tensor_ptr_test.cpp b/extension/tensor/test/tensor_ptr_test.cpp index 70511045a64..0d72be91151 100644 --- a/extension/tensor/test/tensor_ptr_test.cpp +++ b/extension/tensor/test/tensor_ptr_test.cpp @@ -332,6 +332,14 @@ TEST_F(TensorPtrTest, TensorOwningEmptyData) { EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); } +TEST_F(TensorPtrTest, TensorWithInnerZeroSize) { + // A size-0 dim counts as size 1 for the strides outside it, as in PyTorch. + auto tensor = make_tensor_ptr({2, 0, 3}, std::vector{}); + EXPECT_EQ(tensor->strides()[0], 3); + EXPECT_EQ(tensor->strides()[1], 3); + EXPECT_EQ(tensor->strides()[2], 1); +} + TEST_F(TensorPtrTest, TensorDataOnly) { auto tensor = make_tensor_ptr({1.0f, 2.0f, 3.0f, 4.0f}); @@ -1093,6 +1101,59 @@ TEST_F(TensorPtrTest, TensorInvalidDimOrder) { ""); } +TEST_F(TensorPtrTest, ValidateTensorLayoutRejectsInvalidLayouts) { + EXPECT_EQ(validate_tensor_layout({2, 2}, {}, {8, 2}), Error::InvalidArgument); + EXPECT_EQ( + validate_tensor_layout({3, 4}, {1, 0}, {1, 4}), Error::InvalidArgument); + EXPECT_EQ(validate_tensor_layout({4, 5}, {2, 1}, {}), Error::InvalidArgument); + EXPECT_EQ( + validate_tensor_layout({3, 4}, {}, {4, 1, 1}), Error::InvalidArgument); + EXPECT_EQ( + validate_tensor_layout({3, 4}, {0, 1, 2}, {}), Error::InvalidArgument); + EXPECT_EQ(validate_tensor_layout({3, 4}, {1, 1}, {}), Error::InvalidArgument); + EXPECT_EQ(validate_tensor_layout({-1, 4}, {}, {}), Error::InvalidArgument); + EXPECT_EQ( + validate_tensor_layout( + std::vector(17, 1), {}, {}), + Error::InvalidArgument); + const executorch::aten::SizesType large = 1 << 30; + EXPECT_EQ( + validate_tensor_layout({large, large, large}, {}, {}), + Error::InvalidArgument); + // The square of this size overflows StridesType. + const auto huge = executorch::aten::SizesType(1) + << (4 * sizeof(executorch::aten::StridesType)); + EXPECT_EQ( + validate_tensor_layout({0, huge, huge}, {}, {}), Error::InvalidArgument); + EXPECT_EQ( + validate_tensor_layout({2, huge, huge}, {}, {}), Error::InvalidArgument); +} + +TEST_F(TensorPtrTest, ValidateTensorLayoutAcceptsValidLayouts) { + EXPECT_EQ(validate_tensor_layout({}, {}, {}), Error::Ok); + EXPECT_EQ(validate_tensor_layout({2, 3}, {}, {}), Error::Ok); + EXPECT_EQ(validate_tensor_layout({3, 4}, {1, 0}, {1, 3}), Error::Ok); + EXPECT_EQ(validate_tensor_layout({3, 4}, {}, {1, 3}), Error::Ok); + EXPECT_EQ(validate_tensor_layout({1, 4}, {}, {100, 1}), Error::Ok); +} + +TEST_F(TensorPtrTest, TensorWithMaxDims) { + const std::vector sizes(16, 1); + EXPECT_EQ(validate_tensor_layout(sizes, {}, {}), Error::Ok); + EXPECT_EQ(make_tensor_ptr(sizes, {1.f})->dim(), 16); +} + +TEST_F(TensorPtrTest, TensorOverflowingStrides) { + // The square of this size overflows StridesType, but make_tensor_ptr has + // always accepted such a layout. + const auto huge = executorch::aten::SizesType(1) + << (4 * sizeof(executorch::aten::StridesType)); + auto tensor = make_tensor_ptr({0, huge, huge}, std::vector{}); + EXPECT_EQ(tensor->numel(), 0); + EXPECT_EQ(tensor->strides()[1], huge); + EXPECT_EQ(tensor->strides()[2], 1); +} + TEST_F(TensorPtrTest, TensorCustomDeleter) { float data[20] = {4}; auto tensor = make_tensor_ptr({4, 5}, data); diff --git a/extension/wasm/test/unittests.js b/extension/wasm/test/unittests.js index 3d485c2e8b2..d49a6951aa4 100644 --- a/extension/wasm/test/unittests.js +++ b/extension/wasm/test/unittests.js @@ -48,6 +48,7 @@ describe("Tensor", () => { test("fromArray wrong size", () => { expect(() => et.Tensor.fromArray([3, 2], [1, 2, 3, 4])).toThrow(); + expect(() => et.Tensor.fromArray([2], [1, 2, 3])).toThrow("Required 2 elements, given 3"); }); test("full", () => { @@ -95,6 +96,10 @@ describe("Tensor", () => { test("incorrect strides", () => { expect(() => et.Tensor.fromArray([2, 2], [1, 2, 3, 4], et.ScalarType.Float, [1, 1], [2, 1])).toThrow(); }); + + test("strides that overflow", () => { + expect(() => et.Tensor.fromArray([0, 65536, 65536], [])).toThrow("Invalid tensor layout."); + }); }); describe("Module", () => { diff --git a/extension/wasm/wasm_bindings.cpp b/extension/wasm/wasm_bindings.cpp index 5d3cff0cddc..e3a0ebffd40 100644 --- a/extension/wasm/wasm_bindings.cpp +++ b/extension/wasm/wasm_bindings.cpp @@ -11,8 +11,6 @@ #include #include #include -#include -#include #ifdef EXECUTORCH_ENABLE_EVENT_TRACER #include @@ -109,7 +107,7 @@ inline void assert_valid_numel( THROW_IF_ERROR( computed_numel.error(), "Invalid tensor sizes: numel computation failed"); THROW_IF_FALSE( - data.size() >= static_cast(computed_numel.get()), + data.size() == static_cast(computed_numel.get()), "Required %ld elements, given %ld", computed_numel.get(), data.size()); @@ -136,49 +134,15 @@ std::vector convertJSGeneratorToNumberVector(val generator) { return data; } -// make_tensor_ptr() assertions will abort the program if they fail. -// These checks will throw a JS error instead. -void assert_dim_order_and_strides_valid( +// make_tensor_ptr() aborts the program on an invalid layout. This throws a JS +// error instead. +void assert_layout_valid( const std::vector& sizes, - std::vector& dim_order, - std::vector& strides) { - THROW_IF_FALSE( - dim_order.size() == 0 || dim_order.size() == sizes.size(), - "dim_order size must match sizes or be empty."); - THROW_IF_FALSE( - strides.size() == 0 || strides.size() == sizes.size(), - "strides size must match sizes or be empty."); - - if (dim_order.empty()) { - dim_order.resize(sizes.size()); - std::iota(dim_order.begin(), dim_order.end(), 0); - if (!strides.empty()) { - std::sort(dim_order.begin(), dim_order.end(), [&](size_t a, size_t b) { - return strides[a] > strides[b]; - }); - } - } - std::vector computed_strides(sizes.size()); - - auto error = runtime::dim_order_to_stride( - sizes.data(), dim_order.data(), sizes.size(), computed_strides.data()); - THROW_IF_ERROR(error, "Failed to compute strides."); - - if (!strides.empty()) { - for (size_t i = 0; i < sizes.size(); i++) { - THROW_IF_FALSE( - strides[i] == computed_strides[i] || sizes[i] == 1, - "invalid strides for dim %zu: %" ET_PRI_SIZES_AND_STRIDES - "!= %" ET_PRI_SIZES_AND_STRIDES - " while its size is %" ET_PRI_SIZES_AND_STRIDES " != 1", - i, - strides[i], - computed_strides[i], - sizes[i]); - } - } - - strides = std::move(computed_strides); + const std::vector& dim_order, + const std::vector& strides) { + THROW_IF_ERROR( + validate_tensor_layout(sizes, dim_order, strides), + "Invalid tensor layout."); } /** @@ -311,7 +275,7 @@ class ET_EXPERIMENTAL JsTensor { auto strides_vec = convertJSArrayToNumberVector(strides); - assert_dim_order_and_strides_valid(sizes_vec, dim_order_vec, strides_vec); + assert_layout_valid(sizes_vec, dim_order_vec, strides_vec); switch (type) { #define JS_CASE_FROM_ARRAY_VECTOR_TYPE(T, NAME) \ case ScalarType::NAME: { \ @@ -363,7 +327,7 @@ class ET_EXPERIMENTAL JsTensor { auto strides_vec = convertJSArrayToNumberVector(strides); - assert_dim_order_and_strides_valid(sizes_vec, dim_order_vec, strides_vec); + assert_layout_valid(sizes_vec, dim_order_vec, strides_vec); switch (type) { #define JS_CASE_FROM_ITER_VECTOR_TYPE(T, NAME) \