opr_impl.cpp 48.5 KB
Newer Older
1 2 3 4
/**
 * \file dnn/src/arm_common/elemwise_multi_type/opr_impl.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8 9 10 11 12 13
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
 */

#include "./opr_impl.h"
14
#include "kernels.h"
15 16 17
#include "src/common/elemwise_multi_type/kern_defs.cuh"
#include "src/naive/handle.h"

18
#include "src/arm_common/elemwise_helper/elemwise_op.h"
19 20 21 22 23 24 25
#include "src/arm_common/simd_macro/marm_neon.h"

namespace {

using namespace megdnn;

template <int k>
M
Megvii Engine Team 已提交
26 27
void neon_round_shr_saturate_int16_static_k(
        const int16_t* a_ptr, size_t size, int8_t* dst_ptr) {
28 29 30 31 32 33 34 35 36 37 38 39 40 41
    static_assert(k >= 1 && k <= 8, "Shift offset out of range");
    size_t i = 0;
    int16x8_t x0, x1, f0, f1;
    for (; i + 15 < size; i += 16, a_ptr += 16, dst_ptr += 16) {
        x0 = vld1q_s16(a_ptr);
        x1 = vld1q_s16(a_ptr + 8);
        f0 = vshrq_n_s16(x0, 15);
        f1 = vshrq_n_s16(x1, 15);
        x0 = vqaddq_s16(x0, f0);
        x1 = vqaddq_s16(x1, f1);
        vst1_s8(dst_ptr, vqrshrn_n_s16(x0, k));
        vst1_s8(dst_ptr + 8, vqrshrn_n_s16(x1, k));
    }
    for (; i < size; i++, a_ptr++, dst_ptr++) {
M
Megvii Engine Team 已提交
42
        *dst_ptr = megdnn::elemwise_multi_type::round_shr_saturate<int16_t, int8_t>(
43 44 45 46 47 48
                *a_ptr, k);
    }
}

}  // namespace

