megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

/**
 * \file dnn/src/cuda/local_share/forward/batched_matmul.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#include "./algo.h"
#include "src/cuda/local_share/im2col.cuh"
#include "src/cuda/local_share/opr_impl.h"

using namespace megdnn;
using namespace cuda;

bool LocalShareForwardImpl::AlgoBatchedMatMul::is_available(
        const SizeArgs& args) const {
    bool available = true;
    auto&& param = args.opr->param();
    using Param = LocalShare::Param;
    using Format = Param::Format;
    // NCHW format
    available &= param.format == Format::NCHW;
    // only support float
    auto src_dtype = args.src_layout.dtype, filter_dtype = args.filter_layout.dtype,
         dst_dtype = args.dst_layout.dtype;
    available &= (src_dtype == filter_dtype) && (src_dtype == dst_dtype) &&
                 (src_dtype == dtype::Float32());
    // do not support dilate conv
    size_t dh = param.dilate_h, dw = param.dilate_w;
    available &= (dh == 1 && dw == 1);
    return available;
}

WorkspaceBundle LocalShareForwardImpl::AlgoBatchedMatMul::get_workspace_bundle(
        dt_byte* raw_ptr, const SizeArgs& args) const {
    auto&& param = args.opr->param();
    unpack_local_share_params(
            args.src_layout, args.filter_layout, args.dst_layout, param);
    using Param = LocalShare::Param;
    using Sparse = Param::Sparse;
    size_t groups = 1;
    if (param.sparse == Sparse::GROUP) {
        groups = args.filter_layout.shape[0];
    }
    size_t icpg = ci / groups, ocpg = co / groups;
    size_t ws_im2col = n * ci * ho * wo * fh * fw * args.src_layout.dtype.size();
    size_t ws_posttranspose = n * co * ho * wo * args.dst_layout.dtype.size();
    auto&& matmul_opr = args.opr->handle()->create_operator<BatchedMatrixMul>();
    TensorLayout A{
            {groups * sgh * sgw, ho / sgh * wo / sgw * n, icpg * fh * fw},
            dtype::Float32()};
    TensorLayout B{{groups * sgh * sgw, icpg * fh * fw, ocpg}, dtype::Float32()};
    TensorLayout C{
            {groups * sgh * sgw, ho / sgh * wo / sgw * n, ocpg}, dtype::Float32()};
    size_t ws_matmul = matmul_opr->get_workspace_in_bytes(A, B, C);
    WorkspaceBundle ws{raw_ptr, {ws_im2col, ws_matmul, ws_posttranspose}};
    return ws;
}

size_t LocalShareForwardImpl::AlgoBatchedMatMul::get_workspace_in_bytes(
        const SizeArgs& args) const {
    return get_workspace_bundle(nullptr, args).total_size_in_bytes();
}

void LocalShareForwardImpl::AlgoBatchedMatMul::exec(const ExecArgs& args) const {
    auto&& param = args.opr->param();
    unpack_local_share_params(
            args.src_layout, args.filter_layout, args.dst_layout, param);
    using Param = LocalShare::Param;
    using Sparse = Param::Sparse;
    size_t groups = 1;
    if (param.sparse == Sparse::GROUP) {
        groups = args.filter_layout.shape[0];
    }
    size_t icpg = ci / groups, ocpg = co / groups;
    local_share::Param kern_param;
    kern_param.n = n, kern_param.co = co, kern_param.ci = ci, kern_param.hi = hi,
    kern_param.wi = wi, kern_param.ph = ph, kern_param.pw = pw,
    kern_param.grp_ho = ho / sgh, kern_param.grp_wo = wo / sgw, kern_param.sgh = sgh,
    kern_param.sgw = sgw;

    auto ws = get_workspace_bundle(args.workspace.raw_ptr, args);
    auto ws_im2col = ws.get(0);
    auto ws_matmul = ws.get(1);
    auto ws_posttranspose = ws.get(2);
    auto&& stream = cuda_stream(args.opr->handle());
    local_share::_do_local_share_im2col(
            args.src_tensor->ptr<dt_float32>(),
            reinterpret_cast<dt_float32*>(ws_im2col), fh, fw, sh, sw, groups,
            kern_param, stream);

    auto&& matmul_opr = args.opr->handle()->create_operator<BatchedMatrixMul>();
    TensorLayout A{
            {groups * sgh * sgw, ho / sgh * wo / sgw * n, icpg * fh * fw},
            dtype::Float32()};
    TensorLayout B{{groups * sgh * sgw, icpg * fh * fw, ocpg}, dtype::Float32()};
    TensorLayout C{
            {groups * sgh * sgw, ho / sgh * wo / sgw * n, ocpg}, dtype::Float32()};
    TensorND ts_A{ws_im2col, A};
    TensorND ts_B{args.filter_tensor->raw_ptr(), B};
    TensorND ts_C{ws_posttranspose, C};
    Workspace ws_wrapper;
    ws_wrapper.raw_ptr = reinterpret_cast<dt_byte*>(ws_matmul);
    ws_wrapper.size = ws.get_size(1);
    matmul_opr->exec(ts_A, ts_B, ts_C, ws_wrapper);

    {
        TensorLayout C1{
                {n, groups, ocpg, sgh, ho / sgh, sgw, wo / sgw}, dtype::Float32()};
        C1.stride[0] = ho / sgh * wo / sgw * ocpg;
        C1.stride[1] = n * ho * wo * ocpg;
        C1.stride[2] = 1;
        C1.stride[3] = n * ho / sgh * wo * ocpg;
        C1.stride[4] = wo / sgw * ocpg;
        C1.stride[5] = n * ho / sgh * wo / sgw * ocpg;
        C1.stride[6] = ocpg;
        TensorLayout C2 = args.dst_layout;
        TensorND ts_C1{ws_posttranspose, C1};
        TensorND ts_C2{args.dst_tensor->raw_ptr(), C2};
        auto&& relayout_opr = args.opr->handle()->create_operator<Relayout>();
        relayout_opr->exec(ts_C1, ts_C2);
    }
}

// vim: syntax=cpp.doxygen