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_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme.h"
11
12			#include <stddef.h>
13			#include <stdint.h>
14			#include <string.h>
15
16			#include "kai/kai_common.h"
17
18			enum {
19			NR = 2,
20			KR = 2,
21			MAX_N_STEP = NR * ((KAI_SME_VEC_LENGTH_MAX_BYTES / sizeof(uint16_t)) / KR),
22			};
23
24			typedef struct {
25			const void* bias_ptr;
26			size_t width;
27			size_t height;
28			size_t k_chunk_count;
29			size_t in_stride;
30			size_t out_stride;
31			const void* in;
32			void* out;
33			const void* pad_row;
34			} KernelArgs;
35
36			static const size_t kai_num_bytes_input = sizeof(uint16_t);
37			static const size_t kai_num_bytes_output = sizeof(uint16_t);
38			static const size_t kai_num_bytes_bias = sizeof(uint16_t);
39
40			void kai_kernel_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(const KernelArgs* args_ptr);
41
42		113520	size_t kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(void) {
43		113520	return NR * kai_get_sme_vector_length_u16() / KR;
44			}
45
46		11352	size_t kai_get_rhs_offset_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n_idx) {
47		−	KAI_ASSUME(n_idx % kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme() == 0);
48
49		11352	return n_idx * kai_num_bytes_input;
50			}
51
52		11352	size_t kai_get_bias_offset_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n_idx) {
53		11352	return n_idx * kai_num_bytes_bias;
54			}
55
56		34056	static size_t kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(
57			size_t k_chunk_count, size_t k_chunk_length) {
58		68112	return kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme() *
59		34056	(kai_num_bytes_bias + k_chunk_count * kai_roundup(k_chunk_length, KR) * kai_num_bytes_output);
60			}
61
62		22704	size_t kai_get_rhs_packed_offset_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(
63			size_t n_idx, size_t k_chunk_count, size_t k_chunk_length) {
64		−	KAI_ASSUME(n_idx % kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme() == 0);
65
66		22704	const size_t block_idx = n_idx / kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme();
67		68112	return block_idx *
68		22704	kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(k_chunk_count, k_chunk_length);
69		22704	}
70
71		11352	size_t kai_get_rhs_packed_size_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(
72			size_t n, size_t k_chunk_count, size_t k_chunk_length) {
73		11352	const size_t n_nr_blocks = kai_roundup(n, kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme());
74		22704	return kai_get_rhs_packed_offset_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(
75		11352	n_nr_blocks, k_chunk_count, k_chunk_length);
76		11352	}
77
78		11352	void kai_run_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(
79			size_t n, size_t k_chunk_count, size_t k_chunk_length, size_t rhs_stride_row, const void* rhs, const void* bias,
80			void* rhs_packed) {
81		−	KAI_ASSUME(rhs != NULL);
82		−	KAI_ASSUME(bias != NULL);
83		−	KAI_ASSUME(rhs_packed != NULL);
84
85		−	KAI_ASSERT(kai_get_n_step_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme() <= MAX_N_STEP);
86			static const uint16_t pad_row[MAX_N_STEP] = {0};
87
88		11352	KernelArgs args;
89		11352	args.bias_ptr = bias;
90		11352	args.height = k_chunk_length;
91		11352	args.width = n;
92		11352	args.in = rhs;
93		11352	args.out = rhs_packed;
94		11352	args.k_chunk_count = k_chunk_count;
95		11352	args.in_stride = rhs_stride_row;
96		11352	args.out_stride =
97		11352	kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(k_chunk_count, k_chunk_length);
98		11352	args.pad_row = pad_row;
99
100		11352	kai_kernel_rhs_imatmul_pack_kxn_x16p2vlx2b_x16_x16_sme(&args);
101		11352	}
102
103			#endif // Architectural features check.
104