| 24 | } |
| 25 | |
| 26 | MatrixMulForwardImpl::Algorithm* MatrixMulForwardImpl::get_algorithm_heuristic( |
| 27 | const TensorLayout& A, const TensorLayout& B, const TensorLayout& C, |
| 28 | size_t workspace_limit_in_bytes, const AlgoAttribute& positive_attr, |
| 29 | const AlgoAttribute& negative_attr) { |
| 30 | AlgoBase::SizeArgs args{this, A, B, C}; |
| 31 | if (sm_algo_pack.cublas.is_available_attribute( |
| 32 | args, positive_attr, negative_attr, workspace_limit_in_bytes)) { |
| 33 | return &sm_algo_pack.cublas; |
| 34 | } |
| 35 | #if CUDA_VERSION >= 10010 |
| 36 | if (sm_algo_pack.cublas_lt.is_available_attribute( |
| 37 | args, positive_attr, negative_attr, workspace_limit_in_bytes)) { |
| 38 | return &sm_algo_pack.cublas_lt; |
| 39 | } |
| 40 | #endif |
| 41 | |
| 42 | #if CUDA_VERSION >= 10000 |
| 43 | if (sm_algo_pack.wmma_uint4x4x32.is_available_attribute( |
| 44 | args, positive_attr, negative_attr, workspace_limit_in_bytes)) { |
| 45 | return &sm_algo_pack.wmma_uint4x4x32; |
| 46 | } |
| 47 | #endif |
| 48 | |
| 49 | return megdnn::get_algo_match_attribute<MatrixMulForwardImpl>( |
| 50 | sm_algo_pack.all_algos, args, workspace_limit_in_bytes, |
| 51 | "matrix mul forward", positive_attr, negative_attr); |
| 52 | } |
| 53 | |
| 54 | size_t MatrixMulForwardImpl::get_workspace_in_bytes( |
| 55 | const TensorLayout& A, const TensorLayout& B, const TensorLayout& C) { |
no test coverage detected