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

Method get_algorithm_heuristic

dnn/src/cuda/conv_bias/opr_impl.cpp:44–238  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

42}
43
44ConvBiasForward::Algorithm* ConvBiasForwardImpl::get_algorithm_heuristic(
45 const TensorLayout& src, const TensorLayout& filter, const TensorLayout& bias,
46 const TensorLayout& z, const TensorLayout& dst, size_t workspace_limit_in_bytes,
47 const AlgoAttribute& positive_attr, const AlgoAttribute& negative_attr) {
48 using namespace conv_bias;
49 AlgoBase::SizeArgs args{this, src, filter, bias, z, dst};
50#if CUDNN_VERSION >= 8020 && 0 // FIXME(hc): need fix
51 if (sm_algo_pack.cudnn_conv_v8.is_available_attribute(
52 args, positive_attr, negative_attr, workspace_limit_in_bytes)) {
53 return &sm_algo_pack.cudnn_conv_v8;
54 }
55 if (sm_algo_pack.cudnn_conv_bias_activation_v8.is_available_attribute(
56 args, positive_attr, negative_attr, workspace_limit_in_bytes)) {
57 return &sm_algo_pack.cudnn_conv_bias_activation_v8;
58 }
59#endif
60
61 auto dst_layout = *args.dst_layout;
62 if (dst_layout.dtype.enumv() != args.bias_layout->dtype.enumv()) {
63 dst_layout.dtype = DType();
64 args.opr->check_or_deduce_dtype_fwd(
65 args.src_layout->dtype, args.filter_layout->dtype, dst_layout.dtype);
66 }
67 auto conv_args = args;
68
69 auto cudnn_conv_bias_act_from_enum_wrapper =
70 [](cudnnConvolutionFwdAlgo_t algo) -> AlgoBase* {
71 return sm_algo_pack.cudnn_conv_bias_act_from_enum(algo);
72 };
73
74 auto cudnn_conv_from_enum_wrapper =
75 [](cudnnConvolutionFwdAlgo_t algo) -> AlgoBase* {
76 return sm_algo_pack.cudnn_conv_from_enum(algo);
77 };
78
79 auto get_cudnn_algo =
80 [this, &conv_args, &args, workspace_limit_in_bytes, positive_attr,
81 negative_attr](
82 const thin_function<AlgoBase*(cudnnConvolutionFwdAlgo_t)>& cb)
83 -> AlgoBase* {
84 auto cudnn_handle = cuda::cudnn_handle(this->handle());
85 CUDNNForwardDescs desc;
86 conv_args.init_conv_desc(desc);
87#if CUDNN_MAJOR >= 7
88 int max_count = 0;
89 cudnn_check(
90 cudnnGetConvolutionForwardAlgorithmMaxCount(cudnn_handle, &max_count));
91 SmallVector<cudnnConvolutionFwdAlgoPerf_t> algo_perf(max_count);
92 int ret_count = 0;
93 cudnn_check(cudnnGetConvolutionForwardAlgorithm_v7(
94 cudnn_handle, desc.src_desc.desc, desc.filter_desc.desc,
95 desc.conv_desc.conv_desc, desc.dst_desc.desc, max_count, &ret_count,
96 algo_perf.data()));
97 for (int i = 0; i < ret_count; ++i) {
98 auto conv_bias_algo = cb(algo_perf[i].algo);
99 if (conv_bias_algo->is_available_attribute(
100 args, positive_attr, negative_attr, workspace_limit_in_bytes)) {
101 return conv_bias_algo;

Callers

nothing calls this directly

Calls 11

DTypeClass · 0.85
cudnn_handleFunction · 0.85
is_cudnn_supportedFunction · 0.85
cudnn_conv_from_enumMethod · 0.80
init_conv_descMethod · 0.80
cbFunction · 0.50
enumvMethod · 0.45
handleMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected