algos.cpp 39.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
/**
 * \file dnn/src/fallback/conv_bias/im2col/algos.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
 *
 * 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 "src/fallback/conv_bias/im2col/algos.h"
13
#include "src/fallback/conv_bias/im2col/factory.h"
14 15 16 17 18
#include "megdnn/opr_param_defs.h"
#include "src/common/opr_delegate.h"
#include "src/fallback/conv_bias/common.h"
#include "src/fallback/conv_bias/opr_impl.h"
#include "src/naive/convolution/helper.h"
19

20
#include "midout.h"
21

22 23 24 25
MIDOUT_DECL(megdnn_fallback_im2col)

using namespace megdnn;
using namespace fallback;
26
using namespace im2col;
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49

/*======================== AlgoIm2col=======================*/
/*!
 *  *\brief The index of all parts workspace in im2col workspace bundel
 *  *Through witch can convenient get the needed ptr
 */
struct Im2colBundelIndex {
    static constexpr size_t BUNDLE_THREAD_INDEX = 2_z;
};

using Pack_Mode=fallback::MatrixMulImpl::AlgoBase::PackMode;
/*!
 * *\brief Im2colKerns collects all the im2col kerns in it
 */

template <Pack_Mode packmode>
class Im2colKerns;

template <>
class Im2colKerns<Pack_Mode::DEFAULT> {
public:
    //! conv kernel
    static void kerns(
50
            const WorkspaceBundle& bundle, WorkspaceBundle bundle_thread,
51 52
            const ConvBiasImpl::NCBKernParam& param,
            fallback::MatrixMulImpl::KernSizeParam matmul_kernsize_param,
53
            const fallback::MatrixMulImpl::AlgoBase* matmul_algo,
54 55
            const fallback::MatrixMulImpl::AlgoBase::MatmulDescription&
                    matmul_desc,
56
            StrategyParam strategyparam,
57
            fallback::ConvBiasImpl::NCBKernIndex ncb_index,
58 59
            size_t ohw_tile_size, StrategyBase* im2colstrategy) {
        size_t OC = param.filter_meta.ocpg;
60
        size_t output_block_size = std::min(
61 62
                ohw_tile_size,
                strategyparam.ohw - ncb_index.ndrange_id[2] * ohw_tile_size);
63
        size_t output_block_oc_size = std::min(
64 65 66 67 68 69 70 71 72 73 74 75 76 77
                strategyparam.oc_tile_size,
                OC - ncb_index.ndrange_id[3] * strategyparam.oc_tile_size);

        strategyparam.batch_id = ncb_index.ndrange_id[0];
        strategyparam.group_id = ncb_index.ndrange_id[1];
        strategyparam.oc_cur_index =
                ncb_index.ndrange_id[3] *
                strategyparam.oc_tile_size;
        strategyparam.oc_end_index = strategyparam.oc_cur_index +
                                     output_block_oc_size;
        strategyparam.ohw_cur_index =
                ncb_index.ndrange_id[2] * ohw_tile_size;
        strategyparam.output_block_oc_size = output_block_oc_size;
        strategyparam.output_block_size = output_block_size;
78

79 80 81 82
        bundle_thread.set(
                static_cast<int8_t*>(
                        bundle.get(Im2colBundelIndex::BUNDLE_THREAD_INDEX)) +
                bundle_thread.total_size_in_bytes() * ncb_index.thread_id);
83 84 85 86
        fallback::MatrixMulImpl::KernParam matmul_param;
        static_cast<fallback::MatrixMulImpl::KernSizeParam&>(matmul_param) =
                matmul_kernsize_param;

87 88 89
        //! 1.Im2col
        im2colstrategy->exec_im2col(bundle, bundle_thread, strategyparam, param,
                                    matmul_param, matmul_algo);
90

91 92
        //! 2.packb and matmul compute
        im2colstrategy->exec_matmul(param, strategyparam, bundle, bundle_thread,
93 94
                                    matmul_param, matmul_algo, ncb_index,
                                    matmul_desc);
95

96 97 98
        //! 3.postprocess and copy dst if need
        im2colstrategy->exec_postprocess(param, strategyparam, bundle_thread);
    }
99

100 101
    WorkspaceBundle get_thread_bundle(
            const fallback::ConvBiasImpl::NCBKernSizeParam& param,
102 103
            const fallback::MatrixMulImpl::KernSizeParam& im2col_kern_param,
            const MatrixMulImpl::AlgoBase* matmul_algo, size_t ohw_tile_size,
104 105 106
            size_t oc_tile_size) {
        size_t IC = param.filter_meta.icpg, FH = param.filter_meta.spatial[0],
               FW = param.filter_meta.spatial[1];
107
        size_t pack_oc_size = pack_size(param.filter_meta.format);
108 109 110 111 112
        size_t im2col = 0, packb = 0, bias_temp = 0;
        bool default_pack = matmul_algo->packmode() == Pack_Mode::DEFAULT;
        megdnn_assert(default_pack, "only support default packa");
        size_t im2col_dst_size =
                IC * FH * FW * ohw_tile_size * sizeof(param.src_type);
113 114
        size_t matmul_dst_size = pack_oc_size * oc_tile_size * ohw_tile_size *
                                 sizeof(param.bias_type);
115 116 117 118 119 120 121 122
        //! matmul_dst and im2col_dst use the same memory
        WorkspaceBundle wb = matmul_algo->get_bundle(im2col_kern_param);
        packb = wb.get_size(1);
        im2col = std::max(im2col_dst_size, matmul_dst_size);
        if (param.bias_mode == megdnn::BiasMode::BIAS) {
            bias_temp = oc_tile_size * ohw_tile_size * sizeof(param.bias_type);
        }
        return {nullptr, {packb, im2col, bias_temp}};
123 124 125 126 127 128 129 130
    }
};

