MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / get_algorithm_heuristic

Method get_algorithm_heuristic

dnn/src/cuda/convolution3d/opr_impl.cpp:21–87  ·  view source on GitHub ↗

============== Convolution3DForwardImpl ============== */

Source from the content-addressed store, hash-verified

19
20/* ============== Convolution3DForwardImpl ============== */
21Convolution3DForwardImpl::Algorithm* Convolution3DForwardImpl::get_algorithm_heuristic(
22 const TensorLayout& src, const TensorLayout& filter, const TensorLayout& dst,
23 size_t workspace_limit_in_bytes, const AlgoAttribute& positive_attr,
24 const AlgoAttribute& negative_attr) {
25 AlgoBase::SizeArgs args(this, src, filter, dst);
26
27#if CUDNN_MAJOR < 7 || (CUDNN_MAJOR == 7 && CUDNN_MINOR < 5)
28 if (args.filter_meta.group > 1) {
29 // prefer special chanwise impl since as the group conv of cudnn whose
30 // version is lower than v7.5.0 is still slower than our implementation
31 // in many channel-wise cases
32 if (sm_algo_pack.chanwise.is_available_attribute(
33 args, positive_attr, negative_attr, workspace_limit_in_bytes)) {
34 return &sm_algo_pack.chanwise;
35 }
36 }
37#endif
38
39 auto prefer_1x1x1 = [&args, positive_attr, negative_attr,
40 workspace_limit_in_bytes]() {
41 const size_t MAX_BATCH_SIZE_FOR_1x1x1_MAT_ALGO = 4;
42 size_t batch_size = args.src_layout->shape[0];
43 if (batch_size > MAX_BATCH_SIZE_FOR_1x1x1_MAT_ALGO) {
44 return false;
45 }
46 return sm_algo_pack.a1x1x1.is_available_attribute(
47 args, positive_attr, negative_attr, workspace_limit_in_bytes);
48 };
49
50 auto get_cudnn_algo = [this, &args, workspace_limit_in_bytes, positive_attr,
51 negative_attr]() -> Convolution3DForwardImpl::AlgoBase* {
52 auto cudnn_handle = cuda::cudnn_handle(this->handle());
53 cudnnConvolutionFwdAlgo_t algo;
54 CUDNNForwardDescs desc;
55 args.init_desc(desc);
56
57 bool got = cudnn_get_convolution_fwd_algo_helper(
58 cudnn_handle, desc.src_desc.desc, desc.filter_desc.desc,
59 desc.conv_desc.desc, desc.dst_desc.desc, workspace_limit_in_bytes,
60 &algo, positive_attr, negative_attr);
61 if (got) {
62 return static_cast<AlgoBase*>(
63 megdnn::get_algo_match_attribute<Convolution3DForwardImpl>(
64 sm_algo_pack.cudnn_from_enum(algo), positive_attr,
65 negative_attr));
66 } else {
67 return nullptr;
68 }
69 };
70 if (prefer_1x1x1()) {
71 return &sm_algo_pack.a1x1x1;
72 }
73 if (is_cudnn_supported(args)) {
74 if (auto algo = get_cudnn_algo())
75 return algo;
76 }
77
78 if (args.filter_meta.group > 1 &&

Callers

nothing calls this directly

Calls 6

cudnn_handleFunction · 0.85
is_cudnn_supportedFunction · 0.85
handleMethod · 0.45
init_descMethod · 0.45

Tested by

no test coverage detected