49 50
using namespace elemwise;

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
namespace megdnn {
namespace arm_common {

template <typename stype>
void ElemwiseMultiTypeImpl::neon_round_shr_saturate_bcast_scalar(
        const stype* a_ptr, int8_t k, size_t size, dt_int8* dst_ptr) {
    MEGDNN_MARK_USED_VAR(a_ptr);
    MEGDNN_MARK_USED_VAR(k);
    MEGDNN_MARK_USED_VAR(size);
    MEGDNN_MARK_USED_VAR(dst_ptr);
    megdnn_throw(
            "ElemwiseMultiType (mode=ROUND_SHR_SATURATE) only supports int8, "
            "int16 and int32 on ARM");
}

template <>
void ElemwiseMultiTypeImpl::neon_round_shr_saturate_bcast_scalar<int8_t>(
        const int8_t* a_ptr, int8_t k, size_t size, dt_int8* dst_ptr) {
    size_t i = 0;
    const int8x16_t shift_vec = vdupq_n_s8(-k);
    int8x16_t x0, x1, f0, f1;
    for (; i + 31 < size; i += 32, a_ptr += 32, dst_ptr += 32) {
        x0 = vld1q_s8(a_ptr);
        x1 = vld1q_s8(a_ptr + 16);
        f0 = vshrq_n_s8(x0, 7);
        f1 = vshrq_n_s8(x1, 7);
        x0 = vqaddq_s8(x0, f0);
        x1 = vqaddq_s8(x1, f1);
        vst1q_s8(dst_ptr, vrshlq_s8(x0, shift_vec));
        vst1q_s8(dst_ptr + 16, vrshlq_s8(x1, shift_vec));
    }
    for (; i < size; i++, a_ptr++, dst_ptr++) {
M
Megvii Engine Team 已提交
83
        *dst_ptr = elemwise_multi_type::round_shr_saturate<int8_t, int8_t>(*a_ptr, k);
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
    }
}

template <>
void ElemwiseMultiTypeImpl::neon_round_shr_saturate_bcast_scalar<int16_t>(
        const int16_t* a_ptr, int8_t k, size_t size, dt_int8* dst_ptr) {
    // vqrshrn_n_s16 is significantly faster than vrshlq_s16 + vqmovn_s16, but
    // it requires that shift offset is known at compile time.
    switch (k) {
#define DISPATCH(i)                                                      \
    case i:                                                              \
        neon_round_shr_saturate_int16_static_k<i>(a_ptr, size, dst_ptr); \
        return;
        DISPATCH(1)
        DISPATCH(2)
        DISPATCH(3)
        DISPATCH(4)
        DISPATCH(5)
        DISPATCH(6)
        DISPATCH(7)
        DISPATCH(8)
#undef DISPATCH
        default:
            break;
    }

    size_t i = 0;
    const int16x8_t shift_vec = vdupq_n_s16(-k);
    int16x8_t x0, x1, f0, f1;
    for (; i + 15 < size; i += 16, a_ptr += 16, dst_ptr += 16) {
        x0 = vld1q_s16(a_ptr);
        x1 = vld1q_s16(a_ptr + 8);
        f0 = vshrq_n_s16(x0, 15);
        f1 = vshrq_n_s16(x1, 15);
        x0 = vqaddq_s16(x0, f0);
        x1 = vqaddq_s16(x1, f1);
        vst1_s8(dst_ptr, vqmovn_s16(vrshlq_s16(x0, shift_vec)));
        vst1_s8(dst_ptr + 8, vqmovn_s16(vrshlq_s16(x1, shift_vec)));
    }
    for (; i < size; i++, a_ptr++, dst_ptr++) {
M
Megvii Engine Team 已提交
124
        *dst_ptr = elemwise_multi_type::round_shr_saturate<int16_t, int8_t>(*a_ptr, k);
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141
    }
}

template <>
void ElemwiseMultiTypeImpl::neon_round_shr_saturate_bcast_scalar<int32_t>(
        const int32_t* a_ptr, int8_t k, size_t size, dt_int8* dst_ptr) {
    size_t i = 0;
    const int32x4_t shift_vec = vdupq_n_s32(-k);
    int32x4_t x0, x1, f0, f1;
    int8x8_t o0;
    for (; i + 7 < size; i += 8, a_ptr += 8, dst_ptr += 8) {
        x0 = vld1q_s32(a_ptr);
        x1 = vld1q_s32(a_ptr + 4);
        f0 = vshrq_n_s32(x0, 31);
        f1 = vshrq_n_s32(x1, 31);
        x0 = vqaddq_s32(x0, f0);
        x1 = vqaddq_s32(x1, f1);
M
Megvii Engine Team 已提交
142 143 144
        o0 = vqmovn_s16(vcombine_s16(
                vqmovn_s32(vrshlq_s32(x0, shift_vec)),
                vqmovn_s32(vrshlq_s32(x1, shift_vec))));
145 146 147
        vst1_s8(dst_ptr, o0);
    }
    for (; i < size; i++, a_ptr++, dst_ptr++) {
M
Megvii Engine Team 已提交
148
        *dst_ptr = elemwise_multi_type::round_shr_saturate<int32_t, int8_t>(*a_ptr, k);
149 150 151 152 153
    }
}

template <typename ctype>
void ElemwiseMultiTypeImpl::dispatch_round_shr_saturate_iXxi8xi8_bcast_scalar(
154
        const ElemwiseOpParamN<2>& param, const TensorND& dst) {
155 156
    auto k = param[1].ptr<dt_int8>()[0];
    size_t size = param.size;
157
    auto src = param[0];
158

159 160
    MEGDNN_DISPATCH_CPU_KERN_OPR(neon_round_shr_saturate_bcast_scalar(
            src.ptr<ctype>(), k, size, static_cast<dt_int8*>(dst.raw_ptr())));
161 162 163
}

void ElemwiseMultiTypeImpl::on_round_shr_saturate_iXxi8xi8(
164
        const ElemwiseOpParamN<2>& param, const TensorND& dst) {
165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184
    if (is_vector(param[0].layout) && is_broadcasted_scalar(param[1].layout)) {
        switch (param[0].layout.dtype.enumv()) {
#define cb(t)                                                     \
    case DTypeTrait<t>::enumv:                                    \
        return dispatch_round_shr_saturate_iXxi8xi8_bcast_scalar< \
                DTypeTrait<t>::ctype>(param, dst);
            MEGDNN_FOREACH_COMPUTING_DTYPE_INT(cb)
#undef cb
            default:
                megdnn_throw(
                        "ElemwiseMultiType (mode=ROUND_SHR_SATURATE) only "
                        "supports int8, int16 and int32 on ARM");
        }
    }

    fallback::ElemwiseMultiTypeImpl::on_round_shr_saturate_iXxi8xi8(param, dst);
}

void neon_fuse_add_rmulh_round_shr_saturate_bcast_1c11_int16(
        size_t batch_size, size_t channel_size, size_t channel_stride,
M
Megvii Engine Team 已提交
185 186
        const int16_t* x_ptr, const int16_t* b_ptr, const int16_t M, const int offset,
        const int8_t minv, const int8_t maxv, size_t size, int8_t* dst_ptr) {
187 188 189 190 191 192 193 194 195 196 197 198
    MEGDNN_MARK_USED_VAR(size);
    const int16x8_t shift_vec = vdupq_n_s16(-offset);
    const int16x8_t M_vec = vdupq_n_s16(M);
    const int8x16_t minv_vec = vdupq_n_s8(minv);
    const int8x16_t maxv_vec = vdupq_n_s8(maxv);

    size_t i = 0, b_pos = 0, channel_offset = 0;
    for (size_t batch = 0; batch < batch_size; ++batch) {
        b_pos = 0;
        for (size_t chan = 0; chan < channel_size; ++chan, ++b_pos) {
            auto b_vec = vdupq_n_s16(b_ptr[b_pos]);
            channel_offset += channel_stride;
M
Megvii Engine Team 已提交
199
            for (; i + 15 < channel_offset; i += 16, x_ptr += 16, dst_ptr += 16) {
200 201 202 203 204 205 206 207 208 209 210
                auto x0 = vld1q_s16(x_ptr);
                auto x1 = vld1q_s16(x_ptr + 8);
                x0 = vaddq_s16(x0, b_vec);
                x1 = vaddq_s16(x1, b_vec);
                x0 = vqrdmulhq_s16(x0, M_vec);
                x1 = vqrdmulhq_s16(x1, M_vec);
                // FIXME Theoretically, we should check shift != 0 here,
                auto fixup0 = vshrq_n_s16(x0, 15);
                auto fixup1 = vshrq_n_s16(x1, 15);
                x0 = vqaddq_s16(x0, fixup0);
                x1 = vqaddq_s16(x1, fixup1);
M
Megvii Engine Team 已提交
211 212 213
                auto o0 = vcombine_s8(
                        vqmovn_s16(vrshlq_s16(x0, shift_vec)),
                        vqmovn_s16(vrshlq_s16(x1, shift_vec)));
214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232
                o0 = vminq_s8(o0, maxv_vec);
                o0 = vmaxq_s8(o0, minv_vec);
                vst1q_s8(dst_ptr, o0);
            }
            for (; i + 7 < channel_offset; i += 8, x_ptr += 8, dst_ptr += 8) {
                auto x0 = vld1q_s16(x_ptr);
                x0 = vaddq_s16(x0, b_vec);
                x0 = vqrdmulhq_s16(x0, M_vec);
                // FIXME Theoretically, we should check shift != 0 here,
                auto fixup0 = vshrq_n_s16(x0, 15);
                x0 = vqaddq_s16(x0, fixup0);
                auto o0 = vqmovn_s16(vrshlq_s16(x0, shift_vec));
                o0 = vmin_s8(o0, vget_low_s8(maxv_vec));
                o0 = vmax_s8(o0, vget_low_s8(minv_vec));
                vst1_s8(dst_ptr, o0);
            }
            dt_int16 bias = b_ptr[b_pos];
            for (; i < channel_offset; ++i, ++x_ptr, ++dst_ptr) {
                dt_int16 result = rounding_shift_right_away_from_zero(
M
Megvii Engine Team 已提交
233 234 235
                        round_mulh_saturate<dt_int16>(*x_ptr + bias, M), offset);
                *dst_ptr = static_cast<dt_int8>(
                        std::max<dt_int16>(std::min<dt_int16>(result, maxv), minv));
236 237 238 239 240 241 242
            }
        }
    }
}

void neon_fuse_add_rmulh_round_shr_saturate_bcast_1c11_int32(
        size_t batch_size, size_t channel_size, size_t channel_stride,
M
Megvii Engine Team 已提交
243 244
        const int32_t* x_ptr, const int32_t* b_ptr, const int32_t M, const int offset,
        const int8_t minv, const int8_t maxv, size_t size, int8_t* dst_ptr) {
245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
    MEGDNN_MARK_USED_VAR(size);
    const int32x4_t shift_vec = vdupq_n_s32(-offset);
    const int32x4_t M_vec = vdupq_n_s32(M);
    const int8x8_t minv_vec = vdup_n_s8(minv);
    const int8x8_t maxv_vec = vdup_n_s8(maxv);

    size_t i = 0, b_pos = 0, channel_offset = 0;
    for (size_t batch = 0; batch < batch_size; ++batch) {
        b_pos = 0;
        for (size_t chan = 0; chan < channel_size; ++chan, ++b_pos) {
            int32x4_t b_vec = vdupq_n_s32(b_ptr[b_pos]);
            channel_offset += channel_stride;
            for (; i + 7 < channel_offset; i += 8, x_ptr += 8, dst_ptr += 8) {
                auto x0 = vld1q_s32(x_ptr);
                auto x1 = vld1q_s32(x_ptr + 4);
                x0 = vaddq_s32(x0, b_vec);
                x1 = vaddq_s32(x1, b_vec);
                x0 = vqrdmulhq_s32(x0, M_vec);
                x1 = vqrdmulhq_s32(x1, M_vec);
                // FIXME Theoretically, we should check shift != 0 here,
                auto fixup0 = vshrq_n_s32(x0, 31);
                auto fixup1 = vshrq_n_s32(x1, 31);
                x0 = vqaddq_s32(x0, fixup0);
                x1 = vqaddq_s32(x1, fixup1);
                auto o0 = vqmovn_s32(vrshlq_s32(x0, shift_vec));
                auto o1 = vqmovn_s32(vrshlq_s32(x1, shift_vec));
                auto of = vqmovn_s16(vcombine_s16(o0, o1));
                of = vmin_s8(of, maxv_vec);
                of = vmax_s8(of, minv_vec);
                vst1_s8(dst_ptr, of);
            }
            dt_int32 bias = b_ptr[b_pos];
            for (; i < channel_offset; ++i, ++x_ptr, ++dst_ptr) {
                dt_int32 result = rounding_shift_right_away_from_zero(
M
Megvii Engine Team 已提交
279 280 281
                        round_mulh_saturate<dt_int32>(*x_ptr + bias, M), offset);
                *dst_ptr = static_cast<dt_int8>(
                        std::max<dt_int32>(std::min<dt_int32>(result, maxv), minv));
282 283 284 285 286 287
            }
        }
    }
}

bool ElemwiseMultiTypeImpl::dispatch_fuse_add_rmulh_rshr(
288
        const ElemwiseOpParamN<6>& param, const TensorND& dst) {
289 290 291 292 293 294 295 296 297 298 299
    BroadcastChannelInfo binfo;
    if (is_vector(param[0].layout) &&
        is_broadcasted_channel_like(param[1].layout, binfo) &&
        is_broadcasted_scalar(param[2].layout) &&
        is_broadcasted_scalar(param[3].layout) &&
        is_broadcasted_scalar(param[4].layout) &&
        is_broadcasted_scalar(param[5].layout)) {
        auto offset = param[3].ptr<dt_int8>()[0];
        auto minv = param[4].ptr<dt_int8>()[0];
        auto maxv = param[5].ptr<dt_int8>()[0];
        switch (param[0].layout.dtype.enumv()) {
300 301 302 303 304 305 306 307 308 309 310 311
#define DISPATCH(stype, suffix)                                                      \
    case DTypeTrait<stype>::enumv: {                                                 \
        auto M = param[2].ptr<DTypeTrait<stype>::ctype>()[0];                        \
        auto src0 = param[0];                                                        \
        auto src1 = param[1];                                                        \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                \
                neon_fuse_add_rmulh_round_shr_saturate_bcast_1c11_##suffix(          \
                        binfo.x, binfo.y, binfo.z,                                   \
                        src0.ptr<DTypeTrait<stype>::ctype>(),                        \
                        src1.ptr<DTypeTrait<stype>::ctype>(), M, offset, minv, maxv, \
                        param.size, static_cast<dt_int8*>(dst.raw_ptr())));          \
        break;                                                                       \
312 313 314 315 316 317 318 319 320 321 322 323 324
    }
            DISPATCH(dtype::Int16, int16)
            DISPATCH(dtype::Int32, int32)
            default:
                megdnn_throw("unreachable");
        }
        return true;
    }
    return false;
#undef DISPATCH
}

