Update on "[ET-VK] Implement generic reduction shader + mean, sum, amax, amin"

## Context

Introduce a generic shader to compute reduction along a single dim, and `keepdim = True`. With the generic shader template, `mean`, `sum`, `amin`, and `amax` can be implemented.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/partitioner/supported_ops.py

@@ -105,9 +105,6 @@ def __contains__(self, op):
 ]
 
 NO_DYNAMIC_SHAPE = [
-    # Reduction
-    exir_ops.edge.aten.mean.dim,
-    exir_ops.edge.aten.sum.dim_IntList,
     # Normalization
     exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
     exir_ops.edge.aten.native_layer_norm.default,

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

@@ -109,7 +109,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
     // Iterate over the partial outputs to obtain the overall output
     int group_i = tid.y * NWORKERS;
     accum = shared_vecs[group_i++];
-    for (int i = 1; i < NWORKERS; ++i, group_i++) {
+    for (int i = 1; i < NWORKERS; i++, group_i++) {
       accum = UPDATE_ACCUM(accum, shared_vecs[group_i]);
     }
 
@@ -123,7 +123,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
 
     // Explicitly set padding elements to 0
     if (is_last_texel && nspill > 0) {
-      [[unroll]] for (int i = nspill; i < 4; ++i) {
+      [[unroll]] for (int i = nspill; i < 4; i++) {
         accum[i] = 0;
       }
     }
@@ -165,7 +165,7 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
   // padding elements are ignored
   if (scan_pos[reduce_dim] == tin_limits[reduce_dim] - 1 && nspill > 0) {
     const vec4 intex = load_texel(tin, scan_pos);
-    for (int i = 0; i < nspill; ++i) {
+    for (int i = 0; i < nspill; i++) {
       accum.x = UPDATE_ACCUM(accum.x, intex[i]);
     }
   }
@@ -179,13 +179,13 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
     // Iterate over the partial maximums to obtain the overall maximum
     int group_i = tid.y * NWORKERS;
     accum = shared_vecs[group_i++];
-    for (int i = 1; i < NWORKERS; ++i, group_i++) {
+    for (int i = 1; i < NWORKERS; i++, group_i++) {
       accum = UPDATE_ACCUM(accum, shared_vecs[group_i]);
     }
     // Each element of the texel is itself a partial maximum; iterate over the
     // texel to find the actual maximum
     float accum_final = accum.x;
-    [[unroll]] for (int i = 1; i < 4; ++i) {
+    [[unroll]] for (int i = 1; i < 4; i++) {
       accum_final = UPDATE_ACCUM(accum[i], accum_final);
     }
 

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

@@ -25,8 +25,11 @@ void resize_reduce_node(
   vTensorPtr out = graph->get_tensor(args[0].refs[0]);
   vTensorPtr in = graph->get_tensor(args[1].refs[0]);
 
-  std::vector<int64_t> in_sizes = in->sizes();
-  // out->virtual_resize(in_sizes);
+  int dim = extra_args[0];
+
+  std::vector<int64_t> new_sizes = in->sizes();
+  new_sizes[normalize(dim, new_sizes.size())] = 1;
+  out->virtual_resize(new_sizes);
 }
 
 void add_reduce_node(
@@ -48,12 +51,8 @@ void add_reduce_node(
   // Check that the concat dim is not the reduction dim, if the tensor has a
   // batch dim greater than 1.
   if (graph.dim_of(in) == 4 && graph.size_at<int>(0, in) > 1) {
-    VK_CHECK_COND(
-        graph.concat_dim_of(in) != reduce_dim,
-        "Reduce shader currently does not support concat dim == reduce dim");
-    VK_CHECK_COND(
-        graph.concat_dim_of(out) != reduce_dim,
-        "Reduce shader currently does not support concat dim == reduce dim");
+    VK_CHECK_COND(graph.concat_dim_of(in) != reduce_dim);
+    VK_CHECK_COND(graph.concat_dim_of(out) != reduce_dim);
   }
 
   vkapi::ShaderInfo shader_descriptor;
@@ -97,7 +96,8 @@ void add_reduce_node(
       // Specialization Constants
       {graph.packed_dim_of(out), reduce_dim, group_dim},
       // Resizing Logic
-      resize_reduce_node));
+      resize_reduce_node,
+      {dim}));
 }
 
 #define DEFINE_REDUCE_FN(op_name, out_arg_idx)                           \