Update on "[ET-VK] Implement prepack nodes"

## Context

This diff implements the idea described in the previous diff in this stack. During export, `et_vk.prepack` nodes will be inserted to convert constant tensors to GPU tensor objects. This makes it so that Vulkan operators will not have to account for the possibility that their arguments can potentially be constant tensor data instead of an actual tensor object.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/_passes/insert_prepack_nodes.py

@@ -86,7 +86,7 @@ def is_non_weight_param_tensor(node: torch.fx.Node) -> bool:
             # Set the mem_obj_id to -1 to indicate that this node requires a dedicated
             # memory object. This pass must be executed AFTER the memory planning pass.
             prepack_node.meta["spec"].mem_obj_id = -1
-            node.replace_all_uses_with(prepack_node, lambda x: x != prepack_node)
+            node.replace_all_uses_with(prepack_node, lambda x, y=prepack_node: x != y)
 
     program.graph.eliminate_dead_code()
     return program

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

@@ -51,9 +51,6 @@ void add_binary_op_node(
     const ValueRef alpha,
     const ValueRef out,
     const std::string& op_name) {
-  VK_CHECK_COND(graph.val_is_tensor(in1));
-  VK_CHECK_COND(graph.val_is_tensor(in2));
-
   vTensorPtr t_in1 = graph.get_tensor(in1);
   vTensorPtr t_in2 = graph.get_tensor(in2);
   vTensorPtr t_out = graph.get_tensor(out);

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

@@ -112,7 +112,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) {
@@ -154,7 +154,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;
 }
 
@@ -172,7 +172,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) {
@@ -200,14 +200,15 @@ 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;
 }
 
 void prepack_op(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_standard_prepack_node(graph, args[0], args[1]);
+  return add_prepack_standard_node(graph, args[0], args[1]);
 }
 
 REGISTER_OPERATORS {

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -1520,11 +1520,18 @@ TEST(VulkanComputeGraphTest, test_simple_prepacked_graph) {
   ValueRef c = graph.add_tensor(size_big, vkapi::kFloat);
   ValueRef e = graph.add_tensor(size_big, vkapi::kFloat);
 
+  ValueRef w1_packed = graph.add_tensor(size_small, vkapi::kFloat);
+  ValueRef w2_packed = graph.add_tensor(size_small, vkapi::kFloat);
+
+  auto prepackFn = VK_GET_OP_FN("et_vk.prepack.default");
+  prepackFn(graph, {w1, w1_packed});
+  prepackFn(graph, {w2, w2_packed});
+
   auto addFn = VK_GET_OP_FN("aten.add.Tensor");
-  addFn(graph, {a.value, w1, kDummyValueRef, c});
+  addFn(graph, {a.value, w1_packed, kDummyValueRef, c});
 
   auto mulFn = VK_GET_OP_FN("aten.mul.Tensor");
-  mulFn(graph, {c, w2, e});
+  mulFn(graph, {c, w2_packed, e});
 
   IOValueRef out = {};
   out.value = e;
@@ -2597,8 +2604,7 @@ void test_binary_op(
     std::vector<int64_t> sizes_big,
     std::vector<int64_t> sizes_small,
     vkapi::ScalarType dtype,
-    utils::GPUMemoryLayout memory_layout,
-    bool prepack = true) {
+    utils::GPUMemoryLayout memory_layout) {
   GraphConfig config;
   ComputeGraph graph(config);
 
@@ -2609,12 +2615,7 @@ void test_binary_op(
   // Build graph
 
   IOValueRef arg1 = graph.add_input_tensor(sizes_big, dtype, memory_layout);
-
-  if (prepack) {
-    arg2.value = arg2_w;
-  } else {
-    arg2 = graph.add_input_tensor(sizes_small, dtype, memory_layout);
-  }
+  arg2 = graph.add_input_tensor(sizes_small, dtype, memory_layout);
 
   IOValueRef out;
   out.value = graph.add_tensor(sizes_big, dtype, memory_layout);
@@ -2635,7 +2636,7 @@ void test_binary_op(
 
   for (int i = 1; i < 4; i++) {
     float val_arg1 = i + 1.5;
-    float val_arg2 = prepack ? 2.5f : i - 3.5;
+    float val_arg2 = i - 3.5;
 
     float val_out = val_arg1 + val_arg2;
     if (op_name == "sub") {
@@ -2648,21 +2649,14 @@ void test_binary_op(
       val_out = val_arg1 / val_arg2;
     }
 
-    if (prepack) {
-      execute_graph_and_check_output(graph, {val_arg1}, {val_out});
-    } else {
-      execute_graph_and_check_output(graph, {val_arg1, val_arg2}, {val_out});
-    }
+    execute_graph_and_check_output(graph, {val_arg1, val_arg2}, {val_out});
   }
 }
 
-#define CALL_TEST_FN_FORALL_CONDITIONS(_)                            \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked, false)    \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kHeightPacked, false)   \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kChannelsPacked, false) \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked, true)     \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kHeightPacked, true)    \
-  _(vkapi::kFloat, utils::kTexture3D, utils::kChannelsPacked, true)
+#define CALL_TEST_FN_FORALL_CONDITIONS(_)                   \
+  _(vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked)  \
+  _(vkapi::kFloat, utils::kTexture3D, utils::kHeightPacked) \
+  _(vkapi::kFloat, utils::kTexture3D, utils::kChannelsPacked)
 
 #define CALL_TEST_FN_FOR_W_PACKED(_)                              \
   _(vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked, false) \
@@ -2677,15 +2671,15 @@ void test_binary_op(
   _(vkapi::kFloat, utils::kBuffer, utils::kChannelsPacked, true)
 
 TEST(VulkanComputeGraphOpsTest, add_smoke_test) {
-#define RUN_TESTS(dtype, storage, layout, prepack)                         \
-  test_binary_op("add", {17, 21}, {17, 21}, dtype, layout, prepack);       \
-  test_binary_op("add", {17, 21}, {1, 1}, dtype, layout, prepack);         \
-  test_binary_op("sub", {11, 22}, {11, 22}, dtype, layout, prepack);       \
-  test_binary_op("sub", {11, 22}, {11, 1}, dtype, layout, prepack);        \
-  test_binary_op("add", {7, 17, 17}, {7, 17, 17}, dtype, layout, prepack); \
-  test_binary_op("add", {7, 17, 17}, {7, 1, 17}, dtype, layout, prepack);  \
-  test_binary_op("sub", {9, 9, 7}, {9, 9, 7}, dtype, layout, prepack);     \
-  test_binary_op("sub", {9, 9, 7}, {9, 1, 1}, dtype, layout, prepack);
+#define RUN_TESTS(dtype, storage, layout)                         \
+  test_binary_op("add", {17, 21}, {17, 21}, dtype, layout);       \
+  test_binary_op("add", {17, 21}, {1, 1}, dtype, layout);         \
+  test_binary_op("sub", {11, 22}, {11, 22}, dtype, layout);       \
+  test_binary_op("sub", {11, 22}, {11, 1}, dtype, layout);        \
+  test_binary_op("add", {7, 17, 17}, {7, 17, 17}, dtype, layout); \
+  test_binary_op("add", {7, 17, 17}, {7, 1, 17}, dtype, layout);  \
+  test_binary_op("sub", {9, 9, 7}, {9, 9, 7}, dtype, layout);     \
+  test_binary_op("sub", {9, 9, 7}, {9, 1, 1}, dtype, layout);
 
   CALL_TEST_FN_FORALL_CONDITIONS(RUN_TESTS);