Update on "[ET-VK] Introduce virtual_clone API to support view of view use cases + fix synchronization hazard with view tensors"

## Context

This diff fixes some hazards (not necessarily) bugs with view tensors.

### `virtual_clone` API

Consider the following sequence of calls which may be common in the view of view use case.

```
t1 = graph.add_tensor(...);
// t2 will have the same metadata as t1
t2 = graph.add_tensor_view(t1);
// t3 will also have the same metadata as t2 at this point.
t3 = graph.add_tensor_view(t2);

// t2 metadata will be updated correctly.
t2 = add_transpose_view_node(t1, 0, 1, t2);
// Unfortunately, this node will have an assumption that t3 has the same metadata as t2 to start. However, this is not true.
// As a result, t3 will have incorrect metadata after this node.
t3 = add_transpose_view_node(t2, 1, 2, t3);
```

To address this, the `virtual_clone` API is introduced which will allow view nodes to set the metadata of the output equal to the input before modifying the output.

### WAW synchronization hazards

`vTensorStorage` maintains a `last_access` state which facilitates inserting the correct memory barriers for the underlying `vkImage` or `vkBuffer`. However, when we create a tensor view, `last_access` is not shared between `vTensor` instances that use the same resource.

As as result, writing into a `vTensor` will not update the `last_access` of its views, and vice versa. Therefore, sebsequent accesses of the other tensor that references the same resource will result in a synchronization hazard.

This diff fixes this hazard in a bit of a crude way; if the `vTensor` is a copy, or has copies, then cowardly assume that it has been written to before the current access so that appropriate memory barriers are inserted. This was the selected solution because I thought that adding a map to track last access of tensors that share resources is a bit overkill when the assumption that the underlying resource has been written to before the current access should hold most of the time.

