opr_impl.cpp 35.0 KB
Newer Older
1
/**
2
 g * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
3
 *
4
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
5 6 7
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
8 9
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
10
 */
M
Megvii Engine Team 已提交
11
#include "src/fallback/conv_bias/opr_impl.h"
12 13 14 15 16
#include "src/common/algo_chooser.h"
#include "src/common/metahelper.h"
#include "src/common/opr_delegate.h"
#include "src/common/utils.h"
#include "src/fallback/conv_bias/algos.h"
17
#include "src/fallback/conv_bias/conv1x1/algos.h"
18
#include "src/fallback/conv_bias/conv1x1/algos_conv1x1_gemv.h"
19
#include "src/fallback/conv_bias/gi/fp16/algos.h"
20
#include "src/fallback/conv_bias/gi/fp32/algos.h"
21
#include "src/fallback/conv_bias/im2col/algos.h"
22
#include "src/fallback/convolution/opr_impl.h"
23
#include "src/fallback/general_intrinsic/gi_common.h"
24 25 26
#include "src/naive/convolution/algorithms.h"
#include "src/naive/handle.h"

27 28 29 30 31 32 33 34
#if MEGDNN_X86
#include "src/x86/conv_bias/opr_impl.h"
#elif MEGDNN_AARCH64
#include "src/aarch64/conv_bias/opr_impl.h"
#elif MEGDNN_ARMV7
#include "src/armv7/conv_bias/opr_impl.h"
#endif

35 36 37 38 39
#include <cstring>

using namespace megdnn;
using namespace fallback;

40 41 42 43 44 45 46 47
namespace {

//! TODO: imp is_fallback_exclude_gi_or_naive
bool is_naive(const detail::Algorithm* algo) {
    return algo->handle_type() == Handle::HandleType::NAIVE;
}
}  // anonymous namespace

48
size_t megdnn::fallback::pack_size(param::ConvBias::Format format) {
49
    switch (format) {
50
        case param::ConvBias::Format::NCHW44:
51
        case param::ConvBias::Format::NCHW44_DOT:
52 53 54 55 56 57 58 59 60
        case param::ConvBias::Format::NCHW4:
            return 4_z;
        case param::ConvBias::Format::NCHW88:
            return 8_z;
        default:
            return 1_z;
    }
}

61 62 63 64 65 66 67 68
namespace {
template <typename T>
void incr_ptr(T*& dst, ptrdiff_t delta) {
    dst = reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(dst) + delta);
}

}  // namespace

69 70 71 72 73 74 75 76 77 78 79 80 81
#if MEGDNN_X86
#define SKIP_GEMV()
//! As we haven't direct conv for int8x8x16 yet, if we disable gemv here, it may
//! fallback to naive implementation, which may cause performance very low, so
//! here we just enable im2col for gemv in x86 backend.
//! FIXME: remove it when we add direct conv support for int8x8x16
#else
#define SKIP_GEMV()                                                            \
    if (algo->algoset() == MatrixMulImpl::AlgoBase::AlgoSet::ALGO_TYPE_GEMV) { \
        continue;                                                              \
    }
#endif

82 83 84
class ConvBiasImpl::AlgoPack : NonCopyableObj {
    AlgoNaive algo_naive;
    SmallVector<std::unique_ptr<AlgoBase>> refhold;
85 86
    SmallVector<AlgoBase*> m_all_algos;
    AlgoBase::Mapper m_all_algos_map;
87 88
    SmallVector<fallback::ConvBiasImpl::AlgoBase*> m_gi_winograd_algos;

89 90
    AlgoF32DirectNCHWNCHW44 f32_nchw_nchw44;
    AlgoF32DirectNCHWNCHW44AGENT f32_nchw_nchw44_agent;
91 92 93 94 95 96
    AlgoF32ChannelWiseNCHW44 f32_chanel_wise_nchw44;
    AlgoF32DirectNCHW44 f32_direct_nchw44;