template <>
class Im2colKerns<Pack_Mode::ONLY_PACKA> {
public:
    //! conv kernel
    static void kerns(
131
            const WorkspaceBundle& bundle, WorkspaceBundle bundle_thread,
132 133
            const ConvBiasImpl::NCBKernParam& param,
            fallback::MatrixMulImpl::KernSizeParam matmul_kernsize_param,
134
            const fallback::MatrixMulImpl::AlgoBase* matmul_algo,
135 136
            const fallback::MatrixMulImpl::AlgoBase::MatmulDescription&
                    matmul_desc,
137
            StrategyParam strategyparam,
138
            fallback::ConvBiasImpl::NCBKernIndex ncb_index,
139 140
            size_t ohw_tile_size, StrategyBase* im2colstrategy) {
        size_t OC = param.filter_meta.ocpg;
141
        size_t output_block_size = std::min(
142 143
                ohw_tile_size,
                strategyparam.ohw - ncb_index.ndrange_id[2] * ohw_tile_size);
144
        size_t output_block_oc_size = std::min(
145 146
                strategyparam.oc_tile_size,
                OC - ncb_index.ndrange_id[3] * strategyparam.oc_tile_size);
147

148 149 150 151
        bundle_thread.set(
                static_cast<int8_t*>(
                        bundle.get(Im2colBundelIndex::BUNDLE_THREAD_INDEX)) +
                bundle_thread.total_size_in_bytes() * ncb_index.thread_id);
152 153 154 155 156

        fallback::MatrixMulImpl::KernParam matmul_param;
        static_cast<fallback::MatrixMulImpl::KernSizeParam&>(matmul_param) =
                matmul_kernsize_param;

157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174
        strategyparam.batch_id = ncb_index.ndrange_id[0];
        strategyparam.group_id = ncb_index.ndrange_id[1];
        strategyparam.oc_cur_index =
                ncb_index.ndrange_id[3] *
                strategyparam.oc_tile_size;
        strategyparam.oc_end_index = strategyparam.oc_cur_index +
                                     output_block_oc_size;
        strategyparam.ohw_cur_index =
                ncb_index.ndrange_id[2] * ohw_tile_size;
        strategyparam.output_block_oc_size = output_block_oc_size;
        strategyparam.output_block_size = output_block_size;

        //! 1.Im2col
        im2colstrategy->exec_im2col(bundle, bundle_thread, strategyparam, param,
                                    matmul_param, matmul_algo);

        //! 2.packb and matmul compute
        im2colstrategy->exec_matmul(param, strategyparam, bundle, bundle_thread,
175 176
                                    matmul_param, matmul_algo, ncb_index,
                                    matmul_desc);
177 178 179 180 181 182

        //! 3.postprocess and copy dst if need
        im2colstrategy->exec_postprocess(param, strategyparam, bundle_thread);
    }
    WorkspaceBundle get_thread_bundle(
            const fallback::ConvBiasImpl::NCBKernSizeParam& param,
183 184
            const fallback::MatrixMulImpl::KernSizeParam& im2col_kern_param,
            const MatrixMulImpl::AlgoBase* matmul_algo, size_t ohw_tile_size,
185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203
            size_t oc_tile_size) {
        size_t IC = param.filter_meta.icpg, FH = param.filter_meta.spatial[0],
               FW = param.filter_meta.spatial[1];

        size_t im2col = 0, packb = 0, matmul_dst = 0, bias_temp = 0;
        bool only_packA = matmul_algo->packmode() == Pack_Mode::ONLY_PACKA;
        megdnn_assert(only_packA, "onlysupport onlypackA mode");
        size_t im2col_dst_size =
                IC * FH * FW * ohw_tile_size * sizeof(param.src_type);
        size_t matmul_dst_size =
                oc_tile_size * ohw_tile_size * sizeof(param.bias_type);
        //! matmul_dst and im2col_dst use the same memory
        WorkspaceBundle wb = matmul_algo->get_bundle(im2col_kern_param);
        packb = wb.get_size(1);
        im2col = im2col_dst_size;
        matmul_dst = matmul_dst_size;
        if (param.bias_mode == megdnn::BiasMode::BIAS) {
            bias_temp = oc_tile_size * ohw_tile_size * sizeof(param.bias_type);
        }
204

205
        return {nullptr, {packb, im2col, matmul_dst, bias_temp}};
206 207 208 209 210 211 212 213
    }
};

