131 lines
3.5 KiB
C++
131 lines
3.5 KiB
C++
#ifndef LAYERNORM_H
|
|
#define LAYERNORM_H
|
|
|
|
#include <vector>
|
|
#include <string>
|
|
#include <iostream>
|
|
#include <NvInfer.h>
|
|
#include <memory>
|
|
#include <string.h>
|
|
#include <cstdint>
|
|
#include <stdlib.h>
|
|
|
|
using namespace std;
|
|
|
|
struct welford
|
|
{
|
|
int count = 0;
|
|
double mean = 0.f;
|
|
double M2 = 0.f;
|
|
};
|
|
|
|
namespace nvinfer1{
|
|
class layernorm : public IPluginV2IOExt
|
|
{
|
|
public:
|
|
layernorm();
|
|
layernorm(const void* data, size_t length);
|
|
~layernorm();
|
|
int getNbOutputs() const override
|
|
{
|
|
return 2;
|
|
}
|
|
|
|
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
|
|
|
|
int initialize() override;
|
|
|
|
virtual void terminate() override {};
|
|
|
|
virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0; }
|
|
|
|
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override;
|
|
|
|
virtual size_t getSerializationSize() const override;
|
|
|
|
virtual void serialize(void* buffer) const override;
|
|
|
|
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override {
|
|
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
|
|
}
|
|
|
|
void setPluginNamespace(const char* pluginNamespace) override;
|
|
|
|
const char* getPluginNamespace() const override;
|
|
|
|
const char* getPluginType() const override;
|
|
|
|
const char* getPluginVersion() const override;
|
|
|
|
void destroy() override;
|
|
|
|
IPluginV2IOExt* clone() const override;
|
|
|
|
|
|
|
|
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
|
|
|
|
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
|
|
|
|
bool canBroadcastInputAcrossBatch(int inputIndex) const override;
|
|
|
|
void attachToContext(
|
|
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
|
|
|
|
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
|
|
|
|
void detachFromContext() override;
|
|
|
|
void setInputSize(int s, int l) {
|
|
mInputSize = s;
|
|
Length = l;
|
|
}
|
|
|
|
|
|
private:
|
|
void forwardGpu(const float *const * inputs, float *mean, float *std, cudaStream_t stream, int batchSize = 1);
|
|
int mThreadCount = 256;
|
|
int mInputSize;
|
|
int Length;
|
|
Dims outputDims ;
|
|
const char* mPluginNamespace;
|
|
};
|
|
class layernormCreator : public IPluginCreator
|
|
{
|
|
public:
|
|
layernormCreator();
|
|
|
|
~layernormCreator() override = default;
|
|
|
|
const char* getPluginName() const override;
|
|
|
|
const char* getPluginVersion() const override;
|
|
|
|
const PluginFieldCollection* getFieldNames() override;
|
|
|
|
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
|
|
|
|
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
|
|
|
|
void setPluginNamespace(const char* libNamespace) override
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
const char* getPluginNamespace() const override
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
private:
|
|
std::string mNamespace;
|
|
static PluginFieldCollection mFC;
|
|
static std::vector<PluginField> mPluginAttributes;
|
|
|
|
};
|
|
|
|
REGISTER_TENSORRT_PLUGIN(layernormCreator);
|
|
};
|
|
|
|
#endif // LAYERNORM_H
|