Add validate_tensor_layout to check a tensor layout without aborting, and document the raw-bytes overloads

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<uint8_t>`, pass `ScalarType::Byte`:
+
+```cpp
+std::vector<uint8_t> 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<size_t>());
+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<uint8_t>` 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<Tensor>`, 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.

extension/tensor/tensor_ptr.cpp

@@ -65,23 +65,31 @@ struct Storage final {
   }
 };
 #endif // USE_ATEN_LIB
-} // namespace
 
-TensorPtr make_tensor_ptr(
-    std::vector<executorch::aten::SizesType> sizes,
-    void* data,
-    std::vector<executorch::aten::DimOrderType> dim_order,
-    std::vector<executorch::aten::StridesType> strides,
-    executorch::aten::ScalarType type,
-    executorch::aten::Device device,
-    executorch::aten::TensorShapeDynamism dynamism,
-    std::function<void(void*)> 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<executorch::aten::SizesType>& sizes,
+    std::vector<executorch::aten::DimOrderType>& dim_order,
+    const std::vector<executorch::aten::StridesType>& strides,
+    bool reject_stride_overflow,
+    std::vector<executorch::aten::StridesType>& 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<executorch::aten::StridesType> 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<size_t>(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<executorch::aten::SizesType> sizes,
+    void* data,
+    std::vector<executorch::aten::DimOrderType> dim_order,
+    std::vector<executorch::aten::StridesType> strides,
+    executorch::aten::ScalarType type,
+    executorch::aten::Device device,
+    executorch::aten::TensorShapeDynamism dynamism,
+    std::function<void(void*)> deleter) {
+  std::vector<executorch::aten::StridesType> 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<executorch::aten::SizesType>& sizes,
+    const std::vector<executorch::aten::DimOrderType>& dim_order,
+    const std::vector<executorch::aten::StridesType>& strides) {
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      executorch::aten::safe_numel(sizes.data(), sizes.size()).error());
+  auto mutable_dim_order = dim_order;
+  std::vector<executorch::aten::StridesType> 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

extension/tensor/tensor_ptr.h

@@ -38,6 +38,12 @@ using TensorPtr = std::shared_ptr<executorch::aten::Tensor>;
  * 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<uint8_t>)` 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<executorch::aten::SizesType>& 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<executorch::aten::SizesType>& sizes,
+    const std::vector<executorch::aten::DimOrderType>& dim_order,
+    const std::vector<executorch::aten::StridesType>& strides);
+
 /**
  * Clones a TensorPtr's data onto the given target device, allocating and
  * copying as needed.

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<float>{});
+  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<executorch::aten::SizesType>(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<executorch::aten::SizesType> 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<float>{});
+  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);

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", () => {