KleidiAI Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 98.5% 66 / 7 / 74
Functions: 100.0% 14 / 0 / 14
Branches: 50.0% 1 / 14 / 16
kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.c
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_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
12
13 #include <stddef.h>
14
15 #include "kai/kai_common.h"
16
17 typedef struct {
18 float* dst; // 0
19 const void* lhs_packed; // 0x8
20 const void* rhs_packed; // 0x10
21 size_t dst_stride_row; // 0x18
22 size_t m; // 0x20
23 size_t n; // 0x28
24 size_t k; // 0x30
25 size_t k_internal; // 0x38
26 size_t lhs_stride; // 0x40
27 size_t rhs_stride; // 0x48
28 size_t rhs_row_bytes; // 0x50
29 size_t lhs_end; // 0x58
30 float clamp_min; // 0x60
31 float clamp_max; // 0x64
32 } KernelArgs;
33
34 void kai_kernel_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(KernelArgs* args_ptr);
35
36 // Compute args
37 static const size_t kai_m_step = 1;
38 static const size_t kai_n_step = 4; // multiple of vector length
39 // Packing args
40 static const size_t kai_mr = 1;
41 static const size_t kai_nr = 4; // multiple of vector length
42 static const size_t kai_kr = 4;
43 static const size_t kai_sr = 1;
44 // LHS format args (num. bytes per value, multiplier, zero_point (if asymmetric))
45 static const size_t kai_num_bytes_qvalue_lhs = 1;
46 static const size_t kai_num_bytes_multiplier_lhs = 4;
47 static const size_t kai_num_bytes_zp_lhs = 4;
48 // RHS format args (num. bytes per value, multiplier, zero_point (if asymmetric), and reduction sum (if LHS is
49 // asymmetric))
50 static const size_t kai_num_bytes_qvalue_rhs = 1;
51 static const size_t kai_num_bytes_multiplier_rhs = 4;
52 static const size_t kai_num_bytes_rsum_rhs = 4;
53 // DST format args
54 static const size_t kai_num_bytes_dst_value = 4;
55 // Extra args
56 static const size_t kai_num_bytes_bias = 4;
57 static const size_t kai_k_multiple_of = 32;
58
59 2778 inline static size_t kai_k_roundedup(size_t k) {
60 // Round up k to be a multiple of 32.
61 2778 return kai_roundup(k, kai_k_multiple_of);
62 }
63
64 1157 inline static size_t kai_get_lhs_packed_stride(size_t k) {
65 1157 const size_t k_internal = kai_k_roundedup(k);
66 − KAI_ASSERT((k_internal % kai_k_multiple_of) == 0);
67 1157 const size_t mr = kai_get_mr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot();
68 1157 size_t lhs_packed_stride = mr * ((k_internal * kai_num_bytes_qvalue_lhs) + kai_num_bytes_multiplier_lhs);
69 // Since the LHS matrix is asymmetric with per-row quantization, we must include the
70 // the number of bytes to hold the zero point value
71 1157 lhs_packed_stride += mr * kai_num_bytes_zp_lhs;
72
73 2314 return lhs_packed_stride;
74 1157 }
75
76 1157 inline static size_t kai_get_rhs_packed_stride(size_t k) {
77 1157 const size_t k_internal = kai_k_roundedup(k);
78 − KAI_ASSERT((k_internal % kai_k_multiple_of) == 0);
79
80 1157 const size_t nr = kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot();
81
82 1157 size_t rhs_packed_stride = nr * (k_internal * kai_num_bytes_qvalue_rhs);
83 1157 rhs_packed_stride += nr * kai_num_bytes_multiplier_rhs;
84 // Since the LHS matrix is quantized asymmetric with per-row quantization, we also include
85 // the number of bytes for the reduction sum
86 1157 rhs_packed_stride += nr * kai_num_bytes_rsum_rhs;
87 // Since the bias is packed with the RHS matrix, the stride is adjusted with the number of bytes of the bias
88 1157 rhs_packed_stride += nr * kai_num_bytes_bias;
89
90 2314 return rhs_packed_stride;
91 1157 }
92
93 1848 size_t kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
94 1848 return kai_m_step;
95 }
96
97 1848 size_t kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
98 1848 return kai_n_step * kai_get_sme_vector_length_u8() / kai_kr;
99 }
100
101 1850 size_t kai_get_mr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
102 1850 return kai_mr;
103 }
104
105 3007 size_t kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
106 3007 return kai_nr * kai_get_sme_vector_length_u8() / kai_kr;
107 }
108
109 924 size_t kai_get_kr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
110 924 return kai_kr;
111 }
112
113 924 size_t kai_get_sr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(void) {
114 924 return kai_sr;
115 }
116
117 693 size_t kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(size_t m_idx, size_t k) {
118 − KAI_ASSUME((m_idx % kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot()) == 0);
119
120 693 return (m_idx / kai_mr) * kai_get_lhs_packed_stride(k);
121 }
122
123 693 size_t kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(size_t n_idx, size_t k) {
124 − KAI_ASSUME((n_idx % kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot()) == 0);
125 693 const size_t nr = kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot();
126 1386 return (n_idx / nr) * kai_get_rhs_packed_stride(k);
127 693 }
128
129 462 size_t kai_get_dst_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(
130 size_t m_idx, size_t n_idx, size_t dst_stride) {
131 − KAI_ASSUME((m_idx % kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot()) == 0);
132 − KAI_ASSUME((n_idx % kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot()) == 0);
133
134 462 return (n_idx * kai_num_bytes_dst_value) + m_idx * dst_stride;
135 }
136
137 462 size_t kai_get_dst_size_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(size_t m, size_t n) {
138 462 return m * n * kai_num_bytes_dst_value;
139 }
140
141 464 void kai_run_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(
142 size_t m, //
143 size_t n, //
144 size_t k, //
145 const void* restrict lhs_packed, //
146 const void* restrict rhs_packed, //
147 float* restrict dst, // NOLINT(readability-non-const-parameter)
148 size_t dst_stride_row, //
149 size_t dst_stride_col, //
150 float scalar_min, //
151 float scalar_max) {
152 − KAI_ASSUME(dst_stride_col == sizeof(float));
153
154
1/2
✓ Branch 0 taken 464 times.
✗ Branch 1 not taken.
 464 if (m == 0) {
155 ✗ return;
156 }
157
158 464 const size_t k_internal = kai_k_roundedup(k);
159 464 const size_t lhs_stride = kai_get_lhs_packed_stride(k);
160 464 const size_t rhs_stride = kai_get_rhs_packed_stride(k);
161 464 const size_t nr = kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot();
162
163 464 const size_t rhs_row_bytes = nr * k_internal;
164 464 const size_t lhs_end_ptr = ((size_t)lhs_packed) + (m * lhs_stride);
165
166 464 KernelArgs args;
167
168 464 args.dst = dst;
169 464 args.lhs_packed = lhs_packed;
170 464 args.rhs_packed = rhs_packed;
171 464 args.clamp_max = scalar_max;
172 464 args.clamp_min = scalar_min;
173 464 args.dst_stride_row = dst_stride_row;
174 464 args.m = m;
175 464 args.n = n;
176 464 args.k = k;
177 464 args.k_internal = k_internal;
178 464 args.lhs_stride = lhs_stride;
179 464 args.rhs_stride = rhs_stride;
180 464 args.rhs_row_bytes = rhs_row_bytes;
181 464 args.lhs_end = lhs_end_ptr;
182
183 464 kai_commit_za();
184
185 464 kai_kernel_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot(&args);
186 464 }
187
188 #endif // Architectural features check.
189