KleidiAI Coverage Report

Line	Branch	Exec	Source
1			//
2			// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates <open-source-office@arm.com>
3			//
4			// SPDX-License-Identifier: Apache-2.0
5			//
6
7			#if (!defined(__aarch64__) || !defined(__ARM_FEATURE_SVE2)) && !defined(_M_ARM64)
8			#error This file must be compiled for AArch64, FEAT_SVE2.
9			#else // Architectural features check.
10			#include "kai_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa.h"
11
12			#include <stddef.h>
13			#include <stdint.h>
14
15			#include "kai/kai_common.h"
16
17			typedef struct {
18			const void* A;
19			const void* B;
20			void* C;
21			uint64_t ldcb;
22			uint64_t M;
23			uint64_t N;
24			uint64_t K;
25			uint16_t min;
26			uint16_t max;
27			void* accumulator_buffer;
28			uint64_t flags;
29			} KernelArgs;
30
31			static const size_t kai_mr = 2;
32			static const size_t kai_nr = 2;
33			static const size_t kai_kr = 2;
34			static const size_t kai_sr = 1;
35
36			void kai_kernel_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(KernelArgs* args);
37			uint16_t kai_f16_from_float_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(float value);
38
39			// Returns a constant value specific to this kernel that's relative to vector length
40		200	static size_t kai_get_kernel_vec_length_constant(void) {
41		200	const size_t kernel_vec_length_constant = kai_get_sme_vector_length_u16() / kai_kr;
42		400	return kernel_vec_length_constant;
43		200	}
44
45		67	size_t kai_get_m_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
46		67	return kai_mr * kai_get_kernel_vec_length_constant();
47			}
48
49		82	size_t kai_get_n_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
50		82	return kai_nr * kai_get_kernel_vec_length_constant();
51			}
52
53		17	size_t kai_get_mr_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
54		17	return kai_mr * kai_get_kernel_vec_length_constant();
55			}
56
57		34	size_t kai_get_nr_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
58		34	return kai_nr * kai_get_kernel_vec_length_constant();
59			}
60
61		51	size_t kai_get_kr_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
62		51	return kai_kr;
63			}
64
65		51	size_t kai_get_sr_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) {
66		51	return kai_sr;
67			}
68
69		33	size_t kai_get_lhs_packed_offset_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(size_t m_idx, size_t k) {
70		−	KAI_ASSUME(m_idx % kai_get_m_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0);
71		33	return m_idx * kai_roundup(k, kai_kr) * sizeof(uint16_t);
72			}
73
74		16	static size_t kai_get_rhs_packed_stride_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(size_t k) {
75		32	return kai_get_n_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() *
76		16	(sizeof(uint16_t) + kai_roundup(k, kai_kr) * sizeof(uint16_t));
77			}
78
79		16	size_t kai_get_rhs_packed_offset_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(size_t n_idx, size_t k) {
80		−	KAI_ASSUME(n_idx % kai_get_n_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0);
81		16	const size_t block_idx = n_idx / kai_get_n_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa();
82		32	return block_idx * kai_get_rhs_packed_stride_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(k);
83		16	}
84
85		16	size_t kai_get_dst_offset_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(
86			size_t m_idx, size_t n_idx, size_t dst_stride_row) {
87		−	KAI_ASSUME(m_idx % kai_get_m_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0);
88		−	KAI_ASSUME(n_idx % kai_get_n_step_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0);
89
90		16	return m_idx * dst_stride_row + n_idx * sizeof(uint16_t);
91			}
92
93		16	size_t kai_get_dst_size_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(size_t m, size_t n) {
94		16	return m * n * sizeof(uint16_t);
95			}
96
97		17	void kai_run_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(
98			size_t m, size_t n, size_t k, const void* lhs_packed, const void* rhs_packed, void* dst, size_t dst_stride_row,
99			size_t dst_stride_col, float clamp_min, float clamp_max) {
100		17	KAI_UNUSED(dst_stride_col);
101
102		17	KernelArgs args;
103
104		17	args.A = lhs_packed;
105		17	args.B = rhs_packed;
106		17	args.C = dst;
107		17	args.ldcb = dst_stride_row;
108		17	args.M = m;
109		17	args.N = n;
110		17	args.K = k;
111		17	args.min = kai_f16_from_float_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(clamp_min);
112		17	args.max = kai_f16_from_float_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(clamp_max);
113		17	args.accumulator_buffer = NULL;
114		17	args.flags = 0;
115
116		17	kai_kernel_matmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(&args);
117		17	}
118
119			#endif // Architectural features check.
120