| 79 | } |
| 80 | |
| 81 | WorkspaceBundle MatrixMulForwardImpl::AlgoConv1X1CUDNN::get_workspace_bundle( |
| 82 | void* ptr, const SizeArgs& args) const { |
| 83 | SmallVector<size_t> sizes; |
| 84 | auto conv_opr_ptr = prepare_conv_opr(args); |
| 85 | |
| 86 | size_t m, k, n; |
| 87 | std::tie(m, k, n) = gen_matrixmul_shape(args); |
| 88 | |
| 89 | TensorLayout src_layout({1, k, 1, n}, args.layout_b.dtype); |
| 90 | TensorLayout filter_layout({m, k, 1, 1}, args.layout_a.dtype); |
| 91 | TensorLayout bias_layout(args.layout_a.dtype); |
| 92 | TensorLayout z_layout(args.layout_a.dtype); |
| 93 | TensorLayout dst_layout({1, m, 1, n}, args.layout_c.dtype); |
| 94 | ConvBiasForwardImpl::AlgoBase::SizeArgs conv_size_args( |
| 95 | static_cast<ConvBiasForwardImpl*>(conv_opr_ptr.get()), src_layout, |
| 96 | filter_layout, bias_layout, z_layout, dst_layout); |
| 97 | |
| 98 | sizes.push_back(m_impl->get_workspace_in_bytes(conv_size_args)); |
| 99 | |
| 100 | auto get_trans_layout = [](const TensorLayout& ly) { |
| 101 | size_t m = ly[0], n = ly[1]; |
| 102 | TensorLayout trans{{n, m}, ly.dtype}; |
| 103 | return trans; |
| 104 | }; |
| 105 | if (args.opr->param().transposeA) { |
| 106 | sizes.push_back(get_trans_layout(args.layout_a).span().dist_byte()); |
| 107 | } |
| 108 | if (args.opr->param().transposeB) { |
| 109 | sizes.push_back(get_trans_layout(args.layout_b).span().dist_byte()); |
| 110 | } |
| 111 | |
| 112 | return {ptr, std::move(sizes)}; |
| 113 | } |
| 114 | |
| 115 | size_t MatrixMulForwardImpl::AlgoConv1X1CUDNN::get_workspace_in_bytes( |
| 116 | const SizeArgs& args) const { |
no test coverage detected