Update on "[ET-VK][6/n] aten.view_copy"

aten.view_copy, supporting all packing.

Using ssjia's idea to do a direct lookup.

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

[ghstack-poisoned]

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -9,18 +9,18 @@
 #define divup4(x) ((x + 3) / 4)
 
 // Input: idx is a ivec4 user-level coordinate, sizes is the tensor shape
-// Output: buffer_idx in the continous nchw-buffer.
-#define to_buffer_i(idx, sizes)                      \
-  (idx.x + idx.y* sizes.x + idx.z* sizes.y* sizes.x + \
-      idx.w* sizes.z* sizes.y* sizes.x)
+// Output: buffer_idx in the continuous nchw-buffer.
+#define to_buffer_i(idx, sizes)                          \
+  (idx.x + idx.y * sizes.x + idx.z * sizes.y * sizes.x + \
+   idx.w * sizes.z * sizes.y * sizes.x)
 
 // Inverse of to_buffer_i
-// Input: buffer_idx in the continous nchw-buffer, sizes is the tensor shape
+// Input: buffer_idx in the continuous nchw-buffer, sizes is the tensor shape
 // Output: ivec4 user-level coorindate
-#define from_buffer_i(buf_i, sizes) \
-  ivec4( \
-      buf_i % sizes.x, \
-      (buf_i / (sizes.x)) % sizes.y, \
+#define from_buffer_i(buf_i, sizes)            \
+  ivec4(                                       \
+      buf_i % sizes.x,                         \
+      (buf_i / (sizes.x)) % sizes.y,           \
       (buf_i / (sizes.x * sizes.y)) % sizes.z, \
       (buf_i / (sizes.x * sizes.y * sizes.z)))
 

backends/vulkan/runtime/graph/ops/glsl/view.glsl

@@ -26,50 +26,46 @@ layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
 #define get_packed_stride get_packed_stride_${PACKING}
 
 layout(set = 0, binding = 2) uniform PRECISION restrict OutGpuSizes {
-  uvec4 data;
-}
-out_gpu_sizes;
+  uvec4 out_gpu_sizes;
+};
 
 layout(set = 0, binding = 3) uniform PRECISION restrict OutCpuSizes {
-  uvec4 data;
-}
-out_cpu_sizes;
+  uvec4 out_cpu_sizes;
+};
 
 layout(set = 0, binding = 4) uniform PRECISION restrict InGpuSizes {
-  uvec4 data;
-}
-in_gpu_sizes;
+  uvec4 in_gpu_sizes;
+};
 
 layout(set = 0, binding = 5) uniform PRECISION restrict InCpuSizes {
-  uvec4 data;
-}
-in_cpu_sizes;
+  uvec4 in_cpu_sizes;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 
 void main() {
 	const ivec3 out_pos = ivec3(gl_GlobalInvocationID);
-	const ivec4 out_tensor_idx = to_tensor_idx(out_pos, out_gpu_sizes.data);
+	const ivec4 out_tensor_idx = to_tensor_idx(out_pos, out_gpu_sizes);
 
-  if (all(greaterThanEqual(out_tensor_idx, out_gpu_sizes.data))) {
+  if (all(greaterThanEqual(out_tensor_idx, out_gpu_sizes))) {
     return;
   }
 
   // Assume there is a virtual continous buffer in nchw format. From the output
   // pos, we first calculate the index in the virual buffer, and then calculate
   // the input position from the indx.
 
-  const uint base_index = to_buffer_i(out_tensor_idx, out_cpu_sizes.data);
+  const uint base_index = to_buffer_i(out_tensor_idx, out_cpu_sizes);
   const uvec4 buf_indices =
-    base_index + ivec4(0, 1, 2, 3) * get_packed_stride(out_cpu_sizes.data);
+    base_index + ivec4(0, 1, 2, 3) * get_packed_stride(out_cpu_sizes);
 
   VEC4_T value;
   // Need to look up the 4 values in the output texel separately.
   for (int i=0; i<4; i++) {
-    ivec4 user_coor = from_buffer_i(buf_indices[i], in_cpu_sizes.data);
+    ivec4 user_coor = from_buffer_i(buf_indices[i], in_cpu_sizes);
 
-    ivec4 in_pos_elem = to_texture_pos_elem(user_coor, in_gpu_sizes.data);
+    ivec4 in_pos_elem = to_texture_pos_elem(user_coor, in_gpu_sizes);
 
     VEC4_T intex = VEC4_T(texelFetch(image_in, in_pos_elem.xyz, 0));
 

backends/vulkan/runtime/graph/ops/glsl/view.yaml

@@ -12,4 +12,3 @@ view:
       - VALUE: H_packed
   shader_variants:
     - NAME: view
-

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

@@ -8,30 +8,22 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
 
-#include <executorch/backends/vulkan/runtime/api/api.h>
-
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 
 namespace vkcompute {
 
-void add_view_node(
-    ComputeGraph& graph,
-    ValueRef in,
-    ValueRef size_ref,
-    ValueRef out) {
-  // Note: size_ref is not used here. Since the output tensor's size have been
-  // determined during compilation.
+void add_view_node(ComputeGraph& graph, ValueRef in, ValueRef out) {
   vTensorPtr t_in = graph.get_tensor(in);
   vTensorPtr t_out = graph.get_tensor(out);
-  
+
   std::string kernel_name = "view";
   kernel_name.reserve(kShaderNameReserve);
   add_dtype_suffix(kernel_name, *t_out);
   add_memory_layout_suffix(kernel_name, *t_out);
-  
-  api::utils::uvec3 global_size = t_out->virtual_extents();
+
+  api::utils::uvec3 global_size = t_out->extents();
   api::utils::uvec3 local_size = adaptive_work_group_size(global_size);
 
   graph.execute_nodes().emplace_back(new ExecuteNode(
@@ -40,22 +32,20 @@ void add_view_node(
       global_size,
       local_size,
       {{out, api::MemoryAccessType::WRITE}, {in, api::MemoryAccessType::READ}},
-      {
-        t_out->gpu_sizes_ubo(), 
-        t_out->cpu_sizes_ubo(),
-        t_in->gpu_sizes_ubo(),
-        t_in->cpu_sizes_ubo()}));
+      {t_out->gpu_sizes_ubo(),
+       t_out->cpu_sizes_ubo(),
+       t_in->gpu_sizes_ubo(),
+       t_in->cpu_sizes_ubo()}));
 }
 
-
 void view(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_view_node(graph, args[0], args[1], args[2]);
+  // Note: The second argument size_ref is not used here. Since the output
+  // tensor's size have been determined during compilation.
+  return add_view_node(graph, args[0], args[2]);
 }
 
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.view_copy.default, view);
 }
 
-
-}  // namespace vkcompute
-
+} // namespace vkcompute

