opr_impl.cpp 32.9 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/fp32/algos.h"
20
#include "src/fallback/conv_bias/im2col/algos.h"
21
#include "src/fallback/convolution/opr_impl.h"
22 23 24
#include "src/naive/convolution/algorithms.h"
#include "src/naive/handle.h"

25 26 27 28 29 30 31 32
#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

33 34 35 36 37
#include <cstring>

using namespace megdnn;
using namespace fallback;

38 39 40 41 42 43 44 45
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

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

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

}  // namespace

67 68 69 70 71 72 73 74 75 76 77 78 79
#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

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

87 88
    AlgoF32DirectNCHWNCHW44 f32_nchw_nchw44;
    AlgoF32DirectNCHWNCHW44AGENT f32_nchw_nchw44_agent;
89 90 91 92 93 94
    AlgoF32ChannelWiseNCHW44 f32_chanel_wise_nchw44;
    AlgoF32DirectNCHW44 f32_direct_nchw44;

    AlgoF32Direct f32_direct;
    AlgoF32DirectStride2 f32_direct_stride2;
    AlgoF32DirectStride1 f32_direct_stride1;
95 96 97

public:
    AlgoPack() {
98 99 100 101 102 103 104 105 106 107 108
        //! 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

109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
        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;
            for (uint32_t tile_size : {16, 8, 24, 32}) {
                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());
                refhold.emplace_back(new AlgoFP32WinogradF63_4x4_NCHW44(
                        static_cast<fallback::MatrixMulImpl::AlgoBase*>(algo),
                        tile_size));
                m_gi_winograd_algos.emplace_back(refhold.back().get());
                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());
                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());
            }
        }
        for (auto&& algo : m_gi_winograd_algos) {
            m_all_algos.emplace_back(algo);
        }
        // end fallback gi fp32 algo

177
        refhold.emplace_back(new AlgoConv1x1Gemv());
178
        m_all_algos.emplace_back(refhold.back().get());
179

180 181
        matmul_algos = static_cast<fallback::MatrixMulImpl*>(matmul_opr)
                               ->get_all_packed_algo();
