Validate Method input layouts for non-contiguous tensors

extension/android/executorch_android/consumer-proguard-rules.pro

@@ -32,7 +32,8 @@
 -keep class org.pytorch.executorch.EValue { *; }
 
 # Keep Tensor and its inner classes. The Tensor class has methods and fields
-# accessed from JNI (dtypeJniCode, shape, getRawDataBuffer, nativeNewTensor).
+# accessed from JNI (dtypeJniCode, shape, dimOrder, getRawDataBuffer,
+# nativeNewTensor).
 -keep class org.pytorch.executorch.Tensor { *; }
 -keep class org.pytorch.executorch.Tensor$* { *; }
 

extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt

@@ -7,6 +7,7 @@
  */
 package org.pytorch.executorch
 
+import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.junit.Assert.assertArrayEquals
@@ -639,6 +640,59 @@ class TensorTest {
         .hasMessage("Inconsistent data capacity:4 and shape number elements:2 shape:[1, 2]")
   }
 
+  @Test
+  fun testDimOrder() {
+    val shape = longArrayOf(1, 3, 2, 2)
+    val dimOrder = intArrayOf(0, 2, 3, 1)
+    val data = FloatArray(12) { it.toFloat() }
+
+    assertArrayEquals(dimOrder, Tensor.fromBlob(data, shape, dimOrder).dimOrder())
+
+    val buffer = Tensor.allocateFloatBuffer(data.size)
+    buffer.put(data)
+    assertArrayEquals(dimOrder, Tensor.fromBlob(buffer, shape, dimOrder).dimOrder())
+    assertArrayEquals(intArrayOf(0, 1, 2, 3), Tensor.fromBlob(data, shape).dimOrder())
+
+    assertThatThrownBy { Tensor.fromBlob(data, shape, intArrayOf(0, 1, 2)) }
+        .isInstanceOf(IllegalArgumentException::class.java)
+    assertThatThrownBy { Tensor.fromBlob(data, shape, intArrayOf(0, 1, 2, 2)) }
+        .isInstanceOf(IllegalArgumentException::class.java)
+  }
+
+  @Test
+  fun testDimOrderSerde() {
+    val shape = longArrayOf(1, 3, 2, 2)
+    val dimOrder = intArrayOf(0, 2, 3, 1)
+    val data = FloatArray(12) { it.toFloat() }
+
+    val deser = Tensor.fromByteArray(Tensor.fromBlob(data, shape, dimOrder).toByteArray())
+
+    assertArrayEquals(shape, deser.shape())
+    assertArrayEquals(dimOrder, deser.dimOrder())
+    assertArrayEquals(data, deser.dataAsFloatArray, 0f)
+  }
+
+  @Test
+  fun testSerdeWithoutDimOrder() {
+    val shape = longArrayOf(2, 2)
+    val data = intArrayOf(1, 2, 3, 4)
+    val buffer = ByteBuffer.allocate(1 + 1 + 4 * shape.size + 4 * data.size)
+    buffer.put(DType.INT32.jniCode.toByte())
+    buffer.put(shape.size.toByte())
+    for (s in shape) {
+      buffer.putInt(s.toInt())
+    }
+    for (d in data) {
+      buffer.putInt(d)
+    }
+
+    val deser = Tensor.fromByteArray(buffer.array())
+
+    assertArrayEquals(shape, deser.shape())
+    assertArrayEquals(intArrayOf(0, 1), deser.dimOrder())
+    assertArrayEquals(data, deser.dataAsIntArray)
+  }
+
   @Test
   fun testLongTensorSerde() {
     val data = longArrayOf(1, 2, 3, 4)

extension/android/jni/jni_layer.cpp

@@ -60,8 +60,6 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
 
   static facebook::jni::local_ref<TensorHybrid::javaobject>
   newJTensorFromTensor(const executorch::aten::Tensor& tensor) {
-    // Java wrapper currently only supports contiguous tensors.
-
     const auto scalarType = tensor.scalar_type();
     if (scalar_type_to_java_dtype.count(scalarType) == 0) {
       std::stringstream ss;
@@ -83,6 +81,22 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
     jTensorShape->setRegion(
         0, tensor_shape_vec.size(), tensor_shape_vec.data());
 
+    std::vector<jint> dim_order_vec;
+    dim_order_vec.reserve(tensor_shape_vec.size());
+    const auto tensor_dim_order = tensor.dim_order();
+    if (tensor_dim_order.size() == tensor_shape_vec.size()) {
+      for (const auto& d : tensor_dim_order) {
+        dim_order_vec.push_back(static_cast<jint>(d));
+      }
+    } else {
+      for (size_t i = 0; i < tensor_shape_vec.size(); ++i) {
+        dim_order_vec.push_back(static_cast<jint>(i));
+      }
+    }
+    facebook::jni::local_ref<facebook::jni::JArrayInt> jDimOrder =
+        facebook::jni::make_int_array(dim_order_vec.size());
+    jDimOrder->setRegion(0, dim_order_vec.size(), dim_order_vec.data());
+
     static auto cls = TensorHybrid::javaClassStatic();
     // Note: this is safe as long as the data stored in tensor is valid; the
     // data won't go out of scope as long as the Method for the inference is
@@ -97,10 +111,16 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
         cls->getStaticMethod<facebook::jni::local_ref<TensorHybrid::javaobject>(
             facebook::jni::alias_ref<facebook::jni::JByteBuffer>,
             facebook::jni::alias_ref<jlongArray>,
+            facebook::jni::alias_ref<facebook::jni::JArrayInt>,
             jint,
             facebook::jni::alias_ref<jhybriddata>)>("nativeNewTensor");
     return jMethodNewTensor(
-        cls, jTensorBuffer, jTensorShape, jdtype, makeCxxInstance(tensor));
+        cls,
+        jTensorBuffer,
+        jTensorShape,
+        jDimOrder,
+        jdtype,
+        makeCxxInstance(tensor));
   }
 
   static TensorPtr newTensorFromJTensor(
@@ -112,6 +132,10 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
     static const auto shapeField = cls->getField<jlongArray>("shape");
     auto jshape = jtensor->getFieldValue(shapeField);
 
+    static const auto dimOrderField =
+        cls->getField<facebook::jni::JArrayInt>("dimOrder");
+    auto jdim_order = jtensor->getFieldValue(dimOrderField);
+
     static auto dataBufferMethod = cls->getMethod<
         facebook::jni::local_ref<facebook::jni::JBuffer::javaobject>()>(
         "getRawDataBuffer");
@@ -131,6 +155,29 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
     for (int i = rank - 1; i >= 0; --i) {
       numel *= shapeArr[i];
     }
+
+    std::vector<executorch::aten::DimOrderType> dim_order_vec;
+    dim_order_vec.reserve(rank);
+    if (!jdim_order || jdim_order->size() != rank) {
+      jni_helper::throwExecutorchException(
+          static_cast<uint32_t>(Error::InvalidArgument),
+          "Tensor dim_order size must match shape size");
+      return nullptr;
+    }
+    const auto dimOrderArr = jdim_order->getRegion(0, rank);
+    std::vector<char> seen_dim(rank, 0);
+    for (int i = 0; i < rank; i++) {
+      const jint d = dimOrderArr[i];
+      if (d < 0 || d >= static_cast<jint>(rank) || seen_dim[d] != 0) {
+        jni_helper::throwExecutorchException(
+            static_cast<uint32_t>(Error::InvalidArgument),
+            "Tensor dim_order must be a permutation of shape dimensions");
+        return nullptr;
+      }
+      seen_dim[d] = 1;
+      dim_order_vec.push_back(static_cast<executorch::aten::DimOrderType>(d));
+    }
+
     JNIEnv* jni = facebook::jni::Environment::current();
     if (java_dtype_to_scalar_type.count(jdtype) == 0) {
       std::stringstream ss;
@@ -163,8 +210,12 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
           static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
       return nullptr;
     }
-    return from_blob(
-        jni->GetDirectBufferAddress(jbuffer.get()), shape_vec, scalar_type);
+    return for_blob(
+               jni->GetDirectBufferAddress(jbuffer.get()),
+               shape_vec,
+               scalar_type)
+        .dim_order(std::move(dim_order_vec))
+        .make_tensor_ptr();
   }
 
  private:

runtime/core/portable_type/tensor_impl.h

@@ -175,11 +175,17 @@ class TensorImpl {
 
   /// Returns the order the dimensions are laid out in memory.
   const ArrayRef<DimOrderType> dim_order() const {
+    if (dim_order_ == nullptr) {
+      return {};
+    }
     return ArrayRef<DimOrderType>{dim_order_, static_cast<size_t>(dim_)};
   }
 
   /// Returns the strides of the tensor at each dimension.
   const ArrayRef<StridesType> strides() const {
+    if (strides_ == nullptr) {
+      return {};
+    }
     return ArrayRef<StridesType>{strides_, static_cast<size_t>(dim_)};
   }
 

runtime/executor/method.cpp

@@ -755,6 +755,87 @@ Error Method::parse_values(const NamedDataMap* external_data_map) {
 }
 
 namespace {
+/**
+ * Validates that t_src's physical layout matches t_dst's, which must already
+ * be resized to t_src's sizes. set_input() copies or shares t_src's buffer
+ * under t_dst's layout, so a mismatched source would be silently
+ * reinterpreted. A source with the default contiguous layout is only warned
+ * about, since callers commonly wrap data that is already in the method's
+ * layout in a default tensor just to pass its shape.
+ */
+ET_NODISCARD Error validate_input_tensor_layout(
+    const executorch::aten::Tensor& t_dst,
+    const executorch::aten::Tensor& t_src,
+    size_t input_idx) {
+  const auto dim = t_src.dim();
+  if (dim == 0 || t_src.numel() == 0) {
+    return Error::Ok;
+  }
+
+  auto src_strides = t_src.strides();
+#ifndef USE_ATEN_LIB
+  std::array<executorch::aten::StridesType, kTensorDimensionLimit>
+      derived_strides{};
+  if (src_strides.empty()) {
+    ET_CHECK_OR_RETURN_ERROR(
+        static_cast<size_t>(dim) <= kTensorDimensionLimit,
+        InvalidArgument,
+        "Input %" ET_PRIsize_t " has unsupported rank %zd.",
+        input_idx,
+        dim);
+    std::array<executorch::aten::DimOrderType, kTensorDimensionLimit>
+        default_dim_order{};
+    for (ssize_t i = 0; i < dim; ++i) {
+      default_dim_order[i] = static_cast<executorch::aten::DimOrderType>(i);
+    }
+    const auto src_dim_order = t_src.dim_order();
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        dim_order_to_stride(
+            t_src.sizes().data(),
+            src_dim_order.empty() ? default_dim_order.data()
+                                  : src_dim_order.data(),
+            dim,
+            derived_strides.data()),
+        "Failed to compute strides for input %" ET_PRIsize_t,
+        input_idx);
+    src_strides = {derived_strides.data(), static_cast<size_t>(dim)};
+  }
+#endif
+
+  const auto dst_strides = t_dst.strides();
+  bool matches = true;
+  bool src_is_contiguous = true;
+  int64_t contiguous_stride = 1;
+  for (ssize_t i = dim - 1; i >= 0; --i) {
+    if (t_src.size(i) > 1) {
+      matches = matches && src_strides[i] == dst_strides[i];
+      src_is_contiguous =
+          src_is_contiguous && src_strides[i] == contiguous_stride;
+    }
+    contiguous_stride *= t_src.size(i);
+  }
+
+  if (matches) {
+    return Error::Ok;
+  }
+  if (src_is_contiguous) {
+    ET_LOG(
+        Debug,
+        "Input %" ET_PRIsize_t
+        " has the default contiguous layout but the method expects a "
+        "different dim order; treating its data as already in the expected "
+        "layout. Pass a tensor with the method's dim order to avoid this.",
+        input_idx);
+    return Error::Ok;
+  }
+  ET_LOG(
+      Error,
+      "Input %" ET_PRIsize_t
+      " has strides that do not match the method's expected layout.",
+      input_idx);
+  return Error::InvalidArgument;
+}
+
 /**
  * Private/helper method for populating operator_name from the Operator.
  * operator_name is a char pointer that is already allocated. The size of
@@ -1345,6 +1426,10 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) {
         resize_tensor(t_dst, t_src.sizes()),
         "Error resizing tensor at input %" ET_PRIsize_t,
         input_idx);
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        validate_input_tensor_layout(t_dst, t_src, input_idx),
+        "Error validating tensor layout at input %" ET_PRIsize_t,
+        input_idx);
     auto tensor_meta = this->method_meta().input_tensor_meta(input_idx);
     if (tensor_meta->is_memory_planned()) {
       ET_CHECK_OK_OR_RETURN_ERROR(

runtime/executor/test/CMakeLists.txt

@@ -24,7 +24,7 @@ if(CMAKE_HOST_SYSTEM_NAME STREQUAL "Windows")
       powershell
       ${EXECUTORCH_ROOT}/kernels/test/export_test_model.ps1
       -Modules
-      "\"ModuleAdd,ModuleAddHalf,ModuleAddMul,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleMultipleEntry,ModuleSimpleTrain,ModuleStateful\""
+      "\"ModuleAdd,ModuleAddHalf,ModuleAddMul,ModuleChannelsLast,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleMultipleEntry,ModuleSimpleTrain,ModuleStateful\""
       -outDir
       "${CMAKE_CURRENT_BINARY_DIR}"
       -CondaEnv
@@ -36,7 +36,7 @@ else()
       -m
       test.models.export_program
       --modules
-      "ModuleAdd,ModuleAddHalf,ModuleAddMul,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleMultipleEntry,ModuleSimpleTrain,ModuleStateful"
+      "ModuleAdd,ModuleAddHalf,ModuleAddMul,ModuleChannelsLast,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleMultipleEntry,ModuleSimpleTrain,ModuleStateful"
       --outdir
       "${CMAKE_CURRENT_BINARY_DIR}"
   )
@@ -48,6 +48,7 @@ add_custom_command(
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMul.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+         "${CMAKE_CURRENT_BINARY_DIR}/ModuleChannelsLast.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleDynamicCatUnallocatedIO.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleIndex.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleMultipleEntry.pte"
@@ -103,6 +104,7 @@ set(test_env
     "ET_MODULE_ADD_MUL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMul.pte"
     "ET_MODULE_ADD_MUL_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
     "ET_MODULE_ADD_MUL_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+    "ET_MODULE_CHANNELS_LAST_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleChannelsLast.pte"
     "ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleDynamicCatUnallocatedIO.pte"
     "ET_MODULE_INDEX_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleIndex.pte"
     "ET_MODULE_MULTI_ENTRY_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleMultipleEntry.pte"