wehub-resource-sync
57 lines
2.6 KiB
Plaintext
57 lines
2.6 KiB
Plaintext
#include "Raster.cuh"
|
|
namespace MNN {
|
|
|
|
// Blit don't care offset
|
|
template <typename T>
|
|
__global__ void blit(const T* input, T* output, int sizeZ, int sizeY, int sizeX, int strideZ, int strideY, int strideX,
|
|
int dstStrideZ, int dstStrideY, int dstStrideX) {
|
|
int total = sizeZ * sizeY * sizeX;
|
|
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; i += blockDim.x * gridDim.x) {
|
|
int x = i % sizeX;
|
|
int tmp = i / sizeX;
|
|
int y = tmp % sizeY;
|
|
int z = tmp / sizeY;
|
|
int srcOffset = z * strideZ + y * strideY + x * strideX;
|
|
int dstOffset = z * dstStrideZ + y * dstStrideY + x * dstStrideX;
|
|
output[dstOffset] = input[srcOffset];
|
|
}
|
|
}
|
|
|
|
cudaError_t RasterBlit(nvinfer1::DataType dataType, uint8_t* output, const uint8_t* input, const Tensor::InsideDescribe::Region& reg, int bytes,
|
|
cudaStream_t stream) {
|
|
int count = reg.size[0] * reg.size[1] * reg.size[2];
|
|
|
|
switch (bytes) {
|
|
case 4:
|
|
if (dataType == nvinfer1::DataType::kFLOAT){
|
|
blit<float><<<CAFFE_GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
|
|
(const float*)input, (float*)output, reg.size[0], reg.size[1], reg.size[2], reg.src.stride[0],
|
|
reg.src.stride[1], reg.src.stride[2], reg.dst.stride[0], reg.dst.stride[1], reg.dst.stride[2]);
|
|
}
|
|
break;
|
|
case 2:
|
|
if (dataType == nvinfer1::DataType::kHALF){
|
|
blit<__half><<<CAFFE_GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
|
|
(const __half*)input, (__half*)output, reg.size[0], reg.size[1], reg.size[2], reg.src.stride[0],
|
|
reg.src.stride[1], reg.src.stride[2], reg.dst.stride[0], reg.dst.stride[1], reg.dst.stride[2]);
|
|
}else{
|
|
blit<int16_t><<<CAFFE_GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
|
|
(const int16_t*)input, (int16_t*)output, reg.size[0], reg.size[1], reg.size[2], reg.src.stride[0],
|
|
reg.src.stride[1], reg.src.stride[2], reg.dst.stride[0], reg.dst.stride[1], reg.dst.stride[2]);
|
|
}
|
|
|
|
break;
|
|
case 1:
|
|
blit<int8_t><<<CAFFE_GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
|
|
(const int8_t*)input, (int8_t*)output, reg.size[0], reg.size[1], reg.size[2], reg.src.stride[0],
|
|
reg.src.stride[1], reg.src.stride[2], reg.dst.stride[0], reg.dst.stride[1], reg.dst.stride[2]);
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
|
|
return cudaPeekAtLastError();
|
|
}
|
|
|
|
} // namespace MNN
|