onednn-src 0.1.13

/*******************************************************************************
* Copyright 2026 Intel Corporation
* Copyright 2025 FUJITSU LIMITED
* Copyright 2025-2026 Arm Ltd. and affiliates
*
* 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.
*******************************************************************************/

#include <cassert>

#include "common/c_types_map.hpp"
#include "common/dnnl_thread.hpp"
#include "common/memory_tracking.hpp"
#include "common/type_helpers.hpp"
#include "common/utils.hpp"
#include "cpu/cpu_primitive.hpp"
#include "cpu/matmul/matmul_utils.hpp"
#include "cpu/scale_utils.hpp"

#include "cpu/aarch64/injectors/jit_uni_eltwise_injector.hpp"
#include "cpu/aarch64/jit_generator.hpp"
#include "cpu/aarch64/matmul/jit_int8_kernel_types.hpp"
#include "cpu/aarch64/matmul/jit_int8_matmul.hpp"
#include "cpu/aarch64/matmul/jit_int8_matmul_utils.hpp"

#define GET_OFF(field) (uint32_t) offsetof(call_params_t, field)

#define LDR_IMM(reg, addr, off) \
    { \
        const uint64_t IMM12_MASK = ~uint64_t(0xfff); \
        if (((off) & IMM12_MASK) == 0) { \
            ldr(reg, ptr(addr, off)); \
        } else { \
            add_imm(X_DEFAULT_ADDR, addr, off, X_TMP_0); \
            ldr(reg, ptr(X_DEFAULT_ADDR)); \
        } \
    }

