| 74 | } |
| 75 | |
| 76 | size_t MatrixMul::get_workspace_size_bytes( |
| 77 | const TensorShapeArray& input_shapes, |
| 78 | const TensorShapeArray& output_shapes) const { |
| 79 | // we may change transepose param in the impl, so get the max possible |
| 80 | // workspace by trying all cases |
| 81 | // current implementation in megdnn guarantees that workspaces in different |
| 82 | // cases are on the same order of magnitude |
| 83 | auto mo = megdnn_opr(); |
| 84 | auto&& tparam = mo->param(); |
| 85 | size_t a, b, c, d; |
| 86 | mgb_assert(input_shapes.size() == 2 && output_shapes.size() == 1); |
| 87 | TensorLayout i0(input_shapes[0], input(0)->dtype()), |
| 88 | i1(input_shapes[1], input(1)->dtype()), |
| 89 | out(output_shapes[0], output(0)->dtype()); |
| 90 | |
| 91 | auto transpose = [](TensorLayout& dst, bool& param) { |
| 92 | std::swap(dst.shape[0], dst.shape[1]); |
| 93 | dst.stride[0] = dst[1]; |
| 94 | param ^= 1; |
| 95 | }; |
| 96 | MGB_TRY { |
| 97 | megdnn_opr()->execution_policy() = {}; |
| 98 | a = AlgoChooser<megdnn::MatrixMul>::setup_algo( |
| 99 | {i0, i1, out}, megdnn_opr(), this); |
| 100 | //! Here we just want to save the execution policy got from setup_algo, |
| 101 | //! while change the delaration of get_workspace_in_bytes may cause |
| 102 | //! many changes. |
| 103 | const_cast<MatrixMul*>(this) |
| 104 | ->m_cadidate_execution_policies[get_mask_from_matmul(tparam)] = |
| 105 | megdnn_opr()->execution_policy(); |
| 106 | megdnn_opr()->execution_policy() = {}; |
| 107 | transpose(i0, tparam.transposeA); |
| 108 | b = AlgoChooser<megdnn::MatrixMul>::setup_algo( |
| 109 | {i0, i1, out}, megdnn_opr(), this); |
| 110 | const_cast<MatrixMul*>(this) |
| 111 | ->m_cadidate_execution_policies[get_mask_from_matmul(tparam)] = |
| 112 | megdnn_opr()->execution_policy(); |
| 113 | megdnn_opr()->execution_policy() = {}; |
| 114 | transpose(i1, tparam.transposeB); |
| 115 | c = AlgoChooser<megdnn::MatrixMul>::setup_algo( |
| 116 | {i0, i1, out}, megdnn_opr(), this); |
| 117 | const_cast<MatrixMul*>(this) |
| 118 | ->m_cadidate_execution_policies[get_mask_from_matmul(tparam)] = |
| 119 | megdnn_opr()->execution_policy(); |
| 120 | megdnn_opr()->execution_policy() = {}; |
| 121 | transpose(i0, tparam.transposeA); |
| 122 | d = AlgoChooser<megdnn::MatrixMul>::setup_algo( |
| 123 | {i0, i1, out}, megdnn_opr(), this); |
| 124 | const_cast<MatrixMul*>(this) |
| 125 | ->m_cadidate_execution_policies[get_mask_from_matmul(tparam)] = |
| 126 | megdnn_opr()->execution_policy(); |
| 127 | megdnn_opr()->execution_policy() = {}; |
| 128 | } |
| 129 | MGB_FINALLY({ tparam = this->param(); }); |
| 130 | return std::max(std::max(a, b), std::max(c, d)); |
| 131 | } |
nothing calls this directly
no test coverage detected