diff --git a/ynnpack/kernels/dot/BUILD b/ynnpack/kernels/dot/BUILD index 14f8821d604..aa6d599963c 100644 --- a/ynnpack/kernels/dot/BUILD +++ b/ynnpack/kernels/dot/BUILD @@ -201,6 +201,16 @@ ynn_generate_srcs( ], ) +ynn_generate_srcs( + name = "x86_fp32_k16", + compatible_with = _COMPATIBLE_WITH, + generator = "//ynnpack/kernels/dot/generator:x86_fp32_k16", + output_srcs = [ + "x86_avx512_fp32_k16.cc", + "x86_avx512_fp32_k16.inc", + ], +) + ynn_generate_srcs( name = "x86_fp64", compatible_with = _COMPATIBLE_WITH, @@ -403,6 +413,7 @@ ynn_cc_library( "x86_avx512_bf16_bf16_fp32_k1.inc", "x86_avx512_fp16_fp16_fp32.inc", "x86_avx512_fp32.inc", + "x86_avx512_fp32_k16.inc", "x86_avx512_fp32_k2.inc", "x86_avx512_fp32_k4.inc", "x86_avx512_fp64.inc", @@ -491,6 +502,7 @@ ynn_cc_library( "x86_avx512_fp64.cc", "x86_avx512_fp32_k2.cc", "x86_avx512_fp32_k4.cc", + "x86_avx512_fp32_k16.cc", "x86_avx512_int8_int8_int32.cc", "x86_avx512_int8_int8_int32_k1.cc", "x86_avx512_int8_int8_int32_k16.cc", diff --git a/ynnpack/kernels/dot/generator/BUILD b/ynnpack/kernels/dot/generator/BUILD index d9e5cae90f7..f1c57302694 100644 --- a/ynnpack/kernels/dot/generator/BUILD +++ b/ynnpack/kernels/dot/generator/BUILD @@ -90,6 +90,7 @@ py_library( "x86_fp32_k2", "x86_fp32_k4", "x86_fp32_k8", + "x86_fp32_k16", "x86_fp64", "x86_int8_int8_int32", "x86_int8_int8_int32_k1", diff --git a/ynnpack/kernels/dot/generator/x86_fp32_k16.py b/ynnpack/kernels/dot/generator/x86_fp32_k16.py new file mode 100644 index 00000000000..abbc6277e43 --- /dev/null +++ b/ynnpack/kernels/dot/generator/x86_fp32_k16.py @@ -0,0 +1,85 @@ +# Copyright 2026 Google LLC +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""fp32 x86 AVX512 dot kernel specialized for n == 1.""" + +# pylint: disable=missing-class-docstring +# pylint: disable=invalid-name + +from ynnpack.kernels.dot.generator.dot_base import dot_base +from ynnpack.kernels.dot.generator.dot_base import generate_dot_kernels + + +class x86_avx512_fp32_k16(dot_base): + """fp32 AVX512 dot kernel specialized for n == 1.""" + + def __init__(self): + super().__init__("avx512", "fp32") + self.a_type = "float" + self.b_type = "float" + self.c_type = "float" + self.tile_shape = (1, 1, 16) + + def header(self): + return ( + """\ +#include +""" + + super().header() + + """ +namespace { + +YNN_INTRINSIC float horizontal_sum(__m512 x) { + return _mm512_reduce_add_ps(x); +} + +} // namespace +""" + ) + + def init_c_tile(self, i, j): + return f"__m512 c_{i}_{j} = _mm512_setzero_ps();\n" + + def load_a_tile(self, i, k): + return f"__m512 a_{i}_{k} = _mm512_loadu_ps({self.a_ptr(i, k)});\n" + + def load_b_tile(self, k, j): + return f"__m512 b_{k}_{j} = _mm512_loadu_ps({self.b_ptr(k, j)});\n" + + def product(self, i, j, k): + return f"c_{i}_{j} = _mm512_fmadd_ps(a_{i}_{k}, b_{k}_{j}, c_{i}_{j});\n" + + def finalize_c_tile(self, i, j): + # Horizontal-sum the 16-wide accumulator (which holds 16 independent + # partial sums, one per SIMD lane, interleaved across k) down to the + # single scalar dot product for this row. + return f"float cs_{i}_{j} = horizontal_sum(c_{i}_{j});\n" + + def add_c_tile(self, i, j): + return f"cs_{i}_{j} += *{self.c_in_ptr(i, j)};\n" + + def store_c_tile(self, i, j): + return f"*{self.c_out_ptr(i, j)} = cs_{i}_{j};\n" + + # n is always 1 (tile_shape[1] == block_shape[1] == 1), so the "tail" case + # (n < tile_shape[1]) never actually triggers at runtime; these just need to + # be valid, and reusing the main-case codegen is correct since there's only + # ever one possible n. + def add_c_tile_tail(self, i, j, n): + return self.add_c_tile(i, j) + + def store_c_tile_tail(self, i, j, n): + return self.store_c_tile(i, j) + + +generate_dot_kernels( + x86_avx512_fp32_k16(), + [ + (1, 1, 16), + (2, 1, 16), + (4, 1, 16), + (8, 1, 16), + ], +) diff --git a/ynnpack/kernels/dot/get_dot_kernel_test.cc b/ynnpack/kernels/dot/get_dot_kernel_test.cc index 0133d357719..6ab0010b0d5 100644 --- a/ynnpack/kernels/dot/get_dot_kernel_test.cc +++ b/ynnpack/kernels/dot/get_dot_kernel_test.cc @@ -120,15 +120,15 @@ TEST(get_dot_kernel, small_n) { return get_dot_kernel_name(fp32, {large_shape, 8, large_shape}, arch_flags); }; ASSERT_EQ(fp32_x1(arch_flags_avx2_fma3), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); - ASSERT_EQ(fp32_x1(arch_flags_avx512), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); + ASSERT_EQ(fp32_x1(arch_flags_avx512), "dot_fp32_8x1x16_1x1x16_avx512"); ASSERT_EQ(fp32_x2(arch_flags_sse2), "dot_fp32_8x4x1_1x4x1_sse2"); ASSERT_EQ(fp32_x2(arch_flags_avx2), "dot_fp32_8x4x2_1x4x2_avx2"); ASSERT_EQ(fp32_x2(arch_flags_avx2_fma3), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); - ASSERT_EQ(fp32_x2(arch_flags_avx512), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); + ASSERT_EQ(fp32_x2(arch_flags_avx512), "dot_fp32_8x1x16_1x1x16_avx512"); ASSERT_EQ(fp32_x3(arch_flags_sse2), "dot_fp32_8x4x1_1x4x1_sse2"); ASSERT_EQ(fp32_x3(arch_flags_avx2), "dot_fp32_8x4x2_1x4x2_avx2"); ASSERT_EQ(fp32_x3(arch_flags_avx2_fma3), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); - ASSERT_EQ(fp32_x3(arch_flags_avx512), "dot_fp32_8x4x4_1x4x4_avx512"); + ASSERT_EQ(fp32_x3(arch_flags_avx512), "dot_fp32_8x1x16_1x1x16_avx512"); ASSERT_EQ(fp32_x4(arch_flags_sse2), "dot_fp32_8x4x1_1x4x1_sse2"); ASSERT_EQ(fp32_x4(arch_flags_avx2), "dot_fp32_8x4x2_1x4x2_avx2"); ASSERT_EQ(fp32_x4(arch_flags_avx2_fma3), "dot_fp32_8x1x8_1x1x8_avx2_fma3"); diff --git a/ynnpack/kernels/dot/kernels.inc b/ynnpack/kernels/dot/kernels.inc index 641968cae2b..6cc42864493 100644 --- a/ynnpack/kernels/dot/kernels.inc +++ b/ynnpack/kernels/dot/kernels.inc @@ -46,6 +46,7 @@ YNN_DOT_KERNEL(arch_flag::amxint8, dot_uint8_int8_int32_32x32x64_16x16x4_amxint8 #include "ynnpack/kernels/dot/x86_avx512_fp32.inc" #include "ynnpack/kernels/dot/x86_avx512_fp32_k2.inc" #include "ynnpack/kernels/dot/x86_avx512_fp32_k4.inc" +#include "ynnpack/kernels/dot/x86_avx512_fp32_k16.inc" #include "ynnpack/kernels/dot/x86_avx512_fp64.inc" #include "ynnpack/kernels/dot/x86_avx512_int8_int8_int32.inc" #include "ynnpack/kernels/dot/x86_avx512_int8_int8_int32_k1.inc"