#define VCHECK_BG(f, msg, ...) \
    VCHECK(primitive, create, dispatch, brgemm_matmul, f, msg, ##__VA_ARGS__);

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

using namespace Xbyak_aarch64;
using namespace dnnl::impl::cpu::matmul;
using namespace dnnl::impl::format_tag;
using namespace dnnl::impl::memory_tracking::names;
using namespace dnnl::impl::utils;

using namespace nstl;

using namespace data_type;

template <cpu_isa_t isa>
struct jit_int8_matmul_kernel_t : public jit_generator_t {

    DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_int8_matmul_kernel_t)

    XReg reg_param = abi_param1;
    XReg reg_a = x3;
    XReg reg_b = x4;
    XReg reg_c = x5;
    XReg reg_aux_a = x6;
    XReg reg_aux_b = x7;
    XReg reg_aux_c = x8;
    XReg reg_aux_a1 = x9;
    XReg reg_zp_aux_b_buf = x10;
    XReg reg_aux_c1 = x11;
    XReg reg_ld_loop = x12;
    XReg reg_rd_loop = x13;
    XReg reg_bd_loop = x14;
    XReg reg_tmp = x15;
    XReg reg_tmp_1 = x16;
    XReg reg_bias = x17;
    XReg reg_zp_a = x18;

    XReg reg_scales = x20;
    XReg reg_aux_scales = x24; //used X_TMP_1
    XReg reg_na = x25; //used X_TMP_2
    XReg reg_zp_b = x26; //used X_TMP_3
    XReg reg_zp_aux_b = x27; //used X_TMP_4
    PReg prd_ld = p1;
    PReg prd_st = p2;
    PReg prd_b = p3;
    PReg prd_8 = p4;
    PReg prd_zp_b_tl = p5;
    XReg reg_zp_val_c = x2;

    XReg reg_zp_val_a = reg_scales;
    XReg reg_zp_val_b = reg_bias;

    call_params_t inp;

    void operator()(const call_params_t *p) {
        return jit_generator_t::operator()(p);
    }

    ZReg loadb(int ld) { return ZReg(ld + 1); }
    ZReg acc(int bd, int ld) const {
        return ZReg(bd * brg_.ld_block + ld + brg_.ld_block + 1);
    }
    void zero_regs() {
        for (int a = 0; a < brg_.bd_block / 2; a++)
            for (int b = 0; b < brg_.ld_block; b++)
                eor(acc(a, b).d, acc(a, b).d, acc(a, b).d);
    }
    void store_regs(int bdb, int ldb, int tail) {
        // temporary vector registers used by the epilogue / zero-point paths.
        // these overlap with some B-load registers in the compute loop
        // but are only used after the compute loop is done.
        const ZReg zp_b_val = ZReg(6);
        const ZReg zp_b_prod0 = ZReg(4);
        const ZReg zp_b_prod1 = ZReg(5);

        for (int a = 0; a < bdb; a++) {
            for (int b = 0; b < ldb; b++) {
                if (brg_.is_s8 || brg_.is_u8_s8)
                    scvtf(acc(a, b).s, P_ALL_ONE, acc(a, b).s);
                else
                    ucvtf(acc(a, b).s, P_ALL_ONE, acc(a, b).s);
            }
        }

        for (int a = 0; a < bdb; a++) {
            for (int b = 0; b < ldb; b += 2) {
                if (b + 1 < ldb) {
                    uzp1(z31.d, acc(a, b).d, acc(a, b + 1).d);
                    uzp2(acc(a, b + 1).d, acc(a, b).d, acc(a, b + 1).d);
                    mov(acc(a, b).d, z31.d);
                } else {
                    uzp1(z31.d, acc(a, b).d, acc(a, b).d);
                    uzp2(acc(a, b + 1).d, acc(a, b).d, acc(a, b).d);
                    mov(acc(a, b).d, z31.d);
                }
            }
        }

        if (brg_.zp_type_a != jit_int8_broadcast_t::none) {
            for (int b = 0; b < ldb; b += 2) {
                PReg p = (brg_.is_n_tail && b >= ldb - 2) ? prd_b : P_ALL_ONE;
                ld1w(z31.s, p, ptr(reg_zp_a, b / 2, MUL_VL));
                for (int a = 0; a < bdb; a++) {
                    fsub(acc(a, b).s, acc(a, b).s, z31.s);
                    fsub(acc(a, b + 1).s, acc(a, b + 1).s, z31.s);
                }
            }
        }

        if (brg_.zp_type_b != jit_int8_broadcast_t::none) {
            int ao = 0;
            if (brg_.is_zp_b_int8) {
                // Weights zero-point is per-tensor; when provided as an int8/u8
                // runtime buffer, it contains a single scalar value.
                ldrb(W_TMP_0, ptr(reg_zp_aux_b_buf));
                if (brg_.zp_b_dt == data_type::s8) sxtb(W_TMP_0, W_TMP_0);
                dup(zp_b_val.s, W_TMP_0);
                if (brg_.zp_b_dt == data_type::u8) {
                    ucvtf(zp_b_val.s, P_ALL_ONE, zp_b_val.s);
                } else {
                    scvtf(zp_b_val.s, P_ALL_ONE, zp_b_val.s);
                }
            }
            for (int a = 0; a < bdb; a++) {
                ld1rw(z31.s, P_ALL_ONE, ptr(reg_zp_aux_b, ao * 4));
                ld1rw(z0.s, P_ALL_ONE, ptr(reg_zp_aux_b, (ao + 1) * 4));
                for (int b = 0; b < ldb; b += 2) {
                    if (brg_.is_zp_b_int8) {
                        fmul(zp_b_prod0.s, z31.s, zp_b_val.s);
                        fmul(zp_b_prod1.s, z0.s, zp_b_val.s);
                        fsub(acc(a, b).s, acc(a, b).s, zp_b_prod0.s);
                        fsub(acc(a, b + 1).s, acc(a, b + 1).s, zp_b_prod1.s);
                    } else {
                        fsub(acc(a, b).s, acc(a, b).s, z31.s);
                        fsub(acc(a, b + 1).s, acc(a, b + 1).s, z0.s);
                    }
                }
                ao += 2;
            }
        }

        if (brg_.with_scales) {
            for (int b = 0; b < ldb; b += 2) {
                PReg p = (brg_.is_n_tail && b >= ldb - 2) ? prd_b : P_ALL_ONE;
                if (brg_.is_oc_scales) {
                    ld1w(z31.s, p, ptr(reg_scales, b / 2, MUL_VL));
                } else {
                    ld1w(z31.s, p, ptr(reg_scales));
                }

                for (int a = 0; a < bdb; a++) {
                    fmul(acc(a, b).s, acc(a, b).s, z31.s);
                    fmul(acc(a, b + 1).s, acc(a, b + 1).s, z31.s);
                }
            }
        }

        if (brg_.is_bias) {
            for (int b = 0; b < ldb; b += 2) {
                PReg p = (brg_.is_n_tail && b >= ldb - 2) ? prd_b : P_ALL_ONE;
                ld1w(z31.s, p, ptr(reg_bias, b / 2, MUL_VL));
                for (int a = 0; a < bdb; a++) {
                    fadd(acc(a, b).s, acc(a, b).s, z31.s);
                    fadd(acc(a, b + 1).s, acc(a, b + 1).s, z31.s);
                }
            }
        }

        if (!eltwise_injectors_.empty()) {
            for (int a = 0; a < bdb; a++) {
                for (int b = 0; b < ldb; b += 2) {
                    // apply element-wise op to every vector register
                    // whose index is in [start, end)
                    int start_idx = acc(a, b).getIdx();
                    int end_idx = acc(a, b + 1).getIdx() + 1;
                    for (int i = 0; i < (int)eltwise_injectors_.size(); i++) {
                        eltwise_injectors_[i]->compute_vector_range(
                                start_idx, end_idx);
                    }
                }
            }
        }

        if (brg_.with_dst_scales) {
            ld1rw(z31.s, P_ALL_ONE, ptr(reg_aux_scales));
            for (int b = 0; b < ldb; b += 2) {
                for (int a = 0; a < bdb; a++) {
                    fmul(acc(a, b).s, acc(a, b).s, z31.s);
                    fmul(acc(a, b + 1).s, acc(a, b + 1).s, z31.s);
                }
            }
        }

        if (brg_.zp_type_c != jit_int8_broadcast_t::none) {
            LDR_IMM(reg_zp_val_c, reg_param, GET_OFF(dst_zero_point));
            ldr(W_TMP_0, ptr(reg_zp_val_c));
            dup(z0.s, W_TMP_0);
            scvtf(z0.s, P_ALL_ONE, z0.s);
            for (int b = 0; b < ldb; b += 2) {
                for (int a = 0; a < bdb; a++) {
                    fadd(acc(a, b).s, acc(a, b).s, z0.s);
                    fadd(acc(a, b + 1).s, acc(a, b + 1).s, z0.s);
                }
            }
        }

        mov(reg_tmp, reg_aux_c);
        add_imm(reg_tmp_1, reg_aux_c, brg_.N * brg_.dst_dt_sz, X_TMP_0);
        for (int a = 0; a < bdb; a++) {
            for (int b = 0; b < ldb; b += 2) {
                PReg p = (brg_.is_n_tail && b >= ldb - 2) ? prd_st : P_ALL_ONE;
                int vl = b / 2;
                st1w(acc(a, b).s, p, ptr(reg_tmp, vl, MUL_VL));
                if (a >= bdb - 1 && brg_.is_m_tail) {
                    if (brg_.m_tail % 2 == 0)
                        st1w(acc(a, b + 1).s, p, ptr(reg_tmp_1, vl, MUL_VL));
                } else {
                    st1w(acc(a, b + 1).s, p, ptr(reg_tmp_1, vl, MUL_VL));
                }
            }
            add_imm(reg_tmp, reg_tmp, 2 * brg_.N * brg_.dst_dt_sz, X_TMP_0);
            add_imm(reg_tmp_1, reg_tmp_1, 2 * brg_.N * brg_.dst_dt_sz, X_TMP_0);
        }
    }

    void microkernel(int rdb, int bdb, int ldb, int tail) {
        int a_off = 0, rd, ld, bd;
        mov(reg_tmp, reg_aux_b);
        for (rd = 0; rd < rdb; rd++) {
            int ao = 0;

            // For SVE128 and no N-tail, use NEON loads (v/z regs overlap) to
            // reduce predicated SVE load overhead in the hot loop.
            if (cpu_isa_traits<isa>::vlen == 16 && !brg_.is_n_tail) {
                const int bytes_per_vec = cpu_isa_traits<isa>::vlen;
                for (ld = 0; ld + 1 < ldb; ld += 2) {
                    QReg q0(loadb(ld).getIdx());
                    QReg q1(loadb(ld + 1).getIdx());
                    ldp(q0, q1, ptr(reg_tmp, ld * bytes_per_vec));
                }
                if (ld < ldb) {
                    QReg q(loadb(ld).getIdx());
                    ldr(q, ptr(reg_tmp, ld * bytes_per_vec));
                }
            } else {
                for (ld = 0; ld < ldb; ld++) {
                    PReg p = (brg_.is_n_tail && ld == ldb - 1) ? prd_ld
                                                               : P_ALL_ONE;
                    ld1b(loadb(ld).b, p, ptr(reg_tmp, ld, MUL_VL));
                }
            }
            for (bd = 0; bd < bdb;) {
                const int a_inc = brg_.m_blk * 2;
                const int32_t a_off_bytes = a_off + ao;
                if (cpu_isa_traits<isa>::vlen == 16 && bd + 1 < bdb) {
                    // Load two consecutive A blocks (each 16 bytes for SVE128)
                    // and compute two rows to amortize load overhead.
                    if ((a_off_bytes & 0xf) == 0 && a_off_bytes >= -1024
                            && a_off_bytes <= 1008) {
                        ldp(QReg(0), QReg(31), ptr(reg_aux_a, a_off_bytes));
                    } else {
                        add_imm(X_DEFAULT_ADDR, reg_aux_a, a_off_bytes,
                                X_TMP_0);
                        ldp(QReg(0), QReg(31), ptr(X_DEFAULT_ADDR));
                    }
                    ao += 2 * a_inc;

                    for (ld = 0; ld < ldb; ld++) {
                        if (brg_.is_s8) {
                            smmla(acc(bd, ld).s, z0.b, loadb(ld).b);
                            smmla(acc(bd + 1, ld).s, z31.b, loadb(ld).b);
                        } else if (brg_.is_u8_s8) {
                            usmmla(acc(bd, ld).s, z0.b, loadb(ld).b);
                            usmmla(acc(bd + 1, ld).s, z31.b, loadb(ld).b);
                        } else {
                            ummla(acc(bd, ld).s, z0.b, loadb(ld).b);
                            ummla(acc(bd + 1, ld).s, z31.b, loadb(ld).b);
                        }
                    }
                    bd += 2;
                } else {
                    if (cpu_isa_traits<isa>::vlen == 16) {
                        ldr(QReg(0), ptr(reg_aux_a, a_off_bytes));
                    } else {
                        add_imm(X_DEFAULT_ADDR, reg_aux_a, a_off_bytes,
                                X_TMP_0);
                        ld1rqb(z0.b, P_ALL_ONE, ptr(X_DEFAULT_ADDR));
                    }
                    ao += a_inc;

                    for (ld = 0; ld < ldb; ld++) {
                        if (brg_.is_s8)
                            smmla(acc(bd, ld).s, z0.b, loadb(ld).b);
                        else if (brg_.is_u8_s8)
                            usmmla(acc(bd, ld).s, z0.b, loadb(ld).b);
                        else
                            ummla(acc(bd, ld).s, z0.b, loadb(ld).b);
                    }
                    bd += 1;
                }
            }
            a_off += brg_.m_blk * brg_.k_blk;
            add_imm(reg_tmp, reg_tmp, brg_.k_blk * brg_.n_blk * brg_.ld_block,
                    X_TMP_0);
        }
    }

    void loop_k(int bdb, int ldb, int tail) {
        zero_regs();
        mov(reg_aux_a, reg_aux_a1);
        mov(reg_aux_b, reg_b);
        if (k_full_blks > 0) {
            mov(reg_rd_loop, k_full_blks);

            asm_do_while(reg_rd_loop, [&]() {
                microkernel(brg_.rd_block, bdb, ldb, tail);
                add_imm(reg_aux_a, reg_aux_a,
                        brg_.m_blk * brg_.k_blk * brg_.rd_block, X_TMP_0);
                add_imm(reg_aux_b, reg_aux_b,
                        brg_.k_blk * brg_.n_blk * brg_.ld_block * brg_.rd_block,
                        X_TMP_0);
            });
        }
        if (k_tail_blk > 0) {
            microkernel(k_tail_blk, bdb, ldb, tail);
            add_imm(reg_aux_a, reg_aux_a, brg_.m_blk * brg_.k_blk * k_tail_blk,
                    X_TMP_0);
            add_imm(reg_aux_b, reg_aux_b,
                    brg_.k_blk * brg_.n_blk * brg_.ld_block * k_tail_blk,
                    X_TMP_0);
        }
        if (k_residual_blk > 0) { microkernel(1, bdb, ldb, tail); }
        store_regs(bdb, ldb, tail);
    }

    void loop_k_zp(int bdb, int ldb, int is_a, int is_b) {
        const ZReg zp_b_sum0 = ZReg(3);
        const ZReg zp_b_sum1 = ZReg(4);
        const ZReg zp_b_sum2 = ZReg(5);
        const ZReg zp_b_sum3 = ZReg(6);

        eor(zp_b_sum0.d, zp_b_sum0.d, zp_b_sum0.d);
        eor(zp_b_sum1.d, zp_b_sum1.d, zp_b_sum1.d);
        if (cpu_isa_traits<isa>::vlen == 16) {
            eor(zp_b_sum2.d, zp_b_sum2.d, zp_b_sum2.d);
            eor(zp_b_sum3.d, zp_b_sum3.d, zp_b_sum3.d);
        }
        for (int i = 0; i < 6; i++)
            eor(acc(2, i).d, acc(2, i).d, acc(2, i).d);
        mov(reg_aux_a, reg_aux_a1);
        mov(reg_aux_b, reg_b);
        if (k_full_blks > 0) {
            mov(reg_rd_loop, k_full_blks);

            asm_do_while(reg_rd_loop, [&]() {
                zp_comp(brg_.rd_block, bdb, ldb, is_a, is_b);
                add_imm(reg_aux_a, reg_aux_a,
                        brg_.m_blk * brg_.k_blk * brg_.rd_block, X_TMP_0);
                add_imm(reg_aux_b, reg_aux_b,
                        brg_.k_blk * brg_.n_blk * brg_.ld_block * brg_.rd_block,
                        X_TMP_0);
            });
        }
        if (k_tail_blk > 0) {
            zp_comp(k_tail_blk, bdb, ldb, is_a, is_b);
            add_imm(reg_aux_a, reg_aux_a, brg_.m_blk * brg_.k_blk * k_tail_blk,
                    X_TMP_0);
            add_imm(reg_aux_b, reg_aux_b,
                    brg_.k_blk * brg_.n_blk * brg_.ld_block * k_tail_blk,
                    X_TMP_0);
        }
        if (k_residual_blk > 0) { zp_comp(1, bdb, ldb, is_a, is_b); }

        if (brg_.zp_type_b != jit_int8_broadcast_t::none && is_b == 1) {
            if (cpu_isa_traits<isa>::vlen == 16) {
                uzp1(zp_b_sum0.d, zp_b_sum0.d, zp_b_sum1.d);
                uzp1(zp_b_sum2.d, zp_b_sum2.d, zp_b_sum3.d);
                scvtf(zp_b_sum0.s, P_ALL_ONE, zp_b_sum0.s);
                scvtf(zp_b_sum2.s, P_ALL_ONE, zp_b_sum2.s);
            } else {
                uzp1(zp_b_sum0.d, zp_b_sum0.d, zp_b_sum1.d);
                scvtf(zp_b_sum0.s, P_ALL_ONE, zp_b_sum0.s);
            }
            if (!brg_.is_zp_b_int8) {
                ldr(W_TMP_0, ptr(reg_zp_val_b));
                dup(z0.s, W_TMP_0);
                scvtf(z0.s, P_ALL_ONE, z0.s);
                fmul(zp_b_sum0.s, P_ALL_ONE, z0.s);
                if (cpu_isa_traits<isa>::vlen == 16) {
                    fmul(zp_b_sum2.s, P_ALL_ONE, z0.s);
                }
            } else {
                if (brg_.zp_type_a != jit_int8_broadcast_t::none) {
                    ldr(W_TMP_0, ptr(reg_zp_val_a));
                    dup(z0.s, W_TMP_0);
                    mov_imm(W_TMP_0, brg_.K);
                    dup(z1.s, W_TMP_0);
                    scvtf(z0.s, P_ALL_ONE, z0.s);
                    scvtf(z1.s, P_ALL_ONE, z1.s);
                    fmul(z0.s, z1.s, z0.s);
                    fsub(zp_b_sum0.s, zp_b_sum0.s, z0.s);
                    if (cpu_isa_traits<isa>::vlen == 16) {
                        fsub(zp_b_sum2.s, zp_b_sum2.s, z0.s);
                    }
                }
            }
            st1w(zp_b_sum0.s, P_ALL_ONE, ptr(reg_zp_b));
            if (cpu_isa_traits<isa>::vlen == 16) {
                st1w(zp_b_sum2.s, P_ALL_ONE, ptr(reg_zp_b, 1, MUL_VL));
            }
        }

        if ((brg_.zp_type_a != jit_int8_broadcast_t::none) && is_a == 1) {
            ldr(W_TMP_0, ptr(reg_zp_val_a));
            dup(z2.s, W_TMP_0);
            scvtf(z2.s, P_ALL_ONE, z2.s);
            uzp1(acc(2, 0).d, acc(2, 0).d, acc(2, 1).d);
            if (brg_.ld_block > 2) uzp1(acc(2, 2).d, acc(2, 2).d, acc(2, 3).d);
            if (brg_.ld_block > 4) uzp1(acc(2, 4).d, acc(2, 4).d, acc(2, 5).d);

            scvtf(acc(2, 0).s, P_ALL_ONE, acc(2, 0).s);
            if (brg_.ld_block > 2) scvtf(acc(2, 2).s, P_ALL_ONE, acc(2, 2).s);
            if (brg_.ld_block > 4) scvtf(acc(2, 4).s, P_ALL_ONE, acc(2, 4).s);
            if (brg_.zp_type_b != jit_int8_broadcast_t::none
                    && !brg_.is_zp_b_int8) {
                ldr(W_TMP_0, ptr(reg_zp_val_b));
                dup(z0.s, W_TMP_0);
                mov_imm(W_TMP_0, brg_.K);
                dup(z1.s, W_TMP_0);
                scvtf(z0.s, P_ALL_ONE, z0.s);
                scvtf(z1.s, P_ALL_ONE, z1.s);
                fmul(z0.s, z1.s, z0.s);
                fsub(acc(2, 0).s, acc(2, 0).s, z0.s);
                if (brg_.ld_block > 2) fsub(acc(2, 2).s, acc(2, 2).s, z0.s);
                if (brg_.ld_block > 4) fsub(acc(2, 4).s, acc(2, 4).s, z0.s);
            }
            fmul(acc(2, 0).s, P_ALL_ONE, z2.s);
            if (brg_.ld_block > 2) fmul(acc(2, 2).s, P_ALL_ONE, z2.s);
            if (brg_.ld_block > 4) fmul(acc(2, 4).s, P_ALL_ONE, z2.s);

            st1w(acc(2, 0).s, P_ALL_ONE, ptr(reg_zp_a));
            if (brg_.ld_block > 2)
                st1w(acc(2, 2).s, P_ALL_ONE, ptr(reg_zp_a, 1, MUL_VL));
            if (brg_.ld_block > 4)
                st1w(acc(2, 4).s, P_ALL_ONE, ptr(reg_zp_a, 2, MUL_VL));
        }
    }

    void han_blk() {
        LDR_IMM(reg_tmp, reg_param, GET_OFF(nb));
        LDR_IMM(reg_na, reg_param, GET_OFF(na));
        ldr(WReg(reg_ld_loop.getIdx()), ptr(reg_tmp));
        mov(reg_aux_a1, reg_a);
        mov(reg_aux_c1, reg_c);
        mov(reg_aux_c, reg_aux_c1);
        mov(reg_zp_aux_b, reg_zp_b);

        asm_do_while(reg_ld_loop, [&]() {
            ldr(WReg(reg_bd_loop.getIdx()), ptr(reg_na));
            asm_do_while(reg_bd_loop, [&]() {
                loop_k(bdb, ldb, 0);
                add_imm(reg_aux_a1, reg_aux_a1,
                        div_up(brg_.K, brg_.k_blk) * brg_.k_blk * brg_.bd_block,
                        X_TMP_0);
                add_imm(reg_aux_c, reg_aux_c,
                        brg_.N * brg_.bd_block * brg_.dst_dt_sz, X_TMP_0);
                add_imm(reg_zp_aux_b, reg_zp_aux_b, brg_.m_blk * brg_.dst_dt_sz,
                        X_TMP_0);
            });

            mov(reg_aux_a1, reg_a);
            mov(reg_zp_aux_b, reg_zp_b);
            add_imm(reg_b, reg_b,
                    (brg_.n_blk * brg_.ld_block) * div_up(brg_.K, brg_.k_blk)
                            * brg_.k_blk,
                    X_TMP_0);
            add_imm(reg_aux_c1, reg_aux_c1,
                    brg_.dst_dt_sz * (brg_.n_blk * brg_.ld_block), X_TMP_0);
            add_imm(reg_zp_a, reg_zp_a,
                    brg_.n_blk * brg_.ld_block * brg_.dst_dt_sz, X_TMP_0);
            if (brg_.is_oc_scales)
                add_imm(reg_scales, reg_scales,
                        brg_.dst_dt_sz * (brg_.n_blk * brg_.ld_block), X_TMP_0);
            add_imm(reg_bias, reg_bias,
                    brg_.dst_dt_sz * (brg_.n_blk * brg_.ld_block), X_TMP_0);
            mov(reg_aux_c, reg_aux_c1);
        });
    }

    void han_blk_zp() {
        LDR_IMM(reg_tmp, reg_param, GET_OFF(nb));
        LDR_IMM(reg_na, reg_param, GET_OFF(na));
        mov(reg_aux_a1, reg_a);
        ldr(WReg(reg_bd_loop.getIdx()), ptr(reg_na));
        ldr(WReg(reg_ld_loop.getIdx()), ptr(reg_tmp));

        if (brg_.zp_type_b != jit_int8_broadcast_t::none) {
            asm_for(reg_bd_loop, reg_bd_loop, [&]() {
                loop_k_zp(bdb, ldb, 0, 1);
                add_imm(reg_aux_a1, reg_aux_a1,
                        div_up(brg_.K, brg_.k_blk) * brg_.k_blk * brg_.bd_block,
                        X_TMP_0);
                add_imm(reg_zp_b, reg_zp_b, brg_.m_blk * brg_.dst_dt_sz,
                        X_TMP_0);
            });
        }
        if (brg_.zp_type_a != jit_int8_broadcast_t::none) {
            asm_for(reg_ld_loop, reg_ld_loop, [&]() {
                loop_k_zp(bdb, ldb, 1, 0);
                add_imm(reg_zp_a, reg_zp_a,
                        brg_.n_blk * brg_.ld_block * brg_.dst_dt_sz, X_TMP_0);
                add_imm(reg_b, reg_b,
                        (brg_.n_blk * brg_.ld_block)
                                * div_up(brg_.K, brg_.k_blk) * brg_.k_blk,
                        X_TMP_0);
            });
        }
    }

    void zp_comp(int rdb, int bdb, int ldb, int is_a, int is_b) {
        const ZReg zp_b_sum0 = ZReg(3);
        const ZReg zp_b_sum1 = ZReg(4);
        const ZReg zp_b_sum2 = ZReg(5);
        const ZReg zp_b_sum3 = ZReg(6);

        dup(z0.b, 1);
        int rd, ld;
        if (brg_.zp_type_b != jit_int8_broadcast_t::none && is_b == 1) {
            mov(reg_tmp, reg_aux_a);
            for (rd = 0; rd < rdb; rd++) {
                if (cpu_isa_traits<isa>::vlen == 16) {
                    ld1b(z1.b, P_ALL_ONE / T_z, ptr(reg_tmp));
                    if (brg_.is_s8) {
                        smmla(zp_b_sum0.s, z0.b, z1.b);
                    } else {
                        ummla(zp_b_sum0.s, z0.b, z1.b);
                    }
                    ld1b(z1.b, P_ALL_ONE / T_z, ptr(reg_tmp, 1, MUL_VL));
                    if (brg_.is_s8) {
                        smmla(zp_b_sum1.s, z0.b, z1.b);
                    } else {
                        ummla(zp_b_sum1.s, z0.b, z1.b);
                    }
                    ld1b(z1.b, P_ALL_ONE / T_z, ptr(reg_tmp, 2, MUL_VL));
                    if (brg_.is_s8) {
                        smmla(zp_b_sum2.s, z0.b, z1.b);
                    } else {
                        ummla(zp_b_sum2.s, z0.b, z1.b);
                    }
                    ld1b(z1.b, P_ALL_ONE / T_z, ptr(reg_tmp, 3, MUL_VL));
                    if (brg_.is_s8) {
                        smmla(zp_b_sum3.s, z0.b, z1.b);
                    } else {
                        ummla(zp_b_sum3.s, z0.b, z1.b);
                    }
                } else {
                    ld1b(z1.b, P_ALL_ONE / T_z, ptr(reg_tmp));
                    ld1b(z2.b, P_ALL_ONE / T_z, ptr(reg_tmp, 1, MUL_VL));
                    if (brg_.is_s8) {
                        smmla(zp_b_sum0.s, z0.b, z1.b);
                        smmla(zp_b_sum1.s, z0.b, z2.b);
                    } else {
                        ummla(zp_b_sum0.s, z0.b, z1.b);
                        ummla(zp_b_sum1.s, z0.b, z2.b);
                    }
                }
                add_imm(reg_tmp, reg_tmp, brg_.k_blk * brg_.m_blk, X_TMP_0);
            }
        }
        if ((brg_.zp_type_a != jit_int8_broadcast_t::none) && is_a == 1) {
            mov(reg_tmp, reg_aux_b);

            for (rd = 0; rd < rdb; rd++) {
                for (ld = 0; ld < ldb; ld++) {
                    PReg p = (brg_.is_n_tail && ld == ldb - 1) ? prd_ld
                                                               : P_ALL_ONE;
                    ld1b(acc(1, ld).b, p, ptr(reg_tmp, ld, MUL_VL));
                }
                add_imm(reg_tmp, reg_tmp,
                        brg_.k_blk * brg_.n_blk * brg_.ld_block, X_TMP_0);
                for (ld = 0; ld < ldb; ld++) {
                    if (brg_.is_s8 || brg_.is_u8_s8) {
                        smmla(acc(2, ld).s, z0.b, acc(1, ld).b);
                    } else {
                        ummla(acc(2, ld).s, z0.b, acc(1, ld).b);
                    }
                }
            }
        }
    }

    void config() {
        const int vlen_bytes = cpu_isa_traits<isa>::vlen;
        const int sv_len = vlen_bytes / brg_.dst_dt_sz;
        int m, pred_st = 0, pred_ld = 0, pred_b = sv_len;
        const int n_full = brg_.n_blk * brg_.ld_block;
        const int n_cols = brg_.is_n_tail ? brg_.n_tail : n_full;
        n_blks = div_up(n_cols, sv_len);
        k_full_blks = brg_.K / (brg_.k_blk * brg_.rd_block);
        m = brg_.K % (brg_.k_blk * brg_.rd_block);
        k_tail_blk = m / brg_.k_blk;
        k_residual_blk = m % brg_.k_blk;
        ldb = (brg_.is_n_tail) ? div_up(brg_.n_tail, brg_.n_blk)
                               : brg_.ld_block;
        bdb = (brg_.is_m_tail) ? div_up(brg_.m_tail, 2) : brg_.bd_block / 2;
        rdb = (brg_.is_k_tail) ? div_up(brg_.k_tail, brg_.k_blk) : 4;

        const int pred_zp_b_tl
                = (n_cols % sv_len == 0) ? sv_len : n_cols % sv_len;
        set_preg(prd_8.b, sv_len, X_TMP_0, X_TMP_1);
        set_preg(prd_zp_b_tl.b, pred_zp_b_tl, X_TMP_0, X_TMP_1);

        if (brg_.is_n_tail) {
            pred_b = (brg_.n_tail % sv_len == 0) ? sv_len
                                                 : (brg_.n_tail % sv_len);
            if (brg_.n_tail % brg_.n_blk == 0) {
                pred_st = (brg_.n_tail % (brg_.n_blk * 2) == 0) ? sv_len
                                                                : sv_len / 2;
                pred_ld = vlen_bytes;
            } else {
                pred_ld = (brg_.n_tail % brg_.n_blk) * brg_.k_blk;
                pred_st = (ldb % 2 == 0)
                        ? (sv_len / 2) + (brg_.n_tail % brg_.n_blk)
                        : (brg_.n_tail % brg_.n_blk);
            }
        }
        set_preg(prd_ld.b, pred_ld, X_TMP_0, X_TMP_1);
        set_preg(prd_st.s, pred_st, X_TMP_0, X_TMP_1);
        set_preg(prd_b.s, pred_b, X_TMP_0, X_TMP_1);
    }

    void generate() override {
        preamble();
        config();

        LDR_IMM(reg_a, reg_param, GET_OFF(src));
        LDR_IMM(reg_b, reg_param, GET_OFF(wei));
        LDR_IMM(reg_c, reg_param, GET_OFF(dst));
        LDR_IMM(reg_zp_b, reg_param, GET_OFF(zp_b_ptr));
        LDR_IMM(reg_zp_a, reg_param, GET_OFF(zp_a_ptr));
        if (brg_.is_zp_cal) {
            LDR_IMM(reg_zp_val_b, reg_param, GET_OFF(wei_zero_point));
            LDR_IMM(reg_zp_val_a, reg_param, GET_OFF(src_zero_point));
            han_blk_zp();
        } else {

            LDR_IMM(reg_bias, reg_param, GET_OFF(bias));
            LDR_IMM(reg_scales, reg_param, GET_OFF(scales));
            LDR_IMM(reg_aux_scales, reg_param, GET_OFF(dst_scales));
            LDR_IMM(reg_zp_aux_b_buf, reg_param, GET_OFF(wei_zero_point_buf));
            han_blk();
        }

        postamble();
        if (!eltwise_injectors_.empty()) {
            for (int i = 0; i < (int)eltwise_injectors_.size(); i++) {
                eltwise_injectors_[i]->prepare_table();
            }
        }
    }

    jit_int8_matmul_kernel_t(
            const brg_int8_t &k, const dnnl::impl::post_ops_t &eltwise = {})
        : brg_(k) {
        for (auto &e : eltwise.entry_) {
            eltwise_injectors_.emplace_back(utils::make_unique<
                    jit_uni_eltwise_injector_t<to_vla_sve(isa)>>(
                    this, e.eltwise));
        }
    }
    ~jit_int8_matmul_kernel_t() override = default;

private:
    brg_int8_t brg_;
    int ldb;
    int bdb;
    int rdb;
    int k_full_blks;
    int k_tail_blk;
    int k_residual_blk;
    int n_blks;
    std::vector<std::unique_ptr<jit_uni_eltwise_injector_t<to_vla_sve(isa)>>>
            eltwise_injectors_;
};

