[ET-VK] Introduce add_tensor_like consuming ValueRef

Pull Request resolved: pytorch/executorch#2835

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.

We can guard against this behavior by passing a `ValueRef` directly, and never having to declare a variable of types `ValueRef&/TensorRef&` in the caller's scope. An example is shown in `Staging.cpp`.

ghstack-source-id: 221175880
@exported-using-ghexport

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

Author: Jorge Pineda

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -132,16 +132,27 @@ ValueRef ComputeGraph::add_tensor(
       sizes, dtype, suggested_storage_type(), memory_layout, shared_object_idx);
 }
 
+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);
+}
+
 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,7 @@ 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 specified properties. The
@@ -191,9 +191,25 @@ class ComputeGraph final {
    */
   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);
 
+  /*
+   * 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,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_like(vref, layout);
   vTensor& t = graph.get_val(v).toTensor();
 
   api::ShaderInfo shader = get_nchw_to_image_shader(t);