PDX fast path for spatial mean.out of channels-last tensors (#24164)

Summary:
A float `mean` over the spatial dims of a channels-last tensor,
`x[N,H,W,C].mean((1,2))` (global average pooling), runs `mean.out`
through NNLib's `xa_nn_mean_f32_f32`. For this reduction (dims {1,2} of
an `[N,H,W,C]` tensor) NNLib uses a strided generic loop that costs ~10
cycles per element, which makes it a hotspot in networks that pool this
way.

Add `mean_nhwc_spatial_simd` in front of the NNLib call. For a float
`[N,H,W,C]` tensor reduced over dims {1,2} (negative dims accepted, any
`keepdim`), it zeroes the output row for each batch element, adds every
position row into it with `PDX_ADD_MXF32` (four channels per vector,
positions in order), then divides by the position count with
`PDX_DIV_MXF32`. The per-channel summation order is the sequential
position order, and the division is an exact IEEE division.

PDX vector accesses must be 16-byte aligned on Fusion G3 (an unaligned
access raises `LoadStoreAlignmentCause`), so the fast path is only taken
when both buffers are 16-byte aligned and `C % 4 == 0`; every other case
falls through to the existing NNLib / portable paths unchanged.

On a Fusion G3 core, `mean.out` on these shapes is ~13x faster, with
unchanged output.

Differential Revision: D127346731

Author: rezaasjd

backends/cadence/fusion_g3/operators/op_mean.cpp

@@ -59,6 +59,72 @@ int prepare_data(
   return num_axis_dims;
 }
 
+namespace {
+
+constexpr size_t kSimdLanes = 4;
+constexpr uintptr_t kSimdAlignmentBytes = 16;
+
+// Spatial mean of a channels-last float tensor [N, H, W, C] over dims {1, 2},
+// the pooling of squeeze-and-excitation blocks. NNLib reduces it with a strided
+// generic loop; here every position row is added into the output row with PDX
+// vectors. PDX vector accesses must be 16-byte aligned, so unaligned tensors
+// and channel counts that are not a multiple of the lane width return false.
+bool mean_nhwc_spatial_simd(
+    const Tensor& in,
+    optional<ArrayRef<int64_t>> dim_list,
+    Tensor& out) {
+  if (!dim_list.has_value() || dim_list.value().size() != 2 || in.dim() != 4) {
+    return false;
+  }
+  int64_t first_dim = dim_list.value()[0];
+  int64_t second_dim = dim_list.value()[1];
+  first_dim += first_dim < 0 ? in.dim() : 0;
+  second_dim += second_dim < 0 ? in.dim() : 0;
+  if (!((first_dim == 1 && second_dim == 2) ||
+        (first_dim == 2 && second_dim == 1))) {
+    return false;
+  }
+
+  const size_t batch = in.size(0);
+  const size_t positions = static_cast<size_t>(in.size(1)) * in.size(2);
+  const size_t channels = in.size(3);
+  if (positions == 0 || channels == 0 || channels % kSimdLanes != 0 ||
+      static_cast<size_t>(out.numel()) != batch * channels) {
+    return false;
+  }
+
+  const float* p_in = in.const_data_ptr<float>();
+  float* p_out = out.mutable_data_ptr<float>();
+  if (reinterpret_cast<uintptr_t>(p_in) % kSimdAlignmentBytes != 0 ||
+      reinterpret_cast<uintptr_t>(p_out) % kSimdAlignmentBytes != 0) {
+    return false;
+  }
+
+  const size_t channel_vectors = channels / kSimdLanes;
+  const float count = static_cast<float>(positions);
+  const xb_vecMxf32 v_zero{0.0f, 0.0f, 0.0f, 0.0f};
+  const xb_vecMxf32 v_count{count, count, count, count};
+  for (size_t n = 0; n < batch; ++n) {
+    auto v_out = reinterpret_cast<xb_vecMxf32*>(p_out + n * channels);
+    for (size_t v = 0; v < channel_vectors; ++v) {
+      v_out[v] = v_zero;
+    }
+    for (size_t p = 0; p < positions; ++p) {
+      auto v_in = reinterpret_cast<const xb_vecMxf32*>(
+          p_in + (n * positions + p) * channels);
+      for (size_t v = 0; v < channel_vectors; ++v) {
+        v_out[v] = PDX_ADD_MXF32(v_out[v], v_in[v]);
+      }
+    }
+    for (size_t v = 0; v < channel_vectors; ++v) {
+      v_out[v] = PDX_DIV_MXF32(v_out[v], v_count);
+    }
+  }
+  return true;
+}
+
+} // namespace
+
 Tensor& mean_out(
     KernelRuntimeContext& ctx,
     const Tensor& in,
@@ -100,6 +166,10 @@ Tensor& mean_out(
   if (in.dim() > kNnlibMaxDim)
     optimized = false;
 
+  if (optimized && mean_nhwc_spatial_simd(in, dim_list, out)) {
+    return out;
+  }
+
   if (optimized) {
     float* __restrict__ p_out = out.mutable_data_ptr<float>();
     const float* __restrict__ p_inp =