template <cpu_isa_t isa>
status_t jit_int8_matmul_t<isa>::pd_t::init(engine_t *engine) {

    const auto src_type = src_md(0)->data_type;
    const auto wei_type = weights_md(0)->data_type;
    const auto dst_type = dst_md(0)->data_type;

    const memory_desc_wrapper src_d(src_md_);
    const memory_desc_wrapper weights_d(weights_md_);
    const memory_desc_wrapper dst_d(dst_md_);
    const memory_desc_wrapper bias_d(bias_md_);

    const bool no_runtime_dims_or_strides
            = !(src_d.has_runtime_dims_or_strides()
                    || weights_d.has_runtime_dims_or_strides());

    VDISPATCH_MATMUL(
            no_runtime_dims_or_strides, VERBOSE_RUNTIMEDIM_UNSUPPORTED);
    VDISPATCH_MATMUL(is_dense_format_kind(), VERBOSE_UNSUPPORTED_SPARSE_CFG);

    bool is_u8 = utils::everyone_is(u8, src_type, wei_type);
    bool is_s8 = utils::everyone_is(s8, src_type, wei_type);
    bool is_u8_s8 = utils::everyone_is(u8, src_type)
            && utils::everyone_is(s8, wei_type);

    int dims = src_d.ndims();

    auto check_attr_scales = [&]() -> bool {
        const std::vector<int> supported_args
                = {DNNL_ARG_SRC, DNNL_ARG_WEIGHTS, DNNL_ARG_DST};
        bool ok = attr_scales_ok(supported_args);
        auto is_src_scl
                = !attr()->scales_.get(DNNL_ARG_SRC).has_default_values();
        auto is_wei_scl
                = !attr()->scales_.get(DNNL_ARG_WEIGHTS).has_default_values();
        auto dst_scl_msk = attr()->scales_.get(DNNL_ARG_DST).get_mask();
        auto wei_scl_msk = attr()->scales_.get(DNNL_ARG_WEIGHTS).get_mask();
        auto src_scl_msk = attr()->scales_.get(DNNL_ARG_SRC).get_mask();

        if (src_scl_msk > 0
                || (wei_scl_msk > 0 && wei_scl_msk != 1 << (dims - 1))
                || dst_scl_msk > 0)
            return false;

        if (is_src_scl && is_wei_scl && wei_scl_msk > 0) {
            // This case requires scratchpad.
            if (is_runtime_value(N())) ok = false;
        }
        return ok;
    };

    auto check_bias = [&]() -> bool {
        if (bias_d.format_any()) {
            if (bias_d.has_runtime_dims_or_strides()) return false;
            status_t status = memory_desc_init_by_strides(bias_md_, nullptr);
            if (status != status::success) return false;
        }

        const auto bia_dt = weights_md(1)->data_type;
        return IMPLICATION(with_bias(), bia_dt == f32 && is_bias_1xN());
    };

    auto init_zp_type = [&](brg_int8_t *brg_) -> bool {
        auto zero_points = attr()->zero_points_;

        auto wt_int8 = zero_points.get_data_type(DNNL_ARG_WEIGHTS) == u8
                || zero_points.get_data_type(DNNL_ARG_WEIGHTS) == s8;
        if (!zero_points.has_default_data_type(DNNL_ARG_SRC)
                || !zero_points.has_default_data_type(DNNL_ARG_DST)
                || (!zero_points.has_default_data_type(DNNL_ARG_WEIGHTS)
                        && !wt_int8))
            return false;

        if (!zero_points.has_default_data_type(DNNL_ARG_WEIGHTS)) {
            switch (zero_points.get_data_type(DNNL_ARG_WEIGHTS)) {
                case u8: {
                    brg_->zp_b_dt = u8;
                    brg_->is_zp_b_int8 = true;
                    break;
                }
                case s8: {
                    brg_->zp_b_dt = s8;
                    brg_->is_zp_b_int8 = true;
                    break;
                }
                case s32: {
                    brg_->is_zp_b_int8 = false;
                    break;
                }
                default: return false;
            }
        }

        if (zero_points.get_mask(DNNL_ARG_SRC) > 0
                || zero_points.get_mask(DNNL_ARG_DST) > 0
                || (zero_points.get_mask(DNNL_ARG_WEIGHTS) > 0
                        && (zero_points.get_mask(DNNL_ARG_WEIGHTS))
                                != (3 << (dims - 2))))
            return false;

        brg_->zp_type_a = zero_points.has_default_values(DNNL_ARG_SRC)
                ? jit_int8_broadcast_t::none
                : jit_int8_broadcast_t::per_tensor;

        brg_->zp_type_b = zero_points.has_default_values(DNNL_ARG_WEIGHTS)
                ? jit_int8_broadcast_t::none
                : jit_int8_broadcast_t::per_tensor;

        brg_->zp_type_c = zero_points.has_default_values(DNNL_ARG_DST)
                ? jit_int8_broadcast_t::none
                : jit_int8_broadcast_t::per_tensor;

        return true;
    };

    VDISPATCH_MATMUL(init_zp_type(&brg_), VERBOSE_UNSUPPORTED_ZP_CFG);

    VDISPATCH_MATMUL(check_bias(), VERBOSE_UNSUPPORTED_BIAS_CFG);

    VDISPATCH_MATMUL(check_attr_scales(), VERBOSE_UNSUPPORTED_SCALES_CFG);

    bool no_post_ops = attr()->post_ops_.has_default_values();

    auto with_eltwise = [this]() -> bool {
        // we only support eltwise post-op
        for (auto &e : this->attr()->post_ops_.entry_) {
            if (e.kind != primitive_kind::eltwise) return false;
        }
        return true;
    };

    const bool problem_dt_correct
            = (is_s8 || is_u8 || is_u8_s8) && utils::everyone_is(f32, dst_type);

    VDISPATCH_MATMUL(problem_dt_correct, VERBOSE_UNSUPPORTED_DT);
    VDISPATCH_MATMUL(no_post_ops || with_eltwise(), VERBOSE_UNSUPPORTED_ATTR);
    VDISPATCH_MATMUL(formats_ok(), VERBOSE_UNSUPPORTED_TAG);
    VDISPATCH_MATMUL(get_sve_length() == cpu_isa_traits<isa>::vlen,
            VERBOSE_UNSUPPORTED_ISA);

    auto is_src_any = src_d.format_kind() == format_kind::any;
    auto is_dst_any = dst_d.format_kind() == format_kind::any;

    matmul_helper_t helper(src_d, weights_d, dst_d);
    brg_.K = helper.K();
    brg_.M = helper.M();
    brg_.N = helper.N();
    brg_.B = batch();

    brg_.n_blk = cpu_isa_traits<isa>::vlen / brg_.k_blk;

    int num_threads = dnnl_get_current_num_threads();

    if (brg_.N <= brg_.n_blk * 2)
        brg_.ld_block = 2;
    else if (brg_.N <= brg_.n_blk * 4)
        brg_.ld_block = 4;
    else
        brg_.ld_block = 6;
    m_block_sz = 32;
    const int micro_n = brg_.n_blk * brg_.ld_block;
    n_block_sz = (brg_.ld_block == 6) ? micro_n : 2 * micro_n;

    // micro_n is small for 128 bit.
    // so N tiles can get too tiny and we spend time in OpenMP overhead.
    // try a coarser N tile (4x/2x/1x) and pick the first that keeps
    // total work more than num_threads.
    if (cpu_isa_traits<isa>::vlen == 16 && brg_.ld_block == 6) {
        // pick the largest N size that still gives enough work.
        int best_n_block_sz = micro_n;
        const int num_m_tiles = div_up(brg_.M, m_block_sz);
        // when M is only a few blocks, starting at 4x can reduce N-parallelism.
        const bool mid_m = (brg_.M > m_block_sz) && (brg_.M <= 4 * m_block_sz);
        const bool small_m = num_m_tiles <= 4;

        const int preferred_factor = (mid_m || small_m) ? 2 : 4;
        for (int f = preferred_factor; f >= 1; f /= 2) {
            const int temp_n_block_sz = f * micro_n;
            const int num_n_tiles = div_up(brg_.N, temp_n_block_sz);
            const int work = (int)brg_.B * num_m_tiles * num_n_tiles;
            if (work >= num_threads) {
                best_n_block_sz = temp_n_block_sz;
                break;
            }
        }
        n_block_sz = std::min<int>(brg_.N, best_n_block_sz);
        // keep the coarsened tile but align it to micro_n so that the last tile
        // is correctly handled by the tail kernel when needed.
        if (n_block_sz % micro_n != 0) {
            n_block_sz = div_up(n_block_sz, micro_n) * micro_n;
        }
    }

    int num_a_blocks = div_up(brg_.M, m_block_sz);
    int num_b_blocks = div_up(brg_.N, n_block_sz);
    mm_parallel_work = brg_.B * num_a_blocks * num_b_blocks;
    if (mm_parallel_work < num_threads && brg_.ld_block == 4)
        n_block_sz = micro_n;
    num_b_blocks = div_up(brg_.N, n_block_sz);
    mm_parallel_work = brg_.B * num_a_blocks * num_b_blocks;

    auto b_tag_2d = format_tag::ab;
    auto b_tag_3d = format_tag::abc;
    auto b_tag_4d = format_tag::abcd;

    switch (micro_n) {
        case 4: {
            b_tag_2d = format_tag::BA4b8a;
            b_tag_3d = format_tag::aCB4c8b;
            b_tag_4d = format_tag::abDC4d8c;
            break;
        }
        case 8: {
            b_tag_2d = format_tag::BA8b8a;
            b_tag_3d = format_tag::aCB8c8b;
            b_tag_4d = format_tag::abDC8d8c;
            break;
        }
        case 12: {
            b_tag_2d = format_tag::BA12b8a;
            b_tag_3d = format_tag::aCB12c8b;
            b_tag_4d = format_tag::abDC12d8c;
            break;
        }
        case 16: {
            b_tag_2d = format_tag::BA16b8a;
            b_tag_3d = format_tag::aCB16c8b;
            b_tag_4d = format_tag::abDC16d8c;
            break;
        }
        case 24: {
            b_tag_2d = format_tag::BA24b8a;
            b_tag_3d = format_tag::aCB24c8b;
            b_tag_4d = format_tag::abDC24d8c;
            break;
        }
        default: return status::unimplemented;
    }

    switch (dims) {
        case 2: {
            if (is_src_any)
                VCHECK_BG(memory_desc_init_by_tag(src_md_, format_tag::ab),
                        VERBOSE_UNSUPPORTED_TAG);
            if (is_dst_any)
                VCHECK_BG(memory_desc_init_by_tag(dst_md_, format_tag::ab),
                        VERBOSE_UNSUPPORTED_TAG);
            if (!weights_d.matches_tag(format_tag::ab)) {
                brg_.b_reo = false;
                VCHECK_BG(memory_desc_init_by_tag(weights_md_, b_tag_2d),
                        VERBOSE_UNSUPPORTED_TAG);
            } else {
                VCHECK_BG(memory_desc_init_by_tag(weights_md_, format_tag::ab),
                        VERBOSE_UNSUPPORTED_TAG);
            }
            break;
        }
        case 3: {
            if (is_src_any)
                VCHECK_BG(memory_desc_init_by_tag(src_md_, format_tag::abc),
                        VERBOSE_UNSUPPORTED_TAG);
            if (is_dst_any)
                VCHECK_BG(memory_desc_init_by_tag(dst_md_, format_tag::abc),
                        VERBOSE_UNSUPPORTED_TAG);
            if (!weights_d.matches_tag(format_tag::abc)) {
                brg_.b_reo = false;
                VCHECK_BG(memory_desc_init_by_tag(weights_md_, b_tag_3d),
                        VERBOSE_UNSUPPORTED_TAG);
            } else {
                VCHECK_BG(memory_desc_init_by_tag(weights_md_, format_tag::abc),
                        VERBOSE_UNSUPPORTED_TAG);
            }
            if (src_d.dims()[0] != weights_d.dims()[0])
                return status::unimplemented;
            break;
        }
        case 4: {
            if (is_src_any)
                VCHECK_BG(memory_desc_init_by_tag(src_md_, format_tag::abcd),
                        VERBOSE_UNSUPPORTED_TAG);
            if (is_dst_any)
                VCHECK_BG(memory_desc_init_by_tag(dst_md_, format_tag::abcd),
                        VERBOSE_UNSUPPORTED_TAG);
            if (!weights_d.matches_tag(format_tag::abcd)) {
                brg_.b_reo = false;
                VCHECK_BG(memory_desc_init_by_tag(weights_md_, b_tag_4d),
                        VERBOSE_UNSUPPORTED_TAG);
            } else {
                VCHECK_BG(
                        memory_desc_init_by_tag(weights_md_, format_tag::abcd),
                        VERBOSE_UNSUPPORTED_TAG);
            }
            if (src_d.dims()[0] != weights_d.dims()[0]
                    || src_d.dims()[1] != weights_d.dims()[1])
                return status::unimplemented;
            break;
        }
        default: return status::unimplemented;
    }

    bool is_scales = !attr()->scales_.get(DNNL_ARG_SRC).has_default_values()
            || !attr()->scales_.get(DNNL_ARG_WEIGHTS).has_default_values();

    bool is_dst_scales
            = !attr()->scales_.get(DNNL_ARG_DST).has_default_values();

    const auto &wei_scales = attr()->scales_.get(DNNL_ARG_WEIGHTS);

    brg_.dst_dt_sz = 4;
    brg_.na = 1;
    brg_.nb = 1;
    brg_.m_tail = brg_.M % brg_.m_blk;
    brg_.k_tail = brg_.K % (brg_.k_blk * brg_.rd_block);
    brg_.n_tail = brg_.N % (brg_.n_blk * brg_.ld_block);
    brg_.is_s8 = is_s8;
    brg_.is_u8_s8 = is_u8_s8;
    brg_.is_bias = with_bias();
    brg_.with_scales = is_scales;
    brg_.with_dst_scales = is_dst_scales;
    brg_.is_oc_scales = wei_scales.get_mask() > 0;
    dyn_.K = brg_.K;
    dyn_.N = brg_.N;
    dyn_.M = brg_.M;
    dyn_.B = brg_.B;
    dyn_.mtail = brg_.m_tail;
    dyn_.m_blk = brg_.m_blk;
    dyn_.k_blk = brg_.k_blk;
    dyn_.n_blk = brg_.n_blk * brg_.ld_block;
    dyn_.ntail = brg_.n_tail;
    dyn_.ktail = dyn_.K % brg_.k_blk;

    auto scratchpad = scratchpad_registry().registrar();
    if (brg_.zp_type_a != jit_int8_broadcast_t::none)
        scratchpad.book(key_brgemm_primitive_zp_comp_a,
                div_up(brg_.N, (brg_.n_blk * brg_.ld_block))
                        * (brg_.n_blk * brg_.ld_block) * brg_.dst_dt_sz
                        * brg_.B,
                sizeof(char));
    if (brg_.zp_type_b != jit_int8_broadcast_t::none)
        scratchpad.book(key_brgemm_primitive_zp_comp_b,
                div_up(brg_.M, brg_.m_blk) * brg_.m_blk * brg_.dst_dt_sz
                        * brg_.B,
                sizeof(char));
    scratchpad.book(key_brgemm_primitive_buffer_a,
            brg_.B * div_up(brg_.M, brg_.m_blk) * div_up(brg_.K, brg_.k_blk)
                    * brg_.m_blk * brg_.k_blk,
            sizeof(char));
    if (brg_.b_reo)
        scratchpad.book(key_gemm_blocked_b,
                brg_.B * div_up(brg_.N, (brg_.n_blk * brg_.ld_block))
                        * (brg_.n_blk * brg_.ld_block)
                        * div_up(brg_.K, brg_.k_blk) * brg_.k_blk,
                sizeof(char));
    book_precomputed_scales(scratchpad, attr()->scales_, N());

    return status::success;
}

