Method exec

dnn/src/cuda/matrix_mul/conv1x1.cpp:120–169

}

void MatrixMulForwardImpl::AlgoConv1X1CUDNN::exec(const ExecArgs& args) const {
    SizeArgs size_args(args.opr, args.layout_a, args.layout_b, args.layout_c);

    auto conv_opr_ptr = prepare_conv_opr(size_args);

    size_t m, k, n;
    std::tie(m, k, n) = gen_matrixmul_shape(size_args);

    auto bundle = get_workspace_bundle(args.workspace.raw_ptr, size_args);
    auto A_dst_tensor = args.tensor_a;
    auto B_dst_tensor = args.tensor_b;
    if (args.opr->param().transposeA || args.opr->param().transposeB) {
        auto trans = args.opr->handle()->create_operator<RelayoutForward>();

        auto trans_tensor = [&](size_t workspace_pos, const TensorND& ori_tensor,
                                TensorND& dst_tensor) {
            TensorLayout dst_layout(
                    {ori_tensor.layout.shape[1], ori_tensor.layout.shape[0]},
                    ori_tensor.layout.dtype);
            dst_tensor = TensorND(bundle.get(workspace_pos), dst_layout);
            TensorND src_tensor(ori_tensor.raw_ptr(), dst_layout);
            src_tensor.layout.stride[0] = ori_tensor.layout.stride[1];
            src_tensor.layout.stride[1] = ori_tensor.layout.stride[0];

            trans->exec(src_tensor, dst_tensor, args.opr->handle());
        };

        if (args.opr->param().transposeA) {
            trans_tensor(1, args.tensor_a, A_dst_tensor);
        }
        if (args.opr->param().transposeB) {
            trans_tensor(bundle.nr_workspace() - 1, args.tensor_b, B_dst_tensor);
        }
    }

    TensorLayout src_layout({1, k, 1, n}, args.layout_b.dtype);
    TensorLayout filter_layout({m, k, 1, 1}, args.layout_a.dtype);
    TensorLayout dst_layout({1, m, 1, n}, args.layout_c.dtype);

    TensorND src(B_dst_tensor.raw_ptr(), src_layout);
    TensorND filter(A_dst_tensor.raw_ptr(), filter_layout);
    TensorND z(nullptr, TensorLayout(src_layout.dtype));
    TensorND bias(nullptr, TensorLayout(src_layout.dtype));
    TensorND dst(args.tensor_c.raw_ptr(), dst_layout);

    ConvBiasForwardImpl::AlgoBase::ExecArgs conv_exec_args(
            static_cast<ConvBiasForwardImpl*>(conv_opr_ptr.get()), src, filter, bias, z,
            dst, bundle.get_workspace(0));
    m_impl->exec(conv_exec_args);
}