Merge branch 'main' into toupstream/bin_node

Author: SaoirseARM

.github/workflows/mlx.yml

@@ -66,7 +66,11 @@ jobs:
         echo "::endgroup::"
 
         echo "::group::Build test runners"
-        ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner -j$(( $(sysctl -n hw.ncpu) - 1 ))
+        ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test -j$(( $(sysctl -n hw.ncpu) - 1 ))
+        echo "::endgroup::"
+
+        echo "::group::Run mutable-state (multi-session) unit test"
+        ./cmake-out/backends/mlx/test/mlx_mutable_state_test
         echo "::endgroup::"
 
         echo "::group::Run op unit tests"

backends/arm/operators/op_tosa_identity.py

@@ -3,42 +3,28 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, List
-
-import torch
 import tosa_serializer as ts
 
-from executorch.backends.arm.operators.node_visitor import (
-    NodeVisitor,
-    register_node_visitor,
-)
-from executorch.backends.arm.operators.operator_validation_utils import (
-    validate_num_inputs,
-    validate_same_dtype,
-    validate_valid_dtype,
+from executorch.backends.arm.operators.node_visitor import register_node_visitor
+from executorch.backends.arm.operators.simple_node_visitor import (
+    SimpleNodeVisitor,
+    SimpleNodeVisitorConfig,
 )
-from executorch.backends.arm.tosa.mapping import TosaArg
 
 
 @register_node_visitor
-class IdentityVisitor(NodeVisitor):
+class IdentityVisitor(SimpleNodeVisitor):
     """Lower the TOSA IDENTITY op."""
 
     target = "tosa.IDENTITY.default"
 
-    def define_node(
-        self,
-        node: torch.fx.Node,
-        tosa_graph: Any,
-        inputs: List[TosaArg],
-        output: TosaArg,
-    ) -> None:
-        validate_num_inputs(self.target, inputs, 1)
-        validate_same_dtype(self.target, [inputs[0], output], ts)
-        validate_valid_dtype(
-            self.target,
-            [inputs[0], output],
-            [
+    @classmethod
+    def get_config(cls) -> SimpleNodeVisitorConfig:
+        return SimpleNodeVisitorConfig(
+            tosa_op=ts.Op.IDENTITY,
+            attr_method="IdentityAttribute",
+            num_inputs=1,
+            input_dtypes=[
                 ts.DType.BOOL,
                 ts.DType.INT8,
                 ts.DType.INT16,
@@ -49,16 +35,4 @@ def define_node(
                 ts.DType.FP8E4M3,
                 ts.DType.FP8E5M2,
             ],
-            self.tosa_spec,
-        )
-
-        attr = ts.TosaSerializerAttribute()
-        attr.IdentityAttribute()
-        self._serialize_operator(
-            node,
-            tosa_graph,
-            ts.Op.IDENTITY,
-            [inputs[0].name],
-            [output.name],
-            attr,
         )

backends/arm/quantizer/quantization_annotator.py

@@ -631,6 +631,16 @@ def _get_fixed_qparams_qspec(
 if _transpose_dimname is not None:
     _one_to_one_shared_input_qspec.add(_transpose_dimname)
 
+for _op in (
+    getattr(torch.ops.aten.moveaxis, "int", None),
+    getattr(torch.ops.aten.moveaxis, "intlist", None),
+    getattr(torch.ops.aten.movedim, "int", None),
+    getattr(torch.ops.aten.movedim, "intlist", None),
+):
+    if _op is not None:
+        _one_to_one_shared_input_qspec.add(_op)
+
+
 _one_to_one_shared_input_or_input_act_qspec: set[OpOverload] = {
     torch.ops.aten.alias.default,
     torch.ops.aten.clone.default,

backends/arm/runtime/VGFSetup.cpp

@@ -778,6 +778,73 @@ static bool find_memory_index(
       memory_type_out);
 }
 
+bool VgfRepr::map_persistent_io_memory() {
+  unmap_persistent_io_memory();
+
+  for (auto& io : IOs) {
+    if (io.memory == VK_NULL_HANDLE) {
+      ET_LOG(Error, "Cannot persistently map null Vulkan IO memory");
+      unmap_persistent_io_memory();
+      return false;
+    }
+
+    void* persistent_memory = nullptr;
+
+    // IO resources may alias the same VkDeviceMemory. Vulkan memory must not be
+    // mapped more than once at the same time, so map each unique memory once
+    // and share the returned pointer across aliased IO entries.
+    // Make sure that memory is HOST_VISIBLE and HOST_COHERENT.
+    bool found_existing_mapping = false;
+    auto mapped_memory_it = std::find_if(
+        persistent_mapped_memories.begin(),
+        persistent_mapped_memories.end(),
+        [&](const auto& mapped_memory) {
+          return mapped_memory.memory == io.memory;
+        });
+
+    if (mapped_memory_it != persistent_mapped_memories.end()) {
+      persistent_memory = mapped_memory_it->data;
+      found_existing_mapping = true;
+    }
+
+    if (!found_existing_mapping) {
+      VkResult result = vkMapMemory(
+          vk_device, io.memory, 0, VK_WHOLE_SIZE, 0, &persistent_memory);
+      if (result != VK_SUCCESS) {
+        ET_LOG(
+            Error,
+            "Failed to persistently map Vulkan IO memory, error %d",
+            result);
+        unmap_persistent_io_memory();
+        return false;
+      }
+
+      persistent_mapped_memories.push_back(PersistentMappedMemory{
+          .memory = io.memory,
+          .data = persistent_memory,
+      });
+    }
+
+    io.persistent_memory = persistent_memory;
+  }
+
+  return true;
+}
+
+void VgfRepr::unmap_persistent_io_memory() {
+  for (const auto& mapped_memory : persistent_mapped_memories) {
+    if (mapped_memory.memory != VK_NULL_HANDLE &&
+        mapped_memory.data != nullptr) {
+      vkUnmapMemory(vk_device, mapped_memory.memory);
+    }
+  }
+  persistent_mapped_memories.clear();
+
+  for (auto& io : IOs) {
+    io.persistent_memory = nullptr;
+  }
+}
+
 VkResult allocate_memory(
     VkPhysicalDevice physical,
     VkDevice device,
@@ -1839,6 +1906,7 @@ bool VgfRepr::process_vgf(
                  VK_NULL_HANDLE,
                  tensor_memory,
                  {0, 0, 0},
+                 nullptr,
                  owns_memory,
                  true,
                  is_in});
@@ -1931,6 +1999,7 @@ bool VgfRepr::process_vgf(
                  VK_NULL_HANDLE,
                  buffer_memory,
                  {0, 0, 0},
+                 nullptr,
                  owns_memory,
                  true,
                  is_in});
@@ -2117,6 +2186,7 @@ bool VgfRepr::process_vgf(
                  image_memory,
                  staging_memory,
                  image_extent,
+                 nullptr,
                  true,
                  owns_image_memory,
                  is_in});
@@ -3433,6 +3503,15 @@ bool VgfRepr::process_vgf(
     vkEndCommandBuffer(vk_execute_cmd);
   }
 
+  {
+    VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_MAP_IO_MEMORY");
+
+    if (!map_persistent_io_memory()) {
+      ET_LOG(Error, "Failed to persistently map VGF IO memory");
+      return false;
+    }
+  }
+
   return true;
 }
 
@@ -3493,6 +3572,8 @@ bool VgfRepr::execute_vgf(executorch::runtime::EventTracer* event_tracer) {
 }
 
 void VgfRepr::free_vgf() {
+  unmap_persistent_io_memory();
+
   if (vk_timestamp_query_pool != VK_NULL_HANDLE) {
     vkDestroyQueryPool(vk_device, vk_timestamp_query_pool, nullptr);
     vk_timestamp_query_pool = VK_NULL_HANDLE;

backends/arm/runtime/VGFSetup.h

@@ -45,11 +45,17 @@ typedef struct IO {
   VkDeviceMemory image_memory;
   VkDeviceMemory memory;
   VkExtent3D image_extent;
+  void* persistent_memory = nullptr;
   bool owns_memory = true;
   bool owns_image_memory = true;
   bool is_input;
 } IO;
 
+typedef struct PersistentMappedMemory {
+  VkDeviceMemory memory = VK_NULL_HANDLE;
+  void* data = nullptr;
+} PersistentMappedMemory;
+
 typedef struct SegmentState {
   int segment_id = -1;
   bool use_data_graph_pipeline = true;
@@ -128,18 +134,19 @@ class VgfRepr {
   std::vector<SegmentState> segments;
   std::vector<ResourceAlloc> extra_allocs;
 
-  bool map_io(IO* io, void** handle) {
-    VkResult result =
-        vkMapMemory(vk_device, io->memory, 0, VK_WHOLE_SIZE, 0, handle);
-    if (result != VK_SUCCESS) {
-      ET_LOG(Error, "Failed to map Vulkan IO memory");
+  // Mapping to persistent IO memory
+  static bool map_io(IO* io, void** handle) {
+    if (io->persistent_memory == nullptr) {
+      ET_LOG(Error, "Vulkan IO memory is not persistently mapped");
       return false;
     }
+    *handle = io->persistent_memory;
     return true;
   }
 
-  void unmap_io(IO* io) {
-    vkUnmapMemory(vk_device, io->memory);
+  // Unmapping to persistent IO memory
+  static void unmap_io(IO* io) {
+    (void)io;
   }
 
   ~VgfRepr() {
@@ -168,6 +175,10 @@ class VgfRepr {
 
   bool init_timestamp_queries();
   void read_timestamp_queries(executorch::runtime::EventTracer* event_tracer);
+
+  std::vector<PersistentMappedMemory> persistent_mapped_memories;
+  bool map_persistent_io_memory();
+  void unmap_persistent_io_memory();
 };
 
 } // namespace vgf