template <cpu_isa_t isa>
status_t jit_int8_matmul_t<isa>::init(engine_t *engine) {

    const auto &b1 = pd()->get_b();
    const auto &d1 = pd()->get_d();

    dyn_vals_t d;
    d.K = d1.K;
    d.M = d1.M;
    d.B = d1.B;
    d.N = d1.N;
    d.mtail = d1.mtail;
    d.ktail = d1.ktail;
    d.ntail = d1.ntail;
    d.k_blk = d1.k_blk;
    d.m_blk = d1.m_blk;
    d.n_blk = d1.n_blk;

    brg_int8_t b;
    b.M = b1.M;
    b.K = b1.K;
    b.N = b1.N;
    b.n_blk = b1.n_blk;
    b.na = b1.na;
    b.nb = b1.nb;
    b.m_tail = b1.m_tail;
    b.n_tail = b1.n_tail;
    b.k_tail = b1.k_tail;
    b.dst_dt_sz = b1.dst_dt_sz;
    b.is_s8 = b1.is_s8;
    b.is_u8_s8 = b1.is_u8_s8;
    b.B = b1.B;
    b.is_bias = b1.is_bias;
    b.zp_type_a = b1.zp_type_a;
    b.zp_type_b = b1.zp_type_b;
    b.zp_type_c = b1.zp_type_c;
    b.is_zp_b_int8 = b1.is_zp_b_int8;
    b.zp_b_dt = b1.zp_b_dt;
    b.with_scales = b1.with_scales;
    b.with_dst_scales = b1.with_dst_scales;
    b.is_oc_scales = b1.is_oc_scales;
    b.b_reo = b1.b_reo;
    b.ld_block = b1.ld_block;

    const bool has_eltwise
            = pd()->attr()->post_ops_.find(primitive_kind_t::dnnl_eltwise) >= 0;

    for (int z = 0; z < 2; z++)
        for (int m = 0; m < 2; m++)
            for (int n = 0; n < 2; n++)
                for (int k = 0; k < 2; k++) {
                    int idx = pd()->get_idx(z, m, k, n, b1);
                    if (idx == -1 || idx > 15) continue;
                    b.is_m_tail = m;
                    b.is_k_tail = k;
                    b.is_n_tail = n;
                    b.is_zp_cal = z;

                    if (has_eltwise) {
                        int8_kernels_[idx] = std::unique_ptr<
                                jit_int8_matmul_kernel_t<isa>> {
                                new jit_int8_matmul_kernel_t<isa>(
                                        b, pd()->attr()->post_ops_)};
                    } else {
                        int8_kernels_[idx] = std::unique_ptr<
                                jit_int8_matmul_kernel_t<isa>> {
                                new jit_int8_matmul_kernel_t<isa>(b)};
                    }

                    if (!int8_kernels_[idx]) return status::runtime_error;

                    CHECK(int8_kernels_[idx]->create_kernel());
                }

    d.reorder_a = 1;
    d.reorder_b = 0;
    reo_ker_a_ = std::unique_ptr<jit_int8_matmul_utils_kernel_t> {
            new jit_int8_matmul_utils_kernel_t(d, isa)};
    CHECK(reo_ker_a_->create_kernel());

    d.reorder_b = 1;
    d.reorder_a = 0;
    reo_ker_b_ = std::unique_ptr<jit_int8_matmul_utils_kernel_t> {
            new jit_int8_matmul_utils_kernel_t(d, isa)};
    CHECK(reo_ker_b_->create_kernel());

    return status::success;
}