182
        for (auto&& algo : matmul_algos) {
183 184 185 186 187 188
#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 已提交
189
            if (algo->algoset() == MatrixMulImpl::AlgoBase::AlgoSet::ALGO_TYPE_GEMV) {
190 191
                continue;
            }
192 193
#endif

194 195
            for (size_t ohw_tile_size : {192, 384, 96, 48, 24}) {
                refhold.emplace_back(new AlgoIm2col(
M
Megvii Engine Team 已提交
196
                        static_cast<MatrixMulImpl::AlgoBase*>(algo), ohw_tile_size));
197
                m_all_algos.emplace_back(refhold.back().get());
198
            }
199
            for (size_t oc_tile_size : {48, 24}) {
200
                refhold.emplace_back(new AlgoConv1x1(
M
Megvii Engine Team 已提交
201
                        static_cast<MatrixMulImpl::AlgoBase*>(algo), oc_tile_size));
202
                m_all_algos.emplace_back(refhold.back().get());
203
            }
204

205
#if 0
206 207 208 209 210
        //! 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)));
211
            m_all_algos.emplace_back(refhold.back().get());
212 213
            refhold.emplace_back(new AlgoWinogradF32_4x4(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
214
            m_all_algos.emplace_back(refhold.back().get());
215 216
            refhold.emplace_back(new AlgoWinogradQS8(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
217
            m_all_algos.emplace_back(refhold.back().get());
218 219
            refhold.emplace_back(new AlgoWinogradQS8_8x8(
                    static_cast<MatrixMulImpl::AlgoBase*>(algo)));
220
            m_all_algos.emplace_back(refhold.back().get());
221 222
#endif
        }
223 224 225 226 227
        m_all_algos.emplace_back(&algo_naive);

        for (auto&& algo : m_all_algos) {
            m_all_algos_map.emplace(algo->info().desc, algo);
        }
228
    }
229 230
    const SmallVector<AlgoBase*>& all_algos() const { return m_all_algos; }
    const AlgoBase::Mapper& all_algos_map() const { return m_all_algos_map; }
231 232
};

233 234 235 236 237 238 239
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();
240
}
241 242 243

SmallVector<ConvBiasImpl::AlgoBase*> ConvBiasImpl::select_algo_type(
        ConvAlgoTypePack target_type) {
M
Megvii Engine Team 已提交
244 245 246
    megdnn_assert(
            nr_type_contain(target_type.data_type),
            "ConvBias algo selection only support one type");
247
    SmallVector<ConvBiasImpl::AlgoBase*> algos;
248
    for (auto&& algo : get_all_packed_algo()) {
249 250 251 252 253 254 255 256 257
        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;
}

258 259 260
bool ConvBiasImpl::is_naive_algo(ConvBiasImpl::Algorithm* algo) {
    return algo == nullptr || strcmp(algo->name(), "DEFAULT") == 0;
}
261

M
Megvii Engine Team 已提交
262 263 264 265 266 267 268 269 270 271 272
#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);
273
    auto&& algo = get_algorithm(fparam, workspace.size);
274
    if (!is_naive_algo(algo) &&
275
        NCB_ALGO_FUNC(get_workspace, algo, fparam) <= workspace.size) {
276 277
        exec_with_ncb_kern(fparam, algo);
    } else {
M
Megvii Engine Team 已提交
278 279
        naive::ConvBiasForwardImpl::exec(
                src, filter, bias, z, dst, preprocessed_filter, workspace);
280 281 282
    }
}

M
Megvii Engine Team 已提交
283 284 285 286 287
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) {
288 289 290
    //! 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 已提交
291 292
    auto fparam =
            make_ncb_kern_param(src, filter, bias, dst, workspace, preprocessed_filter);
293
    //! should not pass workspace_size limit otherwise can not find match algo
294 295
    auto&& algo = get_algorithm(fparam);
    if (!is_naive_algo(algo) &&
M
Megvii Engine Team 已提交
296
        NCB_ALGO_FUNC(get_preprocess_workspace, algo, fparam) <= workspace.size) {
297 298 299
        exec_preprocess_with_ncb_kern(fparam, algo);
    } else {
        naive::ConvBiasForwardImpl::exec_preprocess(
M
Megvii Engine Team 已提交
300 301
                src_layout, filter, bias, z_layout, dst_layout, preprocessed_filter,
                workspace);
302 303 304
    }
}

305
size_t ConvBiasImpl::get_workspace_in_bytes(
M
Megvii Engine Team 已提交
306 307
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst,
308
        const PreprocessedFilter* preprocessed_filter) {
309
    TensorLayoutArray layouts{src, filter, bias, z, dst};
310
    AlgorithmCache::Key key{this->handle(), this->get_opr_type(),
M
Megvii Engine Team 已提交
311 312
                            layouts.data(), layouts.size(),
                            &this->param(), sizeof(this->param())};
313
    auto rst = AlgorithmCache::instance().get(key);
314 315 316 317
    if (rst.policy.algo.valid()) {
        return rst.workspace;
    }

M
Megvii Engine Team 已提交
318
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, preprocessed_filter);
319
    auto&& algo = get_algorithm(fparam);
320
    if (is_naive_algo(algo)) {
321 322
        return naive::ConvBiasForwardImpl::get_workspace_in_bytes(
                src, filter, bias, z, dst, preprocessed_filter);
323
    } else {
324 325 326 327 328
        return NCB_ALGO_FUNC(get_workspace, algo, fparam);
    }
}

size_t ConvBiasImpl::get_preprocess_workspace_in_bytes(
M
Megvii Engine Team 已提交
329 330
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
331
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
332
    auto&& algo = get_algorithm(fparam);
333 334 335 336 337 338 339 340 341
    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 已提交
342 343
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
344
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
345
    auto&& algo = get_algorithm(fparam);
346 347 348 349 350
    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);
351 352 353 354
    }
}

