Add test coverage of YNN_FLAG_CONSISTENT_ARITHMETIC

This doesn't test for consistency (which would be difficult to do), but we should test the correctness of the dot subgraph code with this flag set. Currently, there are few if any branches in the code based on this flag, but that will be changing soon.

This uncovered a bug: we didn't mark any arm kernels as consistent.

PiperOrigin-RevId: 835645726

Author: dsharletg

ynnpack/kernels/dot/generator/arm_bf16_bf16_fp32.py

@@ -11,6 +11,7 @@ def __init__(self):
     super().__init__("neon", "bf16_bf16_fp32", "float", (1, 4, 1))
     self.a_type = "bfloat16"
     self.b_type = "bfloat16"
+    self.flags += ["dot_flag::consistent_arithmetic"]
 
   def header(self):
     return super().header() + """

ynnpack/kernels/dot/generator/arm_fp32.py

@@ -8,6 +8,7 @@ def __init__(self):
     super().__init__("neon", "fp32", "float", (1, 4, 1))
     self.a_type = "float"
     self.b_type = "float"
+    self.flags += ["dot_flag::consistent_arithmetic"]
 
   def header(self):
     return super().header() + """

ynnpack/kernels/dot/generator/arm_int8_int8_int32.py

@@ -8,6 +8,7 @@ def __init__(self, arch, tile_shape):
     super().__init__(arch, "int8_int8_int32", "int32_t", tile_shape)
     self.a_type = "int8_t"
     self.b_type = "int8_t"
+    self.flags += ["dot_flag::consistent_arithmetic"]
 
 
 class arm_neon_int8_int8_int32(arm_int8_int8_int32):

ynnpack/kernels/dot/kernels.inc