Differential Revision: [D63642092](https://our.internmc.facebook.com/intern/diff/D63642092/)

[ghstack-poisoned]

Author: SS-JIA

CMakeLists.txt

@@ -458,22 +458,22 @@ endif()
 add_subdirectory(schema)
 
 #
-# executorch_no_prim_ops: Minimal runtime library
+# executorch_core: Minimal runtime library
 #
 # The bare-minimum runtime library, supporting the Program and Method
 # interfaces. Does not contain any operators, including primitive ops. Does not
 # contain any backends.
 #
 
 # Remove any PAL-definition files from the sources.
-list(FILTER _executorch_no_prim_ops__srcs EXCLUDE REGEX
+list(FILTER _executorch_core__srcs EXCLUDE REGEX
      "runtime/platform/default/[^/]*.cpp$"
 )
 
 # Add the source file that maps to the requested default PAL implementation.
 if(EXECUTORCH_PAL_DEFAULT MATCHES "^(posix|minimal)$")
   message(STATUS "executorch: Using PAL default '${EXECUTORCH_PAL_DEFAULT}'")
-  list(APPEND _executorch_no_prim_ops__srcs
+  list(APPEND _executorch_core__srcs
        "runtime/platform/default/${EXECUTORCH_PAL_DEFAULT}.cpp"
   )
 else()
@@ -483,45 +483,49 @@ else()
   )
 endif()
 
-add_library(executorch_no_prim_ops ${_executorch_no_prim_ops__srcs})
-target_link_libraries(executorch_no_prim_ops PRIVATE program_schema)
+add_library(executorch_core ${_executorch_core__srcs})
+
+# Legacy name alias.
+add_library(executorch_no_prim_ops ALIAS executorch_core)
+
+target_link_libraries(executorch_core PRIVATE program_schema)
 if(EXECUTORCH_USE_DL)
   # Check if dl exists for this toolchain and only then link it.
   find_library(DL_LIBRARY_EXISTS NAMES dl)
   # Check if the library was found
   if(DL_LIBRARY_EXISTS)
-    target_link_libraries(executorch_no_prim_ops PRIVATE dl) # For dladdr()
+    target_link_libraries(executorch_core PRIVATE dl) # For dladdr()
   endif()
 endif()
 target_include_directories(
-  executorch_no_prim_ops PUBLIC ${_common_include_directories}
+  executorch_core PUBLIC ${_common_include_directories}
 )
-target_compile_options(executorch_no_prim_ops PUBLIC ${_common_compile_options})
+target_compile_options(executorch_core PUBLIC ${_common_compile_options})
 if(MAX_KERNEL_NUM)
   target_compile_definitions(
-    executorch_no_prim_ops PRIVATE MAX_KERNEL_NUM=${MAX_KERNEL_NUM}
+    executorch_core PRIVATE MAX_KERNEL_NUM=${MAX_KERNEL_NUM}
   )
 endif()
 
 if(EXECUTORCH_BUILD_PYBIND AND APPLE)
   # shared version
   add_library(
-    executorch_no_prim_ops_shared SHARED ${_executorch_no_prim_ops__srcs}
+    executorch_core_shared SHARED ${_executorch_core__srcs}
   )
-  target_link_libraries(executorch_no_prim_ops_shared PRIVATE program_schema)
+  target_link_libraries(executorch_core_shared PRIVATE program_schema)
   if(DL_LIBRARY_EXISTS)
     # For dladdr()
-    target_link_libraries(executorch_no_prim_ops_shared PRIVATE dl)
+    target_link_libraries(executorch_core_shared PRIVATE dl)
   endif()
   target_include_directories(
-    executorch_no_prim_ops_shared PUBLIC ${_common_include_directories}
+    executorch_core_shared PUBLIC ${_common_include_directories}
   )
   target_compile_options(
-    executorch_no_prim_ops_shared PUBLIC ${_common_compile_options}
+    executorch_core_shared PUBLIC ${_common_compile_options}
   )
   if(MAX_KERNEL_NUM)
     target_compile_definitions(
-      executorch_no_prim_ops_shared PRIVATE MAX_KERNEL_NUM=${MAX_KERNEL_NUM}
+      executorch_core_shared PRIVATE MAX_KERNEL_NUM=${MAX_KERNEL_NUM}
     )
   endif()
 endif()
@@ -534,7 +538,7 @@ endif()
 # any backends.
 #
 add_library(executorch ${_executorch__srcs})
-target_link_libraries(executorch PRIVATE executorch_no_prim_ops)
+target_link_libraries(executorch PRIVATE executorch_core)
 target_include_directories(executorch PUBLIC ${_common_include_directories})
 target_compile_options(executorch PUBLIC ${_common_compile_options})
 target_link_options_shared_lib(executorch)
@@ -570,7 +574,7 @@ endif()
 # Install `executorch` library as well as `executorch-config.cmake` under
 # ${CMAKE_INSTALL_PREFIX}/
 install(
-  TARGETS executorch executorch_no_prim_ops
+  TARGETS executorch executorch_core
   DESTINATION lib
   INCLUDES
   DESTINATION ${_common_include_directories}

backends/apple/coreml/CMakeLists.txt

@@ -134,7 +134,7 @@ target_include_directories(
   coremldelegate PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/runtime/util
 )
 target_include_directories(coremldelegate PRIVATE ${EXECUTORCH_ROOT}/..)
-target_link_libraries(coremldelegate PRIVATE executorch_no_prim_ops)
+target_link_libraries(coremldelegate PRIVATE executorch_core)
 
 if(EXECUTORCH_BUILD_DEVTOOLS)
   target_sources(coremldelegate PRIVATE ${SDK_SOURCES} ${PROTOBUF_SOURCES})
@@ -159,7 +159,7 @@ find_library(SQLITE_LIBRARY sqlite3)
 
 target_link_libraries(
   coremldelegate
-  PRIVATE executorch_no_prim_ops ${ACCELERATE_FRAMEWORK} ${COREML_FRAMEWORK}
+  PRIVATE executorch_core ${ACCELERATE_FRAMEWORK} ${COREML_FRAMEWORK}
           ${FOUNDATION_FRAMEWORK} ${SQLITE_LIBRARY}
 )
 
@@ -176,7 +176,7 @@ target_compile_options(coremldelegate PRIVATE "-fno-exceptions")
 
 if(EXECUTORCH_BUILD_DEVTOOLS)
   target_compile_options(
-    executorch_no_prim_ops PUBLIC -DET_EVENT_TRACER_ENABLED
+    executorch_core PUBLIC -DET_EVENT_TRACER_ENABLED
   )
   target_compile_options(coremldelegate PRIVATE "-frtti")
   target_compile_options(libprotobuf-lite PRIVATE "-frtti")

backends/apple/coreml/runtime/workspace/executorchcoreml.xcodeproj/project.pbxproj

@@ -105,7 +105,7 @@
 		C9E7D7962AB3F9BF00CCAE5D /* KeyValueStoreTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = C9E7D78E2AB3F9BF00CCAE5D /* KeyValueStoreTests.mm */; };
 		C9E7D7A22AB3FBB200CCAE5D /* CoreMLBackendDelegateTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = C9E7D7A12AB3FBB200CCAE5D /* CoreMLBackendDelegateTests.mm */; };
 		C9EC7E1B2BC73B3200A6B166 /* MultiArrayTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = C9EC7E1A2BC73B3200A6B166 /* MultiArrayTests.mm */; };