void ElemwiseMultiTypeImpl::on_fuse_add_rmulh_round_shr_saturate_int16x16x16x8(
325
        const ElemwiseOpParamN<6>& param, const TensorND& dst) {
326 327
    if (dispatch_fuse_add_rmulh_rshr(param, dst))
        return;
M
Megvii Engine Team 已提交
328 329
    fallback::ElemwiseMultiTypeImpl::on_fuse_add_rmulh_round_shr_saturate_int16x16x16x8(
            param, dst);
330 331 332
}

void ElemwiseMultiTypeImpl::on_fuse_add_rmulh_round_shr_saturate_int32x32x32x8(
333
        const ElemwiseOpParamN<6>& param, const TensorND& dst) {
334 335
    if (dispatch_fuse_add_rmulh_rshr(param, dst))
        return;
M
Megvii Engine Team 已提交
336 337
    fallback::ElemwiseMultiTypeImpl::on_fuse_add_rmulh_round_shr_saturate_int32x32x32x8(
            param, dst);
338 339
}

M
Megvii Engine Team 已提交
340 341
void ElemwiseMultiTypeImpl::on_quantized_mode(
        const ElemwiseOpParamN<1>& param, const TensorND& dst, Elemwise::Mode mode) {
342 343 344
    megdnn_assert(param[0].layout.dtype.category() == DTypeCategory::QUANTIZED);
    megdnn_assert(dst.layout.dtype.category() == DTypeCategory::QUANTIZED);

M
Megvii Engine Team 已提交
345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366
#define DISPATCH_MODE(_src_dt, _dst_dt)                                           \
    switch (mode) {                                                               \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::RELU, ReluOp)      \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::H_SWISH, HSwishOp) \
        default:                                                                  \
            break;                                                                \
    }

