| 88 | } |
| 89 | |
| 90 | std::optional<ValueRefList> matrix_mul_grad_rule( |
| 91 | const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad, |
| 92 | CustomBackward& backward) { |
| 93 | auto&& matmul = op.cast_final_safe<MatrixMul>(); |
| 94 | size_t dimA = matmul.dimA; |
| 95 | size_t dimB = matmul.dimB; |
| 96 | auto&& param = matmul.param(); |
| 97 | auto&& policy = matmul.policy(); |
| 98 | mgb_assert(inputs.size() == 2); |
| 99 | std::array<ValueRef, 2> inps, input_shapes; |
| 100 | for (size_t i = 0; i < 2; ++i) { |
| 101 | if (inputs_require_grad[i ^ 1]) { |
| 102 | inps[i] = inputs[i]; |
| 103 | input_shapes[i] = get_shape(inputs[i]); |
| 104 | } |
| 105 | } |
| 106 | auto maker = CustomGradMaker(backward, inputs.size()); |
| 107 | maker.output_size(1).output_captured(0, false); |
| 108 | maker.backward([inps_ = std::move(inps), input_shapes_ = std::move(input_shapes), |
| 109 | param, policy, dimA, dimB](Span<ValueRef> grads) { |
| 110 | mgb_assert(grads.size() == 1); |
| 111 | ValueRef grad = grads[0]; |
| 112 | SmallVector<ValueRef> ret(2); |
| 113 | if (!grad) { |
| 114 | return ret; |
| 115 | } |
| 116 | size_t dimG = std::max(dimA, dimB); |
| 117 | if (inps_[1]) { |
| 118 | if (param.transposeA) { |
| 119 | // A^T(2) @ B(2) = G(2), A'(2) = B'(2) @ G'^T(2) -> MatrixMul |
| 120 | auto&& grad_op = MatrixMul::make( |
| 121 | param.transposeB, true, param.compute_mode, param.format, |
| 122 | policy.strategy, policy.workspace_limit, dimB, dimG); |
| 123 | ret[0] = imperative::apply(*grad_op, inps_[1], grad)[0]; |
| 124 | } else { |
| 125 | // A(>=2) @ B(2) = G(>=2), A'(>=2) = G'(>=2) @ B(2) -> MatrixMul |
| 126 | auto&& grad_op = MatrixMul::make( |
| 127 | false, !param.transposeB, param.compute_mode, param.format, |
| 128 | policy.strategy, policy.workspace_limit, dimG, dimB); |
| 129 | ret[0] = imperative::apply(*grad_op, grad, inps_[1])[0]; |
| 130 | } |
| 131 | } |
| 132 | if (inps_[0]) { |
| 133 | if (param.transposeB) { |
| 134 | // A(>=2) @ B^T(2) = G(>=2), B'(2) = G'^T(>=2) @ A(>=2) -> MatrixMul |
| 135 | // (specialized) |
| 136 | auto&& grad_op = MatrixMul::make( |
| 137 | true, param.transposeA, param.compute_mode, param.format, |
| 138 | policy.strategy, policy.workspace_limit, dimG, dimA); |
| 139 | ret[1] = imperative::apply(*grad_op, grad, inps_[0])[0]; |
| 140 | } else { |
| 141 | // A(>=2) @ B(2) = G(>=2), B'(2) = G'(>=2) @ A(>=2) -> MatrixMul |
| 142 | // (specialized) |
| 143 | auto&& grad_op = MatrixMul::make( |
| 144 | !param.transposeA, false, param.compute_mode, param.format, |
| 145 | policy.strategy, policy.workspace_limit, dimA, dimG); |
| 146 | ret[1] = imperative::apply(*grad_op, inps_[0], grad)[0]; |
| 147 | } |
nothing calls this directly
no test coverage detected