[ET-VK] Performance improvement to few indexing functions. (pytorch#9171)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#9138 by
@trivedivivek
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/61/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/61/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/60/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/trivedivivek/61/orig
@diff-train-skip-merge

---------

Co-authored-by: Vivek Trivedi <[email protected]>

Author: pytorchbot

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -104,16 +104,19 @@ ivec4 tidx_to_4bufi(
 }
 
 ivec4 nchwi_to_tidx(const int nchwi, const ivec4 sizes) {
+  const int nchwi_div_x = nchwi / sizes.x;
+  const int nchwi_div_y = nchwi_div_x / sizes.y;
   return ivec4(
       nchwi % sizes.x,
-      (nchwi / (sizes.x)) % sizes.y,
-      (nchwi / (sizes.x * sizes.y)) % sizes.z,
-      (nchwi / (sizes.x * sizes.y * sizes.z)));
+      nchwi_div_x % sizes.y,
+      nchwi_div_y % sizes.z,
+      nchwi_div_y / sizes.z);
 }
 
 int tidx_to_nchwi(const ivec4 tidx, const ivec4 sizes) {
-  return tidx.w * sizes.x * sizes.y * sizes.z + tidx.z * sizes.x * sizes.y +
-      tidx.y * sizes.x + tidx.x;
+  const int sizes_xy = sizes.x * sizes.y;
+  return tidx.w * sizes_xy * sizes.z + tidx.z * sizes_xy + tidx.y * sizes.x +
+      tidx.x;
 }
 
 // TODO(ssjia): make this function use dim order so that it can work with any
@@ -360,8 +363,8 @@ ivec4 to_texture_elem_pos(ivec4 idx, ivec4 sizes, int packed_dim) {
   //  pos[4] is set to a placeholder value
   ivec4 pos = idx.xyzx;
   pos[BATCH_AXIS] += idx.w * sizes[BATCH_AXIS];
-  pos[packed_dim] /= 4;
-  pos.w = idx[packed_dim] % 4;
+  pos[packed_dim] >>= 2;
+  pos.w = idx[packed_dim] & 0x3;
   return pos;
 }