onednn-src 0.1.13

Source of oneAPI Deep Neural Network Library (oneDNN)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248

/*******************************************************************************
* Copyright 2022 Intel Corporation
* Copyright 2024 FUJITSU LIMITED
* Copyright 2025 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 "common/dnnl_thread.hpp"
#include "cpu/aarch64/cpu_isa_traits.hpp"

#include "cpu/aarch64/matmul/brgemm_matmul_reorders.hpp"

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

status_t brgemm_matmul_copy_reorder_t::pd_t::init(
        engine_t *engine, engine_t *src_engine, engine_t *dst_engine) {
    using namespace status;
    using namespace format_tag;

    status_t status = cpu_reorder_pd_t::init(engine, src_engine, dst_engine);
    if (status != success) return status;

    const memory_desc_wrapper id(src_md_), od(dst_md_);
    const int ndims = id.ndims();

    const auto type_i = id.data_type();
    const auto type_o = od.data_type();
    // TODO: enable support for type_i != type_o cases
    const bool dt_ok = true && type_i == type_o
            && utils::one_of(type_o, data_type::s8, data_type::bf16,
                    data_type::f16, data_type::f32);
    const bool is_f16 = utils::one_of(data_type::f16, type_i, type_o);
    const bool is_s8s8 = type_i == data_type::s8 && type_o == data_type::s8;
    const bool has_adj_scale
            = od.extra().flags & memory_extra_flags::scale_adjust;
    UNUSED(is_f16);
    UNUSED(is_s8s8);
    assert(!(is_f16 || is_s8s8));
    const bool args_ok = true && dt_ok && id.is_dense()
            && utils::one_of(ndims, 2, 3) && !has_adj_scale
            && attr()->has_default_values() && od.is_blocking_desc()
            && !od.has_runtime_dims_or_strides() && !od.has_zero_dim();
    if (!args_ok) return invalid_arguments;

    const auto &dims = id.dims();
    // TODO: enable for itag = {ba, acb}
    format_tag_t itag = id.matches_one_of_tag(ab, abc);
    format_tag_t otag = format_tag::undef;

    const auto vnni_granularity = data_type_vnni_granularity(type_o);
    switch (vnni_granularity) {
        case 4:
            otag = od.matches_one_of_tag(aCB16b64c4b, BA16a64b4a, aCB16b48c4b,
                    BA16a48b4a, aCB16b32c4b, BA16a32b4a, aCB16b16c4b,
                    BA16a16b4a);
            break;
        case 2:
            otag = od.matches_one_of_tag(aCB16b64c2b, BA16a64b2a, aCB16b48c2b,
                    BA16a48b2a, aCB16b32c2b, BA16a32b2a, aCB16b16c2b,
                    BA16a16b2a);
            break;
        case 1:
            otag = od.matches_one_of_tag(aCB16b64c, BA16a64b, aCB16b48c,
                    BA16a48b, aCB16b32c, BA16a32b, aCB16b16c, BA16a16b);
            break;
        default: otag = format_tag::undef;
    }

    if (utils::one_of(format_tag::undef, itag, otag)) return invalid_arguments;

    // initialize all required fields to generate copy_b kernel
    matmul_conf_for_reorder_.wei_tag = itag;
    matmul_conf_for_reorder_.batch = ndims > 2 ? dims[ndims - 3] : 1;
    matmul_conf_for_reorder_.K = dims[ndims - 2];
    matmul_conf_for_reorder_.N = dims[ndims - 1];
    matmul_conf_for_reorder_.wei_n_blk = matmul_conf_for_reorder_.N_blk
            = matmul_conf_for_reorder_.LDB
            = matmul::get_default_n_block(otag, matmul_conf_for_reorder_);
    matmul_conf_for_reorder_.N_tail
            = matmul_conf_for_reorder_.N % matmul_conf_for_reorder_.N_blk;
    matmul_conf_for_reorder_.K_blk = 16 * vnni_granularity;
    matmul_conf_for_reorder_.K_tail
            = matmul_conf_for_reorder_.K % matmul_conf_for_reorder_.K_blk;
    matmul_conf_for_reorder_.src_dt = matmul_conf_for_reorder_.wei_dt = type_o;
    matmul_conf_for_reorder_.a_dt_sz = matmul_conf_for_reorder_.tr_a_dt_sz
            = types::data_type_size(matmul_conf_for_reorder_.src_dt);
    matmul_conf_for_reorder_.b_dt_sz = matmul_conf_for_reorder_.tr_b_dt_sz
            = types::data_type_size(matmul_conf_for_reorder_.wei_dt);
    matmul_conf_for_reorder_.s8s8_comp_b_str = utils::rnd_up(
            matmul_conf_for_reorder_.N, matmul_conf_for_reorder_.wei_n_blk);
    matmul_conf_for_reorder_.s8s8_comp_n_str
            = matmul_conf_for_reorder_.wei_n_blk;
    matmul_conf_for_reorder_.s8s8_compensation_required
            = od.extra().flags & memory_extra_flags::compensation_conv_s8s8;
    const bool req_asymmetric_comp = od.extra().flags
            & memory_extra_flags::compensation_conv_asymmetric_src;
    matmul_conf_for_reorder_.src_zp_type = req_asymmetric_comp
            ? brgemm_broadcast_t::per_tensor
            : brgemm_broadcast_t::none;
    matmul_conf_for_reorder_.has_zero_point_a
            = matmul_conf_for_reorder_.src_zp_type != brgemm_broadcast_t::none;

    // asimd not supported, so we need >sve_128
    if (!mayiuse(sve_128)) return status::unimplemented;
    matmul_conf_for_reorder_.isa = get_max_cpu_isa();

    auto mask_ok = [&](bool check, int mask) {
        return IMPLICATION(
                check, mask == (1 << ndims) - 1 - (1 << (ndims - 2)));
    };

    const bool comp_masks_ok = true
            && mask_ok(matmul_conf_for_reorder_.s8s8_compensation_required,
                    od.extra().compensation_mask)
            && mask_ok(req_asymmetric_comp, od.extra().asymm_compensation_mask);
    if (!comp_masks_ok) return invalid_arguments;

    init_scratchpad();

    return status::success;
}

status_t brgemm_matmul_copy_reorder_t::pd_t::create(reorder_pd_t **reorder_pd,
        engine_t *engine, const primitive_attr_t *attr, engine_t *src_engine,
        const memory_desc_t *src_md, engine_t *dst_engine,
        const memory_desc_t *dst_md) {
    using namespace status;

    auto _pd = std::unique_ptr<pd_t>(new pd_t(
            attr, src_engine->kind(), src_md, dst_engine->kind(), dst_md));
    if (_pd == nullptr) return out_of_memory;
    CHECK(_pd->init(engine, src_engine, dst_engine));
    CHECK(_pd->init_scratchpad_md());
    return safe_ptr_assign<reorder_pd_t>(*reorder_pd, _pd.release());
}

status_t brgemm_matmul_copy_reorder_t::execute_body(
        const exec_ctx_t &ctx) const {
    using namespace utils;

    const auto src = CTX_IN_MEM(const char *, DNNL_ARG_FROM);
    auto dst = CTX_OUT_MEM(char *, DNNL_ARG_TO);
    const memory_desc_wrapper &src_d = pd()->src_md();
    const memory_desc_wrapper &dst_d = pd()->dst_md();
    const auto sdt_sz = types::data_type_size(src_d.data_type());
    const auto type_o = dst_d.data_type();
    const auto ddt_sz = types::data_type_size(type_o);

    const auto &kernel_conf = pd()->matmul_conf_for_reorder_;
    const size_t comp_offset_bytes
            = dst_d.size() - dst_d.additional_buffer_size();
    const size_t s8s8_comp_size_bytes = kernel_conf.s8s8_compensation_required
            ? dst_d.additional_buffer_size(
                      memory_extra_flags::compensation_conv_s8s8)
            : 0;
    const size_t zp_comp_offset_bytes
            = comp_offset_bytes + s8s8_comp_size_bytes;
    int32_t *cp = kernel_conf.s8s8_compensation_required
            ? reinterpret_cast<int32_t *>(dst + comp_offset_bytes)
            : nullptr;
    int32_t *zp = kernel_conf.has_zero_point_a
            ? reinterpret_cast<int32_t *>(dst + zp_comp_offset_bytes)
            : nullptr;

    const int ndims = src_d.ndims();
#define get_blk_off(md, dt_sz, batch, d0, d1) \
    (ndims == 3 ? (dt_sz) * (md).blk_off((batch), (d0), (d1)) \
                : (dt_sz) * (md).blk_off((d0), (d1)))

    parallel_nd(kernel_conf.batch, div_up(kernel_conf.N, kernel_conf.N_blk),
            [&](dim_t batch, dim_t n_blk_idx) {
        const auto n = n_blk_idx * kernel_conf.N_blk;
        const bool is_N_tail = (kernel_conf.N - n < kernel_conf.N_blk);
        auto ker_exec_ctx = matmul::jit_brgemm_matmul_copy_b_t::ctx_t();
        ker_exec_ctx.current_N_blk
                = is_N_tail ? kernel_conf.N_tail : kernel_conf.N_blk;

        assert(!is_runtime_value(kernel_conf.s8s8_comp_b_str));
        const auto comp_offset = batch * kernel_conf.s8s8_comp_b_str
                + n_blk_idx * kernel_conf.s8s8_comp_n_str;

        ker_exec_ctx.zp_a_compensation_ptr = kernel_conf.has_zero_point_a
                ? (void *)&zp[comp_offset]
                : nullptr;
        ker_exec_ctx.compensation_ptr = kernel_conf.s8s8_compensation_required
                ? (void *)&cp[comp_offset]
                : nullptr;

        // required to compute zp compensation
        int tmp_neg_a_zp_val = -1;
        ker_exec_ctx.zp_a_neg_value_ptr = &tmp_neg_a_zp_val;

        int k_blk_idx = 0;
        for (; k_blk_idx < kernel_conf.K / kernel_conf.K_blk; k_blk_idx++) {
            const auto k = k_blk_idx * kernel_conf.K_blk;
            ker_exec_ctx.src
                    = (void *)&src[get_blk_off(src_d, sdt_sz, batch, k, n)];
            ker_exec_ctx.tr_src = (void *)&dst[get_blk_off(
                    dst_d, ddt_sz, batch, k_blk_idx, n_blk_idx)];
            ker_exec_ctx.current_K_start = k;
            ker_exec_ctx.current_K_iters = kernel_conf.K_blk;
            (*kernel_)(&ker_exec_ctx);
        }
        if (kernel_conf.K_tail > 0) {
            const auto k = k_blk_idx * kernel_conf.K_blk;
            ker_exec_ctx.src
                    = (void *)&src[get_blk_off(src_d, sdt_sz, batch, k, n)];
            const auto dst_offset
                    = get_blk_off(dst_d, ddt_sz, batch, k_blk_idx, n_blk_idx);
            ker_exec_ctx.tr_src = (void *)&dst[dst_offset];
            ker_exec_ctx.current_K_start = k;
            ker_exec_ctx.current_K_iters = kernel_conf.K_tail;
            (*kernel_)(&ker_exec_ctx);
            const auto vnni_granularity = data_type_vnni_granularity(type_o);
            const auto dst_zero_out_offset
                    = rnd_up(kernel_conf.K_tail, vnni_granularity)
                    * kernel_conf.N_blk * ddt_sz;
            const auto elems_to_zero
                    = rnd_dn(kernel_conf.K_blk - kernel_conf.K_tail,
                              vnni_granularity)
                    * kernel_conf.N_blk * ddt_sz;
            array_set(&dst[dst_offset + dst_zero_out_offset], 0, elems_to_zero);
        }
    });

#undef get_blk_off

    return status::success;
}

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