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

/**
 * \file src/serialization/impl/flatbuffers_helper.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.
 */

#if MGB_ENABLE_FBS_SERIALIZATION

#include "megbrain/serialization/internal/flatbuffers_helper.h"
#include "megbrain/common.h"

using namespace megdnn;

namespace mgb {
namespace serialization {
namespace fbs {
namespace intl {

megdnn::DTypeEnum convert_dtype_to_megdnn(DTypeEnum fb) {
    switch (fb) {
#define cb(_dt)           \
    case DTypeEnum_##_dt: \
        return megdnn::DTypeEnum::_dt;
        MEGDNN_FOREACH_DTYPE_NAME(cb)
        MEGDNN_FOREACH_PARAMETERIZED_DTYPE(cb)
#undef cb
        default:
            // Float16 may be disabled
            megdnn_trap();
    }
}

DTypeEnum convert_dtype_to_fbs(megdnn::DTypeEnum enumv) {
    switch (enumv) {
#define cb(_dt)                  \
    case megdnn::DTypeEnum::_dt: \
        return DTypeEnum_##_dt;
        MEGDNN_FOREACH_DTYPE_NAME(cb)
        MEGDNN_FOREACH_PARAMETERIZED_DTYPE(cb)
#undef cb
    }
    megdnn_trap();
}

megdnn::DType load_dtype(const fbs::DType* dtype) {
    auto param = dtype->param_as_LinearQuantizationParam();
    switch (dtype->type()) {
#define cb(_dt)           \
    case DTypeEnum_##_dt: \
        return dtype::_dt{};
        MEGDNN_FOREACH_DTYPE_NAME(cb)
#undef cb
        case DTypeEnum_QuantizedS4:
            return dtype::QuantizedS4{param->scale()};
        case DTypeEnum_QuantizedS8:
            return dtype::QuantizedS8{param->scale()};
        case DTypeEnum_QuantizedS16:
            return dtype::QuantizedS16{param->scale()};
        case DTypeEnum_QuantizedS32:
            return dtype::QuantizedS32{param->scale()};
        case DTypeEnum::DTypeEnum_Quantized4Asymm:
            return dtype::Quantized4Asymm{param->scale(), param->zero_point()};
        case DTypeEnum::DTypeEnum_Quantized8Asymm:
            return dtype::Quantized8Asymm{param->scale(), param->zero_point()};
        default:
            // Float16 may be disabled
            megdnn_trap();
    }
    return {};
}

flatbuffers::Offset<fbs::DType> build_dtype(
        flatbuffers::FlatBufferBuilder& builder, megdnn::DType dtype) {
    if (!dtype.valid())
        return {};
    DTypeEnum enumv{};
    switch (dtype.enumv()) {
#define cb(_dt)                  \
    case megdnn::DTypeEnum::_dt: \
        enumv = DTypeEnum_##_dt; \
        break;
        MEGDNN_FOREACH_DTYPE_NAME(cb)
        MEGDNN_FOREACH_PARAMETERIZED_DTYPE(cb)
#undef cb
    }
    DTypeParam param_type = DTypeParam_NONE;
    flatbuffers::Offset<void> param;
    if (dtype.has_param()) {
        switch (dtype.enumv()) {
#define cb(_dt)                  \
    case megdnn::DTypeEnum::_dt: \
        mgb_trap();  // unreachable
            MEGDNN_FOREACH_DTYPE_NAME(cb)
#undef cb
#define CASE_ASYMMETRIC(_dt)                                                           \
    case megdnn::DTypeEnum::_dt: {                                                     \
        auto&& p = dtype.param<dtype::_dt>();                                          \
        param_type = DTypeParam_LinearQuantizationParam;                               \
        param = CreateLinearQuantizationParam(builder, p.scale, p.zero_point).Union(); \
        break;                                                                         \
    }
#define CASE_SYMMETRIC(_dt)                                       \
    case megdnn::DTypeEnum::_dt:                                  \
        param_type = DTypeParam_LinearQuantizationParam;          \
        param = CreateLinearQuantizationParam(                    \
                        builder, dtype.param<dtype::_dt>().scale) \
                        .Union();                                 \
        break;
            CASE_ASYMMETRIC(Quantized4Asymm)
            CASE_ASYMMETRIC(Quantized8Asymm)
            CASE_SYMMETRIC(QuantizedS4)
            CASE_SYMMETRIC(QuantizedS8)
            CASE_SYMMETRIC(QuantizedS16)
            CASE_SYMMETRIC(QuantizedS32)
        }
    }
    DTypeBuilder dt(builder);
    dt.add_type(enumv);
    if (param_type != DTypeParam_NONE) {
        dt.add_param_type(param_type);
        dt.add_param(param);
    }
    return dt.Finish();
}

}  // namespace intl
}  // namespace fbs
}  // namespace serialization
}  // namespace mgb

#endif