| 42 | } |
| 43 | |
| 44 | ConvBiasForward::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; |
nothing calls this directly
no test coverage detected