rebase to pacify GitHub again

[ghstack-poisoned]

Author: swolchok

backends/apple/coreml/runtime/delegate/multiarray.mm

@@ -124,7 +124,10 @@ bool init_bnns_descriptor(BNNSNDArrayDescriptor& bnns_descriptor, const MultiArr
 
 bool copy_using_bnns(const MultiArray& src, MultiArray& dst) {
     if (src.layout().dataType() != dst.layout().dataType()) {
-        return false;
+        // Copying from FP16 to FP32 is supported and this is a common use case
+        if (!(src.layout().dataType() == MultiArray::DataType::Float16 && dst.layout().dataType() == MultiArray::DataType::Float32)) {
+            return false;
+        }
     }
     if (dst.layout().num_bytes() < src.layout().num_bytes()) {
         return false;

backends/arm/quantizer/arm_quantizer.py

@@ -101,18 +101,20 @@ def get_symmetric_quantization_config(
     weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = (
         MinMaxObserver
     )
+
     # Determine the right observer/fake-quant constructor
     if is_qat:
-        # Set plain fake-quant with true min/max
-        weight_observer_or_fake_quant_ctr = FakeQuantize
+        if is_per_channel:
+            weight_observer_or_fake_quant_ctr = PerChannelMinMaxObserver
+        else:
+            # Set plain fake-quant with true min/max
+            weight_observer_or_fake_quant_ctr = FakeQuantize
     else:
         # PTQ: set min/max observer
         weight_observer_or_fake_quant_ctr = (
             PerChannelMinMaxObserver if is_per_channel else MinMaxObserver
         )
 
-    extra_args = {"eps": 2**-12}
-
     weight_quantization_spec = QuantizationSpec(
         dtype=torch.int8,
         quant_min=weight_qmin,

backends/arm/test/misc/test_bn_relu_folding_qat.py

@@ -40,13 +40,17 @@ def forward(self, x: torch.Tensor):
 
 
 models = {
-    "conv_bn_relu": ConvModule(batch_norm=True),
-    "conv_relu": ConvModule(batch_norm=False),
+    # name : (model, is_per_channel)
+    "conv_bn_relu_per_channel": (ConvModule(batch_norm=True), True),
+    "conv_relu_per_channel": (ConvModule(batch_norm=False), True),
+    "conv_bn_relu_per_tensor": (ConvModule(batch_norm=True), False),
+    "conv_relu_per_tensor": (ConvModule(batch_norm=False), False),
 }
 
 
-@common.parametrize("model", models)
-def test_qat_tosa_INT(model: torch.nn.Module):
+@common.parametrize("test_data", models)
+def test_qat_tosa_INT(test_data):
+    model, per_channel = test_data
     pipeline = TosaPipelineINT[input_t1](model, model.test_data, [], [], qtol=1)
     tosa_version = conftest.get_option("tosa_version")
     tosa_profiles = {
@@ -59,7 +63,7 @@ def test_qat_tosa_INT(model: torch.nn.Module):
         Quantize(
             quantizer=quantizer,
             quantization_config=get_symmetric_quantization_config(
-                is_qat=True, is_per_channel=False
+                is_qat=True, is_per_channel=per_channel
             ),
             is_qat=True,
         ),

backends/arm/tosa/dialect/TARGETS

@@ -1,6 +1,36 @@
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 
+python_library(
+    name = "core",
+    srcs = [
+        "lib.py",
+        "ops_registration.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/arm:tosa_specification",
+        "//executorch/exir/dialects:lib",
+    ],
+)
+
+python_library(
+    name = "ops",
+    srcs = glob(["ops/*.py"]),
+    deps = [
+        ":core",
+        "//caffe2:torch",
+        "//executorch/backends/arm:tosa_specification",
+    ],
+)
+
 python_library(
     name = "lib",
-    srcs = glob(["*.py"]),
+    srcs = ["__init__.py"],
+    deps = [
+        ":core",
+        ":ops",
+        "//caffe2:torch",
+        "//executorch/backends/arm:tosa_specification",
+        "//executorch/exir/dialects:lib",
+    ],
 )

backends/vulkan/runtime/VulkanBackend.cpp

@@ -22,6 +22,7 @@
 #include <executorch/runtime/core/event_tracer_hooks_delegate.h>
 #endif // ET_EVENT_TRACER_ENABLED
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/platform/compiler.h>
 #include <executorch/runtime/platform/profiler.h>
 
@@ -47,6 +48,7 @@ using executorch::runtime::Error;
 using executorch::runtime::EValue;
 using executorch::runtime::FreeableBuffer;
 using executorch::runtime::kTensorDimensionLimit;
+using executorch::runtime::NamedDataMap;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
 
@@ -66,14 +68,6 @@ using BytesVector =
     const flatbuffers::Vector<flatbuffers::Offset<vkgraph::VkBytes>>*;
 using UIntVector = const flatbuffers::Vector<uint32_t>*;
 
-const uint8_t* get_constant_data_ptr(
-    VkGraphPtr flatbuffer_graph,
-    const int32_t buffer_idx,
-    const uint8_t* constant_data) {
-  VkBytesPtr constant_bytes = flatbuffer_graph->constants()->Get(buffer_idx);
-  return constant_data + constant_bytes->offset();
-}
-
 vkapi::ScalarType get_scalar_type(const vkgraph::VkDataType& vk_datatype) {
   switch (vk_datatype) {
     case vkgraph::VkDataType::BOOL:
@@ -166,17 +160,22 @@ class GraphBuilder {
   ComputeGraph* compute_graph_;
   VkGraphPtr flatbuffer_;
   const uint8_t* constant_data_;
+  const NamedDataMap* named_data_map_;
+  std::vector<FreeableBuffer> loaded_buffers_from_map_;
 
   std::vector<ValueRef> ref_mapping_;
 
  public:
   explicit GraphBuilder(
       ComputeGraph* compute_graph,
       VkGraphPtr flatbuffer,
-      const uint8_t* constant_data)
+      const uint8_t* constant_data,
+      const NamedDataMap* named_data_map)
       : compute_graph_(compute_graph),
         flatbuffer_(flatbuffer),
         constant_data_(constant_data),
+        named_data_map_(named_data_map),
+        loaded_buffers_from_map_(),
         ref_mapping_() {}
 
   void resize(uint32_t size) {
@@ -212,10 +211,27 @@ class GraphBuilder {
 
     ValueRef ref;
     if (tensor_fb->constant_id() >= 0) {
-      const uint8_t* tensor_data = get_constant_data_ptr(
-          flatbuffer_, tensor_fb->constant_id(), constant_data_);
+      VkBytesPtr constant_bytes =
+          flatbuffer_->constants()->Get(tensor_fb->constant_id());
 
-      ref = compute_graph_->add_tensorref(dims_vector, dtype, tensor_data);
+      if (constant_bytes->named_key() != nullptr &&
+          constant_bytes->offset() == UINT64_MAX &&
+          named_data_map_ != nullptr) {
+        const std::string& data_name = constant_bytes->named_key()->str();
+        Result<FreeableBuffer> buffer =
+            named_data_map_->get_data(data_name.c_str());
+
+        VK_CHECK_COND(
+            buffer.ok(),
+            "Failed to get constant data for key %s from named_data_map. Error code: %u",
+            data_name.c_str(),
+            static_cast<uint32_t>(buffer.error()));
+        ref = compute_graph_->add_tensorref(
+            dims_vector, dtype, std::move(buffer.get()));
+      } else {
+        const uint8_t* tensor_data = constant_data_ + constant_bytes->offset();
+        ref = compute_graph_->add_tensorref(dims_vector, dtype, tensor_data);
+      }
     } else {
       ref = compute_graph_->add_tensor(
           dims_vector,
@@ -479,8 +495,10 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     return true;
   }
 
-  ET_NODISCARD Error
-  compileModel(const void* buffer_pointer, ComputeGraph* compute_graph) const {
+  ET_NODISCARD Error compileModel(
+      const void* buffer_pointer,
+      ComputeGraph* compute_graph,
+      const NamedDataMap* named_data_map) const {
     Result<VulkanDelegateHeader> header =
         VulkanDelegateHeader::parse(buffer_pointer);
 
@@ -506,7 +524,8 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
 
     VkGraphPtr flatbuffer_graph = vkgraph::GetVkGraph(flatbuffer_data);
 
-    GraphBuilder builder(compute_graph, flatbuffer_graph, constant_data);
+    GraphBuilder builder(
+        compute_graph, flatbuffer_graph, constant_data, named_data_map);
 
     builder.build_graph();
 
@@ -532,7 +551,8 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     graph_config.external_adapter = vkapi::set_and_get_external_adapter();
     new (compute_graph) ComputeGraph(graph_config);
 
-    Error err = compileModel(processed->data(), compute_graph);
+    const NamedDataMap* named_data_map = context.get_named_data_map();
+    Error err = compileModel(processed->data(), compute_graph, named_data_map);
 
     // This backend does not need its processed data after compiling the
     // model.

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -897,6 +897,16 @@ VkMemoryRequirements vTensor::get_memory_requirements() const {
   return {};
 }
 
+bool vTensor::memory_is_bound() const {
+  switch (storage_type()) {
+    case utils::kBuffer:
+      return storage_->buffer_.has_memory();
+    case utils::kTexture2D:
+    case utils::kTexture3D:
+      return storage_->image_.has_memory();
+  }
+}
+
 void vTensor::bind_allocation(const vkapi::Allocation& allocation) {
   switch (storage_type()) {
     case utils::kBuffer:
@@ -909,6 +919,18 @@ void vTensor::bind_allocation(const vkapi::Allocation& allocation) {
   }
 }
 
+void vTensor::acquire_allocation(vkapi::Allocation&& allocation) {
+  switch (storage_type()) {
+    case utils::kBuffer:
+      storage_->buffer_.acquire_allocation(std::move(allocation));
+      break;
+    case utils::kTexture2D:
+    case utils::kTexture3D:
+      storage_->image_.acquire_allocation(std::move(allocation));
+      break;
+  }
+}
+
 void vTensor::update_metadata() {
   numel_ = utils::multiply_integers(sizes_);
   strides_ = calculate_strides(sizes_, dim_order_);

backends/vulkan/runtime/api/containers/Tensor.h

@@ -560,6 +560,12 @@ class vTensor final {
    */
   VmaAllocationCreateInfo get_allocation_create_info() const;
 
+  /*
+   * Checks if the tensor's underlying buffer or image resource is bound to a
+   * memory allocation.
+   */
+  bool memory_is_bound() const;
+
   /*
    * Return the VkMemoryRequirements of the underlying resource
    */
@@ -570,6 +576,11 @@ class vTensor final {
    */
   void bind_allocation(const vkapi::Allocation& allocation);
 
+  /*
+   * Binds and acquires a rvalue memory allocation
+   */
+  void acquire_allocation(vkapi::Allocation&& allocation);
+
  private:
   /*
    * Assuming sizes, dim order, or axis mapping was modified, recompute all