template <>
class Im2colKerns<Pack_Mode::NO_PACK> {
public:
    //! conv kernel
    static void kerns(
214
            const WorkspaceBundle& bundle, WorkspaceBundle bundle_thread,
215 216
            const ConvBiasImpl::NCBKernParam& param,
            fallback::MatrixMulImpl::KernSizeParam matmul_kernsize_param,
217
            const fallback::MatrixMulImpl::AlgoBase* matmul_algo,
218 219
            const fallback::MatrixMulImpl::AlgoBase::MatmulDescription&
                    matmul_desc,
220
            StrategyParam strategyparam,
221
            fallback::ConvBiasImpl::NCBKernIndex ncb_index,
222 223
            size_t ohw_tile_size, StrategyBase* im2colstrategy) {
        size_t OC = param.filter_meta.ocpg;
224
        size_t output_block_size = std::min(
225 226
                ohw_tile_size,
                strategyparam.ohw - ncb_index.ndrange_id[2] * ohw_tile_size);
227
        size_t output_block_oc_size = std::min(
228 229 230 231 232 233 234 235 236 237 238 239 240 241
                strategyparam.oc_tile_size,
                OC - ncb_index.ndrange_id[3] * strategyparam.oc_tile_size);

        strategyparam.batch_id = ncb_index.ndrange_id[0];
        strategyparam.group_id = ncb_index.ndrange_id[1];
        strategyparam.oc_cur_index =
                ncb_index.ndrange_id[3] *
                strategyparam.oc_tile_size;
        strategyparam.oc_end_index = strategyparam.oc_cur_index +
                                     output_block_oc_size;
        strategyparam.ohw_cur_index =
                ncb_index.ndrange_id[2] * ohw_tile_size;
        strategyparam.output_block_oc_size = output_block_oc_size;
        strategyparam.output_block_size = output_block_size;
242

243 244 245 246
        bundle_thread.set(
                static_cast<int8_t*>(
                        bundle.get(Im2colBundelIndex::BUNDLE_THREAD_INDEX)) +
                bundle_thread.total_size_in_bytes() * ncb_index.thread_id);
247 248 249 250 251

        fallback::MatrixMulImpl::KernParam matmul_param;
        static_cast<fallback::MatrixMulImpl::KernSizeParam&>(matmul_param) =
                matmul_kernsize_param;

252 253 254
        //! 1.Im2col
        im2colstrategy->exec_im2col(bundle, bundle_thread, strategyparam, param,
                                    matmul_param, matmul_algo);
255

256 257
        //! 2.packb and matmul compute
        im2colstrategy->exec_matmul(param, strategyparam, bundle, bundle_thread,
258 259
                                    matmul_param, matmul_algo, ncb_index,
                                    matmul_desc);
260

261 262 263 264 265
        //! 3.postprocess and copy dst if need
        im2colstrategy->exec_postprocess(param, strategyparam, bundle_thread);
    }
    WorkspaceBundle get_thread_bundle(
            const fallback::ConvBiasImpl::NCBKernSizeParam& param,
266 267
            const fallback::MatrixMulImpl::KernSizeParam& im2col_kern_param,
            const MatrixMulImpl::AlgoBase* matmul_algo, size_t ohw_tile_size,
268 269 270 271
            size_t oc_tile_size) {
        size_t IC = param.filter_meta.icpg, FH = param.filter_meta.spatial[0],
               FW = param.filter_meta.spatial[1];
        size_t ohw = param.osz[0] * param.osz[1];
272

273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294
        size_t im2col = 0, matmul_dst = 0, bias_temp = 0, matmul_compute = 0;
        bool no_pack = matmul_algo->packmode() == Pack_Mode::NO_PACK;
        megdnn_assert(no_pack, "only support no pack");
        bool is_dst_8bit =
                (param.src_type.enumv() == DTypeEnum::QuantizedS8 &&
                 param.dst_type.enumv() == DTypeEnum::QuantizedS8) ||
                (param.src_type.enumv() == DTypeEnum::Quantized8Asymm &&
                 param.dst_type.enumv() == DTypeEnum::Quantized8Asymm);
        size_t im2col_dst_size =
                IC * FH * FW * ohw_tile_size * sizeof(param.src_type);
        size_t matmul_dst_size =
                oc_tile_size * ohw_tile_size * sizeof(param.bias_type);
        im2col = im2col_dst_size;
        if (is_dst_8bit) {
            matmul_dst = matmul_dst_size;
        } else {
            matmul_dst = ohw_tile_size >= ohw ? 0 : matmul_dst_size;
        }
        matmul_compute = matmul_algo->get_workspace(im2col_kern_param);
        if (param.bias_mode == megdnn::BiasMode::BIAS) {
            bias_temp = oc_tile_size * ohw_tile_size * sizeof(param.bias_type);
        }
295

296
        return {nullptr, {im2col, matmul_dst, bias_temp, matmul_compute}};
297 298 299
    }
};

