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

/*******************************************************************************
* Copyright 2020 Intel Corporation
* 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 <cstdlib>
#include <memory>

#include "common/broadcast_strategy.hpp"
#include "common/c_types_map.hpp"
#include "common/verbose.hpp"
#include "cpu/platform.hpp"
#include "cpu/primitive_attr_postops.hpp"
#include "cpu/ref_io_helper.hpp"
#include "cpu/simple_q10n.hpp"

#if DNNL_X64
#include "cpu/x64/jit_gemm_x8s8s32x_convolution_utils.hpp"
#endif

#include "cpu/gemm_x8s8s32x_convolution_utils.hpp"

namespace dnnl {
namespace impl {
namespace cpu {
namespace gemm_x8s8s32x_convolution_utils {

#define VCHECK_PO_BOOL(cond, msg) \
    VCONDCHECK(primitive, create, check, gemm_x8s8s32x, cond, false, msg);

template <typename dst_data_t>
struct ref_pp_ker_t : pp_ker_t {
    ref_pp_ker_t(const convolution_pd_t *pd, const conv_gemm_conf_t &jcp)
        : pp_ker_t(pd, jcp), dst_md_(pd->dst_md()) {}

    using acc_data_t = pp_ker_t::acc_data_t;

    void operator()(void *dst, const acc_data_t *acc, const char *bias,
            const float *scales, float dst_scale, float sum_scale,
            float signed_scale, int g, int mb, size_t start, size_t end,
            const zero_point_call_params_t &zp,
            const void *post_ops_binary_rhs_arg_vec, const void *dst_orig,
            const exec_ctx_t &ctx, const memory_desc_t &dst_md,
            const single_gemm_conv_chunk_desc_t &chunk_desc) const override;

    status_t create_kernel() override {
        if (this->jcp_.with_eltwise || this->jcp_.with_binary) {
            ref_post_ops_
                    = utils::make_unique<ref_post_ops_t>(this->jcp_.post_ops);
            if (!ref_post_ops_) return status::out_of_memory;
            return ref_post_ops_->init(dst_md_);
        }
        return status::success;
    }

private:
    std::unique_ptr<ref_post_ops_t> ref_post_ops_;
    const memory_desc_t *dst_md_;
};

template <typename dst_data_t>
void ref_pp_ker_t<dst_data_t>::operator()(void *void_dst, const acc_data_t *acc,
        const char *bias, const float *scales, float dst_scale, float sum_scale,
        float signed_scale, int g, int mb, size_t start, size_t end,
        const zero_point_call_params_t &zp,
        const void * /* post_ops_binary_rhs_arg_vec */,
        const void * /* dst_orig */, const exec_ctx_t &ctx,
        const memory_desc_t &dst_md,
        const single_gemm_conv_chunk_desc_t &chunk_desc) const {

    if (end <= start) return;

    assert(data_traits_t<dst_data_t>::data_type == jcp_.dst_data_type);

    const lldiv_t dv_start = std::div((long long)start, (long long)jcp_.oc);
    const lldiv_t dv_end = std::div((long long)(end - 1), (long long)jcp_.oc);
    const size_t first_oc = dv_start.rem;
    const size_t last_oc = dv_end.rem;
    const size_t first_os = dv_start.quot;
    const size_t last_os = dv_end.quot;
    const int32_t zp_dst_val = jcp_.zp.dst_exists ? *(zp.dst) : 0;

    ref_post_ops_t::args_t args;
    args.ctx = &ctx;
    args.dst_md = &dst_md;

    for (size_t os = first_os; os <= last_os; os++) {
        const size_t start_oc = (os == first_os) ? first_oc : 0;
        const size_t end_oc = (os == last_os) ? last_oc : jcp_.oc - 1;
        for (size_t oc = start_oc; oc <= end_oc; oc++) {
            const size_t acc_off = os * jcp_.oc + oc;
            const size_t dst_off = os * jcp_.dst_os_stride + oc;

            int32_t data_s32 = acc[acc_off];

            if (jcp_.zp.src_exists) {
                const auto oc_offset = g * jcp_.oc + oc;
                data_s32 += zp.src_comp[oc_offset];
            }

            float data = static_cast<float>(data_s32);

            if (jcp_.signed_input) data *= signed_scale;

            // dequantize data
            data *= scales[(g * jcp_.oc + oc) * jcp_.scale_idx_mult];

            if (jcp_.with_bias) {
                const float b = io::load_float_value(
                        jcp_.bias_data_type, bias, g * jcp_.oc + oc);
                data += b;
            }

            if (jcp_.with_sum)
                data += sum_scale
                        * io::load_float_value(
                                jcp_.sum_data_type, void_dst, dst_off);
            if (jcp_.with_eltwise || jcp_.with_binary) {
                args.l_offset = ((mb * jcp_.ngroups + g) * jcp_.oc + oc)
                                * jcp_.os * jcp_.od
                        + os;
                ref_post_ops_->execute(data, args);
            }

            // quantize data
            if (jcp_.with_dst_scale) data *= dst_scale;
            if (jcp_.zp.dst_exists) data += static_cast<float>(zp_dst_val);

            io::store_float_value(jcp_.dst_data_type, data, void_dst, dst_off);
        }
    }
}

