Update on "[ET-VK] Introduce vTensorPtr to prevent reference invalidation and remove get_val() API"

## Context

Currently when writing operators developers will save a reference to a `vTensor` retrieved from a `ComputeGraph`'s list of `values_` like so:

```
vTensor& vten = graph.get_val(vref).toTensor();
```

However, this is dangerous since if any values are added once the reference has been stored, `values_` which is a `std::vector` may have been resized and therefore have its contents moved, meaning the reference is now invalid.

To protect against this, this changeset introduces the `vTensorPtr` class which is a wrapper around a `vTensor*`.  When constructed, it will increment a counter in the `ComputeGraph` instance, and when destroyed it will decrement the counter. `ComputeGraph` cannot add any values while the counter is not zero.

Since `Value` can be converted to other non-trivial types, this changeset also removes the `get_val` function entirely to guard against unsafe behaviour.

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

[ghstack-poisoned]

Author: SS-JIA

CMakeLists.txt

@@ -144,6 +144,8 @@ option(EXECUTORCH_BUILD_COREML "Build the Core ML backend" OFF)
 
 option(EXECUTORCH_BUILD_CUSTOM "Build the custom kernels" OFF)
 
+option(EXECUTORCH_BUILD_CUSTOM_OPS_AOT "Build the custom ops lib for AOT" OFF)
+
 option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER "Build the Data Loader extension"
        OFF)
 
@@ -185,12 +187,19 @@ cmake_dependent_option(
 cmake_dependent_option(EXECUTORCH_BUILD_CPUINFO "Build cpuinfo library." ON
                        "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF)
 
+if(EXECUTORCH_BUILD_CUSTOM_OPS_AOT)
+  set(EXECUTORCH_BUILD_CUSTOM ON)
+endif()
+
 if(EXECUTORCH_BUILD_CUSTOM)
   set(EXECUTORCH_BUILD_OPTIMIZED ON)
 endif()
 
 if(EXECUTORCH_BUILD_CPUINFO)
   # --- cpuinfo
+  set(ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG
+      ${CMAKE_POSITION_INDEPENDENT_CODE})
+  set(CMAKE_POSITION_INDEPENDENT_CODE ON)
   set(CPUINFO_SOURCE_DIR "backends/xnnpack/third-party/cpuinfo")
   set(CPUINFO_BUILD_TOOLS
       OFF
@@ -212,10 +221,15 @@ if(EXECUTORCH_BUILD_CPUINFO)
       CACHE STRING "")
   set(CLOG_SOURCE_DIR "${CPUINFO_SOURCE_DIR}/deps/clog")
   add_subdirectory("${CPUINFO_SOURCE_DIR}")
+  set(CMAKE_POSITION_INDEPENDENT_CODE
+      ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG})
 endif()
 
 if(EXECUTORCH_BUILD_PTHREADPOOL)
   # --- pthreadpool
+  set(ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG
+      ${CMAKE_POSITION_INDEPENDENT_CODE})
+  set(CMAKE_POSITION_INDEPENDENT_CODE ON)
   set(PTHREADPOOL_SOURCE_DIR "backends/xnnpack/third-party/pthreadpool")
   set(PTHREADPOOL_BUILD_TESTS
       OFF
@@ -235,6 +249,8 @@ if(EXECUTORCH_BUILD_PTHREADPOOL)
        CACHE STRING "")
   endif()
   add_subdirectory("${PTHREADPOOL_SOURCE_DIR}")
+  set(CMAKE_POSITION_INDEPENDENT_CODE
+      ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG})
 endif()
 
 if(NOT PYTHON_EXECUTABLE)
@@ -546,6 +562,9 @@ if(EXECUTORCH_BUILD_PYBIND)
     list(APPEND _dep_libs custom_ops)
   endif()
 