#define DISPATCH_QUANTIZED_MODE(_src_dt, _dst_dt)                                     \
    switch (mode) {                                                                   \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::RELU, ReluOp)          \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::ABS, AbsOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::SIGMOID, SigmoidOp)    \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::EXP, ExpOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::TANH, TanhOp)          \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::FAST_TANH, FastTanhOp) \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::H_SWISH, HSwishOp)     \
        default:                                                                      \
            break;                                                                    \
    }

#define DISPATCH()                                                              \
367 368
    if (param[0].layout.dtype.enumv() == DTypeEnum::Quantized8Asymm &&          \
        dst.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) {               \
M
Megvii Engine Team 已提交
369 370 371 372 373
        DISPATCH_QUANTIZED_MODE(dtype::Quantized8Asymm, dtype::Quantized8Asymm) \
    } else if (                                                                 \
            param[0].layout.dtype.enumv() == DTypeEnum::QuantizedS32 &&         \
            dst.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) {           \
        DISPATCH_MODE(dtype::QuantizedS32, dtype::Quantized8Asymm)              \
374 375 376 377 378
    }

    TensorND src = param[0];

    size_t nr_elems = src.layout.total_nr_elems();
M
Megvii Engine Team 已提交
379 380 381 382 383 384 385 386 387 388
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                            \
    case _mode: {                                                                     \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                        \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                        \
        thin_function<void(const src_ctype*, dst_ctype*, DType, DType, size_t)> run = \
                OpCallerUnary<_op<src_ctype, dst_ctype>, VEC>::run;                   \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                 \
                run(src.ptr<src_ctype>(), dst.ptr<dst_ctype>(), src.layout.dtype,     \
                    dst.layout.dtype, nr_elems));                                     \
        return;                                                                       \
389 390 391 392 393 394 395 396 397 398 399 400
    }

    DISPATCH()

    fallback::ElemwiseMultiTypeImpl::on_quantized_mode(param, dst, mode);

#undef DISPATCH_SINGLE_MODE
#undef DISPATCH
#undef DISPATCH_QUANTIZED_MODE
#undef DISPATCH_MODE
}

M
Megvii Engine Team 已提交
401 402 403 404 405
void ElemwiseMultiTypeImpl::on_quantized_mode(
        const ElemwiseOpParamN<2>& param, const TensorND& dst, Elemwise::Mode mode) {
    megdnn_assert(
            param[0].layout.dtype.enumv() == param[1].layout.dtype.enumv() &&
            param[0].layout.dtype.category() == DTypeCategory::QUANTIZED);
406 407
    megdnn_assert(dst.layout.dtype.category() == DTypeCategory::QUANTIZED);

M
Megvii Engine Team 已提交
408 409 410 411 412 413 414 415 416
#define DISPATCH_MODE(_src_dt, _dst_dt)                                              \
    switch (mode) {                                                                  \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::ADD, AddOp)           \
        DISPATCH_SINGLE_MODE(                                                        \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_RELU, FuseAddReluOp)      \
        DISPATCH_SINGLE_MODE(                                                        \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_H_SWISH, FuseAddHSwishOp) \
        default:                                                                     \
            break;                                                                   \
417 418 419
    }

#if MEGDNN_AARCH64
M
Megvii Engine Team 已提交
420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437
#define DISPATCH_QUANTIZED_MODE(_src_dt, _dst_dt)                                     \
    switch (mode) {                                                                   \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::ADD, AddOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MIN, MinOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MAX, MaxOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::SUB, SubOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MUL, MulOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::TRUE_DIV, TrueDivOp)   \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_RELU, FuseAddReluOp)       \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_SIGMOID, FuseAddSigmoidOp) \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_TANH, FuseAddTanhOp)       \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_H_SWISH, FuseAddHSwishOp)  \
        default:                                                                      \
            break;                                                                    \
438 439
    }