std::vector<ConvBiasImpl::Algorithm*> ConvBiasImpl::get_all_algorithms(
M
Megvii Engine Team 已提交
355 356
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& z, const TensorLayout& dst) {
357
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
358 359
    auto ret = get_all_algorithms_with_ncb(fparam);
    if (ret.empty()) {
M
Megvii Engine Team 已提交
360 361
        return naive::ConvBiasForwardImpl::get_all_algorithms_safe(
                src, filter, bias, z, dst);
362 363 364
    }
    return ret;
}
365
std::vector<ConvBiasImpl::Algorithm*> ConvBiasImpl::get_all_algorithms_safe(
M
Megvii Engine Team 已提交
366 367 368
        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);
369 370
    return ret_safe;
}
371 372

ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_heuristic(
M
Megvii Engine Team 已提交
373 374 375
        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) {
376
    auto fparam = make_ncb_kern_size_param(src, filter, bias, dst, nullptr);
377
    auto result = get_algorithm_heuristic_with_ncb(
378
            fparam, workspace_limit_in_bytes, positive_attr, negative_attr);
379 380
    if (result == nullptr) {
        result = naive::ConvBiasForwardImpl::get_algorithm_heuristic(
M
Megvii Engine Team 已提交
381 382
                src, filter, bias, z, dst, workspace_limit_in_bytes, positive_attr,
                negative_attr);
383 384 385 386
    }
    return result;
}

387 388
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_heuristic_with_ncb(
        const NCBKernSizeParam& param, size_t workspace_limit_in_bytes,
M
Megvii Engine Team 已提交
389
        const AlgoAttribute& positive_attr, const AlgoAttribute& negative_attr) {
390 391 392
    if (ConvBiasImpl::param().format == Param::Format::NHWCD4) {
        return nullptr;
    }
393 394 395 396 397 398
    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) {
399
            bool usable_attribute = static_cast<AlgoBase*>(i)->usable_attribute(
400 401
                    param, AlgoSelectionStrategy::HEURISTIC, positive_attr,
                    negative_attr);
M
Megvii Engine Team 已提交
402 403
            if (usable_attribute && static_cast<AlgoBase*>(i)->get_workspace(param) <=
                                            workspace_limit_in_bytes) {
404 405 406 407 408 409 410 411 412 413 414 415 416
                //! 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;
417 418 419 420 421
        }
    }
    return nullptr;
}

422
ConvBiasImpl::NCBKernSizeParam ConvBiasImpl::make_ncb_kern_size_param(
M
Megvii Engine Team 已提交
423 424
        const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
        const TensorLayout& dst, const PreprocessedFilter* preprocessed_filter) {
425
    auto safe_u32 = [](size_t v) -> uint32_t {
M
Megvii Engine Team 已提交
426 427
        megdnn_assert(
                v <= std::numeric_limits<uint32_t>::max(), "value too large: %zu", v);
428 429 430 431 432 433
        return v;
    };
    size_t spatial_pos;
    if (param().format == Param::Format::NCHW88 ||
        param().format == Param::Format::NCHW8 ||
        param().format == Param::Format::NCHW4 ||
434
        param().format == Param::Format::NCHW44 ||
435
        param().format == Param::Format::NCHW44_DOT ||
436 437
        param().format == Param::Format::NCHW ||
        param().format == Param::Format::NCHW32 ||
438
        param().format == Param::Format::NCHW64) {
439
        spatial_pos = 2;
M
Megvii Engine Team 已提交
440 441 442
    } else if (
            param().format == Param::Format::NHWC ||
            param().format == Param::Format::NHWCD4) {
443 444
        spatial_pos = 1;
    } else {
M
Megvii Engine Team 已提交
445
        megdnn_assert(0, "invalid conv format %d", static_cast<int>(param().format));
446 447
    }
    BiasMode bias_mode;
448 449
    //! 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;
450 451
    if (bias.ndim == 0) {
        bias_mode = BiasMode::NO_BIAS;
452
    } else if (bias.eq_shape(dst) && !dst_only_c) {
453 454 455 456 457 458 459
        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 已提交
460 461 462 463
    static_assert(
            sizeof(CanonizedFilterMeta) == sizeof(ConvolutionImpl::CanonizedFilterMeta),
            "sizeof CanonizedFilterMeta in convolution and conv_bias "
            "should be equal");
464 465
    auto&& fm = check_layout_fwd(src, filter, dst);
    auto& conv_fm = reinterpret_cast<ConvolutionImpl::CanonizedFilterMeta&>(fm);
466

467 468 469 470 471 472 473 474 475 476 477 478 479 480 481
    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,
482 483
             nr_threads,
             reinterpret_cast<const ConvolutionForward::PreprocessedFilter*>(
484 485
                     preprocessed_filter),
             handle()},
486 487 488 489 490 491 492 493
            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,
494 495
        _megdnn_tensor_out dst, _megdnn_workspace workspace,
        const PreprocessedFilter* preprocessed_filter) {
496
    NCBKernParam ret;
M
Megvii Engine Team 已提交
497 498
    static_cast<NCBKernSizeParam&>(ret) = make_ncb_kern_size_param(
            src.layout, filter.layout, bias.layout, dst.layout, preprocessed_filter);
499 500 501 502
    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();
503 504
    ret.workspace_ptr = workspace.raw_ptr;
    ret.workspace_size = workspace.size;
505
    ret.handle = handle();
506 507 508
    return ret;
}

