Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

aarch64: matmul: addition of JIT int8 kernel #2686

Merged
merged 1 commit into from
Mar 20, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion include/oneapi/dnnl/dnnl.hpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*******************************************************************************
* Copyright 2016-2025 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2024-2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -1615,6 +1615,9 @@ struct memory : public handle<dnnl_memory_t> {
BA16a32b4a = dnnl_BA16a32b4a,
BA16a48b4a = dnnl_BA16a48b4a,
BA16a64b4a = dnnl_BA16a64b4a,
BA24b8a = dnnl_BA24b8a,
aCB24c8b = dnnl_aCB24c8b,
abDC24d8c = dnnl_abDC24d8c,
decbA16a = dnnl_decbA16a,
decbA8a = dnnl_decbA8a,
defcbA16a = dnnl_defcbA16a,
Expand Down
5 changes: 4 additions & 1 deletion include/oneapi/dnnl/dnnl_types.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*******************************************************************************
* Copyright 2016-2025 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2024-2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -1047,6 +1047,9 @@ typedef enum {
dnnl_aCBdef8b8c,
dnnl_abdEC16e4c,
dnnl_abDC16d4c,
dnnl_BA24b8a,
dnnl_aCB24c8b,
dnnl_abDC24d8c,

/// Just a sentinel, not real memory format tag. Must be changed after new
/// format tag is added.
Expand Down
5 changes: 4 additions & 1 deletion src/common/c_types_map.hpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*******************************************************************************
* Copyright 2016-2025 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2024-2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -372,6 +372,9 @@ const format_tag_t aCB16b16c = dnnl_aCB16b16c;
const format_tag_t aCB16b32c = dnnl_aCB16b32c;
const format_tag_t aCB16b48c = dnnl_aCB16b48c;
const format_tag_t aCB16b64c = dnnl_aCB16b64c;
const format_tag_t BA24b8a = dnnl_BA24b8a;
const format_tag_t aCB24c8b = dnnl_aCB24c8b;
const format_tag_t abDC24d8c = dnnl_abDC24d8c;
const format_tag_t aCB16b16c2b = dnnl_aCB16b16c2b;
const format_tag_t aCB16b32c2b = dnnl_aCB16b32c2b;
const format_tag_t aCB16b48c2b = dnnl_aCB16b48c2b;
Expand Down
5 changes: 4 additions & 1 deletion src/common/dnnl_debug_autogenerated.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*******************************************************************************
* Copyright 2018-2025 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2024-2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -953,6 +953,9 @@ const char *dnnl_fmt_tag2str(dnnl_format_tag_t v) {
if (v == dnnl_aCBdef8b8c) return "aCBdef8b8c";
if (v == dnnl_abdEC16e4c) return "abdEC16e4c";
if (v == dnnl_abDC16d4c) return "abDC16d4c";
if (v == dnnl_BA24b8a) return "BA24b8a";
if (v == dnnl_aCB24c8b) return "aCB24c8b";
if (v == dnnl_abDC24d8c) return "abDC24d8c";
if (v == dnnl_format_tag_last) return "format_tag_last";
if (v == dnnl_x) return "x";
if (v == dnnl_nc) return "nc";
Expand Down
5 changes: 4 additions & 1 deletion src/common/memory_desc_wrapper.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*******************************************************************************
* Copyright 2016-2025 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2024-2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -202,6 +202,9 @@ status_t memory_desc_wrapper::compute_blocking(
C(BA16a32b, {1, 0}, {16, 32}, {0, 1});
C(BA16a48b, {1, 0}, {16, 48}, {0, 1});
C(BA16a64b, {1, 0}, {16, 64}, {0, 1});
C(BA24b8a, {1, 0}, {24, 8}, {1, 0});
C(aCB24c8b, {0, 2, 1}, {24, 8}, {2, 1});
C(abDC24d8c, {0, 1, 3, 2}, {24, 8}, {3, 2});
C(BA16a16b2a, {1, 0}, {16, 16, 2}, {0, 1, 0});
C(BA16a32b2a, {1, 0}, {16, 32, 2}, {0, 1, 0});
C(BA16a48b2a, {1, 0}, {16, 48, 2}, {0, 1, 0});
Expand Down
94 changes: 94 additions & 0 deletions src/cpu/aarch64/matmul/jit_int8_kernel_types.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
/*******************************************************************************
* Copyright 2025 FUJITSU LIMITED
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#ifndef CPU_AARCH64_JIT_INT8_KERNEL_TYPES_HPP
#define CPU_AARCH64_JIT_INT8_KERNEL_TYPES_HPP

namespace dnnl {
namespace impl {
namespace cpu {
namespace aarch64 {
namespace matmul {

typedef enum {
none = 0,
per_tensor = 1,
per_m = 2,
per_n = 3,
per_k = 4,
} jit_int8_broadcast_t;

struct dyn_vals_t {
int f = 0;
dim_t M = 0;
dim_t K = 0;
dim_t N = 0;
dim_t B = 0;
int is_s8 = 0, is_u8 = 0;
int mtail, ktail, ntail, m_blk, k_blk, n_blk;
int get_min_max = 0, reorder_a = 0, reorder_b = 0, cal_src = 0;
int is_mtail = 0, is_ktail = 0;
};

struct dyn_params_t {
const float *dyn_src;
const int8_t *src;
int8_t *dst;
float *max, *min;
int *nk, *nm, *nn;
int *tl, *mtl, *ntl;
};

struct brg_int8_t {
int M, K, N;
const int m_blk = 8, n_blk = 4, k_blk = 8;
const int ld_block = 6, rd_block = 4, bd_block = 8;
int na, nb;
int m_tail, n_tail, k_tail;
int is_m_tail, is_k_tail, is_n_tail, is_zp_cal;
int dst_dt_sz;
bool is_s8;
bool is_bias;
bool with_scales;
bool with_dst_scales;
bool is_oc_scales;
jit_int8_broadcast_t zp_type_a = jit_int8_broadcast_t::none;
jit_int8_broadcast_t zp_type_b = jit_int8_broadcast_t::none;
jit_int8_broadcast_t zp_type_c = jit_int8_broadcast_t::none;
bool is_zp_b_int8 = false;
bool b_reo = true;
data_type_t zp_b_dt;
dim_t B;
};

struct call_params_t {
const uint8_t *src, *wei;
float *dst;
const float *bias, *scales, *dst_scales;
dim_t M, K, N;
char *buf_B_ptr_;
int *na, *nb;
int32_t *src_zero_point, *wei_zero_point, *dst_zero_point;
const int8_t *wei_zero_point_buf;
float *zp_a_ptr, *zp_b_ptr;
};

} // namespace matmul
} // namespace aarch64
} // namespace cpu
} // namespace impl
} // namespace dnnl
#endif
Loading