Validate kTensorDimensionLimit (#12684)

Summary:

Need to check this limit at deserialization since the ops rely on it

Reviewed By: larryliu0820

Differential Revision: D78675986

Author: JacobSzwejbka

runtime/executor/targets.bzl

@@ -133,6 +133,7 @@ def define_common_targets():
             ],
             deps = [
                 "//executorch/schema:program",
+                "//executorch/runtime/core/exec_aten/util:tensor_dimension_limit"
             ],
             visibility = [
                 "//executorch/runtime/executor/...",

runtime/executor/tensor_parser_aten.cpp

@@ -8,126 +8,257 @@
 
 #include <executorch/runtime/executor/tensor_parser.h>
 
-#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
-#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
-#include <executorch/runtime/core/named_data_map.h>
+#include <executorch/runtime/core/evalue.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
 #include <executorch/schema/program_generated.h>
 
-#include <ATen/ATen.h> // @donotremove @manual=//caffe2/aten:ATen-core
-
 namespace executorch {
-// This file is only used in ATen mode, so we use the runtime_aten namespace.
-namespace runtime {
-namespace aten {
+namespace ET_RUNTIME_NAMESPACE {
 namespace deserialization {
 
+using executorch::aten::ScalarType;
+// Provides access to private Program methods.
+class TensorParser final {
+ public:
+  ET_NODISCARD static Error load_mutable_subsegment_into(
+      const Program* program,
+      size_t mutable_data_segments_index,
+      size_t offset_index,
+      size_t size,
+      void* buffer) {
+    return program->load_mutable_subsegment_into(
+        mutable_data_segments_index, offset_index, size, buffer);
+  }
+};
+
 namespace {
 
-void deleteNothing(void*);
-void deleteNothing(void*) {}
+// Retrieve the buffer specified by the allocation_info
+ET_NODISCARD Result<void*> getMemPlannedPtr(
+    const executorch_flatbuffer::AllocationDetails* allocation_info,
+    size_t nbytes,
+    HierarchicalAllocator* allocator) {
+  // Normal non-constant Tensor. Allocate data using mem_id and offset.
 
+  // TODO(T142455629): make the allocator actually id based and not indexed
+  // based. -1 is a hack to get the memory ids 0 aligned because previously
+  // 0 was reserved
+  const uint32_t memory_id = allocation_info->memory_id() - 1;
+
+  // Originally this field was a single uint32_t, but we need 64 bits for
+  // larger models. To preserve backwards compatibility, the high bits are
+  // managed in a separate uint32_t field.
+  const uint32_t memory_offset_low = allocation_info->memory_offset_low();
+  const uint32_t memory_offset_high = allocation_info->memory_offset_high();
+
+  size_t memory_offset = memory_offset_low;
+  if ((sizeof(size_t) > sizeof(uint32_t)) && (memory_offset_high > 0)) {
+    // The compiler should remove this always-true check on 64-bit systems.
+    ET_CHECK_OR_RETURN_ERROR(
+        sizeof(size_t) >= sizeof(uint64_t),
+        NotSupported,
+        "size_t cannot hold memory offset 0x%08" PRIx32 ".%08" PRIx32,
+        memory_offset_high,
+        memory_offset_low);
+    memory_offset |= static_cast<size_t>(memory_offset_high) << 32;
+  }
+  return allocator->get_offset_address(memory_id, memory_offset, nbytes);
+}
 } // namespace
 
-Result<at::Tensor> parseTensor(
-    const Program* program,
-    MemoryManager* memory_manager,
-    const executorch_flatbuffer::Tensor* s_tensor,
-    const NamedDataMap* named_data_map,
-    Span<NamedData> external_constants) {
-  EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor");
+ET_NODISCARD Result<BoxedEvalueList<executorch::aten::Tensor>> parseTensorList(
+    const flatbuffers::Vector<int32_t>* tensor_indices,
+    EValue* values,
+    size_t values_len,
+    MemoryManager* memory_manager) {
+  EXECUTORCH_SCOPE_PROF("TensorParser::parseTensorList");
 
-  ET_CHECK_OR_RETURN_ERROR(
-      s_tensor->storage_offset() == 0,
-      NotSupported,
-      "Non-zero storage offset %" PRId32 " not supported",
-      s_tensor->storage_offset());
+  auto* tensor_list =
+      memory_manager->method_allocator()
+          ->allocateList<executorch::aten::Tensor>(tensor_indices->size());
+  if (tensor_list == nullptr) {
+    return Error::MemoryAllocationFailed;
+  }
+  auto* evalp_list = memory_manager->method_allocator()->allocateList<EValue*>(
+      tensor_indices->size());
+  if (evalp_list == nullptr) {
+    return Error::MemoryAllocationFailed;
+  }
 
-  // get metadata
-  at::ScalarType type = static_cast<at::ScalarType>(s_tensor->scalar_type());
-  ET_CHECK_OR_RETURN_ERROR(
-      isValid(type),
-      InvalidProgram,
-      "Invalid ScalarType %" PRId8,
-      static_cast<int8_t>(type));
-  auto options = at::CPU(type).options();
+  // For each tensor index look up the corresponding Tensor (which has been
+  // already allocated) and stick it in the list.
+  size_t output_idx = 0;
+  for (int32_t tensor_index : *tensor_indices) {
+    ET_CHECK_OR_RETURN_ERROR(
+        tensor_index >= 0 && static_cast<size_t>(tensor_index) < values_len,
+        InvalidProgram,
+        "Invalid value index %" PRId32 " for TensorList",
+        tensor_index);
 
-  ET_CHECK_OR_RETURN_ERROR(
-      s_tensor->sizes() != nullptr, InvalidProgram, "Missing sizes field");
-  size_t ndim = s_tensor->sizes()->size();
+    // Placement new as the list elements are not initialized, so calling
+    // copy assignment is not defined if it's non trivial.
+    new (&tensor_list[output_idx]) executorch::aten::Tensor(
+        values[static_cast<size_t>(tensor_index)].toTensor());
+    evalp_list[output_idx] = &values[static_cast<size_t>(tensor_index)];
+    output_idx++;
+  }
+
+  return BoxedEvalueList<executorch::aten::Tensor>(
+      evalp_list, tensor_list, tensor_indices->size());
+}
 
+ET_NODISCARD Error validateTensorLayout(
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const TensorLayout& expected_layout) {
   ET_CHECK_OR_RETURN_ERROR(
-      s_tensor->dim_order() != nullptr,
-      InvalidProgram,
-      "Missing dim_order field");
+      static_cast<ScalarType>(s_tensor->scalar_type()) ==
+          expected_layout.scalar_type(),
+      InvalidExternalData,
+      "Scalar type mismatch. Expected %hhd, got %hhd.",
+      static_cast<int8_t>(s_tensor->scalar_type()),
+      static_cast<int8_t>(expected_layout.scalar_type()));
+  int dim = s_tensor->sizes()->size();
   ET_CHECK_OR_RETURN_ERROR(
-      s_tensor->dim_order()->size() == ndim,
-      InvalidProgram,
-      "dim_order size %" PRIu32 " != ndim %zu",
-      s_tensor->dim_order()->size(),
-      ndim);
-
-  // convert int32 in serialization to int64 for aten
-  std::vector<int64_t> sizes(
-      s_tensor->sizes()->begin(), s_tensor->sizes()->end());
-  std::vector<int64_t> strides(ndim);
-  auto status = dim_order_to_stride(
-      s_tensor->sizes()->data(),
-      s_tensor->dim_order()->data(),
-      ndim,
-      strides.data());
+      dim >= 0, InvalidExternalData, "Dim is negative: %d", dim)
   ET_CHECK_OR_RETURN_ERROR(
-      status == Error::Ok,
-      Internal,
-      "dim_order_to_stride returned invalid status");
-
-  // Create a tensor without data first so we can find its expected size before
-  // getting its memory.
-  at::Tensor tensor = at::from_blob(
-      /*data=*/nullptr,
-      sizes,
-      strides,
-      /*storage_offset=*/0,
-      deleteNothing,
-      options);
-
-  if (s_tensor->shape_dynamism() ==
-      executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_UNBOUND) {
-    // Provide fully dynamic tensors with an allocator so they can be resized
-    // within aten kernels.
-    auto impl = tensor.unsafeGetTensorImpl();
-    at::StorageImpl* storage = impl->unsafe_storage().unsafeGetStorageImpl();
-    storage->set_allocator(at::getCPUAllocator());
-    storage->set_resizable(true);
-    storage->set_nbytes(0);
-    impl->set_sizes_contiguous(0);
-    // Leave the data as nullptr since it will be reallocated.
-  } else {
-    // Now that we know how big the tensor is, find and assign its memory.
-    Result<void*> data_ptr = getTensorDataPtr(
-        s_tensor,
-        program,
-        tensor.nbytes(),
-        memory_manager->planned_memory(),
-        named_data_map,
-        external_constants);
-    if (!data_ptr.ok()) {
-      ET_LOG(
-          Error,
-          "getTensorDataPtr() failed: 0x%" PRIx32,
-          static_cast<uint32_t>(data_ptr.error()));
-      return data_ptr.error();
+      static_cast<size_t>(dim) == expected_layout.sizes().size(),
+      InvalidExternalData,
+      "Dim mismatch. Expected %d, got %zu.",
+      dim,
+      expected_layout.sizes().size());
+  for (int i = 0; i < dim; i++) {
+    ET_CHECK_OR_RETURN_ERROR(
+        s_tensor->sizes()->Get(i) == expected_layout.sizes()[i],
+        InvalidExternalData,
+        "Sizes mismatch. Expected %d, got %d for size at index %d.",
+        s_tensor->sizes()->Get(i),
+        expected_layout.sizes()[i],
+        i);
+    ET_CHECK_OR_RETURN_ERROR(
+        s_tensor->dim_order()->Get(i) == expected_layout.dim_order()[i],
+        InvalidExternalData,
+        "Dim order mismatch. Expected %d, got %d for dim at index %d.",
+        s_tensor->dim_order()->Get(i),
+        expected_layout.dim_order()[i],
+        i);
+  }
+  return Error::Ok;
+}
+
+// Check if key exists in entries. If it does, return a pointer to the entry
+// otherwise return a nullptr.
+NamedData* get_data_by_key(const char* key, Span<NamedData> entries) {
+  for (const auto i : c10::irange(entries.size())) {
+    if (strcmp(key, entries[i].key) == 0) {
+      return &entries[i];
     }
-    tensor.unsafeGetTensorImpl()->unsafe_storage().set_data_ptr(
-        at::DataPtr(data_ptr.get(), c10::DeviceType::CPU));
   }
+  return nullptr;
+}
 
-  return tensor;
+ET_NODISCARD Result<void*> getTensorDataPtr(
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const Program* program,
+    size_t nbytes,
+    HierarchicalAllocator* allocator,
+    const NamedDataMap* named_data_map,
+    Span<NamedData> external_constants) {
+  auto data_buffer_idx = s_tensor->data_buffer_idx();
+  const executorch_flatbuffer::AllocationDetails* allocation_info =
+      s_tensor->allocation_info();
+
+  // External tensors.
+  if (s_tensor->extra_tensor_info() != nullptr &&
+      s_tensor->extra_tensor_info()->location() ==
+          executorch_flatbuffer::TensorDataLocation::EXTERNAL) {
+    // Check that fqn is not null.
+    ET_CHECK_OR_RETURN_ERROR(
+        s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr,
+        InvalidExternalData,
+        "Fully qualified name of external tensor is null");
+    const char* fqn =
+        s_tensor->extra_tensor_info()->fully_qualified_name()->c_str();
+
+    // Constant value.
+    if (allocation_info == nullptr) {
+      NamedData* data = get_data_by_key(fqn, external_constants);
+      if (data != nullptr) {
+        return const_cast<void*>(data->buffer.data());
+      }
+      // Should never reach here; these tensors are resolved in
+      // Method::parse_external_constants. Any errors should be caught there.
+      return Error::Internal;
+    } else {
+      // Mutable value.
+      // Look up tensor in named data map.
+      ET_CHECK_OR_RETURN_ERROR(
+          named_data_map != nullptr,
+          InvalidExternalData,
+          "Cannot retrieve external tensor with fqn: %s. The named_data_map is null; most likely no external .ptd file was provided.",
+          fqn);
+      Result<const TensorLayout> tensor_layout_res =
+          named_data_map->get_tensor_layout(fqn);
+      if (!tensor_layout_res.ok()) {
+        return tensor_layout_res.error();
+      }
+      const TensorLayout& tensor_layout = tensor_layout_res.get();
+      Error err = validateTensorLayout(s_tensor, tensor_layout);
+      if (err != Error::Ok) {
+        return err;
+      }
+      // Call load_into.
+      auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator);
+      if (!planned_ptr.ok()) {
+        return planned_ptr.error();
+      }
+      auto load_error =
+          named_data_map->load_data_into(fqn, planned_ptr.get(), nbytes);
+      if (load_error != Error::Ok) {
+        return load_error;
+      }
+
+      return planned_ptr;
+    }
+
+    // Constant, stored in PTE file.
+  } else if (data_buffer_idx > 0 && allocation_info == nullptr) {
+    auto const_data =
+        program->get_constant_buffer_data(data_buffer_idx, nbytes);
+    if (!const_data.ok()) {
+      return const_data.error();
+    }
+
+    // The const_cast is 'ok' here because the program and runtime should
+    // guarantee that this data is never modified.
+    return const_cast<void*>(const_data.get());
+
+    // Memory Planned, with initial state
+  } else if (data_buffer_idx > 0 && allocation_info != nullptr) {
+    auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator);
+    if (!planned_ptr.ok()) {
+      return planned_ptr.error();
+    }
+    auto err = TensorParser::load_mutable_subsegment_into(
+        program, 0, s_tensor->data_buffer_idx(), nbytes, planned_ptr.get());
+
+    if (err != Error::Ok) {
+      return err;
+    }
+    return planned_ptr;
+
+    // Memory planned, no initial state
+  } else if (data_buffer_idx == 0 && allocation_info != nullptr) {
+    return getMemPlannedPtr(allocation_info, nbytes, allocator);
+
+    // Pointer recived at runtime
+  } else { // data_buffer_idx == 0 && allocation_info == nullptr,
+    return nullptr;
+  }
 }
 
 } // namespace deserialization
-} // namespace aten
-} // namespace runtime
+} // namespace ET_RUNTIME_NAMESPACE
 } // namespace executorch

runtime/executor/tensor_parser_portable.cpp

@@ -11,6 +11,7 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_dimension_limit.h>
 #include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
@@ -62,6 +63,13 @@ Result<Tensor> parseTensor(
   const auto serialized_sizes = s_tensor->sizes()->data();
   const auto dim = s_tensor->sizes()->size();
 
+  ET_CHECK_OR_RETURN_ERROR(
+      dim <= kTensorDimensionLimit,
+      InvalidProgram,
+      "Tensor rank too large %" PRIu32 " > %zu",
+      dim,
+      kTensorDimensionLimit)
+
   ET_CHECK_OR_RETURN_ERROR(
       s_tensor->dim_order() != nullptr,
       InvalidProgram,