backends/vulkan/test/op_tests/cases.py

@@ -197,9 +197,9 @@ def get_permute_inputs():
 def get_view_inputs():
     test_suite = VkTestSuite(
         [
-            ((3,4,5), [1, 1, -1]),
-            ((3,4,5), [1, -1, 1]),
-            ((3,4,5), [-1, 1, 1]),
+            ((3, 4, 5), [1, 1, -1]),
+            ((3, 4, 5), [1, -1, 1]),
+            ((3, 4, 5), [-1, 1, 1]),
             ((8, 7, 2, 3), [4, 3, 7, 4]),
             ((8, 7, 2, 3), [7, -1, 2, 1]),
             ((8, 7, 2, 3), [1, 1, 1, -1]),

backends/vulkan/test/op_tests/utils/codegen_base.py

@@ -106,13 +106,15 @@ def gen_case_name(self, inputs: List[Any], prepack: bool = False) -> str:
                     name_str += str(size) + "x"
                 name_str = name_str[:-1]
                 # minus sign is a invalid char for test case. change to "n".
-                name_str = name_str.replace('-', 'n')
+                name_str = name_str.replace("-", "n")
+
             elif isinstance(arg_sizes_or_val, list):
                 for size in arg_sizes_or_val:
                     name_str += str(size) + "c"
                 name_str = name_str[:-1]
                 # minus sign is a invalid char for test case. change to "n".
-                name_str = name_str.replace('-', 'n')
+                name_str = name_str.replace("-", "n")
+
             else:
                 name_str += str(arg_sizes_or_val).replace(".", "p")
         return name_str