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_qsi8cxp2vlx4sb_qs8cx_f32_i32_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 = 4, |
21 |
|
|
MAX_N_STEP = NR * ((KAI_SME_VEC_LENGTH_MAX_BYTES / sizeof(uint8_t)) / KR), |
22 |
|
|
}; |
23 |
|
|
|
24 |
|
|
typedef struct { |
25 |
|
|
const void* bias_ptr; |
26 |
|
|
const void* scale_ptr; |
27 |
|
|
int32_t input_zero_point; |
28 |
|
|
float scale_multiplier; |
29 |
|
|
size_t width; |
30 |
|
|
size_t height; |
31 |
|
|
size_t k_chunk_count; |
32 |
|
|
size_t in_stride; |
33 |
|
|
size_t out_stride; |
34 |
|
|
const void* in; |
35 |
|
|
void* out; |
36 |
|
|
const void* pad_row; |
37 |
|
|
} KernelArgs; |
38 |
|
|
|
39 |
|
|
static const size_t kai_num_bytes_input = sizeof(uint8_t); |
40 |
|
|
static const size_t kai_num_bytes_output = sizeof(uint8_t); |
41 |
|
|
static const size_t kai_num_bytes_bias = sizeof(int32_t); |
42 |
|
|
static const size_t kai_num_bytes_scale = sizeof(float); |
43 |
|
|
|
44 |
|
|
void kai_kernel_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(const KernelArgs* args_ptr); |
45 |
|
|
|
46 |
|
93240 |
size_t kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(void) { |
47 |
|
93240 |
return NR * kai_get_sme_vector_length_u8() / KR; |
48 |
|
|
} |
49 |
|
|
|
50 |
|
9324 |
size_t kai_get_rhs_offset_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(size_t n_idx) { |
51 |
|
− |
KAI_ASSUME(n_idx % kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme() == 0); |
52 |
|
|
|
53 |
|
9324 |
return n_idx * kai_num_bytes_input; |
54 |
|
|
} |
55 |
|
|
|
56 |
|
9324 |
size_t kai_get_bias_offset_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(size_t n_idx) { |
57 |
|
9324 |
return n_idx * kai_num_bytes_bias; |
58 |
|
|
} |
59 |
|
|
|
60 |
|
9324 |
size_t kai_get_scale_offset_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(size_t n_idx) { |
61 |
|
9324 |
return n_idx * kai_num_bytes_scale; |
62 |
|
|
} |
63 |
|
|
|
64 |
|
27972 |
static size_t kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme( |
65 |
|
|
size_t k_chunk_count, size_t k_chunk_length) { |
66 |
|
55944 |
return kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme() * |
67 |
|
27972 |
(kai_num_bytes_bias + k_chunk_count * kai_roundup(k_chunk_length, KR) * kai_num_bytes_output + |
68 |
|
|
kai_num_bytes_scale); |
69 |
|
|
} |
70 |
|
|
|
71 |
|
18648 |
size_t kai_get_rhs_packed_offset_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme( |
72 |
|
|
size_t n_idx, size_t k_chunk_count, size_t k_chunk_length) { |
73 |
|
− |
KAI_ASSUME(n_idx % kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme() == 0); |
74 |
|
|
|
75 |
|
18648 |
const size_t block_idx = n_idx / kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(); |
76 |
|
55944 |
return block_idx * |
77 |
|
18648 |
kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(k_chunk_count, k_chunk_length); |
78 |
|
18648 |
} |
79 |
|
|
|
80 |
|
9324 |
size_t kai_get_rhs_packed_size_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme( |
81 |
|
|
size_t n, size_t k_chunk_count, size_t k_chunk_length) { |
82 |
|
9324 |
const size_t n_nr_blocks = kai_roundup(n, kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme()); |
83 |
|
18648 |
return kai_get_rhs_packed_offset_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme( |
84 |
|
9324 |
n_nr_blocks, k_chunk_count, k_chunk_length); |
85 |
|
9324 |
} |
86 |
|
|
|
87 |
|
9324 |
void kai_run_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme( |
88 |
|
|
size_t n, size_t k_chunk_count, size_t k_chunk_length, size_t rhs_stride_row, const void* rhs, const void* bias, |
89 |
|
|
const void* scale, void* rhs_packed, const struct kai_rhs_pack_qsi8cx_params* params) { |
90 |
|
− |
KAI_ASSUME(rhs != NULL); |
91 |
|
− |
KAI_ASSUME(bias != NULL); |
92 |
|
− |
KAI_ASSUME(scale != NULL); |
93 |
|
− |
KAI_ASSUME(rhs_packed != NULL); |
94 |
|
− |
KAI_ASSUME(params != NULL); |
95 |
|
|
|
96 |
|
− |
KAI_ASSERT(kai_get_n_step_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme() <= MAX_N_STEP); |
97 |
|
|
static const uint8_t pad_row[MAX_N_STEP] = {0}; |
98 |
|
|
|
99 |
|
9324 |
KernelArgs args; |
100 |
|
9324 |
args.bias_ptr = bias; |
101 |
|
9324 |
args.scale_ptr = scale; |
102 |
|
9324 |
args.height = k_chunk_length; |
103 |
|
9324 |
args.width = n; |
104 |
|
9324 |
args.in = rhs; |
105 |
|
9324 |
args.out = rhs_packed; |
106 |
|
9324 |
args.k_chunk_count = k_chunk_count; |
107 |
|
9324 |
args.in_stride = rhs_stride_row; |
108 |
|
9324 |
args.out_stride = |
109 |
|
9324 |
kai_get_rhs_packed_stride_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(k_chunk_count, k_chunk_length); |
110 |
|
9324 |
args.input_zero_point = params->lhs_zero_point; |
111 |
|
9324 |
args.scale_multiplier = params->scale_multiplier; |
112 |
|
9324 |
args.pad_row = pad_row; |
113 |
|
|
|
114 |
|
9324 |
kai_kernel_rhs_imatmul_pack_kxn_qsi8cxp2vlx4sb_qs8cx_f32_i32_sme(&args); |
115 |
|
9324 |
} |
116 |
|
|
|
117 |
|
|
#endif // Architectural features check. |
118 |
|
|
|