300 301 302 303
namespace {
static fallback::MatrixMulImpl::KernSizeParam get_matmul_kern_param(
        const fallback::ConvBiasImpl::NCBKernSizeParam& param,
        size_t ohw_tile_size, size_t oc_tile_size) {
304
    auto format = param::MatrixMul::Format::DEFAULT;
305
    size_t pack_oc_size = pack_size(param.filter_meta.format);
306 307
    if (param.filter_meta.format == param::ConvBias::Format::NCHW44) {
        format = param::MatrixMul::Format::MK4;
308 309
    } else if (param.filter_meta.format ==
               param::ConvBias::Format::NCHW44_DOT) {
310
        format = param::MatrixMul::Format::MK4_DOT;
311
    }
312 313 314 315
    size_t M = oc_tile_size;
    size_t N = ohw_tile_size;
    size_t K = param.filter_meta.icpg * param.filter_meta.spatial[0] *
               param.filter_meta.spatial[1];
316 317
    size_t LDA = pack_oc_size * K, LDB = pack_oc_size * N,
           LDC = N * pack_oc_size;
318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333
    bool is_dst_8bit = (param.src_type.enumv() == DTypeEnum::QuantizedS8 &&
                        param.dst_type.enumv() == DTypeEnum::QuantizedS8) ||
                       (param.src_type.enumv() == DTypeEnum::Quantized8Asymm &&
                        param.dst_type.enumv() == DTypeEnum::Quantized8Asymm);
    return {param.filter_type,
            param.src_type,
            is_dst_8bit ? param.bias_type : param.dst_type,
            M,
            N,
            K,
            LDA,
            LDB,
            LDC,
            false,
            false,
            param::MatrixMul::ComputeMode::DEFAULT,
334
            format};
335 336
}

337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353
static void choice_ohw_oc_block(
        const fallback::ConvBiasImpl::NCBKernSizeParam& param,
        size_t& oc_tile_size, size_t& ohw_tile_size, size_t block_m,
        size_t block_n, const size_t m_ohw_tile_size,
        fallback::MatrixMulImpl::AlgoBase::PackMode pack_mode) {
    //! calculate m_oc_tile_size in choice_ohw_oc_block() fucntion,
    //! when ohw_tile_size < this value ohw_tile_size = ohw
    static constexpr size_t DEFAULT_OHW_MIN_TILE_SIZE = 32;
    //! when nr_threads > 1 and round(ohw,nr_threads)>nr_threads,
    //! oc_tile_size = DEFAULT_OC_TILE_SIZE
    static constexpr size_t DEFAULT_OC_TILE_SIZE = 512;
    //! when oc_tile_size > this value m_oc_tile_size =
    //! DEFAULT_OC_MAX_TILE_SIZE
    static constexpr size_t DEFAULT_OC_MAX_TILE_SIZE = 1024;
    //! when oc_tile_size < this value oc_tile_size =
    //! DEFAULT_OC_MIN_TILE_SIZE the purpose is aligning the calculation
    static constexpr size_t DEFAULT_OC_MIN_TILE_SIZE = 128;
354 355 356
    size_t nr_threads = param.nr_threads;
    size_t OC = param.filter_meta.ocpg;
    size_t ohw = param.osz[0] * param.osz[1];
357 358
    oc_tile_size = DEFAULT_OC_TILE_SIZE;
    ohw_tile_size = m_ohw_tile_size;
359

360 361
    oc_tile_size = std::min(oc_tile_size, OC);
    ohw_tile_size = std::min(ohw_tile_size, ohw);
362 363

    if (nr_threads > 1) {
364 365 366 367 368 369 370 371 372
        if (ohw / ohw_tile_size < nr_threads) {
            ohw_tile_size = round_up(div_ceil(ohw, nr_threads), block_n);
            if (ohw_tile_size < DEFAULT_OHW_MIN_TILE_SIZE) {
                ohw_tile_size = ohw;
                oc_tile_size = round_up(div_ceil(OC, nr_threads), block_m);
                if (oc_tile_size > DEFAULT_OC_MAX_TILE_SIZE) {
                    oc_tile_size = DEFAULT_OC_MAX_TILE_SIZE;
                } else if (oc_tile_size < DEFAULT_OC_MIN_TILE_SIZE) {
                    oc_tile_size = DEFAULT_OC_MIN_TILE_SIZE;
373 374 375 376
                }
            }
        }
    } else {
377 378
        //! in no_pack mode don't do block operation when using single thread
        if (pack_mode == fallback::MatrixMulImpl::AlgoBase::PackMode::NO_PACK) {
379 380
            ohw_tile_size = ohw;
            oc_tile_size = OC;
381 382 383 384
        }
    }
}

385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452
static size_t packA_group_size(
        const MatrixMulImpl::AlgoBase* matmul_algo,
        const fallback::MatrixMulImpl::KernSizeParam& matmul_param,
        const fallback::MatrixMulImpl::AlgoBase::MatmulDescription& matmul_desc,
        size_t packa_parallel_times) {
    if (matmul_desc.packmode ==
        fallback::MatrixMulImpl::AlgoBase::PackMode::DEFAULT) {
        return matmul_algo->get_bundle(matmul_param).get_size(0);
    } else if (matmul_desc.packmode ==
               fallback::MatrixMulImpl::AlgoBase::PackMode::ONLY_PACKA) {
        return packa_parallel_times *
               matmul_algo->get_bundle(matmul_param).get_size(0);
    }
    megdnn_assert(matmul_desc.packmode ==
                  fallback::MatrixMulImpl::AlgoBase::PackMode::NO_PACK);
    //! nopack mode return 0;
    return 0;
}

static WorkspaceBundle get_thread_bundle(
        const fallback::ConvBiasImpl::NCBKernSizeParam& param,
        const MatrixMulImpl::AlgoBase* matmul_algo,
        const fallback::MatrixMulImpl::KernSizeParam& matmul_param,
        const fallback::MatrixMulImpl::AlgoBase::MatmulDescription& matmul_desc,
        size_t oc_tile_size, size_t ohw_tile_size) {
    if (matmul_desc.packmode == Pack_Mode::DEFAULT) {
        MIDOUT_BEGIN(
                megdnn_fallback_im2col,
                midout_iv("ConvBiasImpl::AlgoIm2col::get_bundle_dft"_hash)) {
            Im2colKerns<Pack_Mode::DEFAULT> defaultkern;
            return defaultkern.get_thread_bundle(param, matmul_param,
                                                 matmul_algo, ohw_tile_size,
                                                 oc_tile_size);
        }
        MIDOUT_END();
    } else if (matmul_desc.packmode ==
               fallback::MatrixMulImpl::AlgoBase::PackMode::ONLY_PACKA) {
        MIDOUT_BEGIN(
                megdnn_fallback_im2col,
                midout_iv(
                        "ConvBiasImpl::AlgoIm2col::get_bundle_onlypacka"_hash)) {
            Im2colKerns<Pack_Mode::ONLY_PACKA> onlypackakern;
            return onlypackakern.get_thread_bundle(param, matmul_param,
                                                   matmul_algo, ohw_tile_size,
                                                   oc_tile_size);
        }
        MIDOUT_END();
    } else {
        megdnn_assert(matmul_desc.packmode ==
                      fallback::MatrixMulImpl::AlgoBase::PackMode::NO_PACK);
        MIDOUT_BEGIN(
                megdnn_fallback_im2col,
                midout_iv(
                        "ConvBiasImpl::AlgoIm2col::get_thread_bundle_nopack"_hash)) {
            Im2colKerns<Pack_Mode::NO_PACK> nopackkern;
            return nopackkern.get_thread_bundle(param, matmul_param,
                                                matmul_algo, ohw_tile_size,
                                                oc_tile_size);
        }
        MIDOUT_END();
    }
    return {nullptr, {}};
}

static WorkspaceBundle get_bundle(
        const fallback::ConvBiasImpl::NCBKernSizeParam& param,
        MatrixMulImpl::AlgoBase* matmul_algo, size_t oc_tile_size,
        size_t ohw_tile_size) {
453 454 455 456 457 458 459 460 461 462 463 464 465 466 467
    UNPACK_CONV_F32_NCB_KERN_SIZES(param);
    MEGDNN_MARK_USED_VAR(OC);
    MEGDNN_MARK_USED_VAR(OH);
    MEGDNN_MARK_USED_VAR(OW);
    MEGDNN_MARK_USED_VAR(FH);
    MEGDNN_MARK_USED_VAR(FW);
    MEGDNN_MARK_USED_VAR(SW);
    MEGDNN_MARK_USED_VAR(SH);

    auto IW2 = IH + 2 * PH;
    auto IH2 = IW + 2 * PW;
    bool no_need_pading = (PH == 0 && PW == 0);
    size_t padding = 0, packa_size = 0, packa_group_size = 0;
    size_t nr_threads = param.nr_threads;
    size_t GROUP = param.filter_meta.group;
468 469 470 471 472 473 474 475 476
    fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
            matmul_algo->matmul_description();
    bool default_pack = matmul_desc.packmode == Pack_Mode::DEFAULT;

    //! packmode is default should use oc
    //! packmode is onlypackA should use oc_tile_size
    auto im2col_kern_param = get_matmul_kern_param(
            param, ohw_tile_size, default_pack ? OC : oc_tile_size);
    if (is_enable_filter_preprocess(param)) {
477
        packa_group_size = 0;
478 479 480 481
    } else {
        size_t oc_parallel_times = div_ceil<size_t>(OC, oc_tile_size);
        packa_group_size = packA_group_size(matmul_algo, im2col_kern_param,
                                            matmul_desc, oc_parallel_times);
482
    }
483

484 485 486 487 488 489
    if (no_need_pading) {
        padding = 0;  //! not need  padding
    } else {
        padding = (GROUP * N * IC * IH2 * IW2) *
                  sizeof(param.src_type);  //! for padding
    }
490

491
    packa_size = GROUP * packa_group_size;  //! for packA  size = GROUP * a_size
492

493 494 495
    WorkspaceBundle ws =
            get_thread_bundle(param, matmul_algo, im2col_kern_param,
                              matmul_desc, oc_tile_size, ohw_tile_size);
496 497
    return {nullptr,
            {padding, packa_size, ws.total_size_in_bytes() * nr_threads}};
498 499
}

500 501
}  // namespace

