[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/ComputeGraph.cpp

@@ -285,7 +285,8 @@ ValueRef ComputeGraph::add_tensor_like(
 ValueRef ComputeGraph::add_tensor_like(
     const ValueRef idx,
     const utils::GPUMemoryLayout memory_layout) {
-  return add_tensor(sizes_of(idx), dtype_of(idx), memory_layout);
+  return add_tensor(
+      sizes_of(idx), dtype_of(idx), storage_type_of(idx), memory_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(

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

@@ -18,9 +18,10 @@
 
 namespace vkcompute {
 
-ValueRef prepack_arg(
+ValueRef check_and_prepack_arg(
     ComputeGraph& graph,
     ValueRef arg_ref,
+    const utils::StorageType stype,
     int64_t num_channels,
     const std::string& debug_name) {
   VK_CHECK_COND(
@@ -33,7 +34,7 @@ ValueRef prepack_arg(
   // batch_norm's param are broadcasted on the channel dimension.
   // In this implementation, we pack the weights along the x dimension, and
   // in the shader, we lookup using the along the x.
-  return prepack_if_tensor_ref(graph, arg_ref, utils::kWidthPacked);
+  return prepack_standard(graph, arg_ref, stype, utils::kWidthPacked);
 }
 
 void add_native_batch_norm_node(
@@ -51,22 +52,26 @@ void add_native_batch_norm_node(
   VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor");
   VK_CHECK_COND(out_sizes.size() == 4, "BatchNorm only support 4d tensor");
 
+  // Only the first element of the return value is propagated. The remaining 2
+  // elements are zero-size dummy tensor.
+  ValueRef out_ref = graph.get_value_list(out_tuple_ref)->at(0);
+
+  utils::StorageType stype = graph.storage_type_of(out_ref);
+
   int64_t num_channels = dim_at<kChannel4D>(in_sizes);
 
-  ValueRef arg_weight = prepack_arg(graph, weight_ref, num_channels, "weight");
-  ValueRef arg_bias = prepack_arg(graph, bias_ref, num_channels, "bias");
-  ValueRef arg_mean = prepack_arg(graph, mean_ref, num_channels, "mean");
-  ValueRef arg_var = prepack_arg(graph, var_ref, num_channels, "var");
+  ValueRef arg_weight =
+      check_and_prepack_arg(graph, weight_ref, stype, num_channels, "weight");
+  ValueRef arg_bias =
+      check_and_prepack_arg(graph, bias_ref, stype, num_channels, "bias");
+  ValueRef arg_mean =
+      check_and_prepack_arg(graph, mean_ref, stype, num_channels, "mean");
+  ValueRef arg_var =
+      check_and_prepack_arg(graph, var_ref, stype, num_channels, "var");
   float epsilon = graph.extract_scalar<float>(eps_ref);
 
   vTensorPtr t_in = graph.get_tensor(in_ref);
 
-  // Only the first element of the return value is propagated. The remaining 2
-  // elements are zero-size dummy tensor.
-  const auto out_tuple_val = graph.get_value_list(out_tuple_ref);
-
-  ValueRef out_ref = out_tuple_val->at(0);
-
   VK_CHECK_COND(!graph.val_is_tref(out_ref), "Output should not be tref");
   vTensorPtr t_out = graph.get_tensor(out_ref);
 

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

@@ -51,9 +51,8 @@ void add_binary_op_node(
     const ValueRef alpha,
     const ValueRef out,
     const std::string& op_name) {
-  ValueRef arg1 = prepack_if_tensor_ref(graph, in1);
-  ValueRef arg2 =
-      prepack_if_tensor_ref(graph, in2, graph.estimate_memory_layout_of(arg1));
+  ValueRef arg1 = prepack_standard_like(graph, in1, out, true);
+  ValueRef arg2 = prepack_standard_like(graph, in2, out, true);
 
   vTensorPtr t_in1 = graph.get_tensor(arg1);
   vTensorPtr t_in2 = graph.get_tensor(arg2);

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

@@ -304,7 +304,7 @@ utils::uvec3 create_conv2d_global_wg_size(
 void add_conv2d_node(
     ComputeGraph& graph,
     const ValueRef in,
-    const ValueRef weight,
+    const ValueRef weight_data,
     const ValueRef bias,
     const ValueRef stride,
     const ValueRef padding,
@@ -330,19 +330,18 @@ void add_conv2d_node(
   const int64_t groups_val = graph.get_int(groups);
 
   const Conv2dMethod method =
-      get_conv2d_method(graph, weight, groups_val, transposed_val);
+      get_conv2d_method(graph, weight_data, groups_val, transposed_val);
 
-  ValueRef arg_in = prepack_if_tensor_ref(graph, in);
-  ValueRef arg_weight = prepack_weights(graph, weight, method);
+  ValueRef arg_weight = prepack_weights(graph, weight_data, method);
   ValueRef arg_bias = prepack_biases(
       graph,
       bias,
-      weight,
+      weight_data,
       transposed_val,
       /* storage_type = */ utils::kTexture2D,
       /* memory_layout = */ utils::kWidthPacked);
 
-  vTensorPtr t_in = graph.get_tensor(arg_in);
+  vTensorPtr t_in = graph.get_tensor(in);
   vTensorPtr t_out = graph.get_tensor(out);
   if (t_in->sizes().at(0) > 1) {
     VK_THROW("conv2d: input batch size > 1 is not supported yet!");
@@ -351,20 +350,25 @@ void add_conv2d_node(
 
   Kernel2dParams kernel_params = create_kernel2d_params(
       graph,
-      weight,
+      weight_data,
       /*kernel_size_only = */ false,
       stride,
       padding,
       dilation);
   Conv2dParams extra_params =
-      create_conv2d_params(graph, weight, kernel_params, transposed_val);
+      create_conv2d_params(graph, weight_data, kernel_params, transposed_val);
 
   OutputParams out_params = {out_min_val, out_max_val};
 
   check_conv2d_params(kernel_params, transposed_val);
 
   vkapi::ShaderInfo shader = get_conv2d_shader(
-      graph, *t_out, /*prepack_weights = */ false, method, weight, clamp_out);
+      graph,
+      *t_out,
+      /*prepack_weights = */ false,
+      method,
+      weight_data,
+      clamp_out);
 
   graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
@@ -373,7 +377,7 @@ void add_conv2d_node(
       graph.create_local_wg_size(out),
       // Inputs and Outputs
       {{out, vkapi::MemoryAccessType::WRITE},
-       {{arg_in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
+       {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {
           t_out->logical_limits_ubo(),
@@ -386,7 +390,7 @@ void add_conv2d_node(
       {},
       // Resizing Logic
       resize_conv2d_node,
-      {weight, stride, padding, dilation, transposed, output_padding}));
+      {weight_data, stride, padding, dilation, transposed, output_padding}));
 }
 
 void add_conv1d_node(
@@ -402,9 +406,8 @@ void add_conv1d_node(
     const ValueRef out_max,
     const ValueRef out,
     const bool clamp_out) {
-  ValueRef arg_in = prepack_if_tensor_ref(graph, in);
-  ValueRef arg_weight =
-      prepack_if_tensor_ref(graph, weight, utils::kWidthPacked);
+  ValueRef arg_weight = prepack_standard(
+      graph, weight, graph.storage_type_of(out), utils::kWidthPacked);
   ValueRef arg_bias = prepack_biases(
       graph,
       bias,
@@ -422,7 +425,7 @@ void add_conv1d_node(
     out_max_val = graph.extract_scalar<float>(out_max);
   }
 
-  vTensorPtr t_in = graph.get_tensor(arg_in);
+  vTensorPtr t_in = graph.get_tensor(in);
   vTensorPtr t_weight = graph.get_tensor(arg_weight);
   vTensorPtr t_bias = graph.get_tensor(arg_bias);
   vTensorPtr t_out = graph.get_tensor(out);
@@ -471,7 +474,7 @@ void add_conv1d_node(
       local_size,
       // Inputs and Outputs
       {{out, vkapi::MemoryAccessType::WRITE},
-       {{arg_in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
+       {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {
           t_out->logical_limits_ubo(),

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

@@ -57,9 +57,9 @@ void add_embedding_node(
 }
 
 void embedding(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  ValueRef weight = prepack_if_tensor_ref(graph, args[0]);
-  ValueRef in = prepack_if_tensor_ref(graph, args[1]);
+  ValueRef in = args[1];
   ValueRef out = args[5];
+  ValueRef weight = prepack_standard_like(graph, args[0], out);
 
   add_embedding_node(graph, weight, in, out);
 }

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

@@ -108,9 +108,9 @@ int64_t get_dim_idx(ComputeGraph& graph, ValueRef in, ValueRef dim_ref) {
 }
 
 void index_select(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  ValueRef in = prepack_if_tensor_ref(graph, args[0]);
+  ValueRef in = args[0];
   ValueRef dim_ref = args[1];
-  ValueRef idx = prepack_if_tensor_ref(graph, args[2]);
+  ValueRef idx = args[2];
   ValueRef out = args[3];
 
   const int64_t dim_idx = get_dim_idx(graph, in, dim_ref);

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

@@ -94,8 +94,11 @@ void add_addmm_naive_node(
     const ValueRef out,
     const Params& params,
     const ValueRef mat2_is_transposed) {
-  ValueRef self = prepack_if_tensor_ref(graph, self_data, utils::kWidthPacked);
-  ValueRef mat2 = prepack_if_tensor_ref(graph, mat2_data, utils::kHeightPacked);
+  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);
 
   std::string kernel_name =
       graph.get_bool(mat2_is_transposed) ? "linear_naive" : "addmm_naive";
@@ -145,9 +148,11 @@ void add_addmm_optimized_node(
     const ValueRef out,
     const Params& params,
     const ValueRef mat2_is_transposed) {
+  utils::StorageType stype = graph.storage_type_of(out);
   ValueRef self =
-      prepack_if_tensor_ref(graph, self_data, utils::kChannelsPacked);
-  ValueRef mat2 = prepack_if_tensor_ref(graph, mat2_data, utils::kHeightPacked);
+      prepack_standard(graph, self_data, stype, utils::kChannelsPacked, true);
+  ValueRef mat2 =
+      prepack_standard(graph, mat2_data, stype, utils::kHeightPacked, true);
 
   // Ensure mat1 is width packed
   ValueRef mat1_W_packed = graph.add_tensor_like(mat1, utils::kWidthPacked);
@@ -276,8 +281,8 @@ void linear(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   ValueRef weight_data = args.at(1);
   ValueRef bias = args.at(2);
   ValueRef out = args.at(3);
-  ValueRef weight =
-      prepack_if_tensor_ref(graph, weight_data, utils::kWidthPacked);
+  ValueRef weight = prepack_standard(
+      graph, weight_data, graph.storage_type_of(out), utils::kWidthPacked);
   ValueRef mat2_is_transposed = graph.add_scalar(true);
   if (graph.val_is_none(bias)) {
     return add_matmul_node(graph, input, weight, out, mat2_is_transposed);

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

@@ -62,7 +62,8 @@ void add_matmul_naive_buffer_node(
     const ValueRef mat2_data,
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
-  ValueRef mat2 = prepack_if_tensor_ref(graph, mat2_data, utils::kHeightPacked);
+  ValueRef mat2 = prepack_standard(
+      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
 
   std::string kernel_name = "matmul_naive_buffer";
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
@@ -103,7 +104,8 @@ void add_matmul_naive_texture3d_node(
     const ValueRef mat2_data,
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
-  ValueRef mat2 = prepack_if_tensor_ref(graph, mat2_data, utils::kHeightPacked);
+  ValueRef mat2 = prepack_standard(
+      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
 
   std::string kernel_name = graph.get_bool(mat2_is_transposed)
       ? "matmul_transposed_naive"
@@ -146,7 +148,8 @@ void add_matmul_optimized_node(
     const ValueRef mat2_data,
     const ValueRef out,
     const ValueRef mat2_is_transposed) {
-  ValueRef mat2 = prepack_if_tensor_ref(graph, mat2_data, utils::kHeightPacked);
+  ValueRef mat2 = prepack_standard(
+      graph, mat2_data, graph.storage_type_of(out), utils::kHeightPacked, true);
 
   // Ensure mat1 is width packed
   ValueRef mat1_W_packed = graph.add_tensor_like(mat1, utils::kWidthPacked);

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

@@ -57,8 +57,8 @@ void add_native_layer_norm_node(
     ComputeGraph& graph,
     const ValueRef in,
     const ValueRef normalized_shape,
-    const ValueRef weight,
-    const ValueRef bias,
+    const ValueRef weight_data,
+    const ValueRef bias_data,
     const ValueRef eps,
     const ValueRef out) {
   const auto normalized_shape_dim =
@@ -67,19 +67,16 @@ void add_native_layer_norm_node(
     VK_THROW("native_layer_norm only supports normalized_shape with dim == 1");
   }
 
-  if (graph.val_is_none(weight)) {
+  if (graph.val_is_none(weight_data)) {
     VK_THROW("native_layer_norm requires weight to be non-None");
   }
 
-  if (graph.val_is_none(bias)) {
+  if (graph.val_is_none(bias_data)) {
     VK_THROW("native_layer_norm requires bias to be non-None");
   }
 
-  ValueRef arg_in = prepack_if_tensor_ref(graph, in);
-  ValueRef arg_weight = prepack_if_tensor_ref(
-      graph, weight, graph.estimate_memory_layout_of(arg_in));
-  ValueRef arg_bias = prepack_if_tensor_ref(
-      graph, bias, graph.estimate_memory_layout_of(arg_in));
+  ValueRef arg_weight = prepack_standard_like(graph, weight_data, in);
+  ValueRef arg_bias = prepack_standard_like(graph, bias_data, in);
 
   const auto out_val = graph.get_value_list(out);
   vTensorPtr t_out = graph.get_tensor(out_val->at(0));
@@ -107,7 +104,7 @@ void add_native_layer_norm_node(
       // Inputs and Outputs
       {{{out_val->at(0), out_val->at(1), out_val->at(2)},
         vkapi::MemoryAccessType::WRITE},
-       {{arg_in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
+       {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {t_out->logical_limits_ubo(),
        t_out->sizes_ubo(),

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

@@ -71,8 +71,7 @@ void add_max_pool2d_node(
     const ValueRef dilation,
     const ValueRef ceil_mode,
     const ValueRef out) {
-  ValueRef arg = prepack_if_tensor_ref(graph, in);
-  vTensorPtr t_in = graph.get_tensor(arg);
+  vTensorPtr t_in = graph.get_tensor(in);
 
   const auto out_val = graph.get_value_list(out);
   vTensorPtr t_out = graph.get_tensor(out_val->at(0));
@@ -100,7 +99,7 @@ void add_max_pool2d_node(
       local_size,
       // Inputs and Outputs
       {{{out_val->at(0), out_val->at(1)}, vkapi::MemoryAccessType::WRITE},
-       {arg, vkapi::MemoryAccessType::READ}},
+       {in, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {
           t_out->logical_limits_ubo(),
@@ -149,8 +148,7 @@ void add_avg_pool2d_node(
     const ValueRef count_include_pad,
     const ValueRef divisor_override,
     const ValueRef out) {
-  ValueRef arg = prepack_if_tensor_ref(graph, in);
-  vTensorPtr t_in = graph.get_tensor(arg);
+  vTensorPtr t_in = graph.get_tensor(in);
   vTensorPtr t_out = graph.get_tensor(out);
 
   check_pool2d_args(*t_in, *t_out);
@@ -174,7 +172,7 @@ void add_avg_pool2d_node(
       local_size,
       // Inputs and Outputs
       {{out, vkapi::MemoryAccessType::WRITE},
-       {arg, vkapi::MemoryAccessType::READ}},
+       {in, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {t_out->logical_limits_ubo(),
        t_in->sizes_ubo(),