#else
M
Megvii Engine Team 已提交
440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456
#define DISPATCH_QUANTIZED_MODE(_src_dt, _dst_dt)                                     \
    switch (mode) {                                                                   \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::ADD, AddOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MIN, MinOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MAX, MaxOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::SUB, SubOp)            \
        DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, Elemwise::Mode::MUL, MulOp)            \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_RELU, FuseAddReluOp)       \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_SIGMOID, FuseAddSigmoidOp) \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_TANH, FuseAddTanhOp)       \
        DISPATCH_SINGLE_MODE(                                                         \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_ADD_H_SWISH, FuseAddHSwishOp)  \
        default:                                                                      \
            break;                                                                    \
457 458 459
    }
#endif

M
Megvii Engine Team 已提交
460 461
#define DISPATCH()                                                              \
    if (param[0].layout.dtype.enumv() == DTypeEnum::QuantizedS32 &&             \
462
        dst.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) {               \
M
Megvii Engine Team 已提交
463 464 465 466 467
        DISPATCH_MODE(dtype::QuantizedS32, dtype::Quantized8Asymm)              \
    } else if (                                                                 \
            param[0].layout.dtype.enumv() == DTypeEnum::Quantized8Asymm &&      \
            dst.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) {           \
        DISPATCH_QUANTIZED_MODE(dtype::Quantized8Asymm, dtype::Quantized8Asymm) \
468 469 470 471 472 473 474 475
    }

    TensorND src0 = param[0];
    TensorND src1 = param[1];

    //! VEC + VEC
    if (is_vector(src0.layout) && is_vector(src1.layout)) {
        size_t nr_elems = src0.layout.total_nr_elems();
M
Megvii Engine Team 已提交
476 477 478 479 480 481 482 483 484 485 486 487
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                           \
    case _mode: {                                                                    \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                       \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                       \
        thin_function<void(                                                          \
                const src_ctype*, const src_ctype*, dst_ctype*, DType, DType, DType, \
                size_t)>                                                             \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, VEC_VEC>::run;       \
        MEGDNN_DISPATCH_CPU_KERN_OPR(run(                                            \
                src0.ptr<src_ctype>(), src1.ptr<src_ctype>(), dst.ptr<dst_ctype>(),  \
                src0.layout.dtype, src1.layout.dtype, dst.layout.dtype, nr_elems));  \
        return;                                                                      \
488 489 490 491 492 493 494 495 496
    }

        DISPATCH()

#undef DISPATCH_SINGLE_MODE
    }

    //! VEC + SCALAR
    {
M
Megvii Engine Team 已提交
497
        bool normal_case = is_vector(src0.layout) && is_broadcasted_scalar(src1.layout);
498 499 500 501 502
        bool swap_case = false;
        bool commutable = false;
        if (mode != Elemwise::Mode::SUB && mode != Elemwise::Mode::TRUE_DIV)
            commutable = true;
        if (!normal_case && commutable) {
M
Megvii Engine Team 已提交
503
            swap_case = is_vector(src1.layout) && is_broadcasted_scalar(src0.layout);
504 505 506 507 508
        }
        if (normal_case || swap_case) {
            auto &lhs = src0, &rhs = src1;
            if (swap_case)
                std::swap(lhs, rhs);
M
Megvii Engine Team 已提交
509 510 511 512 513 514 515 516 517 518 519 520 521
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                          \
    case _mode: {                                                                   \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                      \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                      \
        thin_function<void(                                                         \
                const src_ctype*, const src_ctype, dst_ctype*, DType, DType, DType, \
                size_t)>                                                            \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, VEC_SCALAR>::run;   \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                               \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>()[0],                \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,     \
                    dst.layout.dtype, src0.layout.total_nr_elems()));               \
        return;                                                                     \
522 523 524 525 526 527 528 529 530 531
    }

            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }

        //! SCALAR + VEC
        if (!commutable && is_vector(src1.layout) &&
            is_broadcasted_scalar(src0.layout)) {
M
Megvii Engine Team 已提交
532 533 534 535 536 537 538 539 540 541 542 543 544
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                          \
    case _mode: {                                                                   \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                      \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                      \
        thin_function<void(                                                         \
                const src_ctype, const src_ctype*, dst_ctype*, DType, DType, DType, \
                size_t)>                                                            \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, SCALAR_VEC>::run;   \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                               \
                run(src0.ptr<src_ctype>()[0], src1.ptr<src_ctype>(),                \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,     \
                    dst.layout.dtype, src1.layout.total_nr_elems()));               \
        return;                                                                     \
545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569
    }

            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }
    }

    //! VEC + BCAST101
    {
        BroadcastChannelInfo binfo;
        bool normal_case = is_vector(src0.layout) &&
                           is_broadcasted_channel_like(src1.layout, binfo);
        bool swap_case = false;
        bool commutable = false;
        if (mode != Elemwise::Mode::SUB && mode != Elemwise::Mode::TRUE_DIV)
            commutable = true;
        if (!normal_case && commutable) {
            swap_case = is_vector(src1.layout) &&
                        is_broadcasted_channel_like(src0.layout, binfo);
        }
        if (normal_case || swap_case) {
            auto &lhs = src0, &rhs = src1;
            if (swap_case)
                std::swap(lhs, rhs);
M
Megvii Engine Team 已提交
570 571 572 573 574 575 576 577 578 579 580 581 582
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                           \
    case _mode: {                                                                    \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                       \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                       \
        thin_function<void(                                                          \
                const src_ctype*, const src_ctype*, dst_ctype*, DType, DType, DType, \
                size_t, size_t, size_t)>                                             \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, VEC_BCAST101>::run;  \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                    \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,      \
                    dst.layout.dtype, binfo.x, binfo.y, binfo.z));                   \
        return;                                                                      \
583 584 585 586 587 588 589 590 591 592
    }

            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }

        //! BCAST101 + VEC : only for SUB or TRUE_DIV
        if (!commutable && is_vector(src1.layout) &&
            is_broadcasted_channel_like(src0.layout, binfo)) {
M
Megvii Engine Team 已提交
593 594 595 596 597 598 599 600 601 602 603 604 605
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                           \
    case _mode: {                                                                    \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                       \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                       \
        thin_function<void(                                                          \
                const src_ctype*, const src_ctype*, dst_ctype*, DType, DType, DType, \
                size_t, size_t, size_t)>                                             \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, BCAST101_VEC>::run;  \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                    \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,      \
                    dst.layout.dtype, binfo.x, binfo.y, binfo.z));                   \
        return;                                                                      \