    AlgoF32Direct f32_direct;
    AlgoF32DirectStride2 f32_direct_stride2;
    AlgoF32DirectStride1 f32_direct_stride1;
97 98 99

public:
    AlgoPack() {
100 101 102 103 104 105 106 107 108 109 110
        //! fallback gi fp32 algo
        //! now f32_nchw_nchw44_agent is fast than f32_nchw_nchw44
        //! on x86 and rvv platform, so we adjust heuristic order.
#if MEGDNN_AARCH64 || MEGDNN_ARMV7
        m_all_algos.emplace_back(&f32_nchw_nchw44);
        m_all_algos.emplace_back(&f32_nchw_nchw44_agent);
#else
        m_all_algos.emplace_back(&f32_nchw_nchw44_agent);
        m_all_algos.emplace_back(&f32_nchw_nchw44);
#endif

111 112 113 114 115 116 117 118 119 120 121 122 123 124 125
        m_all_algos.emplace_back(&f32_chanel_wise_nchw44);
        m_all_algos.emplace_back(&f32_direct_nchw44);
        m_all_algos.emplace_back(&f32_direct_stride1);
        m_all_algos.emplace_back(&f32_direct_stride2);
        m_all_algos.emplace_back(&f32_direct);

        static CpuOprDelegationStorage<2> storage;
        auto matmul_opr = storage.get<MatrixMul, 0>();
        using MatmulFormat = param::MatrixMul::Format;
        auto&& matmul_algos =
                static_cast<fallback::MatrixMulImpl*>(matmul_opr)
                        ->select_algo_type({AlgoDataType::FLOAT32, MatmulFormat::MK4});
        for (auto&& algo : matmul_algos) {
            if (is_naive(algo))
                continue;
126
            for (uint32_t tile_size : {16, 8, 24, 32, 48, 68}) {
127 128 129 130 131 132 133 134
                refhold.emplace_back(new AlgoFP32WinogradF23_4x4(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
                refhold.emplace_back(new AlgoFP32WinogradF63_4x4(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
135 136 137 138
                refhold.emplace_back(new AlgoFP32WinogradF43_4x4(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
139 140 141 142
                refhold.emplace_back(new AlgoFP32WinogradF63_4x4_NCHW44(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
143 144 145 146
                refhold.emplace_back(new AlgoFP32WinogradF43_4x4_NCHW44(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
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
                refhold.emplace_back(new AlgoFP32WinogradF23_4x4_NCHW44(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
//! uncomment this when low precision mode is done
#if 0
                refhold.emplace_back(new AlgoFP32WinogradF73_4x4_NCHW44(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
#endif
            }
        }

        matmul_algos = static_cast<fallback::MatrixMulImpl*>(matmul_opr)
                               ->select_algo_type(
                                       {AlgoDataType::FLOAT32, MatmulFormat::DEFAULT});
        for (auto&& algo : matmul_algos) {
            if (is_naive(algo))
                continue;
            for (uint32_t tile_size : {16, 8, 24, 32}) {
                refhold.emplace_back(new AlgoFP32WinogradF63(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
172 173 174 175
                refhold.emplace_back(new AlgoFP32WinogradF43(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
176 177 178 179 180 181 182 183 184 185
                refhold.emplace_back(new AlgoFP32WinogradF54(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
                refhold.emplace_back(new AlgoFP32WinogradF45(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
            }
        }
186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201
        // end fallback gi fp32 algo

#if defined(GI_SUPPORT_F16)
        //! fallback gi fp16 algo
        matmul_algos =
                static_cast<fallback::MatrixMulImpl*>(matmul_opr)
                        ->select_algo_type({AlgoDataType::FLOAT16, MatmulFormat::MK8});
        for (auto&& algo : matmul_algos) {
            if (is_naive(algo))
                continue;

            for (uint32_t tile_size : {68, 16, 8, 24, 32, 48}) {
                refhold.emplace_back(new AlgoFP16WinogradF43_8x8_NCHW88(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
202 203 204 205 206

                refhold.emplace_back(new AlgoFP16WinogradF23_8x8_NCHW88(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
207 208 209 210 211

                refhold.emplace_back(new AlgoFP16WinogradF63_8x8_NCHW88(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
212 213 214
            }
        }
#endif
215 216 217 218
        for (auto&& algo : m_gi_winograd_algos) {
            m_all_algos.emplace_back(algo);
        }

219
        refhold.emplace_back(new AlgoConv1x1Gemv());
220
        m_all_algos.emplace_back(refhold.back().get());
221

222 223
        matmul_algos = static_cast<fallback::MatrixMulImpl*>(matmul_opr)
                               ->get_all_packed_algo();
224
        for (auto&& algo : matmul_algos) {
225 226 227 228 229 230
#if MEGDNN_X86
//! As we haven't direct conv for int8x8x16 yet, if we disable gemv here, it may
//! fallback to naive implementation, which may cause performance very low, so
//! here we just enable im2col for gemv in x86 backend.
//! FIXME: remove it when we add direct conv support for int8x8x16
#else
M
Megvii Engine Team 已提交
231
            if (algo->algoset() == MatrixMulImpl::AlgoBase::AlgoSet::ALGO_TYPE_GEMV) {
232 233
                continue;
            }
234 235
#endif

236 237
            for (size_t ohw_tile_size : {192, 384, 96, 48, 24}) {
                refhold.emplace_back(new AlgoIm2col(
M
Megvii Engine Team 已提交
238
                        static_cast<MatrixMulImpl::AlgoBase*>(algo), ohw_tile_size));
239
                m_all_algos.emplace_back(refhold.back().get());
240
            }
241
            for (size_t oc_tile_size : {48, 24}) {
242
                refhold.emplace_back(new AlgoConv1x1(
M
Megvii Engine Team 已提交
243
                        static_cast<MatrixMulImpl::AlgoBase*>(algo), oc_tile_size));
244
                m_all_algos.emplace_back(refhold.back().get());
245
            }
246

247
#if 0
248 249 250 251 252
        //! As these algos maybe very slow, it will make fastrun search slow, so
        //! we disable it, but for the test of strategyhelper, we just keep it.
        //! FIXME: I do not know a better way to do it.
            refhold.emplace_back(new AlgoWinogradF32(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
253
            m_all_algos.emplace_back(refhold.back().get());
254 255
            refhold.emplace_back(new AlgoWinogradF32_4x4(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
256
            m_all_algos.emplace_back(refhold.back().get());
257 258
            refhold.emplace_back(new AlgoWinogradQS8(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
259
            m_all_algos.emplace_back(refhold.back().get());
260 261
            refhold.emplace_back(new AlgoWinogradQS8_8x8(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
262
            m_all_algos.emplace_back(refhold.back().get());
263 264
#endif
        }
265 266 267 268 269
        m_all_algos.emplace_back(&algo_naive);

        for (auto&& algo : m_all_algos) {
            m_all_algos_map.emplace(algo->info().desc, algo);
        }
270
    }
271 272
    const SmallVector<AlgoBase*>& all_algos() const { return m_all_algos; }
    const AlgoBase::Mapper& all_algos_map() const { return m_all_algos_map; }
273 274
};

275 276 277 278 279 280 281
const ConvBiasImpl::AlgoPack& ConvBiasImpl::algo_pack() {
    static AlgoPack algo_pack;
    return algo_pack;
}

SmallVector<ConvBiasImpl::AlgoBase*> ConvBiasImpl::get_all_packed_algo() {
    return algo_pack().all_algos();
282
}
283 284 285

SmallVector<ConvBiasImpl::AlgoBase*> ConvBiasImpl::select_algo_type(
        ConvAlgoTypePack target_type) {
M
Megvii Engine Team 已提交
286 287 288
    megdnn_assert(
            nr_type_contain(target_type.data_type),
            "ConvBias algo selection only support one type");
289
    SmallVector<ConvBiasImpl::AlgoBase*> algos;
290
    for (auto&& algo : get_all_packed_algo()) {
291 292 293 294 295 296 297 298 299
        auto algo_type = algo->get_algo_type();
        if (contain_data_type(algo_type.data_type, target_type.data_type) &&
            algo_type.algo_category == target_type.algo_category) {
            algos.push_back(algo);
        }
    }
    return algos;
}

300 301 302
bool ConvBiasImpl::is_naive_algo(ConvBiasImpl::Algorithm* algo) {
    return algo == nullptr || strcmp(algo->name(), "DEFAULT") == 0;
}
303

M
Megvii Engine Team 已提交
304 305 306 307 308 309 310 311 312 313 314
#define NCB_ALGO_FUNC(name, algo, param) static_cast<AlgoBase*>(algo)->name(param)

void ConvBiasImpl::exec(
        _megdnn_tensor_in src, _megdnn_tensor_in filter, _megdnn_tensor_in bias,
        _megdnn_tensor_in z, _megdnn_tensor_out dst,
        const PreprocessedFilter* preprocessed_filter, _megdnn_workspace workspace) {
    check_exec(
            src.layout, filter.layout, bias.layout, z.layout, dst.layout,
            workspace.size, preprocessed_filter);
    auto fparam =
            make_ncb_kern_param(src, filter, bias, dst, workspace, preprocessed_filter);
315
    auto&& algo = get_algorithm(fparam, workspace.size);
316
    if (!is_naive_algo(algo) &&
317
        NCB_ALGO_FUNC(get_workspace, algo, fparam) <= workspace.size) {
318 319
        exec_with_ncb_kern(fparam, algo);
    } else {
M
Megvii Engine Team 已提交
320 321
        naive::ConvBiasForwardImpl::exec(
                src, filter, bias, z, dst, preprocessed_filter, workspace);
322 323 324
    }
}

M
Megvii Engine Team 已提交
325 326 327 328 329
void ConvBiasImpl::exec_preprocess(
        const TensorLayout& src_layout, _megdnn_tensor_in filter,
        _megdnn_tensor_in bias, const TensorLayout& z_layout,
        const TensorLayout& dst_layout, PreprocessedFilter* preprocessed_filter,
        _megdnn_workspace workspace) {
330 331 332
    //! exec_preprocess currently only support preprocess weights and bias
    //! before exec, src/dst/z will be ignored, just set to nullptr
    TensorND src{nullptr, src_layout}, dst{nullptr, dst_layout};
M
Megvii Engine Team 已提交
333 334
    auto fparam =
            make_ncb_kern_param(src, filter, bias, dst, workspace, preprocessed_filter);
335
    //! should not pass workspace_size limit otherwise can not find match algo
336 337
    auto&& algo = get_algorithm(fparam);
    if (!is_naive_algo(algo) &&
M
Megvii Engine Team 已提交
338
        NCB_ALGO_FUNC(get_preprocess_workspace, algo, fparam) <= workspace.size) {
339 340 341
        exec_preprocess_with_ncb_kern(fparam, algo);
    } else {
        naive::ConvBiasForwardImpl::exec_preprocess(
M
Megvii Engine Team 已提交
342 343
                src_layout, filter, bias, z_layout, dst_layout, preprocessed_filter,
                workspace);
344 345 346
    }
}

347
size_t ConvBiasImpl::get_workspace_in_bytes(
M
Megvii Engine Team 已提交
348 349
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst,
350
        const PreprocessedFilter* preprocessed_filter) {
351
    TensorLayoutArray layouts{src, filter, bias, z, dst};
352
    AlgorithmCache::Key key{this->handle(), this->get_opr_type(),
M
Megvii Engine Team 已提交
353 354
                            layouts.data(), layouts.size(),
                            &this->param(), sizeof(this->param())};
355
    auto rst = AlgorithmCache::instance().get(key);
356 357 358 359
    if (rst.policy.algo.valid()) {
        return rst.workspace;
    }

M
Megvii Engine Team 已提交
360
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, preprocessed_filter);
361
    auto&& algo = get_algorithm(fparam);
362
    if (is_naive_algo(algo)) {
363 364
        return naive::ConvBiasForwardImpl::get_workspace_in_bytes(
                src, filter, bias, z, dst, preprocessed_filter);
365
    } else {
366 367 368 369 370
        return NCB_ALGO_FUNC(get_workspace, algo, fparam);
    }
}

size_t ConvBiasImpl::get_preprocess_workspace_in_bytes(
M
Megvii Engine Team 已提交
371 372
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
373
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
374
    auto&& algo = get_algorithm(fparam);
375 376 377 378 379 380 381 382 383
    if (is_naive_algo(algo)) {
        return naive::ConvBiasForwardImpl::get_preprocess_workspace_in_bytes(
                src, filter, bias, z, dst);
    } else {
        return NCB_ALGO_FUNC(get_preprocess_workspace, algo, fparam);
    }
}

SmallVector<TensorLayout> ConvBiasImpl::deduce_preprocessed_filter_layout(
M
Megvii Engine Team 已提交
384 385
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
386
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
387
    auto&& algo = get_algorithm(fparam);
388 389 390 391 392
    if (is_naive_algo(algo)) {
        return naive::ConvBiasForwardImpl::deduce_preprocessed_filter_layout(
                src, filter, bias, z, dst);
    } else {
        return NCB_ALGO_FUNC(deduce_preprocessed_filter_layout, algo, fparam);
393 394 395 396
    }
}

std::vector<ConvBiasImpl::Algorithm*> ConvBiasImpl::get_all_algorithms(
M
Megvii Engine Team 已提交
397 398
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
399
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
400 401
    auto ret = get_all_algorithms_with_ncb(fparam);
    if (ret.empty()) {
M
Megvii Engine Team 已提交
402 403
        return naive::ConvBiasForwardImpl::get_all_algorithms_safe(
                src, filter, bias, z, dst);
404 405 406
    }
    return ret;
}
407
std::vector<ConvBiasImpl::Algorithm*> ConvBiasImpl::get_all_algorithms_safe(
M
Megvii Engine Team 已提交
408 409 410
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
    auto ret_safe = ConvBiasImpl::get_all_algorithms(src, filter, bias, z, dst);
411 412
    return ret_safe;
}
413 414

ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_heuristic(
M
Megvii Engine Team 已提交
415 416 417
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst, size_t workspace_limit_in_bytes,
        const AlgoAttribute& positive_attr, const AlgoAttribute& negative_attr) {
418
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
419
    auto result = get_algorithm_heuristic_with_ncb(
420
            fparam, workspace_limit_in_bytes, positive_attr, negative_attr);
421 422
    if (result == nullptr) {
        result = naive::ConvBiasForwardImpl::get_algorithm_heuristic(
M
Megvii Engine Team 已提交
423 424
                src, filter, bias, z, dst, workspace_limit_in_bytes, positive_attr,
                negative_attr);
425 426 427 428
    }
    return result;
}

429 430
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_heuristic_with_ncb(
        const NCBKernSizeParam& param, size_t workspace_limit_in_bytes,
M
Megvii Engine Team 已提交
431
        const AlgoAttribute& positive_attr, const AlgoAttribute& negative_attr) {
432 433 434
    if (ConvBiasImpl::param().format == Param::Format::NHWCD4) {
        return nullptr;
    }
435 436 437 438 439 440
    auto algo_data_type = param.deduce_algo_data_type();
    auto suggest_category_order = suggest_algo_category_order(param);
    for (auto category : suggest_category_order) {
        auto&& origin_algos = select_algo_type({algo_data_type, category});
        ConvBiasImpl::Algorithm* heuristic_algo = nullptr;
        for (auto i : origin_algos) {
441
            bool usable_attribute = static_cast<AlgoBase*>(i)->usable_attribute(
442 443
                    param, AlgoSelectionStrategy::HEURISTIC, positive_attr,
                    negative_attr);
M
Megvii Engine Team 已提交
444 445
            if (usable_attribute && static_cast<AlgoBase*>(i)->get_workspace(param) <=
                                            workspace_limit_in_bytes) {
446 447 448 449 450 451 452 453 454 455 456 457 458
                //! store the first usable algo if no prefer algo, choose it as
                //! the target algo
                if (!heuristic_algo) {
                    heuristic_algo = i;
                }
                //! choose the first prefer algo
                if (i->is_preferred(param)) {
                    return i;
                }
            }
        }
        if (heuristic_algo) {
            return heuristic_algo;
459 460 461 462 463
        }
    }
    return nullptr;
}

464
ConvBiasImpl::NCBKernSizeParam ConvBiasImpl::make_ncb_kern_size_param(
M
Megvii Engine Team 已提交
465 466
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& dst, const PreprocessedFilter* preprocessed_filter) {
467
    auto safe_u32 = [](size_t v) -> uint32_t {
M
Megvii Engine Team 已提交
468 469
        megdnn_assert(
                v <= std::numeric_limits<uint32_t>::max(), "value too large: %zu", v);
470 471 472 473 474 475
        return v;
    };
    size_t spatial_pos;
    if (param().format == Param::Format::NCHW88 ||
        param().format == Param::Format::NCHW8 ||
        param().format == Param::Format::NCHW4 ||
476
        param().format == Param::Format::NCHW44 ||
477
        param().format == Param::Format::NCHW44_DOT ||
478 479
        param().format == Param::Format::NCHW ||
        param().format == Param::Format::NCHW32 ||
480
        param().format == Param::Format::NCHW64) {
481
        spatial_pos = 2;
M
Megvii Engine Team 已提交
482 483 484
    } else if (
            param().format == Param::Format::NHWC ||
            param().format == Param::Format::NHWCD4) {
485 486
        spatial_pos = 1;
    } else {
M
Megvii Engine Team 已提交
487
        megdnn_assert(0, "invalid conv format %d", static_cast<int>(param().format));
488 489
    }
    BiasMode bias_mode;
490 491
    //! dst only channel BIAS is viewed as BROADCAST_CHANNEL_BIAS
    bool dst_only_c = dst[0] == 1 && dst[spatial_pos] == 1 && dst[spatial_pos + 1] == 1;
492 493
    if (bias.ndim == 0) {
        bias_mode = BiasMode::NO_BIAS;
494
    } else if (bias.eq_shape(dst) && !dst_only_c) {
495 496 497 498 499 500 501
        bias_mode = BiasMode::BIAS;
    } else {
        //! just check the ndim, the detail shape check is in check_exec
        megdnn_assert(bias.ndim == dst.ndim);
        bias_mode = BiasMode::BROADCAST_CHANNEL_BIAS;
    }

M
Megvii Engine Team 已提交
502 503 504 505
    static_assert(
            sizeof(CanonizedFilterMeta) == sizeof(ConvolutionImpl::CanonizedFilterMeta),
            "sizeof CanonizedFilterMeta in convolution and conv_bias "
            "should be equal");
506 507
    auto&& fm = check_layout_fwd(src, filter, dst);
    auto& conv_fm = reinterpret_cast<ConvolutionImpl::CanonizedFilterMeta&>(fm);
508

509 510 511 512 513 514 515 516 517 518 519 520 521 522 523
    size_t nr_threads = static_cast<naive::HandleImpl*>(handle())
                                ->megcore_dispatcher()
                                ->nr_threads();
    return {{safe_u32(src[0]),
             {{safe_u32(src[spatial_pos]), safe_u32(src[spatial_pos + 1])}},
             {{safe_u32(dst[spatial_pos]), safe_u32(dst[spatial_pos + 1])}},
             conv_fm,
             src.dtype,
             filter.dtype,
             dst.dtype,
             src.stride[0],
             dst.stride[0],
             {src.stride[0], src.stride[1], src.stride[2], src.stride[3]},
             {dst.stride[0], dst.stride[1], dst.stride[2], dst.stride[3]},
             param().compute_mode,
524 525
             nr_threads,
             reinterpret_cast<const ConvolutionForward::PreprocessedFilter*>(
526 527
                     preprocessed_filter),
             handle()},
528 529 530 531 532 533 534 535
            bias.dtype,
            bias.stride[0],
            bias_mode,
            param().nonlineMode};
}

ConvBiasImpl::NCBKernParam ConvBiasImpl::make_ncb_kern_param(
        _megdnn_tensor_in src, _megdnn_tensor_in filter, _megdnn_tensor_in bias,
536 537
        _megdnn_tensor_out dst, _megdnn_workspace workspace,
        const PreprocessedFilter* preprocessed_filter) {
538
    NCBKernParam ret;
M
Megvii Engine Team 已提交
539 540
    static_cast<NCBKernSizeParam&>(ret) = make_ncb_kern_size_param(
            src.layout, filter.layout, bias.layout, dst.layout, preprocessed_filter);
541 542 543 544
    ret.src_ptr = src.get_ref_ptr();
    ret.filter_ptr = filter.get_ref_ptr();
    ret.bias_ptr = bias.get_ref_ptr();
    ret.dst_ptr = dst.get_ref_ptr();
545 546
    ret.workspace_ptr = workspace.raw_ptr;
    ret.workspace_size = workspace.size;
547
    ret.handle = handle();
548 549 550
    return ret;
}

M
Megvii Engine Team 已提交
551 552
void ConvBiasImpl::exec_with_ncb_kern(
        const NCBKernParam& param, ConvBiasImpl::Algorithm* algo) {
553
    auto&& ncb_kerns = NCB_ALGO_FUNC(dispatch_kerns, algo, param);
554
    for (auto&& kernel : ncb_kerns) {
555
        auto run = [kernel, param](size_t index, size_t thread_id) {
556
            CpuNDRange ndrange_id(kernel.global_size, index);
557
            kernel.kern(param, {thread_id, ndrange_id});
558 559 560 561 562 563
        };
        static_cast<naive::HandleImpl*>(handle())->dispatch_kern(
                run, kernel.global_size.total_size());
    }
}

564 565
void ConvBiasImpl::exec_preprocess_with_ncb_kern(
        const NCBKernParam& param, ConvBiasImpl::Algorithm* algo) {
566
    auto&& ncb_kerns = NCB_ALGO_FUNC(dispatch_preprocess_kerns, algo, param);
567 568 569 570 571 572 573 574
    for (auto&& kernel : ncb_kerns) {
        auto run = [kernel, param](size_t index, size_t thread_id) {
            CpuNDRange ndrange_id(kernel.global_size, index);
            kernel.kern(param, {thread_id, ndrange_id});
        };
        static_cast<naive::HandleImpl*>(handle())->dispatch_kern(
                run, kernel.global_size.total_size());
    }
575 576 577 578 579 580 581
}

std::vector<ConvBiasImpl::Algorithm*> ConvBiasImpl::get_all_algorithms_with_ncb(
        const NCBKernSizeParam& param) {
    MEGDNN_MARK_USED_VAR(param);
    std::vector<Algorithm*> algos;
    std::vector<Algorithm*> prefer_algos;
582
    for (auto&& algo : get_all_packed_algo()) {
583 584
        if (algo->usable(param, AlgoSelectionStrategy::FULL_RUN)) {
            if (algo->is_preferred(param)) {
585 586 587 588 589 590 591 592 593 594 595
                prefer_algos.push_back(algo);
            } else {
                algos.push_back(algo);
            }
        }
    }
    //! Prefer algo inserted from begin
    algos.insert(algos.begin(), prefer_algos.begin(), prefer_algos.end());
    return algos;
}

596 597
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_from_desc(
        const AlgorithmDesc& desc) {
598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634
    if (!desc.valid()) {
        return nullptr;
    } else {
        switch (desc.handle_type) {
            case Handle::HandleType::FALLBACK: {
                const auto& map = algo_pack().all_algos_map();
                megdnn_assert(map.find(desc) != map.end());
                return map.at(desc);
            };

#if MEGDNN_X86
            case Handle::HandleType::X86:
                return x86::ConvBiasImpl::get_algo_from_desc(desc);
#elif MEGDNN_AARCH64 || MEGDNN_ARMV7
            case Handle::HandleType::ARM_COMMON:
                return arm_common::ConvBiasImpl::get_algo_from_desc(desc);
#if MEGDNN_AARCH64
            case Handle::HandleType::AARCH64:
                return aarch64::ConvBiasImpl::get_algo_from_desc(desc);
#else
            case Handle::HandleType::ARMV7:
                return armv7::ConvBiasImpl::get_algo_from_desc(desc);
#endif
#endif
            case Handle::HandleType::NAIVE: {
                auto algo = static_cast<naive::HandleImpl*>(handle())
                                    ->default_conv_bias_fwd_algo();
                megdnn_assert(algo->info().desc == desc);
                return algo;
            }
            default:
                megdnn_throw("Unknown handle type");
                return nullptr;
        }
    }
}

635 636
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm(
        const NCBKernSizeParam& param, size_t workspace_size) {
637 638 639
    if (ConvBiasImpl::param().format == Param::Format::NHWCD4) {
        return nullptr;
    }
640
    if (auto algo = get_algorithm_from_desc(execution_policy().algo)) {
641
        return algo;
642 643 644
    }
    if (!m_prev_selected_algo ||
        memcmp(&m_prev_selected_algo_sizep, &param, sizeof(NCBKernSizeParam))) {
645
        m_prev_selected_algo = get_algorithm_heuristic_with_ncb(
M
Megvii Engine Team 已提交
646
                param, workspace_size, AlgoAttribute::DEFAULT, AlgoAttribute::DEFAULT);
647 648 649 650 651
        m_prev_selected_algo_sizep = param;
    }
    return m_prev_selected_algo;
}

652 653 654 655 656 657 658
SmallVector<AlgoCategory> ConvBiasImpl::suggest_algo_category_order(
        const NCBKernSizeParam& param) const {
    auto IC = param.filter_meta.icpg;
    auto OC = param.filter_meta.ocpg;
    auto FH = param.filter_meta.spatial[0];
    auto FW = param.filter_meta.spatial[1];
    //! TODO: now winograd only support in fast-run
659

660 661 662 663 664 665 666 667 668
    //! im2col + matmul
    bool im2col_prefer = (IC >= 32 || OC >= 32);
    //! quantized algo use matmul when direct algo is unusable
    if (param.src_type.category() == DTypeCategory::QUANTIZED) {
        im2col_prefer = is_matmul_quantized_prefer(param);
    }
    //! conv1x1
    im2col_prefer |= (FH == 1 && FW == 1);
    if (im2col_prefer) {
M
Megvii Engine Team 已提交
669
        return {AlgoCategory::IM2COL, AlgoCategory::DIRECT, AlgoCategory::NAIVE};
670
    } else {
M
Megvii Engine Team 已提交
671
        return {AlgoCategory::DIRECT, AlgoCategory::IM2COL, AlgoCategory::NAIVE};
672 673 674
    }
}

675 676 677 678 679
const char* ConvBiasImpl::get_algorithm_set_name() const {
    // fallback version 0
    return "F0";
}

680
namespace megdnn {
681
namespace fallback {
682

683
size_t ConvBiasImpl::NCBKernParam::src_offset(
M
Megvii Engine Team 已提交
684 685
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
686
    size_t batch_offset = batch_id * inp_bs * src_type.size();
M
Megvii Engine Team 已提交
687 688 689 690
    size_t group_offset = group_pack_size * group_pack_id * filter_meta.icpg * isz[0] *
                          isz[1] * src_type.size();
    size_t channel_offset =
            channel_pack_size * channel_pack_id * isz[0] * isz[1] * src_type.size();
691
    return (batch_offset + group_offset + channel_offset);
692 693 694
}

template <typename T>
695 696 697 698 699 700 701 702 703 704 705
const T* ConvBiasImpl::NCBKernParam::src(
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
    return reinterpret_cast<T*>(
            reinterpret_cast<ptrdiff_t>(src_ptr.get_ptr()) +
            src_offset(
                    batch_id, group_pack_id, channel_pack_id, group_pack_size,
                    channel_pack_size));
}

size_t ConvBiasImpl::NCBKernParam::filter_offset(
M
Megvii Engine Team 已提交
706
        size_t group_pack_id, size_t pack_group_size) const {
707 708 709
    size_t group_offset = 0_z;
    switch (filter_meta.format) {
        case Param::Format::NCHW: {
710
            group_offset = pack_group_size * group_pack_id * filter_meta.icpg *
711 712 713 714 715 716 717 718 719
                           filter_meta.ocpg * filter_meta.spatial[0] *
                           filter_meta.spatial[1] * filter_type.size();
            break;
        }
        case Param::Format::NCHW88: {
            size_t group = filter_meta.group;
            size_t icpg = filter_meta.icpg;
            size_t ocpg = filter_meta.ocpg;
            //! four format of weight layout
720 721 722
            //! 1. {oc/8, ic/8, fh, fw, 8, 8},
            //! 2. {g, oc/8, ic/8, fh, fw, 8, 8},
            //! 3. {g/8, fh, fw, 1, 1, 8}, 4. {oc/8, fh, fw, ic, 8}
M
Megvii Engine Team 已提交
723 724 725 726 727
            megdnn_assert(
                    (icpg % 8 == 0 && ocpg % 8 == 0) ||
                            (group % 8 == 0 && icpg == 1 && ocpg == 1 &&
                             pack_group_size > 1) ||
                            (group == 1 && ocpg % 8 == 0),
728
                    "The filter shape is not right of nchw88");
729 730 731 732 733 734
            group_offset = pack_group_size * group_pack_id * filter_meta.icpg *
                           filter_meta.ocpg * filter_meta.spatial[0] *
                           filter_meta.spatial[1] * filter_type.size();

            break;
        }
735
        case Param::Format::NCHW44_DOT:
736 737 738 739 740 741 742
        case Param::Format::NCHW44: {
            size_t group = filter_meta.group;
            size_t icpg = filter_meta.icpg;
            size_t ocpg = filter_meta.ocpg;
            //! four format of weight layout
            //! 1. {oc/4, ic/4, fh, fw, 4, 4},
            //! 2. {g, oc/4, ic/4, fh, fw, 4, 4},
743 744
            //! 3. {g/4, fh, fw, 1, 1, 4},
            //! 4. {oc/4, fh, fw, ic, 4}
M
Megvii Engine Team 已提交
745 746 747 748 749
            megdnn_assert(
                    (icpg % 4 == 0 && ocpg % 4 == 0) ||
                            (group % 4 == 0 && icpg == 1 && ocpg == 1 &&
                             pack_group_size > 1) ||
                            (group == 1 && ocpg % 4 == 0),
750
                    "The filter shape is not right of nchw44");
751
            group_offset = pack_group_size * group_pack_id * filter_meta.icpg *
752 753 754 755 756 757
                           filter_meta.ocpg * filter_meta.spatial[0] *
                           filter_meta.spatial[1] * filter_type.size();

            break;
        }
        default:
758
            megdnn_assert(0, "other filter format is not support yet");
759
    }
760
    return group_offset;
761 762 763
}

template <typename T>
764 765 766 767 768 769 770 771
const T* ConvBiasImpl::NCBKernParam::filter(
        size_t group_pack_id, size_t pack_group_size) const {
    size_t group_offset = filter_offset(group_pack_id, pack_group_size);
    return reinterpret_cast<T*>(
            reinterpret_cast<ptrdiff_t>(filter_ptr.get_ptr()) + group_offset);
}

size_t ConvBiasImpl::NCBKernParam::bias_offset(
M
Megvii Engine Team 已提交
772 773
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
774 775
    size_t batch_offset = 0_z;
    size_t group_offset = 0_z;
776
    size_t channel_offset = 0_z;
777 778
    if (bias_mode == BiasMode::BIAS) {
        batch_offset = batch_id * bias_bs * bias_type.size();
M
Megvii Engine Team 已提交
779 780
        group_offset = group_pack_size * group_pack_id * filter_meta.ocpg * osz[0] *
                       osz[1] * bias_type.size();
781 782
        channel_offset = channel_pack_size * channel_pack_id * osz[0] * osz[1] *
                         bias_type.size();
783
    } else if (bias_mode == BiasMode::BROADCAST_CHANNEL_BIAS) {
M
Megvii Engine Team 已提交
784 785
        group_offset =
                group_pack_size * group_pack_id * filter_meta.ocpg * bias_type.size();
786
        channel_offset = channel_pack_size * channel_pack_id * bias_type.size();
787
    }
788
    return (batch_offset + group_offset + channel_offset);
789 790 791
}

template <typename T>
792 793 794 795 796 797 798 799 800 801 802
const T* ConvBiasImpl::NCBKernParam::bias(
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
    return reinterpret_cast<T*>(
            reinterpret_cast<ptrdiff_t>(bias_ptr.get_ptr()) +
            bias_offset(
                    batch_id, group_pack_id, channel_pack_id, group_pack_size,
                    channel_pack_size));
}

size_t ConvBiasImpl::NCBKernParam::dst_offset(
M
Megvii Engine Team 已提交
803 804
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
805
    size_t batch_offset = batch_id * out_bs * dst_type.size();
M
Megvii Engine Team 已提交
806 807 808 809
    size_t group_offset = group_pack_size * group_pack_id * filter_meta.ocpg * osz[0] *
                          osz[1] * dst_type.size();
    size_t channel_offset =
            channel_pack_size * channel_pack_id * osz[0] * osz[1] * dst_type.size();
810 811 812 813 814 815 816
    return (batch_offset + group_offset + channel_offset);
}

template <typename T>
T* ConvBiasImpl::NCBKernParam::dst(
        size_t batch_id, size_t group_pack_id, size_t channel_pack_id,
        size_t group_pack_size, size_t channel_pack_size) const {
M
Megvii Engine Team 已提交
817
    return reinterpret_cast<T*>(
818 819 820 821
            reinterpret_cast<ptrdiff_t>(dst_ptr.get_ptr()) +
            dst_offset(
                    batch_id, group_pack_id, channel_pack_id, group_pack_size,
                    channel_pack_size));
822 823
}

824 825 826 827 828 829 830 831 832 833 834 835
#define INST(T)                                                      \
    template const T* ConvBiasImpl::NCBKernParam::src<T>(            \
            size_t batch_id, size_t group_id, size_t channel_id,     \
            size_t group_pack_size, size_t channel_pack_size) const; \
    template const T* ConvBiasImpl::NCBKernParam::bias<T>(           \
            size_t batch_id, size_t group_id, size_t channel_id,     \
            size_t group_pack_size, size_t channel_pack_size) const; \
    template const T* ConvBiasImpl::NCBKernParam::filter<T>(         \
            size_t group_id, size_t group_pack_size) const;          \
    template T* ConvBiasImpl::NCBKernParam::dst<T>(                  \
            size_t batch_id, size_t group_id, size_t channel_id,     \
            size_t group_pack_size, size_t channel_pack_size) const;
836 837 838 839

#define INST_DT(d) INST(DTypeTrait<d>::ctype)

MEGDNN_FOREACH_COMPUTING_DTYPE(INST_DT)
840
INST(void)
841 842 843 844 845
#undef INST
#undef INST_DT
}  // namespace fallback
}  // namespace megdnn

846
// vim: syntax=cpp.doxygen