502
size_t ConvBiasImpl::AlgoIm2col::get_workspace(
503
        const NCBKernSizeParam& p) const {
504
    MIDOUT_BEGIN(megdnn_fallback_im2col, 0, 0) {
505 506 507 508 509 510 511 512
        fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
                m_matmul_algo->matmul_description();
        size_t oc_tile_size = 0, ohw_tile_size = 0;
        choice_ohw_oc_block(p, oc_tile_size, ohw_tile_size,
                            matmul_desc.innerblocksize.m, matmul_desc.innerblocksize.n,
                            m_ohw_tile_size, matmul_desc.packmode);
        return get_bundle(p, m_matmul_algo, oc_tile_size, ohw_tile_size)
                .total_size_in_bytes();
513 514 515 516 517 518
    }
    MIDOUT_END();
    return 0;
}

SmallVector<ConvBiasImpl::NCBKern> ConvBiasImpl::AlgoIm2col::dispatch_kerns(
519
        const NCBKernSizeParam& param) const {
520
    MIDOUT_BEGIN(megdnn_fallback_im2col, 0, 1) {
521 522 523 524 525 526 527
        UNPACK_CONV_F32_NCB_KERN_SIZES(param);
        MEGDNN_MARK_USED_VAR(SH);
        MEGDNN_MARK_USED_VAR(SW);
        MEGDNN_MARK_USED_VAR(IH);
        MEGDNN_MARK_USED_VAR(IW);
        MEGDNN_MARK_USED_VAR(FH);
        MEGDNN_MARK_USED_VAR(FW);
528
        size_t oc_tile_size = 0, ohw_tile_size = 0;
529
        size_t ohw = OH * OW;
530 531
        size_t GROUP = param.filter_meta.group;
        bool need_padding = (PH != 0 || PW != 0);
532

533
        fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
534 535
                m_matmul_algo->matmul_description();

536 537 538 539
        bool default_pack = matmul_desc.packmode == Pack_Mode::DEFAULT;
        bool no_pack = matmul_desc.packmode == Pack_Mode::NO_PACK;
        bool only_packA = matmul_desc.packmode == Pack_Mode::ONLY_PACKA;
        bool enable_filter_preprocess = is_enable_filter_preprocess(param);
540
        choice_ohw_oc_block(param, oc_tile_size, ohw_tile_size,
541 542 543
                            matmul_desc.innerblocksize.m,
                            matmul_desc.innerblocksize.n, m_ohw_tile_size,
                            matmul_desc.packmode);
544

545
        WorkspaceBundle bundle = get_bundle(param,m_matmul_algo,oc_tile_size,ohw_tile_size);
546 547
        size_t ohw_parallel_times = div_ceil(ohw, ohw_tile_size);
        size_t oc_parallel_times = div_ceil<size_t>(OC, oc_tile_size);
548
        size_t packa_parallel_times = 0;
549
        size_t pack_oc_size = pack_size(param.filter_meta.format);
550

551
        if (only_packA) {
552
            packa_parallel_times = div_ceil<size_t>(OC, oc_tile_size);
553
        } else if (default_pack) {
554 555
            packa_parallel_times =
                    div_ceil<size_t>(OC, matmul_desc.innerblocksize.m);
556 557 558
        }

        auto matmul_param = get_matmul_kern_param(
559
                param, ohw_tile_size, default_pack ? OC : oc_tile_size);
560

561 562 563
        WorkspaceBundle bundle_thread =
                get_thread_bundle(param, m_matmul_algo, matmul_param,
                                  matmul_desc, oc_tile_size, ohw_tile_size);
564 565 566 567 568 569 570
        StrategyParam strategyparam;
        strategyparam.ohw = ohw;
        strategyparam.is_dst_8bit =
                (param.src_type.enumv() == DTypeEnum::QuantizedS8 &&
                 param.dst_type.enumv() == DTypeEnum::QuantizedS8) ||
                (param.src_type.enumv() == DTypeEnum::Quantized8Asymm &&
                 param.dst_type.enumv() == DTypeEnum::Quantized8Asymm);
571
        strategyparam.is_ohw_size_bigger = (ohw_tile_size >= ohw);
572 573
        strategyparam.skip_copy_dst =
                strategyparam.is_ohw_size_bigger && !strategyparam.is_dst_8bit;
574
        strategyparam.oc_tile_size = oc_tile_size;
575
        strategyparam.pack_oc_size = pack_oc_size;
576 577 578
        strategyparam.enable_filter_preprocess = enable_filter_preprocess;
        strategyparam.packA_group_size = packA_group_size(
                m_matmul_algo, matmul_param, matmul_desc, packa_parallel_times);
579

580 581 582 583
        SmallVector<ConvBiasImpl::NCBKern> ret_kern;
        MIDOUT_BEGIN(
                megdnn_fallback_im2col,
                midout_iv("ConvBiasImpl::AlgoIm2col::dispatch_kerns"_hash)) {
584 585 586 587
            StrategyBase* im2colstrategy =
                    Factory::get_im2col_strategy(param, m_matmul_algo);
            auto kern_padding = [bundle, im2colstrategy,
                                 pack_oc_size = pack_oc_size](
588
                                        const NCBKernParam& param,
589 590
                                        const NCBKernIndex& ncb_index) mutable {
                bundle.set(param.workspace_ptr);
591 592
                im2colstrategy->copy_padding_kern(bundle, param, ncb_index,
                                                  pack_oc_size);
593 594 595
            };

            auto kern_packA = [bundle, matmul_algo = m_matmul_algo,
596
                               matmul_param, im2colstrategy,
597 598
                               strategyparam = strategyparam,
                               matmul_desc = matmul_desc](
599 600 601
                                      const NCBKernParam& param,
                                      const NCBKernIndex& ncb_index) mutable {
                bundle.set(param.workspace_ptr);
602 603 604 605

                im2colstrategy->packA_kern(bundle, param, matmul_param,
                                           matmul_algo, ncb_index, matmul_desc,
                                           strategyparam);
606 607
            };
            if (default_pack) {
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 635 636 637 638 639
                MIDOUT_BEGIN(
                        megdnn_fallback_im2col,
                        midout_iv(
                                "ConvBiasImpl::AlgoIm2col::dispatch_kerns_default_pack"_hash)) {
                    auto kern_compute_default =
                            [bundle, bundle_thread, matmul_param,
                             matmul_algo = m_matmul_algo,
                             ohw_tile_size = ohw_tile_size,
                             strategyparam = strategyparam,
                             matmul_desc = matmul_desc, im2colstrategy](
                                    const NCBKernParam& param,
                                    const NCBKernIndex& ncb_index) mutable {
                                bundle.set(param.workspace_ptr);
                                Im2colKerns<Pack_Mode::DEFAULT>::kerns(
                                        bundle, bundle_thread, param,
                                        matmul_param, matmul_algo, matmul_desc,
                                        strategyparam, ncb_index, ohw_tile_size,
                                        im2colstrategy);
                            };
                    if (!enable_filter_preprocess) {
                        ret_kern.push_back(
                                {kern_packA, {GROUP, packa_parallel_times}});
                    }
                    if (need_padding) {
                        ret_kern.push_back(
                                {kern_padding,
                                 {param.n, GROUP, IC / pack_oc_size}});
                    }
                    ret_kern.push_back({kern_compute_default,
                                        {N, GROUP, ohw_parallel_times,
                                         oc_parallel_times}});
                    return ret_kern;
640
                }
641 642
                MIDOUT_END();
                return {};
643
            } else if (only_packA) {
644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674
                MIDOUT_BEGIN(
                        megdnn_fallback_im2col,
                        midout_iv(
                                "ConvBiasImpl::AlgoIm2col::dispatch_kerns_onlypacka"_hash)) {
                    auto kern_compute_onlypackA =
                            [bundle, bundle_thread, matmul_param,
                             matmul_algo = m_matmul_algo,
                             strategyparam = strategyparam,
                             ohw_tile_size = ohw_tile_size,
                             matmul_desc = matmul_desc, im2colstrategy](
                                    const NCBKernParam& param,
                                    const NCBKernIndex& ncb_index) mutable {
                                bundle.set(param.workspace_ptr);
                                Im2colKerns<Pack_Mode::ONLY_PACKA>::kerns(
                                        bundle, bundle_thread, param,
                                        matmul_param, matmul_algo, matmul_desc,
                                        strategyparam, ncb_index, ohw_tile_size,
                                        im2colstrategy);
                            };
                    if (!enable_filter_preprocess) {
                        ret_kern.push_back(
                                {kern_packA, {GROUP, packa_parallel_times}});
                    }
                    if (need_padding) {
                        ret_kern.push_back(
                                {kern_padding, {param.n, GROUP, IC}});
                    }
                    ret_kern.push_back({kern_compute_onlypackA,
                                        {N, GROUP, ohw_parallel_times,
                                         oc_parallel_times}});
                    return ret_kern;
675
                }
676 677
                MIDOUT_END();
                return {};
678
            } else if (no_pack) {
679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705
                MIDOUT_BEGIN(
                        megdnn_fallback_im2col,
                        midout_iv(
                                "ConvBiasImpl::AlgoIm2col::dispatch_kerns_no_pack"_hash)) {
                    auto kern_compute_nopack =
                            [bundle, bundle_thread, matmul_param,
                             matmul_algo = m_matmul_algo,
                             strategyparam = strategyparam,
                             ohw_tile_size = ohw_tile_size,
                             matmul_desc = matmul_desc, im2colstrategy](
                                    const NCBKernParam& param,
                                    const NCBKernIndex& ncb_index) mutable {
                                bundle.set(param.workspace_ptr);
                                Im2colKerns<Pack_Mode::NO_PACK>::kerns(
                                        bundle, bundle_thread, param,
                                        matmul_param, matmul_algo, matmul_desc,
                                        strategyparam, ncb_index, ohw_tile_size,
                                        im2colstrategy);
                            };
                    if (need_padding) {
                        ret_kern.push_back(
                                {kern_padding, {param.n, GROUP, IC}});
                    }
                    ret_kern.push_back({kern_compute_nopack,
                                        {N, GROUP, ohw_parallel_times,
                                         oc_parallel_times}});
                    return ret_kern;
706
                }
707 708
                MIDOUT_END();
                return {};
709
            }
710
            return {};
711 712 713
        }
        MIDOUT_END();
        return {};
714 715 716 717 718 719
    }
    MIDOUT_END();
    return {};
}