606 607 608 609 610 611 612 613 614 615 616 617
    }

            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }
    }

    //! VEC + BCAST101x4
    {
        BroadcastChannelInfo binfo;
        if (is_vector(src0.layout) &&
618 619
            (is_broadcastedx_channel_like<4>(src1.layout, binfo) ||
             is_broadcastedx_channel_like<8>(src1.layout, binfo))) {
M
Megvii Engine Team 已提交
620 621 622 623 624 625 626 627 628 629 630 631 632 633 634
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                            \
    case _mode: {                                                                     \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                        \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                        \
        thin_function<void(                                                           \
                const src_ctype*, const src_ctype*, dst_ctype*, DType, DType, DType,  \
                size_t, size_t, size_t, size_t)>                                      \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, VEC_BCAST101xX>::run; \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                 \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                     \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,       \
                    dst.layout.dtype, batch_size, binfo.x, binfo.y, binfo.z));        \
        return;                                                                       \
    }
            size_t batch_size = src0.layout.shape[0] / (binfo.x * binfo.y * binfo.z);
635 636 637 638 639 640 641 642
            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }

        //! BCAST101x + VEC
        if (is_vector(src1.layout) &&
            is_broadcastedx_channel_like<4>(src0.layout, binfo)) {
M
Megvii Engine Team 已提交
643 644 645 646 647 648 649 650 651 652 653 654 655 656 657
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                            \
    case _mode: {                                                                     \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                        \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                        \
        thin_function<void(                                                           \
                const src_ctype*, const src_ctype*, dst_ctype*, DType, DType, DType,  \
                size_t, size_t, size_t, size_t)>                                      \
                run = OpCallerBinary<_op<src_ctype, dst_ctype>, BCAST101xX_VEC>::run; \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                 \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                     \
                    dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,       \
                    dst.layout.dtype, batch_size, binfo.x, binfo.y, binfo.z));        \
        return;                                                                       \
    }
            size_t batch_size = src1.layout.shape[0] / (binfo.x * binfo.y * binfo.z);
658 659 660 661 662 663 664 665 666 667 668 669 670
            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }
    }

    fallback::ElemwiseMultiTypeImpl::on_quantized_mode(param, dst, mode);

#undef DISPATCH_MODE
#undef DISPATCH_QUANTIZED_MODE
#undef DISPATCH
}

M
Megvii Engine Team 已提交
671 672
void ElemwiseMultiTypeImpl::on_quantized_mode(
        const ElemwiseOpParamN<3>& param, const TensorND& dst, Elemwise::Mode mode) {
673 674 675 676 677 678
    megdnn_assert(
            param[0].layout.dtype.enumv() == param[1].layout.dtype.enumv() &&
            param[0].layout.dtype.enumv() == param[2].layout.dtype.enumv() &&
            param[0].layout.dtype.category() == DTypeCategory::QUANTIZED);
    megdnn_assert(dst.layout.dtype.category() == DTypeCategory::QUANTIZED);

M
Megvii Engine Team 已提交
679 680 681 682 683 684
#define DISPATCH_QUANTIZED_MODE(_src_dt, _dst_dt)                               \
    switch (mode) {                                                             \
        DISPATCH_SINGLE_MODE(                                                   \
                _src_dt, _dst_dt, Elemwise::Mode::FUSE_MUL_ADD3, FuseMulAdd3Op) \
        default:                                                                \
            break;                                                              \
685 686
    }

M
Megvii Engine Team 已提交
687
#define DISPATCH()                                                              \
688 689
    if (param[0].layout.dtype.enumv() == DTypeEnum::Quantized8Asymm &&          \
        dst.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) {               \
M
Megvii Engine Team 已提交
690
        DISPATCH_QUANTIZED_MODE(dtype::Quantized8Asymm, dtype::Quantized8Asymm) \
691 692 693 694 695 696 697
    }

    TensorND src0 = param[0];
    TensorND src1 = param[1];
    TensorND src2 = param[2];

    //! VEC + VEC + VEC
M
Megvii Engine Team 已提交
698
    if (is_vector(src0.layout) && is_vector(src1.layout) && is_vector(src2.layout)) {
699
        size_t nr_elems = src0.layout.total_nr_elems();
M
Megvii Engine Team 已提交
700 701 702 703 704 705 706 707 708 709 710 711 712
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                           \
    case _mode: {                                                                    \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                       \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                       \
        thin_function<void(                                                          \
                const src_ctype*, const src_ctype*, const src_ctype*, dst_ctype*,    \
                DType, DType, DType, DType, size_t)>                                 \
                run = OpCallerTernary<_op<src_ctype, dst_ctype>, VEC_VEC_VEC>::run;  \
        MEGDNN_DISPATCH_CPU_KERN_OPR(run(                                            \
                src0.ptr<src_ctype>(), src1.ptr<src_ctype>(), src2.ptr<src_ctype>(), \
                dst.ptr<dst_ctype>(), src0.layout.dtype, src1.layout.dtype,          \
                src2.layout.dtype, dst.layout.dtype, nr_elems));                     \
        return;                                                                      \
713 714 715 716 717 718 719 720 721 722
    }

        DISPATCH()

#undef DISPATCH_SINGLE_MODE
    }

    //! VEC + VEC + SCALAR
    if (is_vector(src0.layout) && is_vector(src1.layout) &&
        is_broadcasted_scalar(src2.layout)) {
M
Megvii Engine Team 已提交
723 724 725 726 727 728 729 730 731 732 733 734 735 736
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                             \
    case _mode: {                                                                      \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                         \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                         \
        thin_function<void(                                                            \
                const src_ctype*, const src_ctype*, const src_ctype, dst_ctype*,       \
                DType, DType, DType, DType, size_t)>                                   \
                run = OpCallerTernary<_op<src_ctype, dst_ctype>, VEC_VEC_SCALAR>::run; \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                  \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                      \
                    src2.ptr<src_ctype>()[0], dst.ptr<dst_ctype>(), src0.layout.dtype, \
                    src1.layout.dtype, src2.layout.dtype, dst.layout.dtype,            \
                    src0.layout.total_nr_elems()));                                    \
        return;                                                                        \
737 738 739 740 741 742 743 744 745 746 747 748 749 750
    }

        DISPATCH()