template <cpu_isa_t isa>
jit_int8_matmul_t<isa>::jit_int8_matmul_t(const pd_t *apd) : primitive_t(apd) {}
template <cpu_isa_t isa>
jit_int8_matmul_t<isa>::~jit_int8_matmul_t() = default;

template <cpu_isa_t isa>
status_t jit_int8_matmul_t<isa>::execute(const exec_ctx_t &ctx) const {
    const auto *weights_b = CTX_IN_MEM(const float *, DNNL_ARG_WEIGHTS);
    const auto *src_b = CTX_IN_MEM(const float *, DNNL_ARG_SRC);
    auto dst = CTX_OUT_MEM(float *, DNNL_ARG_DST);
    const auto *bias = CTX_IN_MEM(const float *, DNNL_ARG_BIAS);

    const int32_t *src_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC);
    const int32_t *wei_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_WEIGHTS);
    const int32_t *dst_zero_points = CTX_IN_MEM(
            const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST);

    DEFINE_ARG_SCALES_BUFFER(src_scales, DNNL_ARG_SRC);
    DEFINE_ARG_SCALES_BUFFER(wei_scales, DNNL_ARG_WEIGHTS);
    DEFINE_ARG_SCALES_BUFFER(dst_scales, DNNL_ARG_DST);

    const auto &b = pd()->get_b();
    const auto &d = pd()->get_d();

    const auto &scratchpad = ctx.get_scratchpad_grantor();

    int num_threads = dnnl_get_current_num_threads();
    char *src = scratchpad.template get<char>(key_brgemm_primitive_buffer_a);
    char *weights = (b.b_reo)
            ? scratchpad.template get<char>(key_gemm_blocked_b)
            : (char *)weights_b;
    char *zp_ptr_a;
    if (b.zp_type_a != jit_int8_broadcast_t::none)
        zp_ptr_a
                = scratchpad.template get<char>(key_brgemm_primitive_zp_comp_a);
    char *zp_ptr_b;
    if (b.zp_type_b != jit_int8_broadcast_t::none)
        zp_ptr_b
                = scratchpad.template get<char>(key_brgemm_primitive_zp_comp_b);
    const float *oscales = precompute_scales(
            scratchpad, src_scales, wei_scales, pd()->N(), pd()->attr());

    const dim_t B = b.B;
    const dim_t M = b.M;
    const dim_t N = b.N;
    const dim_t K = b.K;

    auto reorder_a = [&]() {
        int m_blks = div_up(M, b.m_blk);
        int k_blks = div_up(K, b.k_blk);
        int parallel_work = B * m_blks * k_blks;
        int blk_per_bt = m_blks * k_blks;
        auto tmp_src = src_b;
        // If parallel_work == 1, we limit num threads to 1 as parallel(1, ...)
        // does not create a parallel section. We do not limit number of threads
        // for case 1 < parallel_work_amount_ < dnnl_get_max_threads() to avoid
        // potential overhead on spawning different number of OMP threads from
        // layer to layer.
        int nt = parallel_work > 1 ? num_threads : 1;
        parallel(nt, [&](const int ithr, const int nthr) {
            int start {0}, end {0};
            balance211(parallel_work, nt, ithr, start, end);

            int bt = start / blk_per_bt;
            int bs = start % blk_per_bt;
            int nobl = end - start;
            int nobt = 1;
            int noblf = end - start, nobll;

            if (bs + nobl > blk_per_bt) {
                nobt += div_up(nobl - (blk_per_bt - bs), blk_per_bt);
                noblf = blk_per_bt - bs;
                nobll = (nobl - (blk_per_bt - bs)) % blk_per_bt;
                if (nobll == 0) nobll = blk_per_bt;
            }
            int nob;
            for (int i = 0; i < nobt; i++) {
                nob = (i == 0) ? noblf : ((i == nobt - 1) ? nobll : blk_per_bt);
                bs = start % blk_per_bt;
                int m_blk_src = bs / k_blks;
                int k_blk_src = bs % k_blks;
                int m_blk_dst = bs / k_blks;
                int k_blk_dst = bs % k_blks;

                int k1 = std::min(k_blks - k_blk_src, nob);
                int k_tmp = nob - k1;
                int m1 = (k_tmp > 0) ? k_tmp / k_blks : 0;
                int k2 = (k_tmp > 0) ? k_tmp % k_blks : 0;
                int src_ad = (bt * M * K) + (m_blk_src * b.m_blk * K)
                        + (k_blk_src * b.k_blk);
                int dst_ad = (bt * m_blks * k_blks * b.m_blk * b.k_blk)
                        + (m_blk_dst * k_blks * b.m_blk * b.k_blk)
                        + (k_blk_dst * b.m_blk * b.k_blk);
                int src_new = src_ad, dst_new = dst_ad;

                dyn_params_t k;

                if (k1 > 0) {
                    int a = 1;
                    int mtl = (d.mtail > 0) ? 1 : 0;
                    int tl = (d.ktail > 0) ? 1 : 0;
                    if (k1 + k_blk_src < k_blks) tl = 0;
                    if (1 + m_blk_src < m_blks) mtl = 0;
                    k.src = (int8_t *)tmp_src + src_ad;
                    k.dst = (int8_t *)src + dst_ad;
                    k.nm = &a;
                    k.nk = &k1;
                    k.tl = &tl;
                    k.mtl = &mtl;
                    (*reo_ker_a_)(&k);
                }

                if (m1 > 0) {
                    int mtl = (d.mtail > 0) ? 1 : 0;
                    int tl = (d.ktail > 0) ? 1 : 0;
                    if (1 + m1 + m_blk_src < m_blks) mtl = 0;
                    if (k1 != k_blks) {
                        src_new = src_ad - b.k_blk * (k_blks - k1)
                                + b.m_blk * K;
                    } else {
                        src_new = src_ad + b.m_blk * K;
                    }
                    dst_new = dst_ad + b.m_blk * b.k_blk * k1;
                    k.src = (int8_t *)tmp_src + src_new;
                    k.dst = (int8_t *)src + dst_new;
                    k.nm = &m1;
                    k.nk = &k_blks;
                    k.tl = &tl;
                    k.mtl = &mtl;
                    (*reo_ker_a_)(&k);
                }
                if (k2 > 0) {
                    int a = 1, tl = 0;
                    int mtl = (d.mtail > 0) ? 1 : 0;
                    if (1 + 1 + m1 + m_blk_src < m_blks) mtl = 0;
                    if (m1 < 1) {
                        src_new = src_ad - b.k_blk * (k_blks - k1)
                                + (b.m_blk * K);
                        dst_new = dst_ad + b.m_blk * b.k_blk * k1;
                    } else {
                        src_new += K * m1 * b.m_blk;
                        dst_new += b.m_blk * b.k_blk * k_blks * m1;
                    }
                    k.src = (int8_t *)tmp_src + src_new;
                    k.dst = (int8_t *)src + dst_new;
                    k.nm = &a;
                    k.nk = &k2;
                    k.tl = &tl;
                    k.mtl = &mtl;
                    (*reo_ker_a_)(&k);
                }
                bt++;
                start += nob;
            }
        });
    };

    auto reorder_b = [&]() {
        int k_blks = div_up(K, d.k_blk);
        int n_blks = div_up(N, d.n_blk);
        int parallel_work = B * n_blks * k_blks;
        int blk_per_bt = n_blks * k_blks;
        int nt = parallel_work > 1 ? num_threads : 1;
        parallel(nt, [&](const int ithr, const int nthr) {
            int start {0}, end {0};
            balance211(parallel_work, nt, ithr, start, end);
            int bt = start / blk_per_bt;
            int bs = start % blk_per_bt;
            int nobl = end - start;
            int nobt = 1;
            int noblf = end - start, nobll;

            if (bs + nobl > blk_per_bt) {
                nobt += div_up(nobl - (blk_per_bt - bs), blk_per_bt);
                noblf = blk_per_bt - bs;
                nobll = (nobl - (blk_per_bt - bs)) % blk_per_bt;
                if (nobll == 0) nobll = blk_per_bt;
            }
            int nob;
            for (int i = 0; i < nobt; i++) {
                nob = (i == 0) ? noblf : ((i == nobt - 1) ? nobll : blk_per_bt);
                bs = start % blk_per_bt;
                int n_blk_src = bs / k_blks;
                int k_blk_src = bs % k_blks;
                int n_blk_dst = bs / k_blks;
                int k_blk_dst = bs % k_blks;

                int k1 = std::min(k_blks - k_blk_src, nob);
                int k_tmp = nob - k1;
                int n1 = (k_tmp > 0) ? k_tmp / k_blks : 0;
                int k2 = (k_tmp > 0) ? k_tmp % k_blks : 0;
                int src_ad = (bt * N * K) + (n_blk_src * d.n_blk)
                        + (k_blk_src * d.k_blk * N);
                int dst_ad = (bt * n_blks * k_blks * d.k_blk * d.n_blk)
                        + (n_blk_dst * k_blks * d.k_blk * d.n_blk)
                        + (k_blk_dst * d.k_blk * d.n_blk);
                int src_new = src_ad, dst_new = dst_ad;

                dyn_params_t k;

                if (k1 > 0) {
                    int a = 1;
                    int ntl = (d.ntail > 0) ? 1 : 0;
                    int tl = (d.ktail > 0) ? 1 : 0;

                    if (k1 + k_blk_src < k_blks) tl = 0;
                    if (1 + n_blk_src < n_blks) ntl = 0;
                    k.src = (int8_t *)weights_b + src_ad;
                    k.dst = (int8_t *)weights + dst_ad;
                    k.nn = &a;
                    k.nk = &k1;
                    k.tl = &tl;
                    k.ntl = &ntl;
                    (*reo_ker_b_)(&k);
                }

                if (n1 > 0) {
                    int ntl = (d.ntail > 0) ? 1 : 0;
                    int tl = (d.ktail > 0) ? 1 : 0;
                    if (1 + n1 + n_blk_src < n_blks) ntl = 0;

                    if (k1 != k_blks) {
                        src_new = src_ad - d.k_blk * N * (k_blks - k1)
                                + d.n_blk;
                    } else {
                        src_new = src_ad + d.n_blk;
                    }
                    dst_new = dst_ad + d.k_blk * d.n_blk * k1;
                    k.src = (int8_t *)weights_b + src_new;
                    k.dst = (int8_t *)weights + dst_new;
                    k.nn = &n1;
                    k.nk = &k_blks;
                    k.tl = &tl;
                    k.ntl = &ntl;
                    (*reo_ker_b_)(&k);
                }
                if (k2 > 0) {
                    int a = 1, tl = 0;
                    int ntl = (d.ntail > 0) ? 1 : 0;
                    if (1 + 1 + n1 + n_blk_src < n_blks) ntl = 0;
                    if (n1 < 1) {
                        src_new = src_ad - d.k_blk * N * (k_blks - k1)
                                + d.n_blk;
                        dst_new = dst_ad + d.k_blk * d.n_blk * k1;
                    } else {
                        src_new += n1 * d.n_blk;
                        dst_new += d.k_blk * d.n_blk * k_blks * n1;
                    }
                    k.src = (int8_t *)weights_b + src_new;
                    k.dst = (int8_t *)weights + dst_new;
                    k.nn = &a;
                    k.nk = &k2;
                    k.tl = &tl;
                    k.ntl = &ntl;
                    (*reo_ker_b_)(&k);
                }
                bt++;
                start += nob;
            }
        });
    };

    auto kernel_execute
            = [&](int idx, int na, int nb, int m_blk_adr, int n_blk_adr,
                      int dst_adr, int bias_addr, int scl_addr,
                      int zp_ptr_a_adr, int zp_ptr_b_adr, int zp_b_buf) {
        call_params_t p;
        p.na = &na;
        p.nb = &nb;
        p.src = (uint8_t *)src + m_blk_adr;
        p.wei = (uint8_t *)weights + n_blk_adr;
        p.dst = dst + dst_adr;
        p.bias = (float *)bias + bias_addr;
        p.scales = oscales + scl_addr;
        p.dst_scales = dst_scales;
        p.src_zero_point = src_zero_points;
        if (b.is_zp_b_int8)
            p.wei_zero_point_buf = (const int8_t *)wei_zero_points;
        else
            p.wei_zero_point = wei_zero_points;
        p.dst_zero_point = dst_zero_points;
        p.M = M;
        p.N = N;
        p.K = K;
        p.zp_a_ptr = (float *)zp_ptr_a + zp_ptr_a_adr;
        p.zp_b_ptr = (float *)zp_ptr_b + zp_ptr_b_adr;
        (*int8_kernels_[idx])(&p);
    };

    auto kernel_execute_zp = [&]() {
        int num_a_blocks = div_up(M, b.m_blk);
        int num_b_blocks = div_up(N, (b.n_blk * b.ld_block));
        int ktail = (b.k_tail == 0) ? 0 : 1;
        int parallel_work = B * num_a_blocks;
        int nt = parallel_work > 1 ? num_threads : 1;
        if (b.zp_type_b != jit_int8_broadcast_t::none)
            parallel(nt, [&](const int ithr, const int nthr) {
                int start {0}, end {0};
                balance211(parallel_work, nt, ithr, start, end);
                int batch = start / num_a_blocks;
                int m_st = start % num_a_blocks;
                int m_ed = end - start + m_st;
                int mtail
                        = (m_ed == num_a_blocks) ? ((b.m_tail > 0) ? 1 : 0) : 0;
                int m_blk_adr = (batch
                                        * (num_a_blocks * b.m_blk
                                                * div_up(K, b.k_blk) * b.k_blk))
                        + m_st * b.m_blk * div_up(K, b.k_blk) * b.k_blk;
                int zp_ptr_b_adr
                        = (batch * (num_a_blocks * b.m_blk)) + m_st * b.m_blk;

                int idx = pd()->get_idx(1, 0, ktail, 0, b);
                if (idx < 0) {
                    assert(!"Requested int8 matmul kernel was not created.");
                    return;
                }
                int n_a = m_ed - m_st;
                if (mtail) n_a -= 1;
                kernel_execute(
                        idx, n_a, 0, m_blk_adr, 0, 0, 0, 0, 0, zp_ptr_b_adr, 0);

                if (mtail) {
                    idx = pd()->get_idx(1, mtail, ktail, 0, b);
                    if (idx < 0) {
                        assert(!"Requested int8 matmul kernel was not "
                                "created.");
                        return;
                    }
                    m_blk_adr += n_a * b.m_blk * div_up(K, b.k_blk) * b.k_blk;
                    zp_ptr_b_adr += n_a * b.m_blk;
                    kernel_execute(idx, 1, 0, m_blk_adr, 0, 0, 0, 0, 0,
                            zp_ptr_b_adr, 0);
                }
                start++;
            });

        parallel_work = B * num_b_blocks;
        nt = parallel_work > 1 ? num_threads : 1;
        if (b.zp_type_a != jit_int8_broadcast_t::none)
            parallel(nt, [&](const int ithr, const int nthr) {
                int start {0}, end {0};
                balance211(parallel_work, nt, ithr, start, end);
                int batch = start / num_b_blocks;
                int n_st = start % num_b_blocks;
                int n_ed = n_st + end - start;
                int ntail
                        = (n_ed == num_b_blocks) ? ((b.n_tail > 0) ? 1 : 0) : 0;
                int n_blk_adr = (batch
                                        * (num_b_blocks * (b.n_blk * b.ld_block)
                                                * div_up(K, b.k_blk) * b.k_blk))
                        + n_st * (b.n_blk * b.ld_block) * div_up(K, b.k_blk)
                                * b.k_blk;
                int zp_ptr_a_adr
                        = (batch * num_b_blocks * (b.n_blk * b.ld_block))
                        + n_st * (b.n_blk * b.ld_block);

                int idx = pd()->get_idx(1, 0, ktail, 0, b);
                if (idx < 0) {
                    assert(!"Requested int8 matmul kernel was not created.");
                    return;
                }
                int n_b = n_ed - n_st;
                if (ntail == 1) n_b -= 1;

                kernel_execute(
                        idx, 0, n_b, 0, n_blk_adr, 0, 0, 0, zp_ptr_a_adr, 0, 0);

                if (ntail) {
                    idx = pd()->get_idx(1, 0, ktail, 1, b);
                    if (idx < 0) {
                        assert(!"Requested int8 matmul kernel was not "
                                "created.");
                        return;
                    }
                    n_blk_adr += n_b * (b.n_blk * b.ld_block)
                            * div_up(K, b.k_blk) * b.k_blk;
                    zp_ptr_a_adr += n_b * (b.n_blk * b.ld_block);
                    kernel_execute(idx, 0, 1, 0, n_blk_adr, 0, 0, 0,
                            zp_ptr_a_adr, 0, 0);
                }

                start++;
            });
    };

    if (b.b_reo) reorder_b();

    reorder_a();

    if (b.zp_type_a != jit_int8_broadcast_t::none
            || b.zp_type_b != jit_int8_broadcast_t::none)
        kernel_execute_zp();

    int m_block1, n_block1, m_block1_rs, n_block1_rs, num_a_blocks_act,
            num_b_blocks_act, num_a_blocks, num_b_blocks;
    int m_block_sz = pd()->m_block_sz;
    int n_block_sz = pd()->n_block_sz;

    m_block1 = div_up(m_block_sz, b.m_blk);
    n_block1 = div_up(n_block_sz, (b.n_blk * b.ld_block));
    m_block1_rs = div_up(M % m_block_sz, b.m_blk);
    n_block1_rs = div_up(N % n_block_sz, (b.n_blk * b.ld_block));

    num_a_blocks_act = div_up(M, b.m_blk);
    num_b_blocks_act = div_up(N, (b.n_blk * b.ld_block));
    num_a_blocks = div_up(M, m_block_sz);
    num_b_blocks = div_up(N, n_block_sz);
    int ktail = (b.k_tail == 0) ? 0 : 1;
    int parallel_work = pd()->mm_parallel_work;
    int nt = parallel_work > 1 ? num_threads : 1;
    parallel(nt, [&](const int ithr, const int nthr) {
        int start {0}, end {0};
        balance211(parallel_work, nt, ithr, start, end);
        while (start < end) {
            int batch = start / (num_a_blocks * num_b_blocks);
            int batch_start = start % (num_a_blocks * num_b_blocks);
            int m_block = batch_start % num_a_blocks;
            int n_block = batch_start / num_a_blocks;
            int mtail
                    = (m_block1_rs != 0 && m_block == num_a_blocks - 1) ? 1 : 0;
            int ntail
                    = (n_block1_rs != 0 && n_block == num_b_blocks - 1) ? 1 : 0;
            int dst_adr = (batch * M * N) + m_block * b.m_blk * m_block1 * N
                    + n_block * (b.n_blk * b.ld_block) * n_block1;
            int m_blk_adr = (batch
                                    * (num_a_blocks_act * b.m_blk
                                            * div_up(K, b.k_blk) * b.k_blk))
                    + m_block * b.m_blk * m_block1 * div_up(K, b.k_blk)
                            * b.k_blk;
            int n_blk_adr = (batch
                                    * (num_b_blocks_act * (b.n_blk * b.ld_block)
                                            * div_up(K, b.k_blk) * b.k_blk))
                    + n_block * (b.n_blk * b.ld_block) * n_block1
                            * div_up(K, b.k_blk) * b.k_blk;
            int zp_ptr_a_adr
                    = (batch * (num_b_blocks_act * (b.n_blk * b.ld_block)))
                    + n_block * (b.n_blk * b.ld_block) * n_block1;
            int zp_ptr_b_adr = (batch * (num_a_blocks_act * b.m_blk))
                    + m_block * b.m_blk * m_block1;
            int bias_addr = n_block * (b.n_blk * b.ld_block) * n_block1;
            int zp_b_buf = n_block * (b.n_blk * b.ld_block) * n_block1;
            int scl_addr = (b.is_oc_scales)
                    ? (n_block * (b.n_blk * b.ld_block) * n_block1)
                    : 0;
            int idx = pd()->get_idx(0, 0, ktail, 0, b);
            if (idx < 0) {
                assert(!"Requested int8 matmul kernel was not created.");
                return;
            }
            int n_a = m_block1, n_b = n_block1;
            n_a = (mtail) ? ((b.m_tail) ? m_block1_rs - 1 : m_block1_rs)
                          : m_block1;
            n_b = (ntail) ? ((b.n_tail) ? n_block1_rs - 1 : n_block1_rs)
                          : n_block1;

            if (n_a > 0 && n_b > 0) {

                kernel_execute(idx, n_a, n_b, m_blk_adr, n_blk_adr, dst_adr,
                        bias_addr, scl_addr, zp_ptr_a_adr, zp_ptr_b_adr,
                        zp_b_buf);
            }

            if (mtail && b.m_tail > 0 && n_b > 0) {
                int new_dst_adr = dst_adr + b.m_blk * n_a * N;
                int new_m_blk_adr = m_blk_adr
                        + b.m_blk * n_a * div_up(K, b.k_blk) * b.k_blk;
                int new_zp_ptr_b_adr = zp_ptr_b_adr + b.m_blk * n_a;
                idx = pd()->get_idx(0, 1, ktail, 0, b);
                if (idx < 0) {
                    assert(!"Requested int8 matmul kernel was not created.");
                    return;
                }
                int na = 1;
                kernel_execute(idx, na, n_b, new_m_blk_adr, n_blk_adr,
                        new_dst_adr, bias_addr, scl_addr, zp_ptr_a_adr,
                        new_zp_ptr_b_adr, zp_b_buf);
            }

            if (ntail && b.n_tail > 0 && n_a > 0) {
                int new_dst_adr = dst_adr + (b.n_blk * b.ld_block) * n_b;
                int new_n_blk_adr = n_blk_adr
                        + (b.n_blk * b.ld_block) * n_b * div_up(K, b.k_blk)
                                * b.k_blk;
                int new_zp_b_buf = zp_b_buf + (b.n_blk * b.ld_block) * n_b;
                int new_zp_ptr_a_adr
                        = zp_ptr_a_adr + (b.n_blk * b.ld_block) * n_b;
                int new_bias_addr = bias_addr + (b.n_blk * b.ld_block) * n_b;
                int new_scl_addr = scl_addr
                        + ((b.is_oc_scales) ? ((b.n_blk * b.ld_block) * n_b)
                                            : 0);
                idx = pd()->get_idx(0, 0, ktail, 1, b);
                if (idx < 0) {
                    assert(!"Requested int8 matmul kernel was not created.");
                    return;
                }
                int nb = 1;

                kernel_execute(idx, n_a, nb, m_blk_adr, new_n_blk_adr,
                        new_dst_adr, new_bias_addr, new_scl_addr,
                        new_zp_ptr_a_adr, zp_ptr_b_adr, new_zp_b_buf);
            }

            if (mtail && b.m_tail > 0 && ntail && b.n_tail > 0) {
                int new_dst_adr = dst_adr + (b.n_blk * b.ld_block) * n_b
                        + b.m_blk * n_a * N;
                int new_m_blk_adr = m_blk_adr
                        + b.m_blk * n_a * div_up(K, b.k_blk) * b.k_blk;
                int new_n_blk_adr = n_blk_adr
                        + (b.n_blk * b.ld_block) * n_b * div_up(K, b.k_blk)
                                * b.k_blk;
                int new_zp_b_buf = zp_b_buf + (b.n_blk * b.ld_block) * n_b;
                int new_zp_ptr_a_adr
                        = zp_ptr_a_adr + (b.n_blk * b.ld_block) * n_b;
                int new_zp_ptr_b_adr = zp_ptr_b_adr + b.m_blk * n_a;
                int new_bias_addr = bias_addr + (b.n_blk * b.ld_block) * n_b;
                int new_scl_addr = scl_addr
                        + ((b.is_oc_scales) ? ((b.n_blk * b.ld_block) * n_b)
                                            : 0);
                idx = pd()->get_idx(0, 1, ktail, 1, b);
                if (idx < 0) {
                    assert(!"Requested int8 matmul kernel was not created.");
                    return;
                }
                int nb = 1, na = 1;
                kernel_execute(idx, na, nb, new_m_blk_adr, new_n_blk_adr,
                        new_dst_adr, new_bias_addr, new_scl_addr,
                        new_zp_ptr_a_adr, new_zp_ptr_b_adr, new_zp_b_buf);
            }
            start++;
        }
    });

    return status::success;
}

template struct jit_int8_matmul_t<sve_256>;
template struct jit_int8_matmul_t<sve_128>;

} // namespace matmul
} // namespace aarch64
} // namespace cpu
} // namespace impl
} // namespace dnnl