add int to generated shader variants for view (pytorch#3970)

Summary:
Pull Request resolved: pytorch#3970

Used in joiner model.

Also updating to use the newer shader layout codegen format

Reviewed By: jorgep31415

Differential Revision: D58499751

fbshipit-source-id: fa7fcf0d89e0afba3944c5f9a1b3c7017791a573

Author: nathanaelsee

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

@@ -16,16 +16,11 @@ layout(std430) buffer;
 
 #include "indexing_utils.h"
 
-layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
-layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
+${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
 
-layout(set = 0, binding = 2) uniform PRECISION restrict OutSizes {
-  ivec4 out_sizes;
-};
-
-layout(set = 0, binding = 3) uniform PRECISION restrict InSizes {
-  ivec4 in_sizes;
-};
+${layout_declare_ubo(2, "ivec4", "out_sizes")}
+${layout_declare_ubo(3, "ivec4", "in_sizes")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -52,10 +47,10 @@ void main() {
     if (out_tensor_idx[out_packed_dim]++ < out_sizes[out_packed_dim]) {
       ivec4 user_coor = from_nchw_buffer_i(buf_indices[i], in_sizes);
       ivec4 in_pos_elem = to_texture_elem_pos(user_coor, in_sizes, in_packed_dim);
-      VEC4_T intex = texelFetch(image_in, in_pos_elem.xyz, 0);
+      VEC4_T intex = texelFetch(t_in, in_pos_elem.xyz, 0);
       value[i] = intex[in_pos_elem.w];
     }
   }
 
-  imageStore(image_out, out_pos, value);
+  imageStore(t_out, out_pos, value);
 }

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

@@ -2,9 +2,11 @@ view:
   parameter_names_with_default_values:
     DTYPE: float
     NDIM: 3
+    STORAGE: texture3d
   generate_variant_forall:
     DTYPE:
       - VALUE: half
       - VALUE: float
+      - VALUE: int
   shader_variants:
     - NAME: view

backends/vulkan/test/test_vulkan_delegate.py

@@ -1083,6 +1083,22 @@ def forward(self, x):
             memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
         )
 
+    def test_vulkan_backend_view_int(self):
+        class ViewModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return x.view([-1, x.size(-1)])
+
+        sample_inputs = (torch.randint(size=(3, 6, 2, 7), high=100, dtype=torch.int32),)
+
+        self.lower_module_and_test_output(
+            ViewModule(),
+            sample_inputs,
+            memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        )
+
     def test_vulkan_backend_unsqueeze(self):
         class UnsqueezeModule(torch.nn.Module):
             def __init__(self):