#undef DISPATCH_SINGLE_MODE
    }

    //! BCAST101 + VEC + BCAST101
    {
        BroadcastChannelInfo binfo;
        bool normal_case = is_vector(src1.layout) &&
                           is_broadcasted_channel_like(src0.layout, binfo) &&
                           src0.layout.eq_shape(src2.layout);
        if (normal_case) {
M
Megvii Engine Team 已提交
751 752 753 754 755 756
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                           \
    case _mode: {                                                                    \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                       \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                       \
        thin_function<void(                                                          \
                const src_ctype*, const src_ctype*, const src_ctype*, dst_ctype*,    \
757
                DType, DType, DType, DType, size_t, size_t, size_t, size_t)>         \
M
Megvii Engine Team 已提交
758 759
                run = OpCallerTernary<                                               \
                        _op<src_ctype, dst_ctype>, BCAST101_VEC_BCAST101>::run;      \
760 761 762 763 764
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                                \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                    \
                    src2.ptr<src_ctype>(), dst.ptr<dst_ctype>(), src0.layout.dtype,  \
                    src1.layout.dtype, src2.layout.dtype, dst.layout.dtype, binfo.x, \
                    binfo.y, binfo.z, binfo.y* binfo.z));                            \
M
Megvii Engine Team 已提交
765
        return;                                                                      \
766 767 768 769
    }

            DISPATCH()

770 771 772 773 774 775 776 777
#undef DISPATCH_SINGLE_MODE
        }
    }

    //! VEC + BCAST101x4 + VEC
    {
        BroadcastChannelInfo binfo;
        if (is_vector(src0.layout) &&
778 779
            (is_broadcastedx_channel_like<4>(src1.layout, binfo) ||
             is_broadcastedx_channel_like<8>(src1.layout, binfo)) &&
780
            src0.layout.eq_shape(src2.layout)) {
M
Megvii Engine Team 已提交
781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                          \
    case _mode: {                                                                   \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                      \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                      \
        thin_function<void(                                                         \
                const src_ctype*, const src_ctype*, const src_ctype*, dst_ctype*,   \
                DType, DType, DType, DType, size_t, size_t, size_t, size_t)>        \
                run = OpCallerTernary<                                              \
                        _op<src_ctype, dst_ctype>, VEC_BCAST101xX_VEC>::run;        \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                               \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                   \
                    src2.ptr<src_ctype>(), dst.ptr<dst_ctype>(), src0.layout.dtype, \
                    src1.layout.dtype, src2.layout.dtype, dst.layout.dtype,         \
                    batch_size, binfo.x, binfo.y, binfo.z));                        \
        return;                                                                     \
    }

            size_t batch_size = src0.layout.shape[0] / (binfo.x * binfo.y * binfo.z);
799 800 801 802 803 804 805
            DISPATCH()

#undef DISPATCH_SINGLE_MODE
        }

        //! BCAST101x + VEC +BCAST101x
        if (is_vector(src1.layout) &&
806 807
            (is_broadcastedx_channel_like<4>(src0.layout, binfo) ||
             is_broadcastedx_channel_like<8>(src0.layout, binfo)) &&
808
            src0.layout.eq_shape(src2.layout)) {
M
Megvii Engine Team 已提交
809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826
#define DISPATCH_SINGLE_MODE(_src_dt, _dst_dt, _mode, _op)                          \
    case _mode: {                                                                   \
        using src_ctype = typename DTypeTrait<_src_dt>::ctype;                      \
        using dst_ctype = typename DTypeTrait<_dst_dt>::ctype;                      \
        thin_function<void(                                                         \
                const src_ctype*, const src_ctype*, const src_ctype*, dst_ctype*,   \
                DType, DType, DType, DType, size_t, size_t, size_t, size_t)>        \
                run = OpCallerTernary<                                              \
                        _op<src_ctype, dst_ctype>, BCAST101xX_VEC_BCAST101xX>::run; \
        MEGDNN_DISPATCH_CPU_KERN_OPR(                                               \
                run(src0.ptr<src_ctype>(), src1.ptr<src_ctype>(),                   \
                    src2.ptr<src_ctype>(), dst.ptr<dst_ctype>(), src0.layout.dtype, \
                    src1.layout.dtype, src2.layout.dtype, dst.layout.dtype,         \
                    batch_size, binfo.x, binfo.y, binfo.z));                        \
        return;                                                                     \
    }

            size_t batch_size = src1.layout.shape[0] / (binfo.x * binfo.y * binfo.z);
827 828
            DISPATCH()

829 830 831 832 833 834 835 836 837
#undef DISPATCH_SINGLE_MODE
        }
    }

    fallback::ElemwiseMultiTypeImpl::on_quantized_mode(param, dst, mode);
