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
11			#include "kai_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla.h"
12
13			#include <stddef.h>
14			#include <stdint.h>
15
16			#include "kai/kai_common.h"
17
18			typedef struct {
19			float maxval;
20			float minval;
21			const void* A_ptr;
22			const void* B_ptr;
23			size_t N;
24			size_t K;
25			void* output_ptr;
26			uint64_t flags;
27			} KernelArgs;
28
29			static const size_t kai_m_step = 1;
30			static const size_t kai_nr = 2;
31			static const size_t kai_n_step = 8;
32			static const size_t kai_kr = 1;
33			static const size_t kai_sr = 1;
34
35			void kai_kernel_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(KernelArgs* args_ptr);
36
37		52	size_t kai_get_m_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(void) {
38		52	return kai_m_step;
39			}
40
41		236	size_t kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(void) {
42		236	return kai_n_step * kai_get_sme_vector_length_u32() / kai_kr;
43			}
44
45		46	size_t kai_get_nr_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(void) {
46		46	return kai_nr * kai_get_sme_vector_length_u32() / kai_kr;
47			}
48
49		46	size_t kai_get_kr_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(void) {
50		46	return kai_kr;
51			}
52
53		46	size_t kai_get_sr_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(void) {
54		46	return kai_sr;
55			}
56
57		46	size_t kai_get_lhs_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(size_t m_idx, size_t k) {
58		−	KAI_ASSUME(m_idx == 0);
59
60		46	return m_idx * k;
61			}
62
63		46	static size_t kai_get_rhs_packed_stride_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(size_t k) {
64		92	return kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla() *
65		46	(kai_roundup(k, kai_kr) * sizeof(float) + sizeof(float));
66			}
67
68		46	size_t kai_get_rhs_packed_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(size_t n_idx, size_t k) {
69		−	KAI_ASSUME(n_idx % kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla() == 0);
70
71		46	const size_t block_idx = n_idx / kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla();
72		92	return block_idx * kai_get_rhs_packed_stride_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(k);
73		46	}
74
75		46	size_t kai_get_dst_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(size_t m_idx, size_t n_idx, size_t dst_stride) {
76		−	KAI_ASSUME(m_idx == 0);
77		−	KAI_ASSUME(n_idx % kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla() == 0);
78
79		46	return (m_idx * dst_stride) + (n_idx * sizeof(float));
80			}
81
82		46	size_t kai_get_dst_size_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(size_t m, size_t n) {
83		46	return m * n * sizeof(float);
84			}
85
86		47	void kai_run_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(
87			size_t m, size_t n, size_t k, const void* lhs, size_t lhs_stride, const void* rhs_packed, void* dst,
88			size_t dst_stride_row, size_t dst_stride_col, float clamp_min, float clamp_max) {
89		47	KAI_UNUSED(dst_stride_row);
90		47	KAI_UNUSED(dst_stride_col);
91		47	KAI_UNUSED(lhs_stride);
92		−	KAI_ASSUME(m == 1);
93
94		47	uint64_t flags = 2;
95
96		47	KernelArgs args;
97
98		47	args.maxval = clamp_max;
99		47	args.minval = clamp_min;
100		47	args.A_ptr = lhs;
101		47	args.B_ptr = rhs_packed;
102		47	args.N = n;
103		47	args.K = k;
104		47	args.output_ptr = dst;
105		47	args.flags = flags;
106
107		47	kai_kernel_matmul_clamp_f32_f32_f32p2vlx1b_1x8vl_sme_mla(&args);
108		47	}
109
110			#endif // Architectural features check.
111