KleidiAI Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 29 / 12 / 41
Functions: 100.0% 7 / 0 / 7
Branches: -% 0 / 24 / 24
kai/ukernels/matmul/pack/kai_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme.c
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
11 #include "kai_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme.h"
12
13 #include <stddef.h>
14 #include <stdint.h>
15
16 #include "kai/kai_common.h"
17
18 enum {
19 NR = 2,
20 KR = 1,
21 };
22
23 typedef struct {
24 const void* bias_ptr;
25 size_t width;
26 size_t height;
27 size_t in_stride;
28 size_t out_stride;
29 const void* in;
30 void* out;
31 } KernelArgs;
32
33 static const size_t kai_num_bytes_input = sizeof(uint32_t);
34 static const size_t kai_num_bytes_output = sizeof(uint32_t);
35 static const size_t kai_num_bytes_bias = sizeof(float);
36
37 void kai_kernel_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(const KernelArgs* args_ptr);
38
39 4620 size_t kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(void) {
40 4620 return NR * kai_get_sme_vector_length_u32() / KR;
41 }
42
43 420 size_t kai_get_rhs_offset_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(size_t n_idx) {
44 − KAI_ASSUME(n_idx % kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme() == 0);
45
46 420 return n_idx * kai_num_bytes_input;
47 }
48
49 528 size_t kai_get_bias_offset_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(size_t n_idx) {
50 528 return n_idx * kai_num_bytes_bias;
51 }
52
53 1260 size_t kai_get_rhs_packed_stride_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(size_t k) {
54 2520 return kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme() *
55 1260 (kai_num_bytes_bias + kai_roundup(k, KR) * kai_num_bytes_output);
56 }
57
58 840 size_t kai_get_rhs_packed_offset_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(size_t n_idx, size_t k) {
59 − KAI_ASSUME(n_idx % kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme() == 0);
60
61 840 const size_t block_idx = n_idx / kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme();
62 1680 return block_idx * kai_get_rhs_packed_stride_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(k);
63 840 }
64
65 420 size_t kai_get_rhs_packed_size_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(size_t n, size_t k) {
66 420 const size_t n_nr_blocks = kai_roundup(n, kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme());
67 840 return kai_get_rhs_packed_offset_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(n_nr_blocks, k);
68 420 }
69
70 420 void kai_run_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(
71 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,
72 const void* bias, const void* scale, void* rhs_packed, size_t extra_bytes, const void* params) {
73 − KAI_ASSUME(num_groups == 1);
74 − KAI_ASSUME(nr == kai_get_n_step_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme());
75 − KAI_ASSUME(kr == KR);
76 − KAI_ASSUME(sr == 1);
77 − KAI_ASSUME(rhs != NULL);
78 − KAI_ASSUME(bias != NULL);
79 − KAI_ASSUME(scale == NULL);
80 − KAI_ASSUME(rhs_packed != NULL);
81 − KAI_ASSUME(extra_bytes == 0);
82 − KAI_ASSUME(params == NULL);
83
84 420 KernelArgs args;
85 420 args.bias_ptr = bias;
86 420 args.height = k;
87 420 args.width = n;
88 420 args.in = rhs;
89 420 args.out = rhs_packed;
90 420 args.in_stride = rhs_stride_row;
91 420 args.out_stride = kai_get_rhs_packed_stride_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(args.height);
92
93 420 kai_commit_za();
94
95 420 kai_kernel_rhs_pack_kxn_f32p2vlx1biasf32_f32_f32_sme(&args);
96 420 }
97
98 #endif // Architectural features check.
99