Fix binary_op int variant (#3745)

Summary:
Pull Request resolved: #3745

We were using the wrong `sampler3D` meant for `float`/`half`, when we should use `isampler3D` for `int`. [The codegen has a `Dict` mapping](https://www.internalfb.com/code/fbsource/[ac566fa2c8cb64cceb5d47fe4a9c97241a73341d]/fbcode/executorch/backends/vulkan/runtime/api/gen_vulkan_spv.py?lines=58-71) to templatize this logic. Since we're already here, modernize to the new layout generation functions.

## Before
```
layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
layout(set = 0, binding = 2) uniform PRECISION sampler3D image_other;
```
Input values were read incorrectly and got wild results.
```
ET-VK
tensor([[-2147483648, -2147483648,           0],
        [-2147483648,           0,           0]], dtype=torch.int32)
Eager-mode
tensor([[-123, -131,  182],
        [ -29,   24,  138]], dtype=torch.int32)
````

## After
```
layout(set = 0, binding = 1) uniform PRECISION ${SAMPLER_T[NDIM][DTYPE]} image_in;
layout(set = 0, binding = 2) uniform PRECISION ${SAMPLER_T[NDIM][DTYPE]} image_other;
```
OR modernize to the new layout generation functions, which already use and abstract `SAMPLER_T` to avoid the same future mistake.
```
${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
${layout_declare_tensor(2, "r", "t_other", DTYPE, STORAGE)}
```
ghstack-source-id: 227761995

Reviewed By: nathanaelsee

Differential Revision: D57731470

fbshipit-source-id: f233089e5b4ced6f3d47b90a9c84795cf62bb4de

Author: jorgep31415

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

@@ -19,29 +19,14 @@
 
 layout(std430) buffer;
 
-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(set = 0, binding = 2) uniform PRECISION sampler3D image_other;
-
-layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
-  ivec4 out_sizes;
-};
-
-layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
-  ivec4 in_sizes;
-};
-
-layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
-  ivec4 other_sizes;
-};
-
-layout(set = 0, binding = 6) uniform PRECISION restrict BroadcastParams {
-  ivec2 broadcast_params;
-};
-
-layout(set = 0, binding = 7) uniform PRECISION restrict Alpha {
-  float alpha;
-};
+${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)}
+${layout_declare_tensor(2, "r", "t_other", DTYPE, STORAGE)}
+${layout_declare_ubo(3, "ivec4", "out_sizes")}
+${layout_declare_ubo(4, "ivec4", "in_sizes")}
+${layout_declare_ubo(5, "ivec4", "other_sizes")}
+${layout_declare_ubo(6, "ivec2", "broadcast_params")}
+${layout_declare_ubo(7, "float", "alpha")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -57,13 +42,13 @@ void main() {
 
   ivec4 in_idx = broadcast_indices(idx, in_sizes);
   VEC4_T in_texel = VEC4_T(texelFetch(
-    image_in,
+    t_in,
     to_texture_pos(in_idx, in_sizes, packed_dim),
     0));
 
   ivec4 other_idx = broadcast_indices(idx, other_sizes);
   VEC4_T other_texel = VEC4_T(texelFetch(
-    image_other,
+    t_other,
     to_texture_pos(other_idx, other_sizes, packed_dim),
     0));
 
@@ -75,5 +60,5 @@ void main() {
     other_texel = other_texel.xxxx;
   }
 
-  imageStore(image_out, pos, VEC4_T(op(in_texel, other_texel, alpha)));
+  imageStore(t_out, pos, VEC4_T(op(in_texel, other_texel, alpha)));
 }

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

@@ -10,6 +10,7 @@ binary_op:
     NDIM: 3
     DTYPE: float
     PACKING: C_packed
+    STORAGE: texture3d
   generate_variant_forall:
     DTYPE:
       - VALUE: half

backends/vulkan/test/test_vulkan_delegate.py

@@ -202,6 +202,23 @@ def forward(self, x, y, w):
 
         self.lower_module_and_test_output(add_module, sample_inputs)
 
+    def test_vulkan_backend_add_int(self):
+        class AddIntModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, y):
+                z = x + y
+                return z
+
+        add_int_module = AddIntModule()
+        sample_inputs = (
+            torch.randint(low=-100, high=100, size=(2, 3), dtype=torch.int32),
+            torch.randint(low=-100, high=100, size=(2, 3), dtype=torch.int32),
+        )
+
+        self.lower_module_and_test_output(add_int_module, sample_inputs)
+
     def test_vulkan_backend_zero_dim_tensor(self):
         class ZeroDimModule(torch.nn.Module):
             def __init__(self):