Update on "[vulkan] Distribute weight prepacking along y dimension fo…

…r conv2d"

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

[ghstack-poisoned]

aten/src/ATen/native/vulkan/glsl/conv2d.glsl

@@ -42,12 +42,12 @@ void main() {
 
     vec4 sum = uBias.data[pos.z];
 
-    for (int z = 0; z < uBlock.kernel.z; ++z) {
-      const ivec4 kz = block + 4 * z;
+    for (int z = 0; z < uBlock.kernel.z; z+=4) {
+      const ivec4 kz = block + z;
 
       for (int y = start.y, ky = kstart.y; y < end.y; y += uBlock.dilate.y, ++ky) {
         for (int x = start.x, kx = kstart.x; x < end.x; x += uBlock.dilate.x, ++kx) {
-          const vec4 In = texelFetch(uInput, ivec3(x, y, z), 0);
+          const vec4 In = texelFetch(uInput, ivec3(x, y, z/4), 0);
 
           sum = fma(In.xxxx, texelFetch(uKernel, ivec3(kx, (uBlock.kernel.y*tower) + ky, kz.x), 0), sum);
           sum = fma(In.yyyy, texelFetch(uKernel, ivec3(kx, (uBlock.kernel.y*tower) + ky, kz.y), 0), sum);

aten/src/ATen/native/vulkan/glsl/conv2d_pw.glsl

@@ -37,9 +37,9 @@ void main() {
 
     vec4 sum = uBias.data[pos.z];
 
-    for (int z = 0; z < uBlock.kernel.x; ++z) {
-      const vec4 In = texelFetch(uInput, ivec3(ipos.x, ipos.y, z), 0);
-      const ivec4 kz = block + 4 * z;
+    for (int z = 0; z < uBlock.kernel.x; z+=4) {
+      const vec4 In = texelFetch(uInput, ivec3(ipos.x, ipos.y, z/4), 0);
+      const ivec4 kz = block + z;
 
       sum = fma(In.xxxx, texelFetch(uKernel, ivec3(0, tower, kz.x), 0), sum);
       sum = fma(In.yyyy, texelFetch(uKernel, ivec3(0, tower, kz.y), 0), sum);