From d95d3c4019dc366bea90102d62c4e1b385b6bd18 Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
Date: Mon, 26 May 2025 21:41:04 -0700
Subject: [PATCH] [ET-VK] Storing positions in uint16 to instead of int in
 conv2d pw shader.

This diff modifies the `conv2d_pw_s1p0.glsl` shader to store positions in `uint16` instead of `int`. The changes include adding the necessary extension for explicit arithmetic types, updating the type definitions for `TILE_SIZE_X` and `TILE_SIZE_Y`, and changing the type of the `pos` array.

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

[ghstack-poisoned]
---
 .../runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl   | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl
index 36c7a61eb3d..9d424683077 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl
@@ -8,12 +8,14 @@
 
 #version 450 core
 
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+
 #define PRECISION ${PRECISION}
 
 #define VEC4_T ${texel_type(DTYPE)}
 
-#define TILE_SIZE_X ${TILE_SIZE_X}
-#define TILE_SIZE_Y ${TILE_SIZE_Y}
+#define TILE_SIZE_X ${TILE_SIZE_X}us
+#define TILE_SIZE_Y ${TILE_SIZE_Y}us
 
 #define op(X, A, B) ${OPERATOR}
 
@@ -63,11 +65,11 @@ void main() {
   // +--------+--------+
   // | pos[2] | pos[3] |
   // +--------+--------+
-  int pos[TILE_SIZE_X * TILE_SIZE_Y * 2];
-  for (int y = 0, i = 0; y < TILE_SIZE_Y; ++y) {
-    for (int x = 0; x < TILE_SIZE_X; ++x) {
-      pos[i * 2] = out_pos[0] * TILE_SIZE_X + x;
-      pos[i * 2 + 1] = out_pos[1] * TILE_SIZE_Y + y;
+  uint16_t pos[TILE_SIZE_X * TILE_SIZE_Y * 2];
+  for (uint16_t y = 0us, i = 0us; y < TILE_SIZE_Y; ++y) {
+    for (uint16_t x = 0us; x < TILE_SIZE_X; ++x) {
+      pos[i * 2] = uint16_t(out_pos[0]) * TILE_SIZE_X + x;
+      pos[i * 2 + 1] = uint16_t(out_pos[1]) * TILE_SIZE_Y + y;
       i++;
     }
   }