[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: jorgep31415

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -123,6 +123,13 @@ 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,
@@ -132,16 +139,18 @@ ValueRef ComputeGraph::add_tensor(
       sizes, dtype, suggested_storage_type(), memory_layout, shared_object_idx);
 }
 
+ValueRef ComputeGraph::add_tensor(
+    TensorRef& tref,
+    const api::GPUMemoryLayout memory_layout) {
+  return add_tensor(tref.sizes, tref.dtype, memory_layout);
+}
+
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const api::ScalarType dtype,
     const int64_t shared_object_idx) {
   return add_tensor(
-      sizes,
-      dtype,
-      suggested_storage_type(),
-      suggested_memory_layout(sizes),
-      shared_object_idx);
+      sizes, dtype, suggested_memory_layout(sizes), shared_object_idx);
 }
 
 ValueRef ComputeGraph::add_tensorref(

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -172,7 +172,15 @@ class ComputeGraph final {
       const api::ScalarType dtype,
       const api::StorageType storage_type,
       const api::GPUMemoryLayout memory_layout,
-      const int64_t shared_object_idx);
+      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
@@ -184,14 +192,22 @@ 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
    * `vTensor`.
    */
   ValueRef add_tensor(
       const std::vector<int64_t>& sizes,
-      const api::ScalarType dtype = api::ScalarType::Float,
+      const api::ScalarType dtype,
       const int64_t shared_object_idx = -1);
 
   /*

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

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