-		F24817E52BC655E100E80D98 /* libexecutorch_no_prim_ops.a in Frameworks */ = {isa = PBXBuildFile; fileRef = F24817E42BC655E100E80D98 /* libexecutorch_no_prim_ops.a */; };
+		F24817E52BC655E100E80D98 /* libexecutorch_core.a in Frameworks */ = {isa = PBXBuildFile; fileRef = F24817E42BC655E100E80D98 /* libexecutorch_core.a */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXCopyFilesBuildPhase section */
@@ -310,7 +310,7 @@
 		C9EA3FE52B73EF6300B7D7BD /* Accelerate.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Accelerate.framework; path = System/Library/Frameworks/Accelerate.framework; sourceTree = SDKROOT; };
 		C9EC7E092BC662A300A6B166 /* objc_array_util.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; name = objc_array_util.h; path = ../util/objc_array_util.h; sourceTree = "<group>"; };
 		C9EC7E1A2BC73B3200A6B166 /* MultiArrayTests.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; name = MultiArrayTests.mm; path = ../test/MultiArrayTests.mm; sourceTree = "<group>"; };
-		F24817E42BC655E100E80D98 /* libexecutorch_no_prim_ops.a */ = {isa = PBXFileReference; lastKnownFileType = archive.ar; name = libexecutorch_no_prim_ops.a; path = ../libraries/libexecutorch_no_prim_ops.a; sourceTree = "<group>"; };
+		F24817E42BC655E100E80D98 /* libexecutorch_core.a */ = {isa = PBXFileReference; lastKnownFileType = archive.ar; name = libexecutorch_core.a; path = ../libraries/libexecutorch_core.a; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -319,7 +319,7 @@
 			buildActionMask = 2147483647;
 			files = (
 				C94D510F2ABDF87500AF47FD /* Accelerate.framework in Frameworks */,
-				F24817E52BC655E100E80D98 /* libexecutorch_no_prim_ops.a in Frameworks */,
+				F24817E52BC655E100E80D98 /* libexecutorch_core.a in Frameworks */,
 				C94D510E2ABDF86800AF47FD /* libsqlite3.tbd in Frameworks */,
 				C94D50D92ABD7B2400AF47FD /* CoreML.framework in Frameworks */,
 				C99883862B95AD7D000953A3 /* libprotobuf-lite.a in Frameworks */,
@@ -540,7 +540,7 @@
 				C96560942AABFDCE005F8126 /* libsqlite3.tbd */,
 				C96560922AABF992005F8126 /* CoreML.framework */,
 				C96560902AABF982005F8126 /* Accelerate.framework */,
-				F24817E42BC655E100E80D98 /* libexecutorch_no_prim_ops.a */,
+				F24817E42BC655E100E80D98 /* libexecutorch_core.a */,
 				C965608D2AABF72A005F8126 /* libexecutorch.a */,
 			);
 			name = "Recovered References";

backends/apple/coreml/scripts/build_tests.sh

@@ -59,7 +59,7 @@ cmake --build "$CMAKE_PROTOBUF_BUILD_DIR_PATH"  -j9 -t libprotobuf-lite
 echo "ExecuTorch: Copying libraries"
 mkdir "$LIBRARIES_DIR_PATH"
 cp -f "$CMAKE_EXECUTORCH_BUILD_DIR_PATH/libexecutorch.a" "$LIBRARIES_DIR_PATH"
-cp -f "$CMAKE_EXECUTORCH_BUILD_DIR_PATH/libexecutorch_no_prim_ops.a" "$LIBRARIES_DIR_PATH"
+cp -f "$CMAKE_EXECUTORCH_BUILD_DIR_PATH/libexecutorch_core.a" "$LIBRARIES_DIR_PATH"
 cp -f "$CMAKE_PROTOBUF_BUILD_DIR_PATH/libprotobuf-lite.a" "$LIBRARIES_DIR_PATH"
 
 #Copy ExecuTorch headers

backends/apple/mps/CMakeLists.txt

@@ -75,7 +75,7 @@ target_link_libraries(
   mpsdelegate
   PRIVATE bundled_program
           mps_schema
-          executorch_no_prim_ops
+          executorch_core
           ${FOUNDATION_FRAMEWORK}
           ${METAL_FRAMEWORK}
           ${MPS_FRAMEWORK}

backends/arm/arm_backend.py

@@ -217,7 +217,7 @@ def preprocess(  # noqa: C901
         # const data directly. Path created and data written only in debug builds.
         tosa_graph = ts.TosaSerializer(artifact_path)
         graph_module = ArmPassManager().transform_to_backend_pipeline(
-            graph_module=edge_program.graph_module, compile_spec=compile_spec
+            exported_program=edge_program, compile_spec=compile_spec
         )
 
         node_visitors = get_node_visitors(edge_program)

backends/arm/arm_partitioner.py

@@ -58,6 +58,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.mm.default,
             exir_ops.edge.aten.repeat.default,
             exir_ops.edge.aten.relu.default,
+            exir_ops.edge.aten.rsqrt.default,
             exir_ops.edge.aten._softmax.default,
             exir_ops.edge.aten.slice_copy.Tensor,
             exir_ops.edge.aten.sub.Tensor,

backends/arm/operators/__init__.py

@@ -28,6 +28,7 @@
     op_quant,
     op_relu,
     op_repeat,
+    op_rsqrt,
     op_sigmoid,
     op_slice,
     op_softmax,

backends/arm/operators/op_rsqrt.py

@@ -0,0 +1,70 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# 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 List
+
+import numpy as np
+import serializer.tosa_serializer as ts
+import torch
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_quant_utils import (
+    dequantize_value,
+    get_quant_node_args,
+    QuantArgs,
+    quantize_value,
+)
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class RsqrtVisitor(NodeVisitor):
+    target = "aten.rsqrt.default"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        if is_quant_node:
+            # Assume quantized input is 8 bit.
+            # Create attribute for 8 bit table lookup.
+            input_node = node.all_input_nodes[0]
+            in_quantargs = get_quant_node_args(input_node)
+            output_node = list(node.users)[0]
+            out_quantargs = get_quant_node_args(output_node)
+            table = rsqrt_table_8bit(in_quantargs, out_quantargs)
+            table_attr = ts.TosaSerializerAttribute()
+            table_attr.TableAttribute(table)
+            tosa_graph.addOperator(
+                TosaOp.Op().TABLE, [inputs[0].name], [output.name], table_attr
+            )
+        else:
+            tosa_graph.addOperator(TosaOp.Op().RSQRT, [inputs[0].name], [output.name])
+
+
+def rsqrt_table_8bit(in_quantargs: QuantArgs, out_quantargs: QuantArgs):
+    """
+    Returns a table mapping 256 entries to rqsrt([qmin,qmax])
+    Reference: https://www.mlplatform.org/tosa/tosa_spec.html#_rsqrt
+    """
+
+    def rqsrt(x):
+        # Convert quantized input to floating point rqsrt input space.
+        v = dequantize_value(x, in_quantargs)
+        # Compute rqsrt.
+        v = 1 / np.sqrt(v)
+        # Convert rqsrt output back to quantized space.
+        return quantize_value(v, out_quantargs)
+
+    return [
+        rqsrt(x)
+        for x in np.linspace(in_quantargs.qmin, in_quantargs.qmax, 256, dtype=np.int8)
+    ]

backends/arm/passes/arm_pass_manager.py

@@ -22,6 +22,7 @@
 )
 from executorch.backends.arm.passes.remove_clone_pass import RemoveClonePass
 from executorch.backends.arm.passes.size_adjust_conv2d_pass import SizeAdjustConv2DPass
+from executorch.exir import ExportedProgram
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.pass_manager import PassManager
 
@@ -32,7 +33,7 @@ def _transform(self, graph_module: torch.fx.GraphModule):
         return self(graph_module).graph_module
 
     def transform_to_backend_pipeline(
-        self, graph_module: torch.fx.GraphModule, compile_spec: list[CompileSpec]
+        self, exported_program: ExportedProgram, compile_spec: list[CompileSpec]
     ):
         """Apply passes before transforming program to backend"""
         self.add_pass(SizeAdjustConv2DPass())
@@ -46,4 +47,4 @@ def transform_to_backend_pipeline(
                 if memory_format == "nhwc":
                     self.add_pass(AnnotateChannelsLastDimOrder())
 
-        return self._transform(graph_module)
+        return self._transform(exported_program.graph_module)

backends/arm/quantizer/quantization_annotation/one_to_one_annotator.py

@@ -35,7 +35,11 @@ def _annotate_one_to_one(
     Typical ops are ops implemented with a lookup table.
     """
     annotated_partitions = []
-    one_to_one_ops = (torch.ops.aten.exp.default, torch.ops.aten.log.default)
+    one_to_one_ops = {
+        torch.ops.aten.exp.default,
+        torch.ops.aten.log.default,
+        torch.ops.aten.rsqrt.default,
+    }
     for node in gm.graph.nodes:
         if node.op != "call_function" or node.target not in one_to_one_ops:
             continue

backends/arm/test/common.py

@@ -4,9 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 import os
 import shutil
 import subprocess
+import sys
 import tempfile
 
 import pytest
@@ -37,6 +39,7 @@ def pytest_configure(config):
                 "Tests are run with --arm_run_corstone300 but corstone300 FVP is not installed."
             )
         _enabled_options.append("corstone300")
+    logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
 
 def pytest_collection_modifyitems(config, items):