Repository: microsoft/onnxruntime
File: onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp (C++)

Code before:

#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p16vlx1b_1x16vl_sme2_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p8x1biasf32_6x8x4_neon_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
#include "kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.h"
#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme.h"
#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
#include "mlas.h"
#include "mlasi_kleidiai.h"
#include "kai_ukernel_interface.h"


// Thread-local reusable buffers to reduce allocation overhead across tiles.
struct KaiTlsBuffers {
    std::vector<float> output_tile;
    std::vector<float> bias_zero;
    std::vector<std::byte> rhs_packed;
    std::vector<std::byte> lhs_packed;
    std::vector<float> gemv_lhs_row_tmp;
};
static thread_local KaiTlsBuffers g_kai_tls;

const kai_matmul_clamp_f32_f32p_f32p_ukernel sgemm_gemm = GetKleidiAISGemmUKernel();
const kai_matmul_clamp_f32_f32_f32p_ukernel sgemm_gemv = GetKleidiAISGemvUKernel();


// Helpers for GEMV
/*++
Routine Description:
    Apply alpha/beta scaling to a 1-D vector with arbitrary destination stride.

Arguments:
    src          - Pointer to the temporary A*B results (length L).
    num_elements - Number of elements.
    alpha        - Scale for the computed product (A*B).
    beta         - Scale for the existing C values.
    dst          - Pointer to the destination in C.
    dst_stride   - Stride, in elements, between successive outputs in C.
    allow_memcpy - If true, allows memcpy path when alpha==1, beta==0, and dst_stride==1.

Notes:
    Uses a memcpy path when alpha==1, beta==0, allow_memcpy is true, and dst_stride==1.
--*/
static inline void ApplyAlphaBetaStrided(const float* src, size_t num_elements, float alpha, float beta, float* dst, size_t dst_stride, bool allow_memcpy) {
    if (alpha == 1.0f && beta == 0.0f && allow_memcpy && dst_stride == 1) {
        std::memcpy(dst, src, num_elements * sizeof(float));
        return;
    }
    for (size_t i = 0; i < num_elements; ++i) {

Review comment (suggestion):
```suggestion
const kai_matmul_clamp_f32_f32p_f32p_ukernel& sgemm_gemm = GetKleidiAISGemmUKernel();
const kai_matmul_clamp_f32_f32_f32p_ukernel& sgemm_gemv = GetKleidiAISGemvUKernel();
```

Code after:

#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p16vlx1b_1x16vl_sme2_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p8x1biasf32_6x8x4_neon_mla.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
#include "kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.h"
#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme.h"
#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
#include "mlas.h"
#include "mlasi_kleidiai.h"
#include "kai_ukernel_interface.h"


// Thread-local reusable buffers to reduce allocation overhead across tiles.
struct KaiTlsBuffers {
    std::vector<float> output_tile;
    std::vector<float> bias_zero;
    std::vector<std::byte> rhs_packed;
    std::vector<std::byte> lhs_packed;
    std::vector<float> gemv_lhs_row_tmp;
};
static thread_local KaiTlsBuffers g_kai_tls;

const kai_matmul_clamp_f32_f32p_f32p_ukernel& sgemm_gemm = GetKleidiAISGemmUKernel();
const kai_matmul_clamp_f32_f32_f32p_ukernel& sgemm_gemv = GetKleidiAISGemvUKernel();


// Helpers for GEMV
/*++
Routine Description:
    Apply alpha/beta scaling to a 1-D vector with arbitrary destination stride.

Arguments:
    src          - Pointer to the temporary A*B results (length L).
    num_elements - Number of elements.
    alpha        - Scale for the computed product (A*B).
    beta         - Scale for the existing C values.
    dst          - Pointer to the destination in C.
    dst_stride   - Stride, in elements, between successive outputs in C.
    allow_memcpy - If true, allows memcpy path when alpha==1, beta==0, and dst_stride==1.

Notes:
    Uses a memcpy path when alpha==1, beta==0, allow_memcpy is true, and dst_stride==1.
--*/
static inline void ApplyAlphaBetaStrided(const float* src, size_t num_elements, float alpha, float beta, float* dst, size_t dst_stride, bool allow_memcpy) {
    if (alpha == 1.0f && beta == 0.0f && allow_memcpy && dst_stride == 1) {
        std::memcpy(dst, src, num_elements * sizeof(float));
        return;
    }
    for (size_t i = 0; i < num_elements; ++i) {