// Interface section

pp_ker_t::pp_ker_t(const convolution_pd_t *pd, const conv_gemm_conf_t &jcp)
    : jcp_(jcp) {}

pp_ker_t *pp_ker_t::create(
        const convolution_pd_t *pd, const conv_gemm_conf_t &jcp) {
#if DNNL_X64
    auto *res
            = x64::gemm_x8s8s32x_convolution_utils::jit_pp_ker_create(pd, jcp);
    if (res) return res;
#endif
    switch (pd->dst_md()->data_type) {
        case data_type::f32: return new ref_pp_ker_t<float>(pd, jcp);
        case data_type::bf16: return new ref_pp_ker_t<bfloat16_t>(pd, jcp);
        case data_type::s32: return new ref_pp_ker_t<int32_t>(pd, jcp);
        case data_type::s8: return new ref_pp_ker_t<int8_t>(pd, jcp);
        case data_type::u8: return new ref_pp_ker_t<uint8_t>(pd, jcp);
        default: assert(!"unexpected data type");
    }
    return nullptr;
}

bool post_ops_ok(const post_ops_t &post_ops, const memory_desc_wrapper *dst_d) {
// TODO: Align x64 and non-x64 support-scopes to avoid implementation surprises.
#if DNNL_X64
    return x64::gemm_x8s8s32x_convolution_utils::post_ops_ok(post_ops, dst_d);
#endif

    for (size_t po_index = 0; po_index < post_ops.entry_.size(); ++po_index) {
        const auto &post_op = post_ops.entry_[po_index];

        bool alg_kind_ok = post_op.is_eltwise() || post_op.is_sum()
                || post_op.is_prelu()
                || (post_op.is_binary()
                        && !post_op.is_binary_with_ternary_op());

        VCHECK_PO_BOOL(alg_kind_ok, "unsupported post-op alg kind");

        if (post_op.is_prelu()) {
            VCHECK_PO_BOOL(post_op.prelu.mask <= 3, "unsupported PReLU mask");

        } else if (post_op.is_binary()) {
            // The intent here is to limit the MASK_INPUT parameter of
            // binary attr-post-ops to {0, 1, 2, 3}. However, there is no
            // standard utility function for doing this, so we use the
            // one provided for broadcasting_strategy_t because the
            // logic is the same as it would be for MASK_INPUT.
            //
            // See the discussion in:
            // https://github.com/uxlfoundation/oneDNN/issues/3803
            const auto &bcast_type = get_rhs_arg_broadcasting_strategy(
                    post_op.binary.src1_desc, *dst_d,
                    {broadcasting_strategy_t::scalar,
                            broadcasting_strategy_t::per_mb,
                            broadcasting_strategy_t::per_oc,
                            broadcasting_strategy_t::spatial});

            VCHECK_PO_BOOL(bcast_type != broadcasting_strategy_t::unsupported,
                    "unsupported binary post-op mask")
        } else if (post_op.is_sum()) {
            VCHECK_PO_BOOL(
                    po_index == 0, "unsupported position for sum post-op")
        }
    }

    return true;
}

bool post_ops_ok(const post_ops_t &post_ops, const memory_desc_t *dst_d) {
    const auto dst_md = memory_desc_wrapper(dst_d);
    return post_ops_ok(post_ops, &dst_md);
}

bool mayiuse_jit_pp_kernel(data_type_t dst_dt) noexcept {
#if DNNL_X64
    return x64::gemm_x8s8s32x_convolution_utils::mayiuse_jit_pp_kernel(dst_dt);
#else
    return false;
#endif
}

} // namespace gemm_x8s8s32x_convolution_utils
} // namespace cpu
} // namespace impl
} // namespace dnnl