| 58 | } |
| 59 | |
| 60 | void BatchedMatrixMulForwardImpl::AlgoDefault::exec(const ExecArgs& args) const { |
| 61 | //! As megbrain may modify param when checking all transpose situations, so |
| 62 | //! here we should copy the param when dispatching kern |
| 63 | auto param = args.opr->param(); |
| 64 | auto kern = [args, param]() { |
| 65 | auto N = args.layout_a.shape[0]; |
| 66 | TensorND A_, B_, C_; |
| 67 | A_.reset_ptr(args.tensor_a.raw_ptr()); |
| 68 | A_.layout = args.layout_a.remove_axis(0); |
| 69 | B_.reset_ptr(args.tensor_b.raw_ptr()); |
| 70 | B_.layout = args.layout_b.remove_axis(0); |
| 71 | C_.reset_ptr(args.tensor_c.raw_ptr()); |
| 72 | C_.layout = args.layout_c.remove_axis(0); |
| 73 | |
| 74 | auto Astrd = args.layout_a.dtype.size() * args.layout_a.stride[0], |
| 75 | Bstrd = args.layout_b.dtype.size() * args.layout_b.stride[0], |
| 76 | Cstrd = args.layout_c.dtype.size() * args.layout_c.stride[0]; |
| 77 | |
| 78 | auto advance_ptr = [](TensorND& dest, ptrdiff_t d) { |
| 79 | dest.reset_ptr( |
| 80 | static_cast<void*>(static_cast<dt_byte*>(dest.raw_ptr()) + d)); |
| 81 | }; |
| 82 | |
| 83 | auto opr = inplace_cpu_handle()->create_operator<MatrixMul>(); |
| 84 | opr->param() = param; |
| 85 | rep(n, N) { |
| 86 | opr->exec(A_, B_, C_, args.workspace); |
| 87 | advance_ptr(A_, Astrd); |
| 88 | advance_ptr(B_, Bstrd); |
| 89 | advance_ptr(C_, Cstrd); |
| 90 | } |
| 91 | }; |
| 92 | |
| 93 | static_cast<naive::HandleImpl*>(args.opr->handle())->dispatch_kern(kern); |
| 94 | } |
| 95 | |
| 96 | // vim: syntax=cpp.doxygen |
no test coverage detected