#undef DISPATCH
#undef DISPATCH_QUANTIZED_MODE
}

838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985
void ElemwiseMultiTypeImpl::on_fuse_mul_add3_int16xf32xf32xf32(
        const ElemwiseOpParamN<3>& param, const TensorND& dst) {
    auto &src0 = param[0], &src1 = param[1], &src2 = param[2];
    BroadcastChannelInfo binfo;

    if (is_vector(src0.layout) &&
        is_NHWC_broadcasted_channel_like(src1.layout, binfo) &&
        src1.layout.eq_layout(src2.layout)) {
        // VEC_BCAST111C_BCAST111C
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_int16xf32xf32xf32_vec_bcast111c_bcast111c(
                        binfo.x, binfo.y, binfo.z,
                        static_cast<dt_int16*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_broadcasted_channel_like(src1.layout, binfo) &&
            src1.layout.eq_layout(src2.layout)) {
        // VEC_BCAST101_BCAST101
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_int16xf32xf32xf32_vec_bcast101_bcast101(
                        binfo.x, binfo.y, binfo.z,
                        static_cast<dt_int16*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_vector(src1.layout) &&
            is_vector(src2.layout)) {
        // VEC_VEC_VEC
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_fuse_mul_add3_int16xf32xf32xf32_vec_vec_vec(
                size, static_cast<dt_int16*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()),
                static_cast<dt_float32*>(src2.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_broadcasted_scalar(src1.layout) &&
            is_broadcasted_scalar(src2.layout)) {
        // VEC_SCALAR_SCALAR
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_int16xf32xf32xf32_vec_scaler_scaler(
                        size, static_cast<dt_int16*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    }
    naive::ElemwiseMultiTypeImpl::on_fuse_mul_add3_int16xf32xf32xf32(param, dst);
}

void ElemwiseMultiTypeImpl::on_fuse_mul_add3_uint8xf32xf32xf32(
        const ElemwiseOpParamN<3>& param, const TensorND& dst) {
    auto &src0 = param[0], &src1 = param[1], &src2 = param[2];
    BroadcastChannelInfo binfo;

    if (is_vector(src0.layout) &&
        is_NHWC_broadcasted_channel_like(src1.layout, binfo) &&
        src1.layout.eq_layout(src2.layout)) {
        // VEC_BCAST111C_BCAST111C
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_uint8xf32xf32xf32_vec_bcast111c_bcast111c(
                        binfo.x, binfo.y, binfo.z,
                        static_cast<dt_uint8*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_broadcasted_channel_like(src1.layout, binfo) &&
            src1.layout.eq_layout(src2.layout)) {
        // VEC_BCAST101_BCAST101
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_uint8xf32xf32xf32_vec_bcast101_bcast101(
                        binfo.x, binfo.y, binfo.z,
                        static_cast<dt_uint8*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_vector(src1.layout) &&
            is_vector(src2.layout)) {
        // VEC_VEC_VEC
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_fuse_mul_add3_uint8xf32xf32xf32_vec_vec_vec(
                size, static_cast<dt_uint8*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()),
                static_cast<dt_float32*>(src2.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_broadcasted_scalar(src1.layout) &&
            is_broadcasted_scalar(src2.layout)) {
        // VEC_SCALAR_SCALAR
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(
                neon_fuse_mul_add3_uint8xf32xf32xf32_vec_scaler_scaler(
                        size, static_cast<dt_uint8*>(src0.raw_ptr()),
                        static_cast<dt_float32*>(src1.raw_ptr()),
                        static_cast<dt_float32*>(src2.raw_ptr()),
                        dst.ptr<dt_float32>()));
        return;
    }

    naive::ElemwiseMultiTypeImpl::on_fuse_mul_add3_uint8xf32xf32xf32(param, dst);
}

void ElemwiseMultiTypeImpl::on_mul_int16xf32xf32(
        const ElemwiseOpParamN<2>& param, const TensorND& dst) {
    auto &src0 = param[0], &src1 = param[1];
    BroadcastChannelInfo binfo;

    if (is_vector(src0.layout) &&
        is_NHWC_broadcasted_channel_like(src1.layout, binfo)) {
        // VEC_BCAST111C
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_mul_int16xf32xf32_vec_bcast111c(
                binfo.x, binfo.y, binfo.z, static_cast<dt_int16*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    } else if (
            is_vector(src0.layout) && is_broadcasted_channel_like(src1.layout, binfo)) {
        // VEC_BCAST101
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_mul_int16xf32xf32_vec_bcast101(
                binfo.x, binfo.y, binfo.z, static_cast<dt_int16*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    } else if (is_vector(src0.layout) && is_vector(src1.layout)) {
        // VEC_VEC
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_mul_int16xf32xf32_vec_vec(
                size, static_cast<dt_int16*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    } else if (is_vector(src0.layout) && is_broadcasted_scalar(src1.layout)) {
        auto size = param.size;
        MEGDNN_DISPATCH_CPU_KERN_OPR(neon_mul_int16xf32xf32_vec_scaler(
                size, static_cast<dt_int16*>(src0.raw_ptr()),
                static_cast<dt_float32*>(src1.raw_ptr()), dst.ptr<dt_float32>()));
        return;
    }

    naive::ElemwiseMultiTypeImpl::on_mul_int16xf32xf32(param, dst);
}

986 987 988 989
}  // namespace arm_common
}  // namespace megdnn

// vim: syntax=cpp.doxygen