| 47 | } |
| 48 | |
| 49 | WorkspaceBundle get_bundle(const NCBKernSizeParam& param) { |
| 50 | UNPACK_CONV_F32_NCB_KERN_SIZES(param); |
| 51 | MEGDNN_MARK_USED_VAR(N); |
| 52 | MEGDNN_MARK_USED_VAR(OH); |
| 53 | MEGDNN_MARK_USED_VAR(OW); |
| 54 | bool can_matrix_mul_direct = |
| 55 | (FH == 1 && FW == 1 && SH == 1 && SW == 1 && PH == 0 && PW == 0); |
| 56 | // temp space to store unrolled matrix |
| 57 | // workspace for matrix mul opr |
| 58 | // workspace for relayout opr |
| 59 | size_t part0, part1, part2; |
| 60 | if (can_matrix_mul_direct) { |
| 61 | part0 = 0; |
| 62 | } else { |
| 63 | part0 = (IC * FH * FW * IH * IW) * param.grad_type.size(); |
| 64 | } |
| 65 | part2 = (OC * IC * FH * FW) * param.filter_type.size(); |
| 66 | if (param.filter_meta.format == param::Convolution::Format::NCHW44) { |
| 67 | TensorLayout A_, B_, C_; |
| 68 | A_ = TensorLayout({IC / 4 * FH * FW, OC / 4, 4, 4}, param.filter_type); |
| 69 | B_ = TensorLayout({OC / 4, IH * IW}, param.diff_type); |
| 70 | C_ = TensorLayout({IC / 4 * FH * FW, IH * IW, 4}, param.grad_type); |
| 71 | auto matmul_algo = get_matmul_opr(param); |
| 72 | part1 = matmul_algo->get_workspace_in_bytes(A_, B_, C_); |
| 73 | } else { |
| 74 | TensorLayout A_, B_, C_; |
| 75 | A_ = TensorLayout({IC * FH * FW, OC}, param.filter_type); |
| 76 | B_ = TensorLayout({OC, IH * IW}, param.diff_type); |
| 77 | C_ = TensorLayout({IC * FH * FW, IH * IW}, param.grad_type); |
| 78 | part1 = get_matmul_opr(param)->get_workspace_in_bytes(A_, B_, C_); |
| 79 | } |
| 80 | return {nullptr, {part0, part1, part2}}; |
| 81 | } |
| 82 | |
| 83 | template <typename ftype, typename dtype, typename gtype> |
| 84 | void kern_matmul(const NCBKernParam& param) { |
no test coverage detected