From 58c55c2c5c921fed60e605daf2c4ea68bb2a6829 Mon Sep 17 00:00:00 2001
From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com>
Date: Thu, 2 Jan 2025 12:16:21 -0800
Subject: [PATCH] [ET-VK] Using vec2 to store output positions to reudce shader
 register footprint.

The diff changes the use of `u16vec3` to `u16vec2` to store output positions in the conv2d_pw op. This change is made to reduce the shader register footprint and improve performance.

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

[ghstack-poisoned]
---
 .../vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl     | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl
index 9d1f6c3bd91..af7c22bb5ad 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl
@@ -53,18 +53,18 @@ void main() {
   // +--------+--------+
   // | pos[2] | pos[3] |
   // +--------+--------+
-  u16vec3 pos[TILE_SIZE * TILE_SIZE];
+  u16vec2 pos[TILE_SIZE * TILE_SIZE];
   for (int y = 0, i = 0; y < TILE_SIZE; ++y) {
     for (int x = 0; x < TILE_SIZE; ++x) {
-      pos[i] = u16vec3(
-          gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y, gpos.z);
+      pos[i] = u16vec2(
+          gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y);
       i++;
     }
   }
 
   // If the top left position is out of bounds, then this invocation will have
   // no work to do.
-  if (any(greaterThanEqual(pos[0], out_limits))) {
+  if (any(greaterThanEqual(u16vec3(pos[0], gpos.z), out_limits))) {
     return;
   }
 
@@ -138,8 +138,8 @@ void main() {
   }
 
   for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
-    if (all(lessThan(pos[i], out_limits))) {
-      imageStore(t_out, pos[i], op(sum[i], out_min, out_max));
+    if (all(lessThan(u16vec3(pos[i], gpos.z), out_limits))) {
+      imageStore(t_out, u16vec3(pos[i], gpos.z), op(sum[i], out_min, out_max));
     }
   }
 }