Update on "[ET-VK][Test] aten.max_pool2d_with_indices"

Due to the below issues, we only check equality of the output tensor and not the index tensor.

1. We can't verify index tensors since VK-float16 vs CPU-float32 deltas can change which index in a pool is the maximum. That can yield completely different integers in the index tensor. Hence, we only verify the output tensor not the index tensor.
2. To actually visualize the index tensor, we need to re-construct the int32 values from the int64 values. Since the `torch.int64` index tensor is serialized as `int32` in Vulkan, Python expects int64 but C++ writes to the buffer as though it is for int32. Hence, we must apply some computation to re-construct the tensor. See below for details. A helper function was included in an earlier version of this change, but was removed for conciseness since we aren't checking that index tensor anyway.

For example, if the first and second elements return 16 and 17, we get this value as the first element:
```
73014444048 = 1000100000000000000000000000000010000
```
We must split this int64 into two int32 values, and construct a new tensor accordingly.
```
10001 | 00000000000000000000000000010000
10001 | 10000
17 | 16
```

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

[ghstack-poisoned]

Author: jorgep31415

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

@@ -59,19 +59,11 @@ void main() {
       if ((x >= 0 && x < in_extents.data.x) && (y >= 0 && y < in_extents.data.y)) {
         const vec4 cur_texel = texelFetch(image_in, ivec3(x, y, pos.z), 0);
 
-        const int cur_idx = x + int(in_extents.data.x) * y;
-        if (cur_texel.x > out_texel.x) {
-          idx_texel.x = cur_idx;
-        }
-        if (cur_texel.y > out_texel.y) {
-          idx_texel.y = cur_idx;
-        }
-        if (cur_texel.z > out_texel.z) {
-          idx_texel.z = cur_idx;
-        }
-        if (cur_texel.w > out_texel.w) {
-          idx_texel.w = cur_idx;
-        }
+        // Set idx if value is greatest in the pool; else, keep the existing idx.
+        ivec4 cur_idx = ivec4(x + int(in_extents.data.x) * y);
+        ivec4 mask = ivec4(greaterThan(cur_texel, out_texel));
+        idx_texel = ivec4(mix(idx_texel, cur_idx, mask));
+
         out_texel = max(cur_texel, out_texel);
       }
       else {

backends/vulkan/runtime/graph/ops/impl/Pool.cpp

@@ -11,7 +11,6 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>

backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.cpp

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+
+int64_t calc_out_size(
+    const int64_t in_size,
+    const int64_t kernel,
+    const int64_t stride,
+    const int64_t padding,
+    const int64_t dilation,
+    const bool ceil_mode) {
+  int64_t c = ceil_mode ? stride - 1 : 0;
+  int64_t out_size =
+      (in_size + 2 * padding - dilation * (kernel - 1) - 1 + c) / stride + 1;
+  if (ceil_mode && (out_size - 1) * stride >= in_size + padding) {
+    --out_size;
+  }
+  return out_size;
+}
+
+api::utils::ivec2 normalize_wh(Value& v) {
+  if (v.isInt()) {
+    return api::utils::make_ivec2({v.toInt(), v.toInt()});
+  } else {
+    auto l = v.toIntList();
+    return api::utils::make_ivec2({l.at(1), l.at(0)});
+  }
+}
+
+} // namespace vulkan
+} // namespace native
+} // namespace at

backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h

@@ -31,24 +31,9 @@ int64_t calc_out_size(
     const int64_t stride,
     const int64_t padding,
     const int64_t dilation,
-    const bool ceil_mode) {
-  int64_t c = ceil_mode ? stride - 1 : 0;
-  int64_t out_size =
-      (in_size + 2 * padding - dilation * (kernel - 1) - 1 + c) / stride + 1;
-  if (ceil_mode && (out_size - 1) * stride >= in_size + padding) {
-    --out_size;
-  }
-  return out_size;
-}
-
-api::utils::ivec2 normalize_wh(Value& v) {
-  if (v.isInt()) {
-    return api::utils::make_ivec2({v.toInt(), v.toInt()});
-  } else {
-    auto l = v.toIntList();
-    return api::utils::make_ivec2({l.at(1), l.at(0)});
-  }
-}
+    const bool ceil_mode);
+
+api::utils::ivec2 normalize_wh(Value& v);
 
 } // namespace vulkan
 } // namespace native