aarch64: matmul: addition of JIT int8 kernel

Author: Shreyas-fuj

include/oneapi/dnnl/dnnl.hpp

@@ -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.
@@ -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,

include/oneapi/dnnl/dnnl_types.h

@@ -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.
@@ -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.

src/common/c_types_map.hpp

@@ -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.
@@ -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;

src/common/dnnl_debug_autogenerated.cpp

@@ -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.
@@ -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";

src/common/memory_desc_wrapper.cpp

@@ -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.
@@ -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});

src/cpu/aarch64/matmul/jit_int8_kernel_types.hpp

@@ -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