Add buffer_to_buffer prepacking

Differential Revision: D61665694

Pull Request resolved: #4840

Author: Esteb37

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

@@ -132,6 +132,33 @@ ValueRef prepack(
   return v;
 }
 
+ValueRef prepack_buffer(
+    ComputeGraph& graph,
+    const ValueRef vref,
+    const utils::GPUMemoryLayout layout) {
+  ValueRef v = graph.add_tensor_like(vref, layout);
+
+  vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR("buffer_to_buffer");
+
+  vkapi::ParamsBindList ubos;
+  ubos.append({graph.numel_ubo(v)});
+
+  graph.prepack_nodes().emplace_back(new PrepackNode(
+      graph,
+      shader,
+      graph.create_global_wg_size(v),
+      graph.create_local_wg_size(v),
+      // Input and Outputs
+      vref,
+      v,
+      // Parameter Buffers
+      ubos,
+      // Specialization Constants
+      {}));
+
+  return v;
+}
+
 ValueRef prepack_if_tensor_ref(
     ComputeGraph& graph,
     const ValueRef v,
@@ -143,6 +170,17 @@ ValueRef prepack_if_tensor_ref(
   }
 }
 
+ValueRef prepack_buffer_if_tensor_ref(
+    ComputeGraph& graph,
+    const ValueRef v,
+    const utils::GPUMemoryLayout layout) {
+  if (graph.val_is_tref(v)) {
+    return prepack_buffer(graph, v, layout);
+  } else {
+    return v;
+  }
+}
+
 ValueRef prepack_if_tensor_ref(ComputeGraph& graph, const ValueRef v) {
   if (graph.val_is_tref(v)) {
     utils::GPUMemoryLayout layout =
@@ -153,4 +191,14 @@ ValueRef prepack_if_tensor_ref(ComputeGraph& graph, const ValueRef v) {
   }
 }
 
+ValueRef prepack_buffer_if_tensor_ref(ComputeGraph& graph, const ValueRef v) {
+  if (graph.val_is_tref(v)) {
+    utils::GPUMemoryLayout layout =
+        graph.suggested_memory_layout(graph.get_tref(v)->sizes);
+    return prepack_buffer(graph, v, layout);
+  } else {
+    return v;
+  }
+}
+
 } // namespace vkcompute

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

@@ -29,6 +29,13 @@ ValueRef prepack_if_tensor_ref(
     const ValueRef v,
     const utils::GPUMemoryLayout layout);
 
+ValueRef prepack_buffer_if_tensor_ref(
+    ComputeGraph& graph,
+    const ValueRef v,
+    const utils::GPUMemoryLayout layout);
+
 ValueRef prepack_if_tensor_ref(ComputeGraph& graph, const ValueRef v);
 
+ValueRef prepack_buffer_if_tensor_ref(ComputeGraph& graph, const ValueRef v);
+
 } // namespace vkcompute