KleidiAI Coverage Report

Line	Branch	Exec	Source
1			//
2			// SPDX-FileCopyrightText: Copyright 2024-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_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 in_stride;
29			size_t out_stride;
30			const void* in;
31			void* out;
32			const void* pad_row;
33			} KernelArgs;
34
35			static const size_t kai_num_bytes_input = sizeof(uint16_t);
36			static const size_t kai_num_bytes_output = sizeof(uint16_t);
37			static const size_t kai_num_bytes_bias = sizeof(uint16_t);
38
39			void kai_kernel_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(const KernelArgs* args_ptr);
40
41		1512	size_t kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(void) {
42		1512	return NR * kai_get_sme_vector_length_u16() / KR;
43			}
44
45		126	size_t kai_get_rhs_offset_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n_idx) {
46		−	KAI_ASSUME(n_idx % kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme() == 0);
47
48		126	return n_idx * kai_num_bytes_input;
49			}
50
51		160	size_t kai_get_bias_offset_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n_idx) {
52		160	return n_idx * kai_num_bytes_bias;
53			}
54
55		378	size_t kai_get_rhs_packed_stride_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t k) {
56		756	return kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme() *
57		378	(kai_num_bytes_bias + kai_roundup(k, KR) * kai_num_bytes_output);
58			}
59
60		252	size_t kai_get_rhs_packed_offset_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n_idx, size_t k) {
61		−	KAI_ASSUME(n_idx % kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme() == 0);
62
63		252	const size_t block_idx = n_idx / kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme();
64		504	return block_idx * kai_get_rhs_packed_stride_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(k);
65		252	}
66
67		126	size_t kai_get_rhs_packed_size_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(size_t n, size_t k) {
68		126	const size_t n_nr_blocks = kai_roundup(n, kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme());
69		252	return kai_get_rhs_packed_offset_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(n_nr_blocks, k);
70		126	}
71
72		126	void kai_run_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(
73			size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t rhs_stride_row, const void* rhs,
74			const void* bias, const void* scale, void* rhs_packed, size_t extra_bytes, const void* params) {
75		−	KAI_ASSUME(num_groups == 1);
76		−	KAI_ASSUME(nr == kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme());
77		−	KAI_ASSUME(kr == KR);
78		−	KAI_ASSUME(sr == 1);
79		−	KAI_ASSUME(rhs != NULL);
80		−	KAI_ASSUME(bias != NULL);
81		−	KAI_ASSUME(scale == NULL);
82		−	KAI_ASSUME(rhs_packed != NULL);
83		−	KAI_ASSUME(extra_bytes == 0);
84		−	KAI_ASSUME(params == NULL);
85
86		−	KAI_ASSERT(kai_get_n_step_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme() <= MAX_N_STEP);
87			static const uint16_t pad_row[MAX_N_STEP] = {0};
88
89		126	KernelArgs args;
90		126	args.bias_ptr = bias;
91		126	args.height = k;
92		126	args.width = n;
93		126	args.in = rhs;
94		126	args.out = rhs_packed;
95		126	args.in_stride = rhs_stride_row;
96		126	args.out_stride = kai_get_rhs_packed_stride_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(args.height);
97		126	args.pad_row = pad_row;
98
99		126	kai_kernel_rhs_pack_kxn_x16p2vlx2b_x16_x16_sme(&args);
100		126	}
101
102			#endif // Architectural features check.
103