bool ConvBiasImpl::AlgoIm2col::usable(
720
         const NCBKernSizeParam& param,
721 722
        AlgoSelectionStrategy /*algo_selection_strategy*/) const {
    MIDOUT_BEGIN(megdnn_fallback_im2col, 0, 2) {
723 724 725 726
        auto format = param.filter_meta.format;
        if (format != param::ConvBias::Format::NCHW &&
            format != param::ConvBias::Format::NCHW44_DOT &&
            format != param::ConvBias::Format::NCHW44) {
727 728 729
            return false;
        }

730 731 732 733 734
        if(param.src_type.enumv() != param.filter_type.enumv()) {
            return false;
        }

        if (param.src_type.enumv() != DTypeEnum::Int8 &&
735 736 737 738 739 740 741 742
            param.src_type.enumv() != DTypeEnum::QuantizedS8 &&
            param.src_type.enumv() != DTypeEnum::Quantized8Asymm &&
#if !MEGDNN_DISABLE_FLOAT16
            param.src_type.enumv() != DTypeEnum::Float16 &&
#endif
            param.src_type.enumv() != DTypeEnum::Float32) {
            return false;
        }
743
        //! make sure 8x8x16 and 8x8x32 biasmode is  nobias and nonlineMode is
744 745
        //! identity otherwise return false mean that 8x8x32 and 8x8x16 not
        //! support PostProcess
746 747 748 749 750 751 752
        if (param.dst_type.enumv() == DTypeEnum::Int16 ||
            param.dst_type.enumv() == DTypeEnum::Int32 ||
            param.dst_type.enumv() == DTypeEnum::QuantizedS32) {
            if (param.bias_mode != megdnn::BiasMode::NO_BIAS ||
                param.nonlineMode != megdnn::NonlineMode::IDENTITY) {
                return false;
            }
753
        }
754
        fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
755
                m_matmul_algo->matmul_description();
756 757 758 759 760 761 762
        //! only matmul's packmode is packa or default support weight preprocess
        if (is_enable_filter_preprocess(param) &&
            (matmul_desc.packmode ==
             fallback::MatrixMulImpl::AlgoBase::PackMode::NO_PACK)) {
            return false;
        }

763 764
        if (format == param::ConvBias::Format::NCHW44 ||
            format == param::ConvBias::Format::NCHW44_DOT) {
765
            //! current NCHW44 im2col only support DEFAULT mode matmul
766
            if (matmul_desc.packmode != Pack_Mode::DEFAULT) {
767 768 769 770 771 772 773
                return false;
                //! nchw44 hybird mode and channel wise is not support
            } else if (param.filter_meta.icpg < 4_z ||
                       param.filter_meta.icpg == 1 ||
                       param.filter_meta.ocpg == 1) {
                return false;
            }
774 775
        }

776
        size_t oc_tile_size = 0, ohw_tile_size = 0;
777
        choice_ohw_oc_block(param, oc_tile_size, ohw_tile_size,
778 779 780
                            matmul_desc.innerblocksize.m,
                            matmul_desc.innerblocksize.n, m_ohw_tile_size,
                            matmul_desc.packmode);
781
        fallback::MatrixMulImpl::KernSizeParam matmul_param =
782
                get_matmul_kern_param(param, ohw_tile_size, oc_tile_size);
783 784
        bool matmulusable = m_matmul_algo->usable(matmul_param);
        return matmulusable &&
785 786
               (!(param.filter_meta.spatial[0] ==
                          param.filter_meta.spatial[1] &&
787
                  param.filter_meta.spatial[0] == 1 &&
788 789
                  param.filter_meta.stride[0] == param.filter_meta.stride[1] &&
                  param.filter_meta.stride[0] == 1)) &&
790 791 792 793 794 795 796 797 798
               (param.filter_meta.dilation[0] ==
                        param.filter_meta.dilation[1] &&
                param.filter_meta.dilation[0] == 1) &&
               param.compute_mode == param::ConvBias::ComputeMode::DEFAULT;
    }
    MIDOUT_END();
    return false;
}

