KFI-543 Integrate SME2 kernels matmul_clamp_f16_qai8dxp_qsi8cxp

cmake/gen/neon_aarch64_microkernels.cmake

@@ -16,6 +16,7 @@ SET(PROD_NEON_AARCH64_MICROKERNEL_SRCS
   src/f32-vsqrt/gen/f32-vsqrt-aarch64-neon-sqrt.c
   src/x8-lut/gen/x8-lut-aarch64-neon-tbx128x4-u64.c
   src/x8-packq/x8-packq-aarch64-neon-f32qp8-u2.c
+  src/x8-packq/x8-packq-aarch64-neon-f16qp8-u2.c
   src/x24-transposec/x24-transposec-4x4-aarch64-neon-tbl128.c
   src/x32-transposec/x32-transposec-4x4-aarch64-neon-tbl128.c)
 

cmake/gen/neonsme2_microkernels.cmake

@@ -23,6 +23,8 @@ SET(PROD_NEONSME2_MICROKERNEL_SRCS
   src/qp8-f32-qc4w-gemm/qp8-f32-qc4w-gemm-minmax-16x64c4-neonsme2.c
   src/qp8-f32-qc8w-gemm/qp8-f32-qc8w-gemm-minmax-1x64c4-neonsme2.c
   src/qp8-f32-qc8w-gemm/qp8-f32-qc8w-gemm-minmax-16x64c4-neonsme2.c
+  src/qp8-f16-qc8w-gemm/qp8-f16-qc8w-gemm-minmax-1x64c4-neonsme2.c
+  src/qp8-f16-qc8w-gemm/qp8-f16-qc8w-gemm-minmax-16x64c4-neonsme2.c
   src/x8-pack-lh/x8-packlh-igemm-neonsme2.c
   src/x8-pack-lh/x8-packlh-neonsme2.c
   src/x16-pack-lh/x16-packlh-igemm-neonsme2.c

gen/neon_aarch64_microkernels.bzl

@@ -12,6 +12,7 @@ PROD_NEON_AARCH64_MICROKERNEL_SRCS = [
     "src/f32-vsqrt/gen/f32-vsqrt-aarch64-neon-sqrt.c",
     "src/x8-lut/gen/x8-lut-aarch64-neon-tbx128x4-u64.c",
     "src/x8-packq/x8-packq-aarch64-neon-f32qp8-u2.c",
+    "src/x8-packq/x8-packq-aarch64-neon-f16qp8-u2.c",
     "src/x24-transposec/x24-transposec-4x4-aarch64-neon-tbl128.c",
     "src/x32-transposec/x32-transposec-4x4-aarch64-neon-tbl128.c",
 ]

gen/neonsme2_microkernels.bzl

@@ -19,6 +19,8 @@ PROD_NEONSME2_MICROKERNEL_SRCS = [
     "src/qp8-f32-qc4w-gemm/qp8-f32-qc4w-gemm-minmax-16x64c4-neonsme2.c",
     "src/qp8-f32-qc8w-gemm/qp8-f32-qc8w-gemm-minmax-1x64c4-neonsme2.c",
     "src/qp8-f32-qc8w-gemm/qp8-f32-qc8w-gemm-minmax-16x64c4-neonsme2.c",
+    "src/qp8-f16-qc8w-gemm/qp8-f16-qc8w-gemm-minmax-1x64c4-neonsme2.c",
+    "src/qp8-f16-qc8w-gemm/qp8-f16-qc8w-gemm-minmax-16x64c4-neonsme2.c",
     "src/x8-pack-lh/x8-packlh-igemm-neonsme2.c",
     "src/x8-pack-lh/x8-packlh-neonsme2.c",
     "src/x16-pack-lh/x16-packlh-igemm-neonsme2.c",

src/configs/gemm-config.c

@@ -58,6 +58,7 @@ static struct xnn_gemm_config qdu8_f32_qc2w_gemm_config = {0};
 static struct xnn_gemm_config qd8_f32_qc8w_gemm_config = {0};
 static struct xnn_gemm_config qp8_f32_qc4w_gemm_config = {0};
 static struct xnn_gemm_config qp8_f32_qc8w_gemm_config = {0};
+static struct xnn_gemm_config qp8_f16_qc8w_gemm_config = {0};
 static struct xnn_gemm_config qp8_f32_qb4w_gemm_config = {0};
 static struct xnn_gemm_config qdu8_f32_qc4w_gemm_config = {0};
 static struct xnn_gemm_config qdu8_f16_qc8w_gemm_config = {0};
@@ -94,6 +95,7 @@ XNN_INIT_ONCE_GUARD(qdu8_f32_qc2w_gemm);
 XNN_INIT_ONCE_GUARD(qd8_f32_qc8w_gemm);
 XNN_INIT_ONCE_GUARD(qp8_f32_qc4w_gemm);
 XNN_INIT_ONCE_GUARD(qp8_f32_qc8w_gemm);
+XNN_INIT_ONCE_GUARD(qp8_f16_qc8w_gemm);
 XNN_INIT_ONCE_GUARD(qp8_f32_qb4w_gemm);
 XNN_INIT_ONCE_GUARD(qdu8_f32_qc4w_gemm);
 XNN_INIT_ONCE_GUARD(qdu8_f16_qc8w_gemm);
@@ -3007,6 +3009,44 @@ static void init_qp8_f32_qc8w_gemm_config(void) {
 #endif  // XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
 }
 
+static void init_qp8_f16_qc8w_gemm_config(void) {
+  qp8_f16_qc8w_gemm_config.log2_input_element_size = XNN_LOG2_SIZEOF_INT8_T;
+  qp8_f16_qc8w_gemm_config.log2_filter_element_size = XNN_LOG2_SIZEOF_UINT8_T;
+  qp8_f16_qc8w_gemm_config.log2_filter_element_bit_size =
+      XNN_LOG2_SIZEOF_UINT8_T + 3;
+  qp8_f16_qc8w_gemm_config.bias_element_size = sizeof(float);
+
+#if XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
+  const struct xnn_hardware_config* hardware_config =
+      xnn_init_hardware_config();
+  assert(hardware_config != NULL);
+  if (XNN_ENABLE_ARM_SME2 && (hardware_config->arch_flags & xnn_arch_arm_sme2)) {
+#if XNN_ENABLE_ARM_SME2
+    const size_t mr =
+        xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2_get_mr();
+    const size_t nr =
+        xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2_get_nr();
+    qp8_f16_qc8w_gemm_config.minmax.qp8gemm[XNN_MR_TO_INDEX(1)] =
+        XNN_INIT_HMP_QP8GEMM_UKERNEL(
+            xnn_qp8_f16_qc8w_gemm_minmax_ukernel_1x64c4__neonsme2);
+    qp8_f16_qc8w_gemm_config.minmax.qp8gemm[XNN_MR_TO_INDEX(mr)] =
+        XNN_INIT_HMP_QP8GEMM_UKERNEL(
+            xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2);
+    qp8_f16_qc8w_gemm_config.init.f16 = xnn_init_f16_minmax_scalar_params;
+    qp8_f16_qc8w_gemm_config.pack_weights_and_biases =
+        xnn_pack_kai_qs8_weights_and_biases;
+    qp8_f16_qc8w_gemm_config.packed_stride_weights_and_biases =
+        xnn_packed_stride_kai_qs8_weights_and_biases;
+    qp8_f16_qc8w_gemm_config.mr = mr;
+    qp8_f16_qc8w_gemm_config.nr = nr;
+    qp8_f16_qc8w_gemm_config.log2_kr = 2;
+    qp8_f16_qc8w_gemm_config.mr_packed = mr;
+    assert(qp8_f16_qc8w_gemm_config.mr <= XNN_MAX_MR);
+#endif  // XNN_ENABLE_ARM_SME2
+  }
+#endif  // XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
+}
+
 static void init_qp8_f32_qb4w_gemm_config(void) {
   // Common parameters.
   qp8_f32_qb4w_gemm_config.log2_input_element_size = XNN_LOG2_SIZEOF_INT8_T;
@@ -6595,6 +6635,17 @@ const struct xnn_gemm_config* xnn_init_qp8_f32_qc8w_gemm_config() {
   return NULL;
 }
 
+const struct xnn_gemm_config* xnn_init_qp8_f16_qc8w_gemm_config() {
+  if (xnn_init_hardware_config() == NULL) {
+    return NULL;
+  }
+  XNN_INIT_ONCE(qp8_f16_qc8w_gemm);
+  if (qp8_f16_qc8w_gemm_config.minmax.qp8gemm[0].function[0] != NULL) {
+    return &qp8_f16_qc8w_gemm_config;
+  }
+  return NULL;
+}
+
 const struct xnn_gemm_config* xnn_init_qp8_f32_qb4w_gemm_config() {
   const struct xnn_hardware_config* hardware_config =
       xnn_init_hardware_config();

src/configs/pack-lh-config.c

@@ -16,6 +16,7 @@
 #include "src/xnnpack/packq.h"
 
 static struct xnn_pack_lh_config qp8_pack_lh_config = {0};
+static struct xnn_pack_lh_config qp8_f16_pack_lh_config = {0};
 static struct xnn_pack_lh_config x8_pack_lh_config = {0};
 static struct xnn_pack_lh_config x16_pack_lh_config = {0};
 static struct xnn_pack_lh_config x32_pack_lh_config = {0};
@@ -24,6 +25,7 @@ static struct xnn_pack_lh_config x32_igemm_pack_lh_config = {0};
 static struct xnn_pack_lh_config x16_igemm_pack_lh_config = {0};
 
 XNN_INIT_ONCE_GUARD(qp8_pack_lh);
+XNN_INIT_ONCE_GUARD(qp8_f16_pack_lh);
 XNN_INIT_ONCE_GUARD(x8_pack_lh);
 XNN_INIT_ONCE_GUARD(x16_pack_lh);
 XNN_INIT_ONCE_GUARD(x32_pack_lh);
@@ -57,6 +59,32 @@ const struct xnn_pack_lh_config* xnn_init_qp8_pack_lh_config() {
   return &qp8_pack_lh_config;
 }
 
+static void init_qp8_f16_pack_lh_config(void) {
+#if XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
+  qp8_f16_pack_lh_config.pack_lh_fn =
+      (xnn_pack_lh_ukernel_fn)xnn_x8_packq_f16qp8_ukernel__aarch64_neon_u2;
+  qp8_f16_pack_lh_config.size_fn =
+      (xnn_pack_lh_size_fn)xnn_x8_packq_f32qp8_packed_size;
+  qp8_f16_pack_lh_config.offset_fn =
+      (xnn_pack_lh_offset_fn)xnn_x8_packq_f32qp8_packed_offset;
+  qp8_f16_pack_lh_config.log2_input_element_size = XNN_LOG2_SIZEOF_FLOAT16;
+  qp8_f16_pack_lh_config.log2_packed_element_size = 0;
+#endif  // XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
+}
+
+const struct xnn_pack_lh_config* xnn_init_qp8_f16_pack_lh_config() {
+  const struct xnn_hardware_config* hardware_config =
+      xnn_init_hardware_config();
+  if (hardware_config == NULL) {
+    return NULL;
+  }
+  XNN_INIT_ONCE(qp8_f16_pack_lh);
+  if (qp8_f16_pack_lh_config.pack_lh_fn == NULL) {
+    return NULL;
+  }
+  return &qp8_f16_pack_lh_config;
+}
+
 static void init_x32_pack_lh_config(void) {
 #if XNN_ARCH_ARM64 && XNN_ENABLE_KLEIDIAI
   const struct xnn_hardware_config* hardware_config =

src/operator-run.c

@@ -527,7 +527,8 @@ XNN_NO_SANITIZE_FUNCTION void xnn_compute_hmp_grouped_qp8gemm(struct gemm_contex
       context->qp8_ukernel.function[uarch_index](
           mr_step, nr_block_size, k_scaled, (const void*)(a + a_offset),
           packed_w, (void*)(c + mr_block_start * cm_stride), cm_stride,
-          /*dst_stride_col=*/sizeof(float), context->fused_params);
+          // QP8 kernels can produce f16 or f32 outputs, so use the configured output element size.
+          /*dst_stride_col=*/1 << context->log2_csize, context->fused_params);
     }
     mr_block_size -= mr_step;
     mr_block_start += mr_step;
@@ -597,7 +598,8 @@ XNN_INLINE static XNN_NO_SANITIZE_FUNCTION void compute_hmp_qp8gemm(
           mr_step, nr_block_size, k_scaled,
           (const void*)((uintptr_t)a + a_offset), packed_w,
           (void*)((uintptr_t)c + mr_block_start * cm_stride), cm_stride,
-          /*dst_stride_col=*/sizeof(float), context->fused_params);
+          // QP8 kernels can produce f16 or f32 outputs, so use the configured output element size.
+          /*dst_stride_col=*/1 << context->log2_csize, context->fused_params);
     }
 
     mr_block_size -= mr_step;

src/operators/fingerprint_id.h.inc

@@ -34,6 +34,7 @@ XNN_FINGERPRINT_ID(fully_connected_nc, qd8, f32, qc4w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qdu8, f32, qc4w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qp8, f32, qc4w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qp8, f32, qc8w)
+XNN_FINGERPRINT_ID(fully_connected_nc, qp8, f16, qc8w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qp8, f32, qb4w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qd8, f32, qb4w)
 XNN_FINGERPRINT_ID(fully_connected_nc, qdu8, f32, qb4w)

src/operators/fully-connected-nc.c

@@ -67,6 +67,7 @@ static enum xnn_operator_type get_operator_type(
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qdu8, f32, qc4w);
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qp8, f32, qc4w);
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qp8, f32, qc8w);
+    XNNPACK_FINGERPRINT_TO_OP_TYPE(qp8, f16, qc8w);
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qp8, f32, qb4w);
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qd8, f32, qb4w);
     XNNPACK_FINGERPRINT_TO_OP_TYPE(qdu8, f32, qb4w);
