Method exec

dnn/src/cuda/group_local/forward/opr_impl.cpp:42–87

}  // namespace

void GroupLocalForwardImpl::exec(
        _megdnn_tensor_in src, _megdnn_tensor_in filter, _megdnn_tensor_out dst,
        _megdnn_workspace workspace) {
    megdnn_assert(
            src.layout.dtype == dtype::Float32(),
            "cuda do not support fp16 group local operator");
    check_exec(src.layout, filter.layout, dst.layout, workspace.size);

    auto handle = concrete_handle(this->handle());
    auto G = filter.layout[0];
    auto IH = src.layout.shape[2], IW = src.layout.shape[3], OH = dst.layout.shape[2],
         OW = dst.layout.shape[3];
    if (prefer_inference_kernel(src.layout, filter.layout, dst.layout)) {
        auto N = src.layout.shape[0], ICg = src.layout.shape[1] / G,
             OCg = dst.layout.shape[1] / G;
        auto FH = filter.layout.shape[4], FW = filter.layout.shape[5];
        auto PH = param().pad_h, PW = param().pad_w;
        auto SH = param().stride_h, SW = param().stride_w;
        const float* sptr = src.ptr<dt_float32>();
        const float* fptr = filter.ptr<dt_float32>();
        float* dptr = dst.ptr<dt_float32>();
        float* wptr = workspace.ptr<dt_float32>();
        auto stream = cuda_stream(this->handle());

        group_local::exec(
                sptr, fptr, dptr, wptr, N, ICg, IH, IW, OCg, OH, OW, FH, FW, G, PH, PW,
                SH, SW, stream);
    } else {
        auto&& opr = get_opr(handle, param());
        TensorND src_g = {src.raw_ptr(), prepare_src_dst(src.layout, G)};
        TensorND dst_g = {dst.raw_ptr(), prepare_src_dst(dst.layout, G)};
        TensorND filter_g = {filter.raw_ptr(), prepare_filter(filter.layout)};
        for (size_t g = 0; g < G; ++g) {
            opr->exec(src_g, filter_g, dst_g, workspace);
            incr_refp(
                    src_g.get_ref_ptr(), src_g.layout.stride[1] *
                                                 src_g.layout.shape[1] *
                                                 src_g.layout.dtype.size());
            incr_refp(
                    dst_g.get_ref_ptr(), dst_g.layout.stride[1] *
                                                 dst_g.layout.shape[1] *
                                                 dst_g.layout.dtype.size());
            incr_refp(filter_g.get_ref_ptr(), filter_g.layout.span().dist_byte());
        }
    }
}

GroupLocalForwardImpl::GroupLocalForwardImpl(Handle* handle)
        : GroupLocalForward(handle) {}