@@ -67,33 +67,41 @@ YNN_DOT_KERNEL(arch_flag::amxint8, dot_uint8_int8_int32_16x64x64_16x16x4_amxint8
 
 #ifndef YNN_DISABLE_SME
 #ifdef YNN_ARCH_ARM64_SME2
-YNN_DOT_KERNEL(arch_flag::sme2, dot_fp32_sme2, sme_vl(float{}),
-               4 * sme_vl(float{}), 1, sme_vl(float{}), 1,
-               dot_flag::transpose_a, float, float, float)
-YNN_DOT_KERNEL(arch_flag::sme2, dot_bf16_bf16_fp32_sme2, sme_vl(float{}),
-               4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
-               dot_flag::transpose_a, bfloat16, bfloat16, float)
-YNN_DOT_KERNEL(arch_flag::sme2, dot_fp16_fp16_fp32_sme2, sme_vl(float{}),
-               4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
-               dot_flag::transpose_a, half, half, float)
-YNN_DOT_KERNEL(arch_flag::sme2, dot_int8_int8_int32_sme2, sme_vl(int32_t{}),
-               4 * sme_vl(int32_t{}), 4, sme_vl(int32_t{}), 4,
-               dot_flag::transpose_a, int8_t, int8_t, int32_t)
+YNN_DOT_KERNEL(arch_flag::sme2, dot_fp32_sme2,
+               sme_vl(float{}), 4 * sme_vl(float{}), 1, sme_vl(float{}), 1,
+               dot_flag::transpose_a | dot_flag::consistent_arithmetic,
+               float, float, float)
+YNN_DOT_KERNEL(arch_flag::sme2, dot_bf16_bf16_fp32_sme2,
+               sme_vl(float{}), 4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
+               dot_flag::transpose_a,
+               bfloat16, bfloat16, float)
+YNN_DOT_KERNEL(arch_flag::sme2, dot_fp16_fp16_fp32_sme2,
+               sme_vl(float{}), 4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
+               dot_flag::transpose_a,
+               half, half, float)
+YNN_DOT_KERNEL(arch_flag::sme2, dot_int8_int8_int32_sme2,
+               sme_vl(int32_t{}), 4 * sme_vl(int32_t{}), 4, sme_vl(int32_t{}), 4,
+               dot_flag::transpose_a | dot_flag::consistent_arithmetic,
+               int8_t, int8_t, int32_t)
 #endif  // YNN_ARCH_ARM64_SME2
 
 #ifdef YNN_ARCH_ARM64_SME
-YNN_DOT_KERNEL(arch_flag::sme, dot_fp32_sme, sme_vl(float{}),
-               4 * sme_vl(float{}), 1, sme_vl(float{}), 1,
-               dot_flag::transpose_a, float, float, float)
-YNN_DOT_KERNEL(arch_flag::sme, dot_bf16_bf16_fp32_sme, sme_vl(float{}),
-               4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
-               dot_flag::transpose_a, bfloat16, bfloat16, float)
-YNN_DOT_KERNEL(arch_flag::sme, dot_fp16_fp16_fp32_sme, sme_vl(float{}),
-               4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
-               dot_flag::transpose_a, half, half, float)
-YNN_DOT_KERNEL(arch_flag::sme, dot_int8_int8_int32_sme, sme_vl(int32_t{}),
-               4 * sme_vl(int32_t{}), 4, sme_vl(int32_t{}), 4,
-               dot_flag::transpose_a, int8_t, int8_t, int32_t)
+YNN_DOT_KERNEL(arch_flag::sme, dot_fp32_sme,
+               sme_vl(float{}), 4 * sme_vl(float{}), 1, sme_vl(float{}), 1,
+               dot_flag::transpose_a | dot_flag::consistent_arithmetic,
+               float, float, float)
+YNN_DOT_KERNEL(arch_flag::sme, dot_bf16_bf16_fp32_sme,
+               sme_vl(float{}), 4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
+               dot_flag::transpose_a,
+               bfloat16, bfloat16, float)
+YNN_DOT_KERNEL(arch_flag::sme, dot_fp16_fp16_fp32_sme,
+               sme_vl(float{}), 4 * sme_vl(float{}), 2, sme_vl(float{}), 2,
+               dot_flag::transpose_a,
+               half, half, float)
+YNN_DOT_KERNEL(arch_flag::sme, dot_int8_int8_int32_sme,
+               sme_vl(int32_t{}), 4 * sme_vl(int32_t{}), 4, sme_vl(int32_t{}), 4,
+               dot_flag::transpose_a | dot_flag::consistent_arithmetic,
+               int8_t, int8_t, int32_t)
 #endif  // YNN_ARCH_ARM64_SME
 #endif  // YNN_DISABLE_SME
 

ynnpack/subgraph/test/dot.cc

@@ -157,7 +157,11 @@ void TestStaticB(A, B, C) {
     Tensor<B> b(to_physical_shape<B>(b_shape));
     b.generate([&]() { return b_gen(rng); });
 
-    SubgraphBuilder subgraph(4);
+    uint32_t subgraph_flags = 0;
+    if (random_bool(rng)) {
+      subgraph_flags |= YNN_FLAG_CONSISTENT_ARITHMETIC;
+    }
+    SubgraphBuilder subgraph(4, subgraph_flags);
     const uint32_t a_id = 0;
     const uint32_t b_id = 1;
     const uint32_t output_id = 3;
@@ -331,7 +335,11 @@ void TestDynamicB(A, B, C) {
         std::uniform_int_distribution<size_t>(1, max_k_dims)(rng);
     const size_t output_rank = std::max(a_rank, b_rank) - num_k_dims + 1;
 
-    SubgraphBuilder subgraph(4);
+    uint32_t subgraph_flags = 0;
+    if (random_bool(rng)) {
+      subgraph_flags |= YNN_FLAG_CONSISTENT_ARITHMETIC;
+    }
+    SubgraphBuilder subgraph(4, subgraph_flags);
     const uint32_t a_id = 0;
     const uint32_t b_id = 1;
     const uint32_t output_id = 3;
@@ -512,7 +520,11 @@ void TestStaticShapeDynamicB(A, B, C) {
     std::sort(inv_b_perm.begin(), inv_b_perm.end(),
               [&](int i, int j) { return b_perm[i] < b_perm[j]; });
 
-    SubgraphBuilder subgraph(4);
+    uint32_t subgraph_flags = 0;
+    if (random_bool(rng)) {
+      subgraph_flags |= YNN_FLAG_CONSISTENT_ARITHMETIC;
+    }
+    SubgraphBuilder subgraph(4, subgraph_flags);
     const uint32_t a_id = 0;
     const uint32_t b_id = 1;
     const uint32_t output_id = 3;