+  if(EXECUTORCH_BUILD_CUSTOM_OPS_AOT)
+    list(APPEND _dep_libs custom_ops_aot_lib)
+  endif()
   # compile options for pybind
 
   set(_pybind_compile_options -Wno-deprecated-declarations -fPIC -frtti

backends/qualcomm/CMakeLists.txt

@@ -253,7 +253,7 @@ target_link_libraries(qnn_executorch_backend
     qnn_executorch_header
     qnn_schema
     qnn_manager
-    executorch
+    executorch_no_prim_ops
     qcir_utils
 )
 target_link_libraries(utils

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -37,6 +37,7 @@ namespace vkcompute {
   }
 
 VALUE_PTR_CLASS_IMPL(vTensorPtr, vTensor, Tensor)
+VALUE_PTR_CLASS_IMPL(TensorRefPtr, TensorRef, TensorRef)
 VALUE_PTR_CLASS_IMPL(StagingPtr, api::StorageBuffer, Staging)
 VALUE_PTR_CLASS_IMPL(IntListPtr, std::vector<int64_t>, IntList)
 VALUE_PTR_CLASS_IMPL(DoubleListPtr, std::vector<double>, DoubleList)
@@ -195,18 +196,17 @@ ValueRef ComputeGraph::add_tensor(
 }
 
 ValueRef ComputeGraph::add_tensor_like(
-    const ValueRef vref,
+    const ValueRef idx,
     const api::StorageType storage_type,
     const api::GPUMemoryLayout memory_layout) {
-  TensorRef tref = get_tref(vref);
-  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
+  return add_tensor(
+      get_sizes_of(idx), get_dtype_of(idx), storage_type, memory_layout);
 }
 
 ValueRef ComputeGraph::add_tensor_like(
-    const ValueRef vref,
+    const ValueRef idx,
     const api::GPUMemoryLayout memory_layout) {
-  TensorRef tref = get_tref(vref);
-  return add_tensor(tref.sizes, tref.dtype, memory_layout);
+  return add_tensor(get_sizes_of(idx), get_dtype_of(idx), memory_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -55,6 +55,7 @@ class ComputeGraph;
   };
 
 DECL_VALUE_PTR_CLASS(vTensorPtr, vTensor)
+DECL_VALUE_PTR_CLASS(TensorRefPtr, TensorRef)
 DECL_VALUE_PTR_CLASS(StagingPtr, api::StorageBuffer)
 DECL_VALUE_PTR_CLASS(IntListPtr, std::vector<int64_t>)
 DECL_VALUE_PTR_CLASS(DoubleListPtr, std::vector<double>)
@@ -132,6 +133,7 @@ class ComputeGraph final {
   }
 
   GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(vTensorPtr, tensor, Tensor)
+  GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(TensorRefPtr, tref, TensorRef)
   GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(StagingPtr, staging, Staging)
   GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(IntListPtr, int_list, IntList)
   GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(DoubleListPtr, double_list, DoubleList)
@@ -148,7 +150,6 @@ class ComputeGraph final {
     return values_.at(idx).is##type_name();                         \
   }
 
-  GET_AND_CHECK_VAL_AS_TYPE_FNS(TensorRef, tref, TensorRef)
   GET_AND_CHECK_VAL_AS_TYPE_FNS(int64_t, int, Int)
   GET_AND_CHECK_VAL_AS_TYPE_FNS(double, double, Double)
   GET_AND_CHECK_VAL_AS_TYPE_FNS(bool, bool, Bool)
@@ -392,6 +393,7 @@ class ComputeGraph final {
   //
 
   friend class vTensorPtr;
+  friend class TensorRefPtr;
   friend class StagingPtr;
   friend class IntListPtr;
   friend class DoubleListPtr;

backends/vulkan/runtime/graph/ops/PrepackNode.cpp

@@ -56,11 +56,11 @@ api::StorageBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) {
     return staging;
   }
 
-  TensorRef tref = graph->get_tref(tref_);
-  size_t numel = api::utils::multiply_integers(tref.sizes);
-  api::StorageBuffer staging(graph->context(), tref.dtype, numel);
-  size_t nbytes = numel * api::element_size(tref.dtype);
-  copy_ptr_to_staging(tref.data, staging, nbytes);
+  TensorRefPtr tref = graph->get_tref(tref_);
+  size_t numel = api::utils::multiply_integers(tref->sizes);
+  api::StorageBuffer staging(graph->context(), tref->dtype, numel);
+  size_t nbytes = numel * api::element_size(tref->dtype);
+  copy_ptr_to_staging(tref->data, staging, nbytes);
   return staging;
 }
 

backends/vulkan/runtime/graph/ops/impl/Conv2d.cpp

@@ -17,6 +17,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 
+#include <iostream>
+
 namespace vkcompute {
 
 void resize_conv2d_node(
@@ -35,8 +37,8 @@ void resize_conv2d_node(
     new_out_sizes.at(ndim - 4) = self->sizes().at(ndim - 4);
   }
 
-  TensorRef weight_ref = graph->get_tref(extra_args[0]);
-  const auto& weight_sizes = weight_ref.sizes;
+  TensorRefPtr weight_ref = graph->get_tref(extra_args[0]);
+  const auto& weight_sizes = weight_ref->sizes;
   new_out_sizes.at(ndim - 3) =
       transposed ? weight_sizes.at(ndim - 3) : weight_sizes.at(ndim - 4);
 
@@ -59,11 +61,14 @@ ValueRef prepack_biases(
     const ValueRef vref,
     const ValueRef weight,
     const bool transposed) {
-  TensorRef tref = graph.get_tref(weight);
-  const int64_t out_channels = transposed ? tref.sizes.at(1) : tref.sizes.at(0);
+  auto sizes = graph.get_sizes_of(weight);
+  const int64_t out_channels = transposed ? sizes.at(1) : sizes.at(0);
 
   ValueRef v = graph.add_tensor(
-      {out_channels}, tref.dtype, api::kTexture2D, api::kWidthPacked);
+      {out_channels},
+      graph.get_dtype_of(weight),
+      api::kTexture2D,
+      api::kWidthPacked);
   vTensorPtr t = graph.get_tensor(v);
 
   api::ShaderInfo shader = get_nchw_to_image_shader(*t);
@@ -102,7 +107,7 @@ api::ShaderInfo get_conv2d_shader(
     case Conv2dMethod::Depthwise:
       kernel_name = "conv2d_dw";
       if (!prepack_weights) {
-        const auto weight_sizes = graph.get_tref(weight).sizes;
+        const auto& weight_sizes = graph.get_tref(weight)->sizes;
         if (weight_sizes.at(2) == 3 && weight_sizes.at(3) == 3) {
           kernel_name += "_output_tile_3x3";
         }
@@ -180,12 +185,12 @@ ValueRef prepack_weights(
     ComputeGraph& graph,
     const ValueRef vref,
     const Conv2dMethod method) {
-  const auto original_sizes = graph.get_tref(vref).sizes;
-  const auto final_sizes = get_final_sizes(graph.get_tref(vref).sizes, method);
+  const auto original_sizes = graph.get_sizes_of(vref);
+  const auto final_sizes = get_final_sizes(original_sizes, method);
 
   ValueRef v = graph.add_tensor(
       final_sizes,
-      graph.get_tref(vref).dtype,
+      graph.get_dtype_of(vref),
       api::kTexture2D,
       api::kChannelsPacked);
   vTensorPtr t = graph.get_tensor(v);
@@ -239,7 +244,7 @@ Conv2dParams create_conv2d_params(
       p.kernel_size.data[1] +
           (p.kernel_size.data[1] - 1) * (p.dilation.data[1] - 1),
   });
-  const auto weight_sizes = graph.get_tref(weight).sizes;
+  const auto weight_sizes = graph.get_sizes_of(weight);
   const int32_t in_group_size =
       api::utils::safe_downcast<int32_t>(api::utils::align_up(
           transposed ? weight_sizes.at(0) : weight_sizes.at(1), INT64_C(4)));
@@ -267,7 +272,7 @@ Conv2dMethod get_conv2d_method(
     const ValueRef weight,
     const int64_t groups,
     const bool transposed) {
-  const auto weight_sizes = graph.get_tref(weight).sizes;
+  const auto weight_sizes = graph.get_sizes_of(weight);
   if (!transposed && weight_sizes.at(0) == groups && weight_sizes.at(1) == 1) {
     return Conv2dMethod::Depthwise;
   }

backends/vulkan/runtime/graph/ops/impl/Staging.cpp

@@ -97,7 +97,7 @@ ValueRef prepack_if_tensor_ref(
 ValueRef prepack_if_tensor_ref(ComputeGraph& graph, const ValueRef v) {
   if (graph.val_is_tref(v)) {
     api::GPUMemoryLayout layout =
-        graph.suggested_memory_layout(graph.get_tref(v).sizes);
+        graph.suggested_memory_layout(graph.get_tref(v)->sizes);
     return prepack(graph, v, layout);
   } else {
     return v;

backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.cpp

@@ -21,7 +21,7 @@ api::utils::ivec2 make_ivec2_kernel_size(
   if (kernel_size_only) {
     return make_ivec2_from_list(graph, weight);
   } else {
-    const auto weight_sizes = graph.get_tref(weight).sizes;
+    const auto weight_sizes = graph.get_tref(weight)->sizes;
     return api::utils::make_ivec2({weight_sizes.at(3), weight_sizes.at(2)});
   }
 }

backends/vulkan/test/op_tests/utils/codegen.py

@@ -248,8 +248,8 @@ def virtual_resize(self, ref: ValueRefList) -> str:
         assert ref.src_cpp_type == AT_TENSOR and ref.is_in
         if self.prepack_ref(ref):
             return ""
-        ret_str = f"{self.graph}{self.dot}get_val({ref.name}.value).toTensor()"
-        ret_str += f".virtual_resize({ref.src_cpp_name}.sizes().vec());\n"
+        ret_str = f"{self.graph}{self.dot}get_tensor({ref.name}.value)"
+        ret_str += f"->virtual_resize({ref.src_cpp_name}.sizes().vec());\n"
         return ret_str
 
     def copy_into_staging(self, ref: ValueRefList) -> str:

backends/xnnpack/CMakeLists.txt

@@ -81,7 +81,7 @@ add_library(xnnpack_backend STATIC ${_xnnpack_backend__srcs})
 target_link_libraries(xnnpack_backend
                       PRIVATE
                       ${xnnpack_third_party}
-                      executorch
+                      executorch_no_prim_ops
                       xnnpack_schema)
 
 target_include_directories(xnnpack_backend