Update on "[ET-VK] Clean up prepack API"

## Context

As title, revamp the prepacking API:

* Make the naming more clear; i.e. `prepack_if_tensor_ref` to `prepack_standard` to disambiguate the packing logic that will used.
* Instead of passing through the `v` argument if it is a `Tensor` by default, this functionality must be toggled via the `passthrough` argument. The goal is to  encourage developers to be more explicit about what types they expect the operator arguments to be.
* Consolidate API surface and reduce the number of overloads

Past the API changes, I have also removed a bunch of unnecessary calls to `prepack_if_tensor_ref` throughout the operator implementations. The most common cases were calling it on an input tensor which is not necessary.

## The "big picture" for prepacking

`TensorRef` objects and prepacking are used whenever we are dealing with a Tensor whose data is serialized with the model. However, these "serialized tensors" all belong to one of two categories

* Weight/biases: trained weights and biases that act as the state for a i.e. Convolutional or Linear layer. These tensors are used only within the `nn.Module` that they belong to
* Persistent tensors: tensors whose data just happen to be invariant to the inputs, and their data can be serialized with the model itself. They are treated as regular tensors and may be used in several operators throughout the model. One example is `freqs_sin` and `freqs_cos` in Llama models which are used to calculate rotary positional encodings

For weights and biases, the way that the serialized data should be packed may be dependent on the operator it is used in. However, for persistent tensors they must be packed with the "standard" staging to tensor algorithm since they are the same as regular tensors.

While it is well known which operators expect weight tensors. However, persistent tensors are tricky because they can be used as an argument to any operator. This would mean that every operator needs to account for the possibility that one of their inputs will be a serialized tensor. This is undesirable because it adds an additional layer of indirection when processing operator inputs on top of the fact that every argument is actually a reference to a`Value` object in the graph, which itself is a wrapper. It also makes things complicated for the operator developer. Another downside is that persistent tensors will be packed multiple times, once by each operator that uses it.

To address this, I plan to handle persistent tensors at export time by inserting a `prepack()` operator for them which will cause operators that use the serialized tensor to see a Tensor object instead of a TensorRef object. This will make it so that the only operators that should expect to prepack an argument are tensors that expect a weight argument, and also avoid packing persistent tensors multiple times.

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

[ghstack-poisoned]

Author: SS-JIA

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

@@ -95,10 +95,10 @@ void add_addmm_naive_node(
     const Params& params,
     const ValueRef mat2_is_transposed) {
   utils::StorageType stype = graph.storage_type_of(out);
-  ValueRef self =
-      prepack_standard(graph, self_data, stype, utils::kWidthPacked, true);
-  ValueRef mat2 =
-      prepack_standard(graph, mat2_data, stype, utils::kHeightPacked, true);
+  ValueRef self = prepack_standard(
+      graph, self_data, stype, utils::kWidthPacked, /*passthrough = */ true);
+  ValueRef mat2 = prepack_standard(
+      graph, mat2_data, stype, utils::kHeightPacked, /*passthrough = */ true);
 
   std::string kernel_name =
       graph.get_bool(mat2_is_transposed) ? "linear_naive" : "addmm_naive";
@@ -149,10 +149,10 @@ void add_addmm_optimized_node(
     const Params& params,
     const ValueRef mat2_is_transposed) {
   utils::StorageType stype = graph.storage_type_of(out);
-  ValueRef self =
-      prepack_standard(graph, self_data, stype, utils::kChannelsPacked, true);
-  ValueRef mat2 =
-      prepack_standard(graph, mat2_data, stype, utils::kHeightPacked, true);
+  ValueRef self = prepack_standard(
+      graph, self_data, stype, utils::kChannelsPacked, /*passthrough=*/true);
+  ValueRef mat2 = prepack_standard(
+      graph, mat2_data, stype, utils::kHeightPacked, /*passthrough=*/true);
 
   // Ensure mat1 is width packed
   ValueRef mat1_W_packed = graph.add_tensor_like(mat1, utils::kWidthPacked);

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

@@ -63,7 +63,11 @@ void add_matmul_naive_buffer_node(
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
   ValueRef mat2 = prepack_standard(
-      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
+      graph,
+      mat2_data,
+      graph.storage_type_of(out),
+      utils::kHeightPacked,
+      /*passthrough = */ true);
 
   std::string kernel_name = "matmul_naive_buffer";
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
@@ -105,7 +109,11 @@ void add_matmul_naive_texture3d_node(
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
   ValueRef mat2 = prepack_standard(
-      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
+      graph,
+      mat2_data,
+      graph.storage_type_of(out),
+      utils::kHeightPacked,
+      /*passthrough = */ true);
 
   std::string kernel_name = graph.get_bool(mat2_is_transposed)
       ? "matmul_transposed_naive"
@@ -149,7 +157,11 @@ void add_matmul_optimized_node(
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
   ValueRef mat2 = prepack_standard(
-      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
+      graph,
+      mat2_data,
+      graph.storage_type_of(out),
+      utils::kHeightPacked,
+      /*passthrough = */ true);
 
   // Ensure mat1 is width packed
   ValueRef mat1_W_packed = graph.add_tensor_like(mat1, utils::kWidthPacked);

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

@@ -110,7 +110,7 @@ void add_tensor_to_staging_node(
       {SV(graph.packed_dim_of(in_tensor))}));
 }
 
-void add_standard_prepack_node(
+void add_prepack_standard_node(
     ComputeGraph& graph,
     const ValueRef tensor_data,
     const ValueRef tensor) {
@@ -152,7 +152,7 @@ ValueRef prepack_standard(
   }
   VK_CHECK_COND(graph.val_is_tref(tensor_data));
   ValueRef tensor = graph.add_tensor_like(tensor_data, storage_type, layout);
-  add_standard_prepack_node(graph, tensor_data, tensor);
+  add_prepack_standard_node(graph, tensor_data, tensor);
   return tensor;
 }
 
@@ -170,7 +170,7 @@ ValueRef prepack_standard_like(
       passthrough);
 }
 
-void add_direct_buffer_copy_prepack_node(
+void add_prepack_direct_buffer_copy_node(
     ComputeGraph& graph,
     const ValueRef tensor_data,
     const ValueRef tensor) {
@@ -198,9 +198,10 @@ void add_direct_buffer_copy_prepack_node(
 ValueRef prepack_direct_copy_buffer(
     ComputeGraph& graph,
     const ValueRef tensor_data) {
+  VK_CHECK_COND(graph.val_is_tref(tensor_data));
   ValueRef tensor =
       graph.add_tensor_like(tensor_data, utils::kBuffer, utils::kWidthPacked);
-  add_direct_buffer_copy_prepack_node(graph, tensor_data, tensor);
+  add_prepack_direct_buffer_copy_node(graph, tensor_data, tensor);
   return tensor;
 }