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

/**
 * \file dnn/src/rocm/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/rocm/utils.h"

namespace megdnn {
namespace rocm {

namespace batch_normalization {

void BNTensorDescHolder::setup(const TensorLayout& x, const ParamDim& param_dim) {
    TensorShape xy_shape(x);

    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 = miopenBNPerActivation;
            break;
        case ParamDim::DIM_1CHW:
            bn_mode = miopenBNPerActivation;
            break;
        case ParamDim::DIM_1C11:
            bn_mode = miopenBNSpatial;
            break;
        default:
            megdnn_throw("Unknown param dim type of batch normalization.");
    }
    xy_desc.set(TensorLayout(xy_shape, x.dtype));
    param_desc.set(xy_desc.desc, bn_mode);
}

}  // namespace batch_normalization

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, _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);
    auto handle = concrete_handle(this->handle())->miopen_handle();
    m_tensor_desc.setup(src.layout, m_param.param_dim);

    float alpha = 1.0f, beta = 0.0f;
    switch (m_param.fwd_mode) {
        case param::BN::FwdMode::TRAINING:
            miopen_check(miopenBatchNormalizationForwardTraining(
                    handle, m_tensor_desc.bn_mode, &alpha, &beta,
                    m_tensor_desc.xy_desc.desc,     // xDesc
                    src.raw_ptr(),                  // x
                    m_tensor_desc.xy_desc.desc,     // yDesc
                    dst.raw_ptr(),                  // y
                    m_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()));

            break;
        case param::BN::FwdMode::INFERENCE:
            miopen_check(miopenBatchNormalizationForwardInference(
                    handle, m_tensor_desc.bn_mode, &alpha, &beta,
                    m_tensor_desc.xy_desc.desc, src.raw_ptr(),
                    m_tensor_desc.xy_desc.desc, dst.raw_ptr(),
                    m_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.");
    }
}

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, _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);
    auto handle = concrete_handle(this->handle())->miopen_handle();
    m_tensor_desc.setup(x.layout, m_param.param_dim);

    float alpha = 1.0, beta = 0.0;
    miopen_check(miopenBatchNormalizationBackward(
            handle, m_tensor_desc.bn_mode, &alpha, &beta, &alpha, &beta,
            m_tensor_desc.xy_desc.desc, x.raw_ptr(), m_tensor_desc.xy_desc.desc,
            dy.raw_ptr(), m_tensor_desc.xy_desc.desc, dx.raw_ptr(),
            m_tensor_desc.param_desc.desc, bn_scale.raw_ptr(), d_bn_scale.raw_ptr(),
            d_bn_bias.raw_ptr(), m_param.epsilon, saved_batch_mean.raw_ptr(),
            saved_batch_inv_variance.raw_ptr()));
}

}  // namespace rocm
}  // namespace megdnn

// vim: syntax=cpp.doxygen