#include <cmath>
|
#include <stdio.h>
|
#include <cassert>
|
#include <iostream>
|
#include "chunk.h"
|
#include <cuda_runtime.h>
|
#define NV_CUDA_CHECK(status) \
|
{ \
|
if (status != 0) \
|
{ \
|
std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \
|
<< " at line " << __LINE__ << std::endl; \
|
abort(); \
|
} \
|
}
|
|
|
namespace nvinfer1
|
{
|
Chunk::Chunk()
|
{
|
|
}
|
Chunk::Chunk(const void* buffer, size_t size)
|
{
|
assert(size == sizeof(_n_size_split));
|
_n_size_split = *reinterpret_cast<const int*>(buffer);
|
}
|
Chunk::~Chunk()
|
{
|
|
}
|
int Chunk::getNbOutputs() const
|
{
|
return 2;
|
}
|
|
Dims Chunk::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
|
{
|
assert(nbInputDims == 1);
|
assert(index == 0 || index == 1);
|
return Dims3(inputs[0].d[0] / 2, inputs[0].d[1], inputs[0].d[2]);
|
}
|
|
int Chunk::initialize()
|
{
|
return 0;
|
}
|
|
void Chunk::terminate()
|
{
|
}
|
|
size_t Chunk::getWorkspaceSize(int maxBatchSize) const
|
{
|
return 0;
|
}
|
|
int Chunk::enqueue(int batchSize,
|
const void* const* inputs,
|
void** outputs,
|
void* workspace,
|
cudaStream_t stream)
|
{
|
//batch
|
for (int b = 0; b < batchSize; ++b)
|
{
|
NV_CUDA_CHECK(cudaMemcpy((char*)outputs[0] + b * _n_size_split, (char*)inputs[0] + b * 2 * _n_size_split, _n_size_split, cudaMemcpyDeviceToDevice));
|
NV_CUDA_CHECK(cudaMemcpy((char*)outputs[1] + b * _n_size_split, (char*)inputs[0] + b * 2 * _n_size_split + _n_size_split, _n_size_split, cudaMemcpyDeviceToDevice));
|
}
|
// NV_CUDA_CHECK(cudaMemcpy(outputs[0], inputs[0], _n_size_split, cudaMemcpyDeviceToDevice));
|
// NV_CUDA_CHECK(cudaMemcpy(outputs[1], (void*)((char*)inputs[0] + _n_size_split), _n_size_split, cudaMemcpyDeviceToDevice));
|
return 0;
|
}
|
|
size_t Chunk::getSerializationSize() const
|
{
|
return sizeof(_n_size_split);
|
}
|
|
void Chunk::serialize(void *buffer)const
|
{
|
*reinterpret_cast<int*>(buffer) = _n_size_split;
|
}
|
|
const char* Chunk::getPluginType()const
|
{
|
return "CHUNK_TRT";
|
}
|
const char* Chunk::getPluginVersion() const
|
{
|
return "1.0";
|
}
|
|
void Chunk::destroy()
|
{
|
delete this;
|
}
|
|
void Chunk::setPluginNamespace(const char* pluginNamespace)
|
{
|
_s_plugin_namespace = pluginNamespace;
|
}
|
|
const char* Chunk::getPluginNamespace() const
|
{
|
return _s_plugin_namespace.c_str();
|
}
|
|
DataType Chunk::getOutputDataType(int index,
|
const nvinfer1::DataType* inputTypes,
|
int nbInputs) const
|
{
|
assert(index == 0 || index == 1);
|
return DataType::kFLOAT;
|
}
|
|
bool Chunk::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
|
{
|
return false;
|
}
|
|
bool Chunk::canBroadcastInputAcrossBatch(int inputIndex) const
|
{
|
return false;
|
}
|
|
void Chunk::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) {}
|
|
void Chunk::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput)
|
{
|
_n_size_split = in->dims.d[0] / 2 * in->dims.d[1] * in->dims.d[2] *sizeof(float);
|
}
|
void Chunk::detachFromContext() {}
|
|
// Clone the plugin
|
IPluginV2IOExt* Chunk::clone() const
|
{
|
Chunk *p = new Chunk();
|
p->_n_size_split = _n_size_split;
|
p->setPluginNamespace(_s_plugin_namespace.c_str());
|
return p;
|
}
|
|
//----------------------------
|
PluginFieldCollection ChunkPluginCreator::_fc{};
|
std::vector<PluginField> ChunkPluginCreator::_vec_plugin_attributes;
|
|
ChunkPluginCreator::ChunkPluginCreator()
|
{
|
_vec_plugin_attributes.clear();
|
_fc.nbFields = _vec_plugin_attributes.size();
|
_fc.fields = _vec_plugin_attributes.data();
|
}
|
|
const char* ChunkPluginCreator::getPluginName() const
|
{
|
return "CHUNK_TRT";
|
}
|
|
const char* ChunkPluginCreator::getPluginVersion() const
|
{
|
return "1.0";
|
}
|
|
const PluginFieldCollection* ChunkPluginCreator::getFieldNames()
|
{
|
return &_fc;
|
}
|
|
IPluginV2IOExt* ChunkPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
|
{
|
Chunk* obj = new Chunk();
|
obj->setPluginNamespace(_s_name_space.c_str());
|
return obj;
|
}
|
|
IPluginV2IOExt* ChunkPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
|
{
|
Chunk* obj = new Chunk(serialData,serialLength);
|
obj->setPluginNamespace(_s_name_space.c_str());
|
return obj;
|
}
|
|
void ChunkPluginCreator::setPluginNamespace(const char* libNamespace)
|
{
|
_s_name_space = libNamespace;
|
}
|
|
const char* ChunkPluginCreator::getPluginNamespace() const
|
{
|
return _s_name_space.c_str();
|
}
|
}//namespace nvinfer1
|
