Minor reduction cleanups

- Move neondot kernels to the correct #if guard
- Minor reformatting to make things a little easier to read

PiperOrigin-RevId: 841025925

Author: dsharletg

ynnpack/kernels/reduce/arm_neon.cc

@@ -38,7 +38,7 @@ static f32x4x16 reduce_add(
     float32x4_t hi =
         vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(b.v[i].v), 16));
 
-    a.v[2 * i] += f32x4{lo};
+    a.v[2 * i + 0] += f32x4{lo};
     a.v[2 * i + 1] += f32x4{hi};
   }
 
@@ -54,7 +54,8 @@ static f32x4 reduce_add(
       vreinterpretq_f32_u32(vandq_u32(pairs, vdupq_n_u32(0xFFFF0000)));
 
   a += f32x4{odd};
-  return a += f32x4{even};
+  a += f32x4{even};
+  return a;
 }
 
 static f32x4x16 reduce_add(
@@ -66,7 +67,7 @@ static f32x4x16 reduce_add(
         vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(b.v[i].v), 16));
     float32x4_t hi =
         vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(b.v[i].v), 16));
-    a.v[2 * i].v = vmlaq_f32(a.v[2 * i].v, lo, lo);
+    a.v[2 * i + 0].v = vmlaq_f32(a.v[2 * i + 0].v, lo, lo);
     a.v[2 * i + 1].v = vmlaq_f32(a.v[2 * i + 1].v, hi, hi);
   }
 

ynnpack/kernels/reduce/arm_neonbf16.cc

@@ -36,7 +36,7 @@ static f32x4x16& operator+=(f32x4x16& a, bf16x8x8 b) {
   YNN_UNROLL
   for (size_t i = 0; i < 8; ++i) {
     uint16x8x2_t zipped = vzipq_u16(b.v[i].v, zero);
-    a.v[2 * i].v = vbfdotq_f32(a.v[2 * i].v,
+    a.v[2 * i + 0].v = vbfdotq_f32(a.v[2 * i + 0].v,
         vreinterpretq_bf16_u16(zipped.val[0]), one);
     a.v[2 * i + 1].v = vbfdotq_f32(a.v[2 * i + 1].v,
         vreinterpretq_bf16_u16(zipped.val[1]), one);
@@ -62,7 +62,7 @@ static f32x4x16 reduce_add(
   YNN_UNROLL
   for (size_t i = 0; i < 8; ++i) {
     uint16x8x2_t zipped = vzipq_u16(b.v[i].v, zero);
-    a.v[2 * i].v = vbfdotq_f32(a.v[2 * i].v,
+    a.v[2 * i + 0].v = vbfdotq_f32(a.v[2 * i + 0].v,
         vreinterpretq_bf16_u16(zipped.val[0]),
         vreinterpretq_bf16_u16(zipped.val[0]));
     a.v[2 * i + 1].v = vbfdotq_f32(a.v[2 * i + 1].v,

ynnpack/kernels/reduce/arm_neonfp16arith.cc

@@ -46,7 +46,7 @@ static f32x4x16& operator+=(f32x4x16& a, f16x8x8 b) {
     f32x4 b_1(vcvt_f32_f16(vget_high_f16(
         reinterpret_cast<float16x8_t>(b.v[i].v))));
 
-    a.v[2 * i] += b_0;
+    a.v[2 * i + 0] += b_0;
     a.v[2 * i + 1] += b_1;
   }
 
@@ -75,7 +75,7 @@ static f32x4x16 reduce_add(
     f32x4 b_1(vcvt_f32_f16(vget_high_f16(
         reinterpret_cast<float16x8_t>(b.v[i].v))));
 
-    a.v[2 * i] += b_0 * b_0;
+    a.v[2 * i + 0] += b_0 * b_0;
     a.v[2 * i + 1] += b_1 * b_1;
   }
 

ynnpack/kernels/reduce/sum.inc

@@ -6,11 +6,13 @@ YNN_UNARY_REDUCE_KERNEL(arch_flag::neonbf16, sum_bf16_fp32_neonbf16, bfloat16, f
 #ifdef YNN_ARCH_ARM_NEONFP16ARITH
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neonfp16arith, sum_fp16_fp32_neonfp16arith, half, float)
 #endif  // YNN_ARCH_ARM_NEONFP16ARITH
+#ifdef YNN_ARCH_ARM_NEONDOT
+YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_int8_int32_neondot, int8_t, int32_t)
+YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_uint8_int32_neondot, uint8_t, int32_t)
+#endif  // YNN_ARCH_ARM_NEONDOT
 #ifdef YNN_ARCH_ARM_NEON
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_fp32_neon, float, float)
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_bf16_fp32_neon, bfloat16, float)
-YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_int8_int32_neondot, int8_t, int32_t)
-YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_uint8_int32_neondot, uint8_t, int32_t)
 #endif  // YNN_ARCH_ARM_NEON
 
 #ifdef YNN_ARCH_X86_AVX512BF16

ynnpack/kernels/reduce/sum_squared.inc

@@ -6,11 +6,13 @@ YNN_UNARY_REDUCE_KERNEL(arch_flag::neonbf16, sum_squared_bf16_fp32_neonbf16, bfl
 #ifdef YNN_ARCH_ARM_NEONFP16ARITH
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neonfp16arith, sum_squared_fp16_fp32_neonfp16arith, half, float)
 #endif  // YNN_ARCH_ARM_NEONFP16ARITH
+#ifdef YNN_ARCH_ARM_NEONDOT
+YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_squared_int8_int32_neondot, int8_t, int32_t)
+YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_squared_uint8_int32_neondot, uint8_t, int32_t)
+#endif  // YNN_ARCH_ARM_NEONDOT
 #ifdef YNN_ARCH_ARM_NEON
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_squared_fp32_neon, float, float)
 YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_squared_bf16_fp32_neon, bfloat16, float)
-YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_squared_int8_int32_neondot, int8_t, int32_t)
-YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_squared_uint8_int32_neondot, uint8_t, int32_t)
 #endif  // YNN_ARCH_ARM_NEON
 
 #ifdef YNN_ARCH_X86_AVX512BF16