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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230

/**
 * \file dnn/src/cuda/batch_normalization/opr_impl.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 "./opr_impl.h"

#include "src/cuda/utils.h"

namespace megdnn {
namespace cuda {

namespace batch_normalization {

BNTensorDescHolder::BNTensorDescHolder(
        const TensorLayout& x, const ParamDim& param_dim, const FwdMode& fwd_mode) {
    TensorShape xy_shape(x);
    Format xy_format = Format::NCHW;

    switch (param_dim) {
        case ParamDim::DIM_11HW:
            // xy: N, C, H, W --> (N*C), 1, H, W
            xy_shape.shape[0] = xy_shape.shape[0] * xy_shape.shape[1];
            xy_shape.shape[1] = 1;
            bn_mode = CUDNN_BATCHNORM_PER_ACTIVATION;
            break;
        case ParamDim::DIM_1CHW:
            bn_mode = CUDNN_BATCHNORM_PER_ACTIVATION;
            break;
        case ParamDim::DIM_1C11:
            bn_mode = CUDNN_BATCHNORM_SPATIAL;
            break;
        case ParamDim::DIM_111C:
            bn_mode = CUDNN_BATCHNORM_SPATIAL;
            xy_format = Format::NHWC;
#if CUDNN_VERSION >= 7410
            if (fwd_mode == FwdMode::TRAINING) {
                bn_mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
            }
#endif  // CUDNN_VERSION >= 7400
            break;
        default:
            megdnn_throw("Unknown param dim type of batch normalization.");
    }
    xy_desc.set(TensorLayout(xy_shape, x.dtype), xy_format);
    param_desc.set(xy_desc.desc, bn_mode);
}

size_t get_reserve_size(
        const cudnnHandle_t& handle, const BNTensorDescHolder& tensor_desc) {
#if CUDNN_VERSION >= 7410
    size_t reserve_size;
    cudnn_check(cudnnGetBatchNormalizationTrainingExReserveSpaceSize(
            handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN,
            nullptr,                   // activationDesc
            tensor_desc.xy_desc.desc,  // xDesc
            &reserve_size));
    return reserve_size;
#else
    return 0;
#endif  // CUDNN_VERSION >= 7410
}
}  // namespace batch_normalization

using batch_normalization::BNTensorDescHolder;

size_t BNForwardImpl::get_workspace_in_bytes(
        const TensorLayout& src, const TensorLayout&, const TensorLayout&,
        const TensorLayout&, const TensorLayout&, const TensorLayout&,
        const TensorLayout&, const TensorLayout&, const TensorLayout&) {
#if CUDNN_VERSION >= 7410
    auto handle = cudnn_handle(this->handle());
    BNTensorDescHolder tensor_desc(src, m_param.param_dim, m_param.fwd_mode);

    size_t workspace_size;
    cudnn_check(cudnnGetBatchNormalizationForwardTrainingExWorkspaceSize(
            handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN,
            tensor_desc.xy_desc.desc,     // xDesc
            tensor_desc.xy_desc.desc,     // yDesc
            tensor_desc.xy_desc.desc,     // zDesc
            tensor_desc.param_desc.desc,  // bnScaleBiasMeanVarDesc
            nullptr,                      // activationDesc
            &workspace_size));
    return workspace_size;
#else
    return 0;
#endif  // CUDNN_VERSION >= 7410
}

size_t BNForwardImpl::get_reserve_in_bytes(const TensorLayout& src) {
    BNTensorDescHolder tensor_desc(src, m_param.param_dim, m_param.fwd_mode);
    return batch_normalization::get_reserve_size(
            cudnn_handle(this->handle()), tensor_desc);
}

void BNForwardImpl::exec(
        _megdnn_tensor_in src, _megdnn_tensor_in bn_scale, _megdnn_tensor_in bn_bias,
        _megdnn_tensor_out mean, _megdnn_tensor_out variance,
        _megdnn_tensor_out batch_mean, _megdnn_tensor_out batch_inv_variance,
        _megdnn_tensor_out reserve, _megdnn_tensor_out dst,
        _megdnn_workspace workspace) {
    check_exec(
            src.layout, bn_scale.layout, bn_bias.layout, mean.layout, variance.layout,
            batch_mean.layout, batch_inv_variance.layout, dst.layout, workspace.size,
            reserve.layout.access_bytes());
    auto handle = cudnn_handle(this->handle());
    BNTensorDescHolder tensor_desc(src.layout, m_param.param_dim, m_param.fwd_mode);

    float alpha = 1.0f, beta = 0.0f;
    switch (m_param.fwd_mode) {
        case param::BN::FwdMode::TRAINING:
#if CUDNN_VERSION >= 7410
            cudnn_check(cudnnBatchNormalizationForwardTrainingEx(
                    handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN, &alpha,
                    &beta,                                    // one & zero
                    tensor_desc.xy_desc.desc, src.raw_ptr(),  // xDesc & x
                    nullptr, nullptr,                         // zDesc & z
                    tensor_desc.xy_desc.desc, dst.raw_ptr(),  // yDesc & y
                    tensor_desc.param_desc.desc,              // bnScaleBiasMeanVarDesc
                    bn_scale.raw_ptr(), bn_bias.raw_ptr(), m_param.avg_factor,
                    mean.raw_ptr(), variance.raw_ptr(), m_param.epsilon,
                    batch_mean.raw_ptr(), batch_inv_variance.raw_ptr(), nullptr,
                    workspace.raw_ptr, workspace.size, reserve.raw_ptr(),
                    reserve.layout.access_bytes()));
#else
            cudnn_check(cudnnBatchNormalizationForwardTraining(
                    handle, tensor_desc.bn_mode, &alpha, &beta,
                    tensor_desc.xy_desc.desc, src.raw_ptr(),  // xDesc & x
                    tensor_desc.xy_desc.desc, dst.raw_ptr(),  // yDesc & y
                    tensor_desc.param_desc.desc,              // bnScaleBiasMeanVarDesc
                    bn_scale.raw_ptr(), bn_bias.raw_ptr(), m_param.avg_factor,
                    mean.raw_ptr(), variance.raw_ptr(), m_param.epsilon,
                    batch_mean.raw_ptr(), batch_inv_variance.raw_ptr()));
#endif  // CUDNN_VERSION >= 7410
            break;
        case param::BN::FwdMode::INFERENCE:
            cudnn_check(cudnnBatchNormalizationForwardInference(
                    handle, tensor_desc.bn_mode, &alpha, &beta,
                    tensor_desc.xy_desc.desc, src.raw_ptr(), tensor_desc.xy_desc.desc,
                    dst.raw_ptr(), tensor_desc.param_desc.desc, bn_scale.raw_ptr(),
                    bn_bias.raw_ptr(), mean.raw_ptr(), variance.raw_ptr(),
                    m_param.epsilon));
            break;
        default:
            megdnn_throw("Unknown forward mode type of batch normalization.");
    }
}

size_t BNBackwardImpl::get_workspace_in_bytes(
        const TensorLayout& x, const TensorLayout&, const TensorLayout&,
        const TensorLayout&, const TensorLayout&, const TensorLayout&,
        const TensorLayout&, const TensorLayout&, const TensorLayout&) {
#if CUDNN_VERSION >= 7410
    auto handle = cudnn_handle(this->handle());
    BNTensorDescHolder tensor_desc(x, m_param.param_dim, m_param.fwd_mode);

    size_t workspace_size;
    cudnn_check(cudnnGetBatchNormalizationBackwardExWorkspaceSize(
            handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN,
            tensor_desc.xy_desc.desc,     // xDesc
            tensor_desc.xy_desc.desc,     // yDesc
            tensor_desc.xy_desc.desc,     // dyDesc
            nullptr,                      // dzDesc
            tensor_desc.xy_desc.desc,     // dxDesc
            tensor_desc.param_desc.desc,  // dBnScaleBiasDesc
            nullptr,                      // activationDesc
            &workspace_size));
    return workspace_size;
#else
    return 0;
#endif  // CUDNN_VERSION >= 7410
}

size_t BNBackwardImpl::get_reserve_in_bytes(const TensorLayout& src) {
    BNTensorDescHolder tensor_desc(src, m_param.param_dim, m_param.fwd_mode);
    return batch_normalization::get_reserve_size(
            cudnn_handle(this->handle()), tensor_desc);
}

void BNBackwardImpl::exec(
        _megdnn_tensor_in x, _megdnn_tensor_in dy, _megdnn_tensor_in saved_batch_mean,
        _megdnn_tensor_in saved_batch_inv_variance, _megdnn_tensor_in bn_scale,
        _megdnn_tensor_in reserve, _megdnn_tensor_out d_bn_scale,
        _megdnn_tensor_out d_bn_bias, _megdnn_tensor_out dx,
        _megdnn_workspace workspace) {
    check_exec(
            x.layout, dy.layout, saved_batch_mean.layout,
            saved_batch_inv_variance.layout, bn_scale.layout, d_bn_scale.layout,
            d_bn_bias.layout, dx.layout, workspace.size, reserve.layout.access_bytes());
    auto handle = cudnn_handle(this->handle());
    BNTensorDescHolder tensor_desc(x.layout, m_param.param_dim, m_param.fwd_mode);

    float alpha = 1.0, beta = 0.0;
#if CUDNN_VERSION >= 7410
    cudnn_check(cudnnBatchNormalizationBackwardEx(
            handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN, &alpha, &beta, &alpha,
            &beta, tensor_desc.xy_desc.desc,
            x.raw_ptr(),                                      // xDesc & x
            nullptr, nullptr,                                 // yDesc & y
            tensor_desc.xy_desc.desc, dy.raw_ptr(),           // dyDesc & dy
            nullptr, nullptr,                                 // dzDesc & dz
            tensor_desc.xy_desc.desc, dx.raw_ptr(),           // dxDesc & dx
            tensor_desc.param_desc.desc, bn_scale.raw_ptr(),  // bnScale
            nullptr,                                          // bnBias
            d_bn_scale.raw_ptr(), d_bn_bias.raw_ptr(),        // dScale, dBias
            m_param.epsilon, saved_batch_mean.raw_ptr(),
            saved_batch_inv_variance.raw_ptr(), nullptr, workspace.raw_ptr,
            workspace.size, reserve.raw_ptr(), reserve.layout.access_bytes()));
#else
    cudnn_check(cudnnBatchNormalizationBackward(
            handle, tensor_desc.bn_mode, &alpha, &beta, &alpha, &beta,
            tensor_desc.xy_desc.desc, x.raw_ptr(),            // xDesc & x
            tensor_desc.xy_desc.desc, dy.raw_ptr(),           // dyDesc & dy
            tensor_desc.xy_desc.desc, dx.raw_ptr(),           // dxDesc & dx
            tensor_desc.param_desc.desc, bn_scale.raw_ptr(),  // bnScale
            d_bn_scale.raw_ptr(), d_bn_bias.raw_ptr(),        // dScale, dBias
            m_param.epsilon, saved_batch_mean.raw_ptr(),
            saved_batch_inv_variance.raw_ptr()));
#endif
}

}  // namespace cuda
}  // namespace megdnn

// vim: syntax=cpp.doxygen