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

/**
 * \file dnn/src/cuda/cudnn_wrapper.h
 * 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.
 */
#pragma once

#include <unordered_map>
#include "megdnn/basic_types.h"
#include "megdnn/oprs/nn.h"
#include "src/cuda/cudnn_with_check.h"

namespace megdnn {
namespace cuda {

/*!
 * \brief get compute_type of convolution operations
 */
cudnnDataType_t get_compute_type_fp16(param::Convolution::ComputeMode comp_mode);

class TensorDesc {
public:
    TensorDesc();
    //! default layout is nchw
    void set(
            const TensorLayout& layout,
            const param::Convolution::Format = param::Convolution::Format::NCHW);
    std::string to_string();
    ~TensorDesc();
    cudnnTensorDescriptor_t desc;
};

template <typename Param>
class FilterDesc {
public:
    FilterDesc();
    void set(const typename ConvolutionBase<Param>::CanonizedFilterMeta& meta);
    std::string to_string();
    ~FilterDesc();
    cudnnFilterDescriptor_t desc;
};

class ConvDesc {
public:
    ConvDesc();
    void set(DType data_type, const param::Convolution& param, const size_t nr_group);
    ~ConvDesc();
    cudnnConvolutionDescriptor_t desc;
};

class LRNDesc {
public:
    LRNDesc();
    void set(const param::LRN& param);
    ~LRNDesc();
    cudnnLRNDescriptor_t desc;
};

class BNParamDesc {
public:
    BNParamDesc();
    void set(const cudnnTensorDescriptor_t xDesc, cudnnBatchNormMode_t mode);
    ~BNParamDesc();
    cudnnTensorDescriptor_t desc;
};

// the classes below is used to deal with 3d situations
class Tensor3DDesc {
public:
    Tensor3DDesc();
    //! default layout is NCDHW
    void set(const TensorLayout& layout, bool is_ndhwc = false);
    ~Tensor3DDesc();
    cudnnTensorDescriptor_t desc;
};

class Filter3DDesc {
public:
    Filter3DDesc();
    void set(const Convolution3DBase::CanonizedFilterMeta& meta);
    ~Filter3DDesc();
    cudnnFilterDescriptor_t desc;
};

class Conv3DDesc {
public:
    Conv3DDesc();
    void set(const param::Convolution3D& param, const size_t nr_group);
    ~Conv3DDesc();
    cudnnConvolutionDescriptor_t desc;
};

class CudnnAlgoPack {
public:
    //! algorithm attr
    struct Attr {
        std::string name;
        bool is_reproducible;
        bool accuracy_depend_on_batch;
    };

    static const std::unordered_map<cudnnConvolutionBwdDataAlgo_t, Attr>
    conv_bwd_data_algos();

    static const std::unordered_map<cudnnConvolutionBwdFilterAlgo_t, Attr>
    conv_bwd_flt_algos();

    static const std::unordered_map<cudnnConvolutionFwdAlgo_t, Attr> conv_fwd_algos();

    static const std::unordered_map<cudnnConvolutionBwdDataAlgo_t, Attr>
    conv3d_bwd_data_algos();

    static const std::unordered_map<cudnnConvolutionBwdFilterAlgo_t, Attr>
    conv3d_bwd_flt_algos();

    static const std::unordered_map<cudnnConvolutionFwdAlgo_t, Attr> conv3d_fwd_algos();
};

}  // namespace cuda
}  // namespace megdnn

namespace std {

#define DEF_HASH(_type)                                                \
    template <>                                                        \
    struct hash<_type> {                                               \
        std::size_t operator()(const _type& algo) const {              \
            return std::hash<uint32_t>()(static_cast<uint32_t>(algo)); \
        }                                                              \
    }

DEF_HASH(cudnnConvolutionBwdDataAlgo_t);
DEF_HASH(cudnnConvolutionBwdFilterAlgo_t);
DEF_HASH(cudnnConvolutionFwdAlgo_t);

#undef DEF_HASH
}  // namespace std

// vim: syntax=cpp.doxygen