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_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa.h" |
11 |
|
|
|
12 |
|
|
#include <stddef.h> |
13 |
|
|
#include <stdint.h> |
14 |
|
|
|
15 |
|
|
#include "kai/kai_common.h" |
16 |
|
|
|
17 |
|
|
typedef struct { |
18 |
|
|
const void* A; |
19 |
|
|
const void* B; |
20 |
|
|
void* C; |
21 |
|
|
uint64_t ldcb; |
22 |
|
|
uint64_t M; |
23 |
|
|
uint64_t N; |
24 |
|
|
uint64_t K; |
25 |
|
|
uint16_t min; |
26 |
|
|
uint16_t max; |
27 |
|
|
void* accumulator_buffer; |
28 |
|
|
uint64_t flags; |
29 |
|
|
} KernelArgs; |
30 |
|
|
|
31 |
|
|
static const size_t kai_mr = 2; |
32 |
|
|
static const size_t kai_nr = 2; |
33 |
|
|
static const size_t kai_kr = 2; |
34 |
|
|
|
35 |
|
|
void kai_kernel_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(KernelArgs* args); |
36 |
|
|
uint16_t kai_f16_from_float_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(float value); |
37 |
|
|
|
38 |
|
|
// Returns a constant value specific to this kernel that's relative to vector length |
39 |
|
34056 |
static size_t kai_get_kernel_vec_length_constant(void) { |
40 |
|
34056 |
const size_t kernel_vec_length_constant = kai_get_sme_vector_length_u16() / kai_kr; |
41 |
|
68112 |
return kernel_vec_length_constant; |
42 |
|
34056 |
} |
43 |
|
|
|
44 |
|
11352 |
size_t kai_get_m_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) { |
45 |
|
11352 |
return kai_mr * kai_get_kernel_vec_length_constant(); |
46 |
|
|
} |
47 |
|
|
|
48 |
|
22704 |
size_t kai_get_n_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(void) { |
49 |
|
22704 |
return kai_nr * kai_get_kernel_vec_length_constant(); |
50 |
|
|
} |
51 |
|
|
|
52 |
|
5676 |
size_t kai_get_lhs_packed_offset_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
53 |
|
|
size_t m_idx, size_t k_chunk_count, size_t k_chunk_length) { |
54 |
|
− |
KAI_ASSUME(m_idx % kai_get_m_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0); |
55 |
|
5676 |
return m_idx * k_chunk_count * kai_roundup(k_chunk_length, kai_kr) * sizeof(uint16_t); |
56 |
|
|
} |
57 |
|
|
|
58 |
|
5676 |
static size_t kai_get_rhs_packed_stride_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
59 |
|
|
size_t k_chunk_count, size_t k_chunk_length) { |
60 |
|
11352 |
return kai_get_n_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() * |
61 |
|
5676 |
(sizeof(uint16_t) + k_chunk_count * kai_roundup(k_chunk_length, kai_kr) * sizeof(uint16_t)); |
62 |
|
|
} |
63 |
|
|
|
64 |
|
5676 |
size_t kai_get_rhs_packed_offset_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
65 |
|
|
size_t n_idx, size_t k_chunk_count, size_t k_chunk_length) { |
66 |
|
− |
KAI_ASSUME(n_idx % kai_get_n_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0); |
67 |
|
5676 |
const size_t block_idx = n_idx / kai_get_n_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(); |
68 |
|
17028 |
return block_idx * |
69 |
|
5676 |
kai_get_rhs_packed_stride_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
70 |
|
5676 |
k_chunk_count, k_chunk_length); |
71 |
|
5676 |
} |
72 |
|
|
|
73 |
|
5676 |
size_t kai_get_dst_offset_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
74 |
|
|
size_t m_idx, size_t n_idx, size_t dst_stride_row) { |
75 |
|
− |
KAI_ASSUME(m_idx % kai_get_m_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0); |
76 |
|
− |
KAI_ASSUME(n_idx % kai_get_n_step_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa() == 0); |
77 |
|
|
|
78 |
|
5676 |
return m_idx * dst_stride_row + n_idx * sizeof(uint16_t); |
79 |
|
|
} |
80 |
|
|
|
81 |
|
5676 |
size_t kai_get_dst_size_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(size_t m, size_t n) { |
82 |
|
5676 |
return m * n * sizeof(uint16_t); |
83 |
|
|
} |
84 |
|
|
|
85 |
|
5676 |
void kai_run_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa( |
86 |
|
|
size_t m, size_t n, size_t k_chunk_count, size_t k_chunk_length, const void* lhs_packed, const void* rhs_packed, |
87 |
|
|
void* dst, size_t dst_stride_row, float clamp_min, float clamp_max) { |
88 |
|
5676 |
KernelArgs args; |
89 |
|
|
|
90 |
|
5676 |
args.A = lhs_packed; |
91 |
|
5676 |
args.B = rhs_packed; |
92 |
|
5676 |
args.C = dst; |
93 |
|
5676 |
args.ldcb = dst_stride_row; |
94 |
|
5676 |
args.M = m; |
95 |
|
5676 |
args.N = n; |
96 |
|
5676 |
args.K = k_chunk_count * kai_roundup(k_chunk_length, kai_kr); |
97 |
|
5676 |
args.min = kai_f16_from_float_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(clamp_min); |
98 |
|
5676 |
args.max = kai_f16_from_float_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(clamp_max); |
99 |
|
5676 |
args.accumulator_buffer = NULL; |
100 |
|
5676 |
args.flags = 0; |
101 |
|
|
|
102 |
|
5676 |
kai_kernel_imatmul_clamp_f16_f16p2vlx2_f16p2vlx2b_2vlx2vl_sme_mopa(&args); |
103 |
|
5676 |
} |
104 |
|
|
|
105 |
|
|
#endif // Architectural features check. |
106 |
|
|
|