| 66 | |
| 67 | template <class Conv> |
| 68 | uint64_t eval_conv_computation( |
| 69 | const TensorShape& src_shape, const TensorShape& filter_shape, |
| 70 | const TensorShape& dst_shape, cg::OperatorNodeBase* opr) { |
| 71 | using Param = opr::ConvolutionForward::Param; |
| 72 | auto&& param = opr->cast_final_safe<Conv>().param(); |
| 73 | |
| 74 | if (param.format == Param::Format::NHWCD4) { |
| 75 | size_t fh, fw; |
| 76 | size_t group = 1; |
| 77 | if (param.sparse == Param::Sparse::DENSE) { |
| 78 | fh = filter_shape[1]; |
| 79 | fw = filter_shape[2]; |
| 80 | group = 1; |
| 81 | } else { |
| 82 | // chanwise conv |
| 83 | mgb_assert(param.sparse == Param::Sparse::GROUP); |
| 84 | fh = filter_shape[2]; |
| 85 | fw = filter_shape[3]; |
| 86 | group = filter_shape[0]; |
| 87 | |
| 88 | if (filter_shape.ndim == 5) { |
| 89 | group *= 4; |
| 90 | } |
| 91 | } |
| 92 | return dst_shape.total_nr_elems() * fh * fw * src_shape[2] * 4 / group * 2; |
| 93 | } |
| 94 | auto eval_conv_computation_nchwx = [¶m, &src_shape, &filter_shape, |
| 95 | &dst_shape]() -> uint64_t { |
| 96 | size_t fh, fw; |
| 97 | bool hybird_nchwx = false; |
| 98 | size_t group = 1; |
| 99 | if (param.sparse == Param::Sparse::DENSE) { |
| 100 | //! if nchwxx mode src is nchw output is nchwxx |
| 101 | if (dst_shape.ndim == 5 && src_shape.ndim == 4) { |
| 102 | fh = filter_shape[1]; |
| 103 | fw = filter_shape[2]; |
| 104 | hybird_nchwx = true; |
| 105 | } else { |
| 106 | fh = filter_shape[2]; |
| 107 | fw = filter_shape[3]; |
| 108 | } |
| 109 | group = 1; |
| 110 | } else { |
| 111 | mgb_assert(param.sparse == Param::Sparse::GROUP); |
| 112 | fh = filter_shape[3]; |
| 113 | fw = filter_shape[4]; |
| 114 | group = filter_shape[0]; |
| 115 | } |
| 116 | if (param.format == Param::Format::NCHW88) { |
| 117 | //! if channel wise weight layout is {group/8, FH, FW, 1, 1, 8} |
| 118 | if (filter_shape[1] == 1 && filter_shape[2] == 1) { |
| 119 | group *= 8; |
| 120 | } |
| 121 | size_t computation = |
| 122 | dst_shape.total_nr_elems() * fh * fw * src_shape[1] / group * 2; |
| 123 | return hybird_nchwx ? computation : computation * 8; |
| 124 | } |
| 125 | if (param.format == Param::Format::NCHW44 || |
nothing calls this directly
no test coverage detected