@@ -1157,6 +1158,21 @@ static const struct fc_variant qp8_f32_qc8w_variant = {
     .kernel_scale_element_size = sizeof(float),
 };
 
+static const struct fc_variant qp8_f16_qc8w_variant = {
+    .check_output_bounds = check_output_bounds_f32,
+    .check_kernel_zero_point = UNUSED,
+    .check_block_size = UNUSED,
+    .check_flags = UNUSED,
+    .setup_gemm_ukernels = setup_gemm_ukernels,
+    .setup_params = setup_params_f16,
+    .setup_packing_params = setup_packing_params_qs8_qc8w,
+    .setup_packing_functions = setup_packing_functions_from_gemm_config,
+    .setup_scale_params = setup_scale_params_qs8_qc8w,
+    .fingerprint_constraints = {},
+    .extra_weights_bytes = sizeof(float) * 2,
+    .kernel_scale_element_size = sizeof(float),
+};
+
 static const struct fc_variant qd8_f16_qb4w_variant = {
     .check_output_bounds = check_output_bounds_f32,
     .check_kernel_zero_point = check_kernel_zero_point_is_0_or_8_qu8,
@@ -1482,6 +1498,11 @@ static enum xnn_status setup_variant_and_gemm_config(
       context->gemm_config = xnn_init_qp8_f32_qc8w_gemm_config();
       context->fingerprint_id = xnn_fingerprint_id_fully_connected_nc_qp8_f32_qc8w;
       break;
+    case xnn_operator_type_fully_connected_nc_qp8_f16_qc8w:
+      *variant = &qp8_f16_qc8w_variant;
+      context->gemm_config = xnn_init_qp8_f16_qc8w_gemm_config();
+      context->fingerprint_id = xnn_fingerprint_id_fully_connected_nc_qp8_f16_qc8w;
+      break;
     case xnn_operator_type_fully_connected_nc_qp8_f32_qb4w:
       *variant = &qp8_f32_qb4w_variant;
       context->gemm_config = xnn_init_qp8_f32_qb4w_gemm_config();
@@ -2143,6 +2164,30 @@ enum xnn_status xnn_create_fully_connected_nc_qp8_f32_qc8w(
   return create_fully_connected_nc_helper(&context);
 }
 
+enum xnn_status xnn_create_fully_connected_nc_qp8_f16_qc8w(
+    size_t input_channels, size_t output_channels, size_t input_stride,
+    size_t output_stride, const float* kernel_scale, const void* kernel,
+    const float* bias, float output_min, float output_max, uint32_t flags,
+    xnn_weights_cache_t weights_cache, xnn_operator_t* fully_connected_op_out) {
+  struct fc_context context = {
+      .input_channels = input_channels,
+      .output_channels = output_channels,
+      .input_stride = input_stride,
+      .output_stride = output_stride,
+      .kernel_scale.f32 = kernel_scale,
+      .kernel = kernel,
+      .bias = bias,
+      .output_min = output_min,
+      .output_max = output_max,
+      .flags = flags,
+      .weights_cache = weights_cache,
+      .operator_type = xnn_operator_type_fully_connected_nc_qp8_f16_qc8w,
+      .fully_connected_op_out = fully_connected_op_out,
+      .should_fingerprint = true,
+  };
+  return create_fully_connected_nc_helper(&context);
+}
+
 enum xnn_status xnn_create_fully_connected_nc_qp8_f32_qb4w(
     size_t input_channels, size_t output_channels, size_t input_stride,
     size_t output_stride, size_t block_size, uint8_t kernel_zero_point,
@@ -2909,6 +2954,9 @@ static XNN_NO_SANITIZE_FUNCTION enum xnn_status reshape_fully_connected_nc(
     case xnn_operator_type_fully_connected_nc_qp8_f32_qc8w:
       packed_lh_config = xnn_init_qp8_pack_lh_config();
       break;
+    case xnn_operator_type_fully_connected_nc_qp8_f16_qc8w:
+      packed_lh_config = xnn_init_qp8_f16_pack_lh_config();
+      break;
     case xnn_operator_type_fully_connected_nc_pf16:
       packed_lh_config = xnn_init_x16_pack_lh_config();
       break;
@@ -3463,6 +3511,19 @@ enum xnn_status xnn_reshape_fully_connected_nc_qp8_f32_qc8w(
       threadpool);
 }
 
+enum xnn_status xnn_reshape_fully_connected_nc_qp8_f16_qc8w(
+    xnn_operator_t fully_connected_op, size_t batch_size,
+    size_t* workspace_size, pthreadpool_t threadpool) {
+  return reshape_fully_connected_nc(
+      fully_connected_op, xnn_operator_type_fully_connected_nc_qp8_f16_qc8w,
+      batch_size,
+      /*dynamic_quantization=*/false,
+      /*log2_output_element_size=*/XNN_LOG2_SIZEOF_FLOAT16,
+      &fully_connected_op->params.f16_minmax,
+      sizeof(fully_connected_op->params.f16_minmax), workspace_size,
+      threadpool);
+}
+
 enum xnn_status xnn_reshape_fully_connected_nc_qp8_f32_qb4w(
     xnn_operator_t fully_connected_op, size_t batch_size,
     size_t* workspace_size, pthreadpool_t threadpool) {
@@ -3812,6 +3873,14 @@ enum xnn_status xnn_setup_fully_connected_nc_qp8_f32_qc8w(
       input, output, workspace, /*row_sum=*/NULL, /*quantization_params=*/NULL);
 }
 
+enum xnn_status xnn_setup_fully_connected_nc_qp8_f16_qc8w(
+    xnn_operator_t fully_connected_op, const int8_t* input, float* output,
+    void* workspace) {
+  return setup_fully_connected_nc(
+      fully_connected_op, xnn_operator_type_fully_connected_nc_qp8_f16_qc8w,
+      input, output, workspace, /*row_sum=*/NULL, /*quantization_params=*/NULL);
+}
+
 enum xnn_status xnn_setup_fully_connected_nc_qp8_f32_qb4w(
     xnn_operator_t fully_connected_op, const int8_t* input, float* output,
     void* workspace) {

src/qp8-f16-qc8w-gemm/qp8-f16-qc8w-gemm-minmax-16x64c4-neonsme2.c

@@ -0,0 +1,47 @@
+// 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.
+
+#include <assert.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include "src/xnnpack/math.h"
+#include "src/xnnpack/microparams.h"
+
+#if XNN_ENABLE_KLEIDIAI
+#include "kai/ukernels/matmul/matmul_clamp_f16_qai8dxp_qsi8cxp/kai_matmul_clamp_f16_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
+
+size_t xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2_get_mr() {
+  return
+      kai_get_mr_matmul_clamp_f16_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa();
+}
+
+size_t xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2_get_nr() {
+  return
+      kai_get_nr_matmul_clamp_f16_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa();
+}
+#endif  // XNN_ENABLE_KLEIDIAI
+
+// Wraps the
+// `kai_matmul_clamp_f16_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa` GEMM
+// microkernel with a name that is compatible with our tooling.
+void xnn_qp8_f16_qc8w_gemm_minmax_ukernel_16x64c4__neonsme2(
+    size_t m, size_t n, size_t k, const void* lhs_packed,
+    const void* rhs_packed, float* dst, size_t dst_stride_row,
+    size_t dst_stride_col, struct xnn_f32_minmax_params* minmax_params) {
+#if XNN_ENABLE_KLEIDIAI
+  const struct xnn_f16_minmax_params* f16_minmax_params =
+      (const struct xnn_f16_minmax_params*) minmax_params;
+  kai_run_matmul_clamp_f16_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa(
+      m, n, k, lhs_packed, rhs_packed, dst, dst_stride_row, dst_stride_col,
+      xnn_float16_to_float(f16_minmax_params->scalar.min),
+      xnn_float16_to_float(f16_minmax_params->scalar.max));
+#else
+  assert(
+      "Calling KleidiAI microkernel wrapper, but XNNPACK was compiled without "
+      "`XNN_ENABLE_KLEIDIAI`." &&
+      0);
+#endif  // XNN_ENABLE_KLEIDIAI
+}