799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 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
SmallVector<TensorLayout>
ConvBiasImpl::AlgoIm2col::deduce_preprocessed_filter_layout(
        const NCBKernSizeParam& param) const {
    MIDOUT_BEGIN(
            megdnn_fallback_im2col,
            midout_iv(
                    "ConvBiasImpl::AlgoIm2col::deduce_preprocessed_filter_layout"_hash)) {
        fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
                m_matmul_algo->matmul_description();

        //! only support default_pack and only_packa mode
        if (matmul_desc.packmode == Pack_Mode::NO_PACK) {
            return {};
        }

        size_t GROUP = param.filter_meta.group;
        bool default_pack = matmul_desc.packmode == Pack_Mode::DEFAULT;

        size_t OC = param.filter_meta.ocpg;
        SmallVector<TensorLayout> preprocessed_layouts;
        size_t oc_tile_size = 0, ohw_tile_size = 0;
        choice_ohw_oc_block(param, oc_tile_size, ohw_tile_size,
                            matmul_desc.innerblocksize.m,
                            matmul_desc.innerblocksize.n, m_ohw_tile_size,
                            matmul_desc.packmode);
        auto matmul_param = get_matmul_kern_param(
                param, ohw_tile_size, default_pack ? OC : oc_tile_size);

        size_t packa_parallel_times = div_ceil<size_t>(
                OC, default_pack ? matmul_desc.innerblocksize.m : oc_tile_size);

        size_t packa_group_size = packA_group_size(
                m_matmul_algo, matmul_param, matmul_desc, packa_parallel_times);
        preprocessed_layouts.push_back(
                {{GROUP, packa_group_size}, dtype::Int8()});
        return preprocessed_layouts;
    }
    MIDOUT_END();
    return {};
}

