| 153 | } |
| 154 | |
| 155 | std::optional<ValueRefList> batched_matrix_mul_grad_rule( |
| 156 | const OpDef& op, Span<ValueRef> inputs, Span<bool> inputs_require_grad, |
| 157 | CustomBackward& backward) { |
| 158 | auto&& bmm = op.cast_final_safe<BatchedMatrixMul>(); |
| 159 | size_t dimA = bmm.dimA; |
| 160 | size_t dimB = bmm.dimB; |
| 161 | auto&& param = bmm.param(); |
| 162 | auto&& policy = bmm.policy(); |
| 163 | mgb_assert(inputs.size() == 2); |
| 164 | std::array<ValueRef, 2> inps, input_shapes; |
| 165 | for (size_t i = 0; i < 2; ++i) { |
| 166 | if (inputs_require_grad[i ^ 1]) { |
| 167 | inps[i] = inputs[i]; |
| 168 | input_shapes[i] = get_shape(inputs[i]); |
| 169 | } |
| 170 | } |
| 171 | auto maker = CustomGradMaker(backward, inputs.size()); |
| 172 | maker.output_size(1).output_captured(0, false); |
| 173 | maker.backward([inps_ = std::move(inps), input_shapes_ = std::move(input_shapes), |
| 174 | param, policy, dimA, dimB](Span<ValueRef> grads) { |
| 175 | mgb_assert(grads.size() == 1); |
| 176 | ValueRef grad = grads[0]; |
| 177 | SmallVector<ValueRef> ret(2); |
| 178 | if (!grad) { |
| 179 | return ret; |
| 180 | } |
| 181 | size_t dimG = std::max(dimA, dimB); |
| 182 | if (inps_[1]) { |
| 183 | if (param.transposeA) { |
| 184 | auto&& grad_op = BatchedMatrixMul::make( |
| 185 | param.transposeB, true, param.compute_mode, param.format, |
| 186 | policy.strategy, policy.workspace_limit, dimB, dimG); |
| 187 | ret[0] = imperative::apply(*grad_op, inps_[1], grad)[0]; |
| 188 | } else { |
| 189 | auto&& grad_op = BatchedMatrixMul::make( |
| 190 | false, !param.transposeB, param.compute_mode, param.format, |
| 191 | policy.strategy, policy.workspace_limit, dimG, dimB); |
| 192 | ret[0] = imperative::apply(*grad_op, grad, inps_[1])[0]; |
| 193 | } |
| 194 | if (dimG != dimA) { |
| 195 | ret[0] = reduce_to(ret[0], input_shapes_[0]); |
| 196 | } |
| 197 | } |
| 198 | if (inps_[0]) { |
| 199 | if (param.transposeB) { |
| 200 | auto&& grad_op = BatchedMatrixMul::make( |
| 201 | true, param.transposeA, param.compute_mode, param.format, |
| 202 | policy.strategy, policy.workspace_limit, dimG, dimA); |
| 203 | ret[1] = imperative::apply(*grad_op, grad, inps_[0])[0]; |
| 204 | } else { |
| 205 | auto&& grad_op = BatchedMatrixMul::make( |
| 206 | !param.transposeA, false, param.compute_mode, param.format, |
| 207 | policy.strategy, policy.workspace_limit, dimA, dimG); |
| 208 | ret[1] = imperative::apply(*grad_op, inps_[0], grad)[0]; |
| 209 | } |
| 210 | if (dimG != dimB) { |
| 211 | ret[1] = reduce_to(ret[1], input_shapes_[1]); |
| 212 | } |
nothing calls this directly
no test coverage detected