grid-sample3d-trt-plugin/src/grid_sample_3d_plugin.h at main · IRCAD/grid-sample3d-trt-plugin · GitHub

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
#include <string>
#include <vector>

#include <NvInferPlugin.h>

#include <grid_sample_3d.h>

#ifndef GRID_SAMPLE_3D_PLUGIN
#define GRID_SAMPLE_3D_PLUGIN

using namespace nvinfer1::plugin;

namespace nvinfer1 {

namespace plugin {

class GridSample3DPlugin : public IPluginV2DynamicExt{

public:

    GridSample3DPlugin(const std::string name,
                       size_t inputChannel,
                       size_t inputDepth,
                       size_t inputHeight,
                       size_t inputWidth,
                       size_t gridDepth,
                       size_t gridHeight,
                       size_t gridWidth,
                       bool alignCorners,
                       GridSample3DInterpolationMode interpolationMode,
                       GridSample3DPaddingMode paddingMode,
                       nvinfer1::DataType dataType);

    GridSample3DPlugin(const std::string name,
                       bool alignCorners,
                       GridSample3DInterpolationMode interpolationMode,
                       GridSample3DPaddingMode paddingMode);

    GridSample3DPlugin(const std::string name,
                       const void* buffer,
                       size_t buffer_size);

    GridSample3DPlugin() = delete;
    ~GridSample3DPlugin() override;

    // IPluginV2DynamicExt Methods
    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;

    nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex,
                                            DimsExprs const* inputs,
                                            int32_t nbInputs,
                                            IExprBuilder& exprBuilder) noexcept override;

    bool supportsFormatCombination(int32_t pos,
                                   PluginTensorDesc const* inOut,
                                   int32_t nbInputs,
                                   int32_t nbOutputs) noexcept override;

    void configurePlugin(DynamicPluginTensorDesc const* in,
                         int32_t nbInputs,
                         DynamicPluginTensorDesc const* out,
                         int32_t nbOutputs) noexcept override;

    size_t getWorkspaceSize(PluginTensorDesc const* inputs,
                            int32_t nbInputs,
                            PluginTensorDesc const* outputs,
                            int32_t nbOutputs) const noexcept override;

    int32_t enqueue(PluginTensorDesc const* inputDesc,
                    PluginTensorDesc const* outputDesc,
                    void const* const* inputs,
                    void* const* outputs,
                    void* workspace,
                    cudaStream_t stream) noexcept override;

    // IPluginV2Ext Methods
    nvinfer1::DataType getOutputDataType(int32_t index,
                                         nvinfer1::DataType const* inputTypes,
                                         int32_t nbInputs) const noexcept override;
    void attachToContext(cudnnContext* cudnnContext,
                         cublasContext* cublasContext,
                         IGpuAllocator* gpuAllocator) noexcept override;

    void detachFromContext() noexcept override;

    // IPluginV2 Methods
    const char* getPluginType() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    int32_t getNbOutputs() const noexcept override;

    int32_t initialize() noexcept override;

    void terminate() noexcept override;

    size_t getSerializationSize() const noexcept override;

    void serialize(void* buffer) const noexcept override;

    void destroy() noexcept override;

    void setPluginNamespace(const char* pluginNamespace) noexcept override;

    const char* getPluginNamespace() const noexcept override;

private:

    const std::string mLayerName;
    size_t mBatch;
    size_t mInputChannel, mInputDepth, mInputWidth, mInputHeight;
    size_t mGridDepth, mGridWidth, mGridHeight;
    bool mAlignCorners;
    std::string mNameSpace;
    GridSample3DInterpolationMode mInterpolationMode;
    GridSample3DPaddingMode mPaddingMode;
    nvinfer1::DataType mDataType;

};

class GridSample3DPluginCreator : public IPluginCreator {

public:
    GridSample3DPluginCreator();
    ~GridSample3DPluginCreator() override;

    const char* getPluginName() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    const PluginFieldCollection* getFieldNames() noexcept override;

    IPluginV2* createPlugin(const char* name,
                            const PluginFieldCollection* fc) noexcept override;

    IPluginV2* deserializePlugin(const char* name,
                                 const void* serialData,
                                 size_t serialLength) noexcept override;

    void setPluginNamespace(const char* libNamespace) noexcept override;

    const char* getPluginNamespace() const noexcept override;

private:
    std::string mNamespace;
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
};

} // namespace plugin

} // namespace nvinfer1

#endif // GRID_SAMPLE_3D_PLUGIN