M
Megvii Engine Team 已提交
509 510
void ConvBiasImpl::exec_with_ncb_kern(
        const NCBKernParam& param, ConvBiasImpl::Algorithm* algo) {
511
    auto&& ncb_kerns = NCB_ALGO_FUNC(dispatch_kerns, algo, param);
512
    for (auto&& kernel : ncb_kerns) {
513
        auto run = [kernel, param](size_t index, size_t thread_id) {
514
            CpuNDRange ndrange_id(kernel.global_size, index);
515
            kernel.kern(param, {thread_id, ndrange_id});
516 517 518 519 520 521
        };
        static_cast<naive::HandleImpl*>(handle())->dispatch_kern(
                run, kernel.global_size.total_size());
    }
}

522 523
void ConvBiasImpl::exec_preprocess_with_ncb_kern(
        const NCBKernParam& param, ConvBiasImpl::Algorithm* algo) {
524
    auto&& ncb_kerns = NCB_ALGO_FUNC(dispatch_preprocess_kerns, algo, param);
525 526 527 528 529 530 531 532
    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());
    }
533 534 535 536 537 538 539
}

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;
540
    for (auto&& algo : get_all_packed_algo()) {
541 542
        if (algo->usable(param, AlgoSelectionStrategy::FULL_RUN)) {
            if (algo->is_preferred(param)) {
543 544 545 546 547 548 549 550 551 552 553
                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;
}

554 555
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm_from_desc(
        const AlgorithmDesc& desc) {
556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592
    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;
        }
    }
}

593 594
ConvBiasImpl::Algorithm* ConvBiasImpl::get_algorithm(
        const NCBKernSizeParam& param, size_t workspace_size) {
595 596 597
    if (ConvBiasImpl::param().format == Param::Format::NHWCD4) {
        return nullptr;
    }
598
    if (auto algo = get_algorithm_from_desc(execution_policy().algo)) {
599
        return algo;
600 601 602
    }
    if (!m_prev_selected_algo ||
        memcmp(&m_prev_selected_algo_sizep, &param, sizeof(NCBKernSizeParam))) {
603
        m_prev_selected_algo = get_algorithm_heuristic_with_ncb(
M
Megvii Engine Team 已提交
604
                param, workspace_size, AlgoAttribute::DEFAULT, AlgoAttribute::DEFAULT);
605 606 607 608 609
        m_prev_selected_algo_sizep = param;
    }
    return m_prev_selected_algo;
}

610 611 612 613 614 615 616
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
617

618 619 620 621 622 623 624 625 626
    //! 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 已提交
627
        return {AlgoCategory::IM2COL, AlgoCategory::DIRECT, AlgoCategory::NAIVE};
628
    } else {
M
Megvii Engine Team 已提交
629
        return {AlgoCategory::DIRECT, AlgoCategory::IM2COL, AlgoCategory::NAIVE};
630 631 632
    }
}

633 634 635 636 637
const char* ConvBiasImpl::get_algorithm_set_name() const {
    // fallback version 0
    return "F0";
}

