Update on "[ET-VK] Introduce add_tensor overloads consuming TensorRef"

From ssjia:
> we should always make sure to store references produced from `graph.get_val()` only after any calls to `graph.add_*()` (i.e. modifications to the values list) are made. This is because `graph.values_`, being a `std::vector`, will reallocate with more space and move its contents if the current allocation is not sufficient. This means that if you store a reference then call `graph.add_*()` then the underlying resource the reference points to may have been moved.

I think we can guard against this behavior by passing a `TensorRef` directly, and never having to declare a variable `TensorRef& tref` in the caller's scope. An example is shown in `Staging.cpp`. We could have it consume `ValueRef` for brevity of the passing parameter but IMO it hinders readability.

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

[ghstack-poisoned]

Author: Jorge Pineda

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -123,13 +123,6 @@ ValueRef ComputeGraph::add_tensor(
   return idx;
 }
 
-ValueRef ComputeGraph::add_tensor(
-    TensorRef& tref,
-    const api::StorageType storage_type,
-    const api::GPUMemoryLayout memory_layout) {
-  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
-}
-
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const api::ScalarType dtype,
@@ -139,9 +132,18 @@ ValueRef ComputeGraph::add_tensor(
       sizes, dtype, suggested_storage_type(), memory_layout, shared_object_idx);
 }
 
-ValueRef ComputeGraph::add_tensor(
-    TensorRef& tref,
+ValueRef ComputeGraph::add_tensor_like(
+    const ValueRef vref,
+    const api::StorageType storage_type,
+    const api::GPUMemoryLayout memory_layout) {
+  TensorRef& tref = get_val(vref).toTensorRef();
+  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
+}
+
+ValueRef ComputeGraph::add_tensor_like(
+    const ValueRef vref,
     const api::GPUMemoryLayout memory_layout) {
+  TensorRef& tref = get_val(vref).toTensorRef();
   return add_tensor(tref.sizes, tref.dtype, memory_layout);
 }
 

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -174,14 +174,6 @@ class ComputeGraph final {
       const api::GPUMemoryLayout memory_layout,
       const int64_t shared_object_idx = -1);
 
-  /*
-   * Add a `vTensor` value to the graph with the properties of `tref`.
-   */
-  ValueRef add_tensor(
-      TensorRef& tref,
-      const api::StorageType storage_type,
-      const api::GPUMemoryLayout memory_layout);
-
   /*
    * Add a `vTensor` value to the graph with the specified properties. The
    * suggested storage type will be used to construct the `vTensor`.
@@ -192,14 +184,6 @@ class ComputeGraph final {
       const api::GPUMemoryLayout memory_layout,
       const int64_t shared_object_idx = -1);
 
-  /*
-   * Add a `vTensor` value to the graph with the properties of `tref`. The
-   * suggested storage type will be used to construct the `vTensor`.
-   */
-  ValueRef add_tensor(
-      TensorRef& tref,
-      const api::GPUMemoryLayout memory_layout);
-
   /*
    * Add a `vTensor` value to the graph with the specified properties. The
    * suggested storage type and memory layout will be used to construct the
@@ -210,6 +194,22 @@ class ComputeGraph final {
       const api::ScalarType dtype,
       const int64_t shared_object_idx = -1);
 
+  /*
+   * Add a `vTensor` value to the graph with the properties of `vref`.
+   */
+  ValueRef add_tensor_like(
+      const ValueRef vref,
+      const api::StorageType storage_type,
+      const api::GPUMemoryLayout memory_layout);
+
+  /*
+   * Add a `vTensor` value to the graph with the properties of `vref`. The
+   * suggested storage type will be used to construct the `vTensor`.
+   */
+  ValueRef add_tensor_like(
+      const ValueRef vref,
+      const api::GPUMemoryLayout memory_layout);
+
   /*
    * Add a `TensorRef` value to the graph with the specific properties. A
    * `TensorRef` is a reference to a `vTensor` whose data is stored in an

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

@@ -63,7 +63,7 @@ ValueRef prepack(
     ComputeGraph& graph,
     const ValueRef vref,
     const api::GPUMemoryLayout layout) {
-  ValueRef v = graph.add_tensor(graph.get_val(vref).toTensorRef(), layout);
+  ValueRef v = graph.add_tensor_like(vref, layout);
   vTensor& t = graph.get_val(v).toTensor();
 
   api::ShaderInfo shader = get_nchw_to_image_shader(t);