SmallVector<ConvBiasImpl::NCBKern>
ConvBiasImpl::AlgoIm2col::dispatch_preprocess_kerns(
        const NCBKernSizeParam& param) const {
    MIDOUT_BEGIN(megdnn_fallback_im2col, 0, 3) {
        size_t OC = param.filter_meta.ocpg;
        size_t oc_tile_size = 0, ohw_tile_size = 0;
        size_t GROUP = param.filter_meta.group;
        fallback::MatrixMulImpl::AlgoBase::MatmulDescription matmul_desc =
                m_matmul_algo->matmul_description();
        choice_ohw_oc_block(param, oc_tile_size, ohw_tile_size,
                            matmul_desc.innerblocksize.m,
                            matmul_desc.innerblocksize.n, m_ohw_tile_size,
                            matmul_desc.packmode);
        WorkspaceBundle bundle =
                get_bundle(param, m_matmul_algo, oc_tile_size, ohw_tile_size);

        Pack_Mode packmode = matmul_desc.packmode;
        bool default_pack = packmode == Pack_Mode::DEFAULT;
        bool only_packA = packmode == Pack_Mode::ONLY_PACKA;
        size_t packa_parallel_times = 0;

        if (only_packA) {
            packa_parallel_times = div_ceil<size_t>(OC, oc_tile_size);
        } else if (default_pack) {
            packa_parallel_times =
                    div_ceil<size_t>(OC, matmul_desc.innerblocksize.m);
        } else {
            //! if nopack return null so that OprWeightPreprocessProxy can run
            //! with nopack mode
            return {};
        }
        auto matmul_param = get_matmul_kern_param(
                param, ohw_tile_size, default_pack ? OC : oc_tile_size);

        StrategyParam strategyparam;
        strategyparam.enable_filter_preprocess =
                is_enable_filter_preprocess(param);
        strategyparam.packA_group_size = packA_group_size(
                m_matmul_algo, matmul_param, matmul_desc, packa_parallel_times);
        SmallVector<ConvBiasImpl::NCBKern> ret_kern;
        StrategyBase* im2colstrategy =
                Factory::get_im2col_strategy(param, m_matmul_algo);

        auto kern_packA = [bundle, matmul_algo = m_matmul_algo, matmul_param,
                           im2colstrategy, strategyparam = strategyparam,
                           matmul_desc = matmul_desc](
                                  const NCBKernParam& param,
                                  const NCBKernIndex& ncb_index) mutable {
            bundle.set(param.workspace_ptr);
            im2colstrategy->packA_kern(bundle, param, matmul_param, matmul_algo,
                                       ncb_index, matmul_desc, strategyparam);
        };
        ret_kern.push_back({kern_packA, {GROUP, packa_parallel_times}});
        return ret_kern;
    }
    MIDOUT_END();
    return {};
}

899
// vim: syntax=cpp.doxygen