638
namespace megdnn {
639
namespace fallback {
640

641
size_t ConvBiasImpl::NCBKernParam::src_offset(
M
Megvii Engine Team 已提交
642 643
        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 {
644
    size_t batch_offset = batch_id * inp_bs * src_type.size();
M
Megvii Engine Team 已提交
645 646 647 648
    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();
649
    return (batch_offset + group_offset + channel_offset);
650 651 652
}

template <typename T>
653 654 655 656 657 658 659 660 661 662 663
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 已提交
664
        size_t group_pack_id, size_t pack_group_size) const {
665 666 667
    size_t group_offset = 0_z;
    switch (filter_meta.format) {
        case Param::Format::NCHW: {
668
            group_offset = pack_group_size * group_pack_id * filter_meta.icpg *
669 670 671 672 673 674 675 676 677
                           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
678 679 680
            //! 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 已提交
681 682 683 684 685 686
            megdnn_assert(
                    (icpg % 8 == 0 && ocpg % 8 == 0) ||
                            (group % 8 == 0 && icpg == 1 && ocpg == 1 &&
                             pack_group_size > 1) ||
                            (group == 1 && ocpg % 8 == 0),
                    "The filter shepe is not right of nchw88");
687 688 689 690 691 692
            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;
        }
693
        case Param::Format::NCHW44_DOT:
694 695 696 697 698 699 700
        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},
701 702
            //! 3. {g/4, fh, fw, 1, 1, 4},
            //! 4. {oc/4, fh, fw, ic, 4}
M
Megvii Engine Team 已提交
703 704 705 706 707 708
            megdnn_assert(
                    (icpg % 4 == 0 && ocpg % 4 == 0) ||
                            (group % 4 == 0 && icpg == 1 && ocpg == 1 &&
                             pack_group_size > 1) ||
                            (group == 1 && ocpg % 4 == 0),
                    "The filter shepe is not right of nchw44");
709
            group_offset = pack_group_size * group_pack_id * filter_meta.icpg *
710 711 712 713 714 715
                           filter_meta.ocpg * filter_meta.spatial[0] *
                           filter_meta.spatial[1] * filter_type.size();

            break;
        }
        default:
716
            megdnn_assert(0, "other filter format is not support yet");
717
    }
718
    return group_offset;
719 720 721
}

template <typename T>
722 723 724 725 726 727 728 729
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 已提交
730 731
        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 {
732 733
    size_t batch_offset = 0_z;
    size_t group_offset = 0_z;
734
    size_t channel_offset = 0_z;
735 736
    if (bias_mode == BiasMode::BIAS) {
        batch_offset = batch_id * bias_bs * bias_type.size();
M
Megvii Engine Team 已提交
737 738
        group_offset = group_pack_size * group_pack_id * filter_meta.ocpg * osz[0] *
                       osz[1] * bias_type.size();
739 740
        channel_offset = channel_pack_size * channel_pack_id * osz[0] * osz[1] *
                         bias_type.size();
741
    } else if (bias_mode == BiasMode::BROADCAST_CHANNEL_BIAS) {
M
Megvii Engine Team 已提交
742 743
        group_offset =
                group_pack_size * group_pack_id * filter_meta.ocpg * bias_type.size();
744
        channel_offset = channel_pack_size * channel_pack_id * bias_type.size();
745
    }
746
    return (batch_offset + group_offset + channel_offset);
747 748 749
}

template <typename T>
750 751 752 753 754 755 756 757 758 759 760
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 已提交
761 762
        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 {
763
    size_t batch_offset = batch_id * out_bs * dst_type.size();
M
Megvii Engine Team 已提交
764 765 766 767
    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();
768 769 770 771 772 773 774
    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 已提交
775
    return reinterpret_cast<T*>(
776 777 778 779
            reinterpret_cast<ptrdiff_t>(dst_ptr.get_ptr()) +
            dst_offset(
                    batch_id, group_pack_id, channel_pack_id, group_pack_size,
                    channel_pack_size));
780 781
}

782 783 784 785 786 787 788 789 790 791 792 793
#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;
794 795 796 797

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

MEGDNN_FOREACH_COMPUTING_DTYPE(INST_DT)
798
INST(void)
799 800 801 802 803
#undef INST
#undef INST_DT
}  // namespace fallback
}  // namespace megdnn

804
// vim: syntax=cpp.doxygen