[ET-VK][Ops] aten.var.dim in reduce

Pull Request resolved: #11198

Incorporated variance logic into reduce by adding additional logic

Differential Revision: [D75247432](https://our.internmc.facebook.com/intern/diff/D75247432/)
ghstack-source-id: 286724866

Author: morelos

backends/vulkan/runtime/graph/ops/glsl/var_buffer.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/reduce_buffer.glsl

@@ -31,15 +31,24 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 layout(constant_id = 3) const int reduce_dim = 0;
 
+$if VARIANCE_MODE:
+  #define VARIANCE_MODE
+
 #define NWORKERS 4
 #define MAX_THREADS 16
 
-shared T shared_sum[NWORKERS];
+shared T shared_accum[NWORKERS];
+#ifdef VARIANCE_MODE
 shared T shared_sum_sq[NWORKERS];
 shared int shared_count[NWORKERS];
+#endif
 
 #include "indexing_utils.h"
 
+#define INIT_ACCUM(first_val) ${INIT_ACCUM}
+#define UPDATE_ACCUM(accum, new_val) ${UPDATE_ACCUM}
+#define POSTPROCESS(accum) ${POSTPROCESS}
+
 void main() {
   const ivec4 out_idx = ivec4(
       gl_GlobalInvocationID.x,
@@ -49,9 +58,11 @@ void main() {
 
   const uint tid = gl_LocalInvocationID[reduce_dim];
 
-  shared_sum[tid] = T(0);
+  shared_accum[tid] = T(0);
+#ifdef VARIANCE_MODE
   shared_sum_sq[tid] = T(0);
   shared_count[tid] = 0;
+#endif
   barrier();
 
   const int R = in_sizes[reduce_dim];
@@ -65,9 +76,25 @@ void main() {
   uint len = q + (tid < rem ? 1u : 0u);
   uint base = tid * q + min(tid, rem);
 
-  T sum = T(0);
+  // Get the first value for initializing the accumulator if needed
+  T first_val = T(0);
+  if (R > 0) {
+    ivec4 first_idx = out_idx;
+    first_idx[reduce_dim] = 0;
+
+    if (reduce_dim == 2) {
+      first_idx[reduce_dim + 1] = 0;
+    }
+
+    first_val = in_buf[tidx_to_bufi(first_idx, in_strides)];
+  }
+
+  // Initialize accumulator
+  T accum = INIT_ACCUM(first_val);
+#ifdef VARIANCE_MODE
   T sum_sq = T(0);
   int count = 0;
+#endif
 
   ivec4 in_idx = out_idx;
   for (uint off = 0u; off < len; ++off) {
@@ -83,39 +110,55 @@ void main() {
 
     T v = in_buf[tidx_to_bufi(in_idx, in_strides)];
 
-    sum += v;
+    accum = UPDATE_ACCUM(accum, v);
+
+#ifdef VARIANCE_MODE
     sum_sq += v * v;
     count += 1;
+#endif
   }
 
-  shared_sum[tid] = sum;
+  shared_accum[tid] = accum;
+#ifdef VARIANCE_MODE
   shared_sum_sq[tid] = sum_sq;
   shared_count[tid] = count;
+#endif
   barrier();
 
   if (tid == 0u) {
-    T tot_sum = T(0);
-    T tot_sum_sq = T(0);
-    int tot_count = 0;
+    T result = shared_accum[0];
+
+#ifdef VARIANCE_MODE
+    T tot_sum = shared_accum[0];
+    T tot_sum_sq = shared_sum_sq[0];
+    int tot_count = shared_count[0];
+#endif
 
-    for (uint i = 0; i < N; ++i) {
-      tot_sum += shared_sum[i];
+    for (uint i = 1; i < N; ++i) {
+#ifdef VARIANCE_MODE
+      tot_sum += shared_accum[i];
       tot_sum_sq += shared_sum_sq[i];
       tot_count += shared_count[i];
+#else
+      result = UPDATE_ACCUM(result, shared_accum[i]);
+#endif
     }
 
-    T var;
+#ifdef VARIANCE_MODE
     if (tot_count > 0) {
       T mean = tot_sum / T(tot_count);
-      var = (tot_sum_sq / T(tot_count)) - (mean * mean);
+      result = (tot_sum_sq / T(tot_count)) - (mean * mean);
       if (pc.unbiased != 0 && tot_count > 1) {
-        var *= T(tot_count) / T(tot_count - 1);
+        result *= T(tot_count) / T(tot_count - 1);
       }
-    } else{
+    } else {
       // NaN to match PyTorch behavior
-      var = T(0.0/0.0);
+      result = T(0.0/0.0);
     }
+#else
+    result = POSTPROCESS(result);
+#endif
 
-    out_buf[tidx_to_bufi(out_idx, out_strides)] = var;
+    out_buf[tidx_to_bufi(out_idx, out_strides)] = result;
   }
 }

backends/vulkan/runtime/graph/ops/glsl/reduce_buffer.yaml

@@ -0,0 +1,32 @@
+# 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.
+
+reduce_buffer:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+    INIT_ACCUM: T(0)
+    UPDATE_ACCUM: accum + new_val
+    POSTPROCESS: accum
+    VARIANCE_MODE: false
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: sum_buffer
+    - NAME: mean_buffer
+      POSTPROCESS: (accum / T(in_sizes[reduce_dim]))
+    - NAME: amax_buffer
+      INIT_ACCUM: first_val
+      UPDATE_ACCUM: max(accum, new_val)
+      POSTPROCESS: accum
+    - NAME: amin_buffer
+      INIT_ACCUM: first_val
+      UPDATE_ACCUM: min(accum, new_val)
+      POSTPROCESS: accum
+    - NAME: var_buffer
+      VARIANCE_MODE: true

backends/vulkan/runtime/graph/ops/glsl/reduce.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/reduce_texture3d.glsl

@@ -23,12 +23,19 @@ ${layout_declare_tensor(B, "r", "tin", DTYPE, STORAGE)}
 ${layout_declare_ubo(B, "ivec3", "tin_limits")}
 ${layout_declare_ubo(B, "ivec4", "tin_sizes")}
 
+layout(push_constant) uniform PushConstants {
+  int unbiased;
+} pc;
+
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 layout(constant_id = 3) const int packed_dim = 0;
 layout(constant_id = 4) const int reduce_dim = 0;
 layout(constant_id = 5) const int group_dim = 1;
 
+$if VARIANCE_MODE:
+  #define VARIANCE_MODE
+
 // A more verbose name would be NWORKERS_PER_GROUP. This describes the number of
 // threads that will co-operate to compute one reduction output. There may be
 // multiple groups computing distinct reduction outputs within one work group.
@@ -39,15 +46,29 @@ layout(constant_id = 5) const int group_dim = 1;
 // work group will write into its assigned element in the shared array.
 #define MAX_NTHREADS 16
 
-
 shared vec4 shared_vecs[MAX_NTHREADS];
+// Second accumulator for variance mode - used for sum of values, prev
+// accumulator is used for sum of squares
+shared vec4 shared_sum_sq[MAX_NTHREADS];
+shared int shared_count[MAX_NTHREADS];
 
 #include "indexing_utils.h"
 
 int tid_to_smi(const ivec2 tid) {
   return tid.x + tid.y * NWORKERS;
 }
 
+vec4 calculate_variance(vec4 sum, vec4 sum_sq, int count) {
+  vec4 mean = sum / float(count);
+  vec4 variance = (sum_sq / float(count)) - (mean * mean);
+
+  if ((pc.unbiased != 0) && (count > 1)) {
+    variance = variance * (float(count) / float(count - 1.0));
+  }
+
+  return variance;
+}
+
 /*
  * The functions below compute reduction along a single dimension for a tensor.
  * The shader template generalize reduction by abstracting the initial value of
@@ -92,25 +113,48 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
   scan_pos[reduce_dim] = 0;
   vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos));
 
+#ifdef VARIANCE_MODE
+  vec4 sum_sq = VEC4_T(0);
+  int count = 0;
+#endif
+
   scan_pos[reduce_dim] = tid.x;
   // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of
   // the reduction row
   for (int i = tid.x; i < tin_sizes[reduce_dim];
        i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
-    accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos));
+    vec4 val = load_texel(tin, scan_pos);
+    accum = UPDATE_ACCUM(accum, val);
+#ifdef VARIANCE_MODE
+    sum_sq += val * val;
+    count += 1;
+#endif
   }
   // Write partial output to shared memory and synchronize work group
   shared_vecs[smi] = accum;
+#ifdef VARIANCE_MODE
+  shared_sum_sq[smi] = sum_sq;
+  shared_count[smi] = count;
+#endif
   barrier();
 
   // Since the reduction row is reduced to only one element, only the "main"
   // thread in the group needs aggregate the partial outputs
   if (tid.x == 0) {
     // 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++) {
-      accum = UPDATE_ACCUM(accum, shared_vecs[group_i]);
+    accum = shared_vecs[group_i];
+#ifdef VARIANCE_MODE
+    sum_sq = shared_sum_sq[group_i];
+    count = shared_count[group_i];
+#endif
+    for (int i = 1; i < NWORKERS; i++) {
+      int idx = tid.y * NWORKERS + i;
+      accum = UPDATE_ACCUM(accum, shared_vecs[idx]);
+#ifdef VARIANCE_MODE
+      sum_sq += shared_sum_sq[idx];
+      count += shared_count[idx];
+#endif
     }
 
     // Determine if there are any padding elements in the final texel of the
@@ -121,14 +165,27 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
     const bool is_last_texel =
         scan_pos[packed_dim] == (tin_limits[packed_dim] - 1);
 
+#ifdef VARIANCE_MODE
+    vec4 variance = calculate_variance(accum, sum_sq, count);
+#endif
+
     // Explicitly set padding elements to 0
     if (is_last_texel && nspill > 0) {
       [[unroll]] for (int i = nspill; i < 4; i++) {
+#ifdef VARIANCE_MODE
+        variance[i] = 0;
+#else
         accum[i] = 0;
+#endif
       }
     }
+
     scan_pos[reduce_dim] = tid.x;
+#ifdef VARIANCE_MODE
+    write_texel(tout, scan_pos, variance);
+#else
     write_texel(tout, scan_pos, POSTPROCESS(accum));
+#endif
   }
 }
 
@@ -153,35 +210,78 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
   scan_pos[reduce_dim] = 0;
   vec4 accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x));
 
+#ifdef VARIANCE_MODE
+  vec4 sum_sq = VEC4_T(0);
+  int count = 0;
+#endif
+
   // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of
   // the reduction row
   scan_pos[reduce_dim] = tid.x;
   for (int i = tid.x * 4; i < reduce_len;
        i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
-    accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos));
+    vec4 val = load_texel(tin, scan_pos);
+    accum = UPDATE_ACCUM(accum, val);
+#ifdef VARIANCE_MODE
+    sum_sq += val * val;
+    count += 4; // Each texel has 4 elements
+#endif
   }
   // For the last texel in the dim, if there are padding elements then each
   // element of the texel needs to be processed individually such that the
   // padding elements are ignored
   if (scan_pos[reduce_dim] == tin_limits[reduce_dim] - 1 && nspill > 0) {
-    const vec4 intex = load_texel(tin, scan_pos);
+    const vec4 val = load_texel(tin, scan_pos);
     for (int i = 0; i < nspill; i++) {
-      accum.x = UPDATE_ACCUM(accum.x, intex[i]);
+      accum.x = UPDATE_ACCUM(accum.x, val[i]);
+#ifdef VARIANCE_MODE
+      sum_sq.x += val[i] * val[i];
+      count += 1;
+#endif
     }
   }
   // Write partial output to shared memory and synchronize work group
   shared_vecs[smi] = accum;
+#ifdef VARIANCE_MODE
+  shared_sum_sq[smi] = sum_sq;
+  shared_count[smi] = count;
+#endif
   barrier();
 
   // Since the reduction row is reduced to only one element, only the "main"
   // thread in the group needs aggregate the partial outputs
   if (tid.x == 0) {
     // Iterate over the partial maximums to obtain the overall maximum
     int group_i = tid.y * NWORKERS;
-    accum = shared_vecs[group_i++];
+    accum = shared_vecs[group_i];
+#ifdef VARIANCE_MODE
+    sum_sq = shared_sum_sq[group_i];
+    count = shared_count[group_i];
+#endif
     for (int i = 1; i < NWORKERS; i++, group_i++) {
-      accum = UPDATE_ACCUM(accum, shared_vecs[group_i]);
+      int idx = tid.y * NWORKERS + i;
+      accum = UPDATE_ACCUM(accum, shared_vecs[idx]);
+#ifdef VARIANCE_MODE
+      sum_sq += shared_sum_sq[idx];
+      count += shared_count[idx];
+#endif
     }
+
+#ifdef VARIANCE_MODE
+    float total_sum = accum.x + accum.y + accum.z + accum.w;
+    float total_sum_sq = sum_sq.x + sum_sq.y + sum_sq.z + sum_sq.w;
+    int total_count = count;
+
+    float mean = total_sum / float(total_count);
+    float variance = (total_sum_sq / float(total_count)) - (mean * mean);
+
+    if ((pc.unbiased != 0) && (total_count > 1)) {
+      variance = variance * (float(total_count) / float(total_count - 1.0));
+    }
+
+    scan_pos[reduce_dim] = tid.x;
+    write_texel(tout, scan_pos, vec4(variance, 0, 0, 0));
+#else
     // Each element of the texel is itself a partial maximum; iterate over the
     // texel to find the actual maximum
     float accum_final = accum.x;
@@ -191,6 +291,7 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
 
     scan_pos[reduce_dim] = tid.x;
     write_texel(tout, scan_pos, POSTPROCESS(vec4(accum_final, 0, 0, 0)));
+#endif
   }
 }