wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 26 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
c8a779b1bb0f603c3e8c1b4d00e7ac2446055496
nvidia--tensorrt/plugin/common/kernels/maskRCNNKernels.cu
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

3102 lines
116 KiB
Plaintext
Raw Blame History

/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <assert.h>
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include "common/plugin.h"
#include "maskRCNNKernels.h"
#include <NvInfer.h>
#include <assert.h>
#include <cub/cub.cuh>
#include <iostream>
#include <stdio.h>
#include <thrust/device_ptr.h>
#include <thrust/fill.h>
#define DUBUG_KERNEL 0
#define DUBUG_BATCH 0
#define DEBUG_T 1
#define dMIN(a, b) ((a) < (b) ? (a) : (b))
#define dMAX(a, b) ((a) > (b) ? (a) : (b))
#define dCLAMP(x, xMin, xMax) ((x) > (xMin) ? ((x) < (xMax) ? (x) : (xMax)) : (xMin))
template <typename BoxType>
struct BBoxT
{
BoxType y1, x1, y2, x2;
};
inline __device__ __half mul_fb(const __half & a, const __half & b) {
#if __CUDA_ARCH__ >= 530
return a * b;
#else
return __float2half(__half2float(a) * __half2float(b));
#endif
}
inline __device__ __half add_fb(const __half & a, const half & b) {
#if __CUDA_ARCH__ >= 530
return a + b;
#else
return __float2half(__half2float(a) + __half2float(b));
#endif
}
template <typename DType>
__global__ void argMaxReset_kernel(
int samples, int NClass, const DType* in_scores, const int* maxIdx, DType* out_scores)
{
int idx = threadIdx.x + blockIdx.x * blockDim.x;
int max_idx = samples * NClass;
if (idx >= max_idx)
return;
int sampleIdx = idx / NClass;
int classIdx = idx % NClass;
if (classIdx != maxIdx[sampleIdx])
out_scores[idx] = 0;
else
out_scores[idx] = in_scores[idx];
}
template <typename DType>
struct ScanItem
{
DType data;
int idx;
};
template <typename DType>
struct GreaterItem
{
__host__ __device__ __forceinline__ ScanItem<DType> operator()(
const ScanItem<DType>& a, const ScanItem<DType>& b) const
{
return (a.data > b.data ? a : b);
}
};
template <typename DType>
__global__ void resetMemValue_kernel(void* outPtr, int samples, float val)
{
DType* out = static_cast<DType*>(outPtr);
int loop = gridDim.x * blockDim.x;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < samples; idx += loop)
{
out[idx] = (DType) val;
}
}
template <>
__global__ void resetMemValue_kernel<half>(void* outPtr, int samples, float val)
{
__half* out = static_cast<__half*>(outPtr);
int loop = gridDim.x * blockDim.x;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < samples; idx += loop)
{
out[idx] = __float2half(val);
}
}
// blockDim.x : NClass
// GroupDim.x : sample count
// GroupDim.y : batch N
// outScore : DType[ N * sample * 1 ]
// outLabel : int[ N * sample * 1 ]
// outBbox : int[ N * sample * 4 ]
template <typename DType, typename BoxType, int Threads = 32>
__global__ void argMaxGroup_kernel(int samples, int start_class_id, int NClass, const void* inScorePtr,
const void* inBboxPtr, const void* validSampleCountPtr, void* outScorePtr, void* outLabelPtr, void* outBboxPtr)
{
const DType* inScore = static_cast<const DType*>(inScorePtr);
const BoxType* inBbox = static_cast<const BoxType*>(inBboxPtr);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
DType* outScore = static_cast<DType*>(outScorePtr);
BoxType* outLabel = static_cast<BoxType*>(outLabelPtr);
BoxType* outBbox = static_cast<BoxType*>(outBboxPtr);
const int N = blockIdx.y;
const int validSamples = validSampleCount[N];
typedef ScanItem<DType> ScanItemD;
typedef cub::BlockReduce<ScanItemD, Threads> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
for (int iSample = blockIdx.x; iSample < validSamples; iSample += gridDim.x)
{
int classOffset = (N * samples + iSample) * NClass; // start from [batch, count, class0]
// total IPerThread * blockDim
ScanItemD maxItem = {0.0f, -1};
for (int i = start_class_id; i < NClass; i += Threads)
{
int curIdx = i + threadIdx.x;
ScanItemD item = {0.0f, -1};
if (curIdx < NClass)
{
item.data = inScore[classOffset + curIdx];
item.idx = curIdx;
}
const int validNum = (NClass - i > Threads ? Threads : NClass - i);
ScanItemD aggregate = BlockReduce(temp_storage).Reduce(item, GreaterItem<DType>(), validNum);
__syncthreads();
if (aggregate.data > maxItem.data)
{
maxItem = aggregate;
}
#if DUBUG_KERNEL
if (N == DUBUG_BATCH && threadIdx.x == 0 && iSample < 15 /*&& maxItem.idx >= 32*/)
{
printf("argMaxGroup N:%d, iSample:%d, maxItem(score:%.3f, idx:%d)validReduceNum:%d\n", N, iSample,
(float) maxItem.data, maxItem.idx, validNum);
}
#endif
}
const int dstOffset = N * samples + iSample;
if (threadIdx.x == 0)
{
outScore[dstOffset] = maxItem.data;
outLabel[dstOffset] = (BoxType) maxItem.idx;
outBbox[dstOffset * 4] = inBbox[(classOffset + maxItem.idx) * 4];
outBbox[dstOffset * 4 + 1] = inBbox[(classOffset + maxItem.idx) * 4 + 1];
outBbox[dstOffset * 4 + 2] = inBbox[(classOffset + maxItem.idx) * 4 + 2];
outBbox[dstOffset * 4 + 3] = inBbox[(classOffset + maxItem.idx) * 4 + 3];
}
}
}
struct BlockClassSumPrefix
{
int total;
// Constructor
__device__ BlockClassSumPrefix()
: total(0)
{
}
// Callback operator to be entered by the first warp of threads in the block.
// Thread-0 is responsible for returning a value for seeding the block-wide scan.
__device__ int operator()(int aggregate)
{
int old = total;
total += aggregate;
return old;
}
};
#define LabelShift (2.5f)
#define MinValidScore (0.01f)
#define ScoreShift (1.0f)
template <typename DType>
__device__ __forceinline__ DType getKey(DType score, int lable, int NClass)
{
return (lable < 0 ? (DType) 0 : ((DType)(NClass - lable - 1) * LabelShift + score + ScoreShift));
}
template <typename DType, typename BoxType>
__device__ __forceinline__ void getScoreLable(DType key, int NClass, DType& score, BoxType& lable)
{
int i = key / LabelShift;
score = (key <= ScoreShift ? (DType) 0 : key - (DType) i * LabelShift - ScoreShift);
score = dCLAMP(score, (DType) 0, (DType) 1.0);
lable = (BoxType)(key <= ScoreShift ? -1 : (NClass - i - 1));
}
// blockDim.x : threads
// gridDim.x : batch N
// validSampleCount INPUT : int [N]
// classStartPos OUTPUT: int [N * (Class + 1)], need memset to zero before this kernel
// outScore OUTPUT : DType [N * samples]
// outLabel OUTPUT : int [N * samples]
// outSampleIdx OUTPUT : int [N * samples]
// outValidSampleCount : int [N]
// IPerThread * Threads >= sample-count
#define MaxClassNum 255
template <typename DType, typename BoxType, int Threads = 256, int IPerThread = 4>
__global__ void sortPerClass_kernel(
// int N,
int samples, int NClass, int background, float scoreThreshold, const void* validSampleCountPtr,
const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr, void* classStartPosPtr, void* outScorePtr,
void* outLabelPtr, void* outSampleIdxPtr, void* outValidSampleCountPtr)
{
typedef cub::BlockExchange<DType, Threads, IPerThread> BlockExchangeKey;
typedef cub::BlockExchange<int, Threads, IPerThread> BlockExchangeI;
typedef cub::BlockRadixSort<DType, Threads, IPerThread, int> BlockRadixSort;
typedef cub::BlockScan<int, Threads> BlockScanClass;
__shared__ union
{
typename BlockExchangeKey::TempStorage storageKey;
typename BlockExchangeI::TempStorage storageI;
typename BlockRadixSort::TempStorage storageSort;
typename BlockScanClass::TempStorage storageScan;
} temp_storage;
__shared__ int smemClassCount[MaxClassNum];
assert(NClass < MaxClassNum);
assert(IPerThread * Threads >= samples);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
const DType* inScore = static_cast<const DType*>(inScorePtr);
const BoxType* inLabel = static_cast<const BoxType*>(inLabelPtr);
int* classStartPos = static_cast<int*>(classStartPosPtr);
DType* outScore = static_cast<DType*>(outScorePtr);
BoxType* outLabel = static_cast<BoxType*>(outLabelPtr);
int* outSampleIdx = static_cast<int*>(outSampleIdxPtr);
int* outValidSampleCount = static_cast<int*>(outValidSampleCountPtr);
for (int s = threadIdx.x; s < NClass + 1; s += blockDim.x)
{
smemClassCount[s] = 0;
}
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
DType key[IPerThread];
int iSample[IPerThread];
for (int i = 0; i < IPerThread; ++i)
{
iSample[i] = -1;
key[i] = -1.0f;
int curIdx = i * Threads + threadIdx.x;
if (curIdx < validSamples)
{
int label = (int) (inLabel[blockOffset + curIdx]);
DType score = inScore[blockOffset + curIdx];
if (label != background && label != -1 && score >= scoreThreshold)
{
key[i] = getKey(score, label, NClass);
iSample[i] = curIdx;
}
}
}
BlockExchangeKey(temp_storage.storageKey).StripedToBlocked(key);
__syncthreads();
BlockExchangeI(temp_storage.storageI).StripedToBlocked(iSample);
__syncthreads();
BlockRadixSort(temp_storage.storageSort).SortDescendingBlockedToStriped(key, iSample);
__syncthreads();
// store Idx
cub::StoreDirectStriped<Threads>(threadIdx.x, outSampleIdx + blockOffset, iSample, validSamples);
BoxType lable[IPerThread];
DType score[IPerThread];
#pragma unroll
for (int i = 0; i < IPerThread; ++i)
{
getScoreLable(key[i], NClass, score[i], lable[i]);
}
cub::StoreDirectStriped<Threads>(threadIdx.x, outScore + blockOffset, score, validSamples);
cub::StoreDirectStriped<Threads>(threadIdx.x, outLabel + blockOffset, lable, validSamples);
// final
for (int i = 0; i < IPerThread; ++i)
{
if (lable[i] >= (BoxType) 0)
{
atomicAdd(&smemClassCount[(int) lable[i]], 1);
}
}
__syncthreads();
int classBlockOffset = N * (NClass + 1); // Exclusive-sum, 1st is 0, last is final sum
#if DUBUG_KERNEL
if (N == DUBUG_BATCH && threadIdx.x == 0)
{
printf("sortPerClass(N:%d) final count of each label, valid samples:%d\n", N, validSamples);
for (int k = 0; k < NClass; ++k)
{
if (smemClassCount[k] > 0)
printf("Batch:%d, L:%d, count:%d, \n", N, k, smemClassCount[k]);
}
}
__syncthreads();
#endif
BlockClassSumPrefix sumPrefix;
for (int s = 0; s < NClass; s += blockDim.x)
{ // s start from block
int iClassSamples = 0;
int iClass = s + threadIdx.x;
if (iClass < NClass)
{
iClassSamples = smemClassCount[iClass];
}
BlockScanClass(temp_storage.storageScan).ExclusiveSum(iClassSamples, iClassSamples, sumPrefix);
__syncthreads();
if (iClass < NClass)
{
classStartPos[classBlockOffset + iClass] = iClassSamples;
}
}
if (threadIdx.x == 0)
{
classStartPos[classBlockOffset + NClass] = sumPrefix.total;
assert(sumPrefix.total <= validSamples); // background data removed.
outValidSampleCount[N] = sumPrefix.total;
#if DUBUG_KERNEL
if (N == DUBUG_BATCH)
printf("After sortPerClass, batch:%d valid samples total:%d\n", N, sumPrefix.total);
#endif
}
}
template <int Threads = 256, int IPerThread = 4>
__global__ void sortPerClass_kernel_half(
// int N,
int samples, int NClass, int background, float scoreThreshold, const void* validSampleCountPtr,
const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr, void* classStartPosPtr, void* outScorePtr,
void* outLabelPtr, void* outSampleIdxPtr, void* outValidSampleCountPtr)
{
typedef cub::BlockExchange<float, Threads, IPerThread> BlockExchangeKey;
typedef cub::BlockExchange<int, Threads, IPerThread> BlockExchangeI;
typedef cub::BlockRadixSort<float, Threads, IPerThread, int> BlockRadixSort;
typedef cub::BlockScan<int, Threads> BlockScanClass;
__shared__ union
{
typename BlockExchangeKey::TempStorage storageKey;
typename BlockExchangeI::TempStorage storageI;
typename BlockRadixSort::TempStorage storageSort;
typename BlockScanClass::TempStorage storageScan;
} temp_storage;
__shared__ int smemClassCount[MaxClassNum];
assert(NClass < MaxClassNum);
assert(IPerThread * Threads >= samples);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
const __half* inScore = static_cast<const __half*>(inScorePtr);
const __half* inLabel = static_cast<const __half*>(inLabelPtr);
int* classStartPos = static_cast<int*>(classStartPosPtr);
__half* outScore = static_cast<__half*>(outScorePtr);
__half* outLabel = static_cast<__half*>(outLabelPtr);
int* outSampleIdx = static_cast<int*>(outSampleIdxPtr);
int* outValidSampleCount = static_cast<int*>(outValidSampleCountPtr);
for (int s = threadIdx.x; s < NClass + 1; s += blockDim.x)
{
smemClassCount[s] = 0;
}
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
float key[IPerThread];
int iSample[IPerThread];
for (int i = 0; i < IPerThread; ++i)
{
iSample[i] = -1;
key[i] = -1.0f;
int curIdx = i * Threads + threadIdx.x;
if (curIdx < validSamples)
{
int label = __half2int_rd(inLabel[blockOffset + curIdx]);
float score = __half2float(inScore[blockOffset + curIdx]);
if (label != background && label != -1 && score >= scoreThreshold)
{
key[i] = getKey<float>(score, label, NClass);
iSample[i] = curIdx;
}
}
}
BlockExchangeKey(temp_storage.storageKey).StripedToBlocked(key);
__syncthreads();
BlockExchangeI(temp_storage.storageI).StripedToBlocked(iSample);
__syncthreads();
BlockRadixSort(temp_storage.storageSort).SortDescendingBlockedToStriped(key, iSample);
__syncthreads();
// store Idx
cub::StoreDirectStriped<Threads>(threadIdx.x, outSampleIdx + blockOffset, iSample, validSamples);
__half lable[IPerThread];
__half score[IPerThread];
for (int i = 0; i < IPerThread; ++i)
{
float label_float;
float score_float;
getScoreLable<float>(key[i], NClass, score_float, label_float);
lable[i] = __float2half(label_float);
score[i] = __float2half(score_float);
}
cub::StoreDirectStriped<Threads>(threadIdx.x, outScore + blockOffset, score, validSamples);
cub::StoreDirectStriped<Threads>(threadIdx.x, outLabel + blockOffset, lable, validSamples);
// final
for (int i = 0; i < IPerThread; ++i)
{
if (__half2float(lable[i]) >= 0)
{
atomicAdd(&smemClassCount[__half2int_rd(lable[i])], 1);
}
}
__syncthreads();
int classBlockOffset = N * (NClass + 1); // Exclusive-sum, 1st is 0, last is final sum
#if DUBUG_KERNEL
if (N == DUBUG_BATCH && threadIdx.x == 0)
{
printf("sortPerClass(N:%d) final count of each label, valid samples:%d\n", N, validSamples);
for (int k = 0; k < NClass; ++k)
{
if (smemClassCount[k] > 0)
printf("Batch:%d, L:%d, count:%d, \n", N, k, smemClassCount[k]);
}
}
__syncthreads();
#endif
BlockClassSumPrefix sumPrefix;
for (int s = 0; s < NClass; s += blockDim.x)
{ // s start from block
int iClassSamples = 0;
int iClass = s + threadIdx.x;
if (iClass < NClass)
{
iClassSamples = smemClassCount[iClass];
}
BlockScanClass(temp_storage.storageScan).ExclusiveSum(iClassSamples, iClassSamples, sumPrefix);
__syncthreads();
if (iClass < NClass)
{
classStartPos[classBlockOffset + iClass] = iClassSamples;
}
}
if (threadIdx.x == 0)
{
classStartPos[classBlockOffset + NClass] = sumPrefix.total;
assert(sumPrefix.total <= validSamples); // background data removed.
outValidSampleCount[N] = sumPrefix.total;
#if DUBUG_KERNEL
if (N == DUBUG_BATCH)
printf("After sortPerClass, batch:%d valid samples total:%d\n", N, sumPrefix.total);
#endif
}
}
template <typename DType>
__device__ __forceinline__ BBoxT<DType> readBbox(const BBoxT<DType>* inBbox, int idx)
{
BBoxT<DType> ret = ((BBoxT<DType>*) (inBbox))[idx];
return ret;
}
template <typename DType>
__device__ __forceinline__ DType boxIoU(const BBoxT<DType>& a, const BBoxT<DType>& b)
{
BBoxT<DType> overlap = {
dMAX(a.y1, b.y1), dMAX(a.x1, b.x1), dMIN(a.y2, b.y2), dMIN(a.x2, b.x2),
};
DType oW = overlap.x2 - overlap.x1;
DType oH = overlap.y2 - overlap.y1;
if (oW < (DType) 0 || oH < (DType) 0)
return (DType) 0;
DType oA = oW * oH;
return (oA / ((a.y2 - a.y1) * (a.x2 - a.x1) + (b.y2 - b.y1) * (b.x2 - b.x1) - oA));
}
// PerClassNMS
// gridDim.x : batch-N
// blockDim.x : Threads
// ItemsPerThreads : = divUp(samples, Threads)
// outFlagSamples OUT: int [N * samples]
template <typename DType, typename BoxType, int Threads = 256, int ItemsPerThreads = 4>
__global__ void PerClassNMS_kernel(
// int N,
int samples, int NClass, const float nmsThreshold, const void* validSampleCountPtr,
// const void *inScorePtr,
const void* inLabelPtr, const void* inBboxPtr, const void* inBboxRefIdxPtr, const void* classStartsPtr,
void* outFlagSamplesPtr)
{
typedef BBoxT<BoxType> BBox;
__shared__ struct
{
BBox refBox[MaxClassNum];
int endIdx[MaxClassNum];
int refIdx[MaxClassNum + 1];
bool markSamples[Threads * ItemsPerThreads];
int done;
} smemClasses;
assert(NClass + 1 < MaxClassNum);
assert(samples <= Threads * ItemsPerThreads);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
// const DType *inScore = static_cast<const DType *>(inScorePtr);
const BoxType* inLabel = static_cast<const BoxType*>(inLabelPtr);
const BBox* inBbox = static_cast<const BBox*>(inBboxPtr);
const int* inBboxRefIdx = static_cast<const int*>(inBboxRefIdxPtr);
const int* classStarts = static_cast<const int*>(classStartsPtr);
int* outFlagSamples = static_cast<int*>(outFlagSamplesPtr);
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
if (threadIdx.x == 0)
{
smemClasses.done = 0;
}
BBox curBox[ItemsPerThreads];
int label[ItemsPerThreads];
#pragma unroll
for (int ite = 0; ite * blockDim.x < validSamples; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
if (curIdx < validSamples)
{
label[ite] = (int) inLabel[blockOffset + curIdx];
curBox[ite] = readBbox(inBbox, blockOffset + inBboxRefIdx[blockOffset + curIdx]);
}
else
{
label[ite] = -1;
}
smemClasses.markSamples[curIdx] = (label[ite] < 0 ? false : true);
}
int classBlockOffset = N * (NClass + 1);
for (int i = threadIdx.x; i < NClass + 1; i += blockDim.x)
{
int refIdx = classStarts[classBlockOffset + i];
smemClasses.refIdx[i] = refIdx;
smemClasses.refBox[i] = readBbox(inBbox, blockOffset + inBboxRefIdx[blockOffset + refIdx]);
}
__syncthreads();
for (int i = threadIdx.x; i < NClass; i += blockDim.x)
{
int endIdx = smemClasses.refIdx[i + 1];
smemClasses.endIdx[i] = endIdx;
if (endIdx == smemClasses.refIdx[i])
{
atomicAdd(&smemClasses.done, 1);
}
}
__syncthreads();
#if DUBUG_KERNEL
// print info
if (N == DUBUG_BATCH && threadIdx.x == 0)
{
printf("batch:%d, before starting NMS, done count:%d\n", N, smemClasses.done);
printf("batch:%d, Total num:%d, startPos:\n", N, validSamples);
for (int k = 0; k < NClass; ++k)
{
if (smemClasses.refIdx[k] != smemClasses.endIdx[k])
{
printf("Batch:%d, label:%d [%d : %d], check ref-label:%d\n", N, k, smemClasses.refIdx[k],
smemClasses.endIdx[k], (int) inLabel[blockOffset + smemClasses.refIdx[k]]);
}
}
printf("\n");
}
__syncthreads();
#endif
// class done to check stop point
while (smemClasses.done < NClass)
{
for (int ite = 0; ite * blockDim.x < validSamples; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
int refIdx = -1;
int endIdx = -1;
if (curIdx < validSamples && smemClasses.markSamples[curIdx])
{
if (label[ite] >= 0)
{
refIdx = smemClasses.refIdx[label[ite]];
endIdx = smemClasses.endIdx[label[ite]];
if (curIdx > refIdx && curIdx < endIdx)
{
BBox refBox = smemClasses.refBox[label[ite]];
if (boxIoU(refBox, curBox[ite]) > (DType) nmsThreshold)
{
smemClasses.markSamples[curIdx] = false;
}
}
}
}
}
__syncthreads();
// push refIdx/refBox forward to next mark
// only the refIdx thread to push itself. other threads idle
for (int i = threadIdx.x; i < NClass; i += blockDim.x)
{
int refIdx = smemClasses.refIdx[i];
int endIdx = smemClasses.endIdx[i];
if (refIdx < endIdx)
{
do
{
++refIdx;
} while (refIdx < endIdx && smemClasses.markSamples[refIdx] == false);
smemClasses.refIdx[i] = refIdx;
if (refIdx < endIdx)
{
smemClasses.refBox[i] = readBbox(inBbox, blockOffset + inBboxRefIdx[blockOffset + refIdx]);
}
else
{
atomicAdd(&smemClasses.done, 1);
}
}
}
__syncthreads();
}
// no need to write all data out
for (int segment = 0; segment < validSamples; segment += blockDim.x)
{
int curIdx = segment + threadIdx.x;
if (curIdx < validSamples)
{
outFlagSamples[blockOffset + curIdx] = (smemClasses.markSamples[curIdx] ? 1 : 0);
}
}
}
template <int Threads = 256, int ItemsPerThreads = 4>
__global__ void PerClassNMS_half_kernel(
// int N,
int samples, int NClass, const float nmsThreshold, const void* validSampleCountPtr,
// const void *inScorePtr,
const void* inLabelPtr, const void* inBboxPtr, const void* inBboxRefIdxPtr, const void* classStartsPtr,
void* outFlagSamplesPtr)
{
typedef BBoxT<__half> BBox;
__shared__ struct
{
BBox refBox[MaxClassNum];
int endIdx[MaxClassNum];
int refIdx[MaxClassNum + 1];
bool markSamples[Threads * ItemsPerThreads];
int done;
} smemClasses;
assert(NClass + 1 < MaxClassNum);
assert(samples <= Threads * ItemsPerThreads);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
// const DType *inScore = static_cast<const DType *>(inScorePtr);
const __half* inLabel = static_cast<const __half*>(inLabelPtr);
const BBox* inBbox = static_cast<const BBox*>(inBboxPtr);
const int* inBboxRefIdx = static_cast<const int*>(inBboxRefIdxPtr);
const int* classStarts = static_cast<const int*>(classStartsPtr);
int* outFlagSamples = static_cast<int*>(outFlagSamplesPtr);
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
if (threadIdx.x == 0)
{
smemClasses.done = 0;
}
BBox curBox[ItemsPerThreads];
int label[ItemsPerThreads];
#pragma unroll
for (int ite = 0; ite * blockDim.x < validSamples; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
if (curIdx < validSamples)
{
label[ite] = __half2int_rd(inLabel[blockOffset + curIdx]);
curBox[ite] = readBbox<__half>(inBbox, blockOffset + inBboxRefIdx[blockOffset + curIdx]);
}
else
{
label[ite] = -1;
}
smemClasses.markSamples[curIdx] = (label[ite] < 0 ? false : true);
}
int classBlockOffset = N * (NClass + 1);
for (int i = threadIdx.x; i < NClass + 1; i += blockDim.x)
{
int refIdx = classStarts[classBlockOffset + i];
smemClasses.refIdx[i] = refIdx;
smemClasses.refBox[i] = readBbox<__half>(inBbox, blockOffset + inBboxRefIdx[blockOffset + refIdx]);
}
__syncthreads();
for (int i = threadIdx.x; i < NClass; i += blockDim.x)
{
int endIdx = smemClasses.refIdx[i + 1];
smemClasses.endIdx[i] = endIdx;
if (endIdx == smemClasses.refIdx[i])
{
atomicAdd(&smemClasses.done, 1);
}
}
__syncthreads();
#if DUBUG_KERNEL
// print info
if (N == DUBUG_BATCH && threadIdx.x == 0)
{
printf("batch:%d, before starting NMS, done count:%d\n", N, smemClasses.done);
printf("batch:%d, Total num:%d, startPos:\n", N, validSamples);
for (int k = 0; k < NClass; ++k)
{
if (smemClasses.refIdx[k] != smemClasses.endIdx[k])
{
printf("Batch:%d, label:%d [%d : %d], check ref-label:%d\n", N, k, smemClasses.refIdx[k],
smemClasses.endIdx[k], (int) inLabel[blockOffset + smemClasses.refIdx[k]]);
}
}
printf("\n");
}
__syncthreads();
#endif
// class done to check stop point
while (smemClasses.done < NClass)
{
for (int ite = 0; ite * blockDim.x < validSamples; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
int refIdx = -1;
int endIdx = -1;
if (curIdx < validSamples && smemClasses.markSamples[curIdx])
{
if (label[ite] >= 0)
{
refIdx = smemClasses.refIdx[label[ite]];
endIdx = smemClasses.endIdx[label[ite]];
if (curIdx > refIdx && curIdx < endIdx)
{
BBox refBox_half = smemClasses.refBox[label[ite]];
BBox curBox_half = curBox[ite];
BBoxT<float> refBox;
BBoxT<float> curBox_float;
refBox.y1 = __half2float(refBox_half.y1);
refBox.x1 = __half2float(refBox_half.x1);
refBox.y2 = __half2float(refBox_half.y2);
refBox.x2 = __half2float(refBox_half.x2);
curBox_float.y1 = __half2float(curBox_half.y1);
curBox_float.x1 = __half2float(curBox_half.x1);
curBox_float.y2 = __half2float(curBox_half.y2);
curBox_float.x2 = __half2float(curBox_half.x2);
if (boxIoU<float>(refBox, curBox_float) > nmsThreshold)
{
smemClasses.markSamples[curIdx] = false;
}
}
}
}
}
__syncthreads();
// push refIdx/refBox forward to next mark
// only the refIdx thread to push itself. other threads idle
for (int i = threadIdx.x; i < NClass; i += blockDim.x)
{
int refIdx = smemClasses.refIdx[i];
int endIdx = smemClasses.endIdx[i];
if (refIdx < endIdx)
{
do
{
++refIdx;
} while (refIdx < endIdx && smemClasses.markSamples[refIdx] == false);
smemClasses.refIdx[i] = refIdx;
if (refIdx < endIdx)
{
smemClasses.refBox[i] = readBbox<__half>(inBbox, blockOffset + inBboxRefIdx[blockOffset + refIdx]);
}
else
{
atomicAdd(&smemClasses.done, 1);
}
}
}
__syncthreads();
}
// no need to write all data out
for (int segment = 0; segment < validSamples; segment += blockDim.x)
{
int curIdx = segment + threadIdx.x;
if (curIdx < validSamples)
{
outFlagSamples[blockOffset + curIdx] = (smemClasses.markSamples[curIdx] ? 1 : 0);
}
}
}
// TopKGather
// gridDim.x : batch-N
// blockDim.x : Threads
// ItemsPerThreads : = divUp(samples, Threads)
// outDetectionCount : int [N], must be set 0 before kernel
#define MaxItemsPerThreads 8
template <typename DType, typename BoxType, int Threads = 256>
__global__ void TopKGatherProposal_kernel(int samples, int keepTopK, const void* validSampleCountPtr,
const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr, const void* inBboxRefIdxPtr,
const void* inFlagSamplesPtr, void* outBboxPtr)
{
typedef BBoxT<BoxType> BBox;
typedef cub::BlockRadixSort<DType, Threads, 1, int> BlockRadixSort1;
typedef cub::BlockRadixSort<DType, Threads, 2, int> BlockRadixSort2;
typedef cub::BlockRadixSort<DType, Threads, 3, int> BlockRadixSort3;
typedef cub::BlockRadixSort<DType, Threads, 4, int> BlockRadixSort4;
typedef cub::BlockRadixSort<DType, Threads, 5, int> BlockRadixSort5;
typedef cub::BlockRadixSort<DType, Threads, 6, int> BlockRadixSort6;
typedef cub::BlockRadixSort<DType, Threads, 7, int> BlockRadixSort7;
typedef cub::BlockRadixSort<DType, Threads, 8, int> BlockRadixSort8;
__shared__ union
{
typename BlockRadixSort8::TempStorage sort8;
typename BlockRadixSort7::TempStorage sort7;
typename BlockRadixSort6::TempStorage sort6;
typename BlockRadixSort5::TempStorage sort5;
typename BlockRadixSort4::TempStorage sort4;
typename BlockRadixSort3::TempStorage sort3;
typename BlockRadixSort2::TempStorage sort2;
typename BlockRadixSort1::TempStorage sort1;
} temp_storage;
assert(MaxItemsPerThreads * Threads >= samples);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
const DType* inScore = static_cast<const DType*>(inScorePtr);
const BBox* inBbox = static_cast<const BBox*>(inBboxPtr);
const int* inBboxRefIdx = static_cast<const int*>(inBboxRefIdxPtr);
const int* inFlagSamples = static_cast<const int*>(inFlagSamplesPtr);
BBox* outBbox = static_cast<BBox*>(outBboxPtr);
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
int finalTopK = dMIN(keepTopK, validSamples);
int idx[MaxItemsPerThreads];
DType score[MaxItemsPerThreads];
int totalItems = (validSamples + (blockDim.x - 1)) / blockDim.x;
for (int ite = 0; ite < totalItems; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
if (curIdx < validSamples && inFlagSamples[blockOffset + curIdx])
{
idx[ite] = curIdx;
score[ite] = inScore[blockOffset + curIdx];
}
else
{
idx[ite] = -1;
score[ite] = 0.0f;
}
}
switch (totalItems)
{
case 0: break;
case 1:
BlockRadixSort1(temp_storage.sort1).SortDescendingBlockedToStriped((DType(&)[1]) score, (int(&)[1]) idx);
break;
case 2:
BlockRadixSort2(temp_storage.sort2).SortDescendingBlockedToStriped((DType(&)[2]) score, (int(&)[2]) idx);
break;
case 3:
BlockRadixSort3(temp_storage.sort3).SortDescendingBlockedToStriped((DType(&)[3]) score, (int(&)[3]) idx);
break;
case 4:
BlockRadixSort4(temp_storage.sort4).SortDescendingBlockedToStriped((DType(&)[4]) score, (int(&)[4]) idx);
break;
case 5:
BlockRadixSort5(temp_storage.sort5).SortDescendingBlockedToStriped((DType(&)[5]) score, (int(&)[5]) idx);
break;
case 6:
BlockRadixSort6(temp_storage.sort6).SortDescendingBlockedToStriped((DType(&)[6]) score, (int(&)[6]) idx);
break;
case 7:
BlockRadixSort7(temp_storage.sort7).SortDescendingBlockedToStriped((DType(&)[7]) score, (int(&)[7]) idx);
break;
case 8:
BlockRadixSort8(temp_storage.sort8).SortDescendingBlockedToStriped((DType(&)[8]) score, (int(&)[8]) idx);
break;
default: assert(false);
}
__syncthreads();
int outBlockOffset = N * keepTopK;
int topkItems = (keepTopK + (Threads - 1)) / Threads;
for (int i = 0; i < topkItems; ++i)
{
int curI = i * blockDim.x + threadIdx.x;
if (curI < keepTopK)
{
BBox oB = {(BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f};
if (curI < finalTopK && idx[i] >= 0 && float(score[i]) > MinValidScore)
{
oB = ((BBox*) inBbox)[blockOffset + inBboxRefIdx[blockOffset + idx[i]]];
}
((BBox*) outBbox)[outBlockOffset + curI] = oB;
}
}
}
#define MaxItemsPerThreads 8
template <typename DType, typename BoxType, int Threads = 256>
__global__ void TopKGather_kernel(int samples, int keepTopK, const void* validSampleCountPtr, const void* inScorePtr,
const void* inLabelPtr, const void* inBboxPtr, const void* inBboxRefIdxPtr, const void* inFlagSamplesPtr,
void* outDetectionPtr)
{
typedef BBoxT<BoxType> BBox;
typedef cub::BlockRadixSort<DType, Threads, 1, int> BlockRadixSort1;
typedef cub::BlockRadixSort<DType, Threads, 2, int> BlockRadixSort2;
typedef cub::BlockRadixSort<DType, Threads, 3, int> BlockRadixSort3;
typedef cub::BlockRadixSort<DType, Threads, 4, int> BlockRadixSort4;
typedef cub::BlockRadixSort<DType, Threads, 5, int> BlockRadixSort5;
typedef cub::BlockRadixSort<DType, Threads, 6, int> BlockRadixSort6;
typedef cub::BlockRadixSort<DType, Threads, 7, int> BlockRadixSort7;
typedef cub::BlockRadixSort<DType, Threads, 8, int> BlockRadixSort8;
__shared__ union
{
typename BlockRadixSort8::TempStorage sort8;
typename BlockRadixSort7::TempStorage sort7;
typename BlockRadixSort6::TempStorage sort6;
typename BlockRadixSort5::TempStorage sort5;
typename BlockRadixSort4::TempStorage sort4;
typename BlockRadixSort3::TempStorage sort3;
typename BlockRadixSort2::TempStorage sort2;
typename BlockRadixSort1::TempStorage sort1;
} temp_storage;
assert(MaxItemsPerThreads * Threads >= samples);
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
const DType* inScore = static_cast<const DType*>(inScorePtr);
const BoxType* inLabel = static_cast<const BoxType*>(inLabelPtr); // InLabel keeps INT32
const BBox* inBbox = static_cast<const BBox*>(inBboxPtr);
const int* inBboxRefIdx = static_cast<const int*>(inBboxRefIdxPtr);
const int* inFlagSamples = static_cast<const int*>(inFlagSamplesPtr);
DType* outDetections = static_cast<DType*>(outDetectionPtr);
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
int finalTopK = dMIN(keepTopK, validSamples);
int idx[MaxItemsPerThreads];
DType score[MaxItemsPerThreads];
int totalItems = (validSamples + (blockDim.x - 1)) / blockDim.x;
for (int ite = 0; ite < totalItems; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
if (curIdx < validSamples && inFlagSamples[blockOffset + curIdx])
{
idx[ite] = curIdx;
score[ite] = inScore[blockOffset + curIdx];
}
else
{
idx[ite] = -1;
score[ite] = 0.0f;
}
}
switch (totalItems)
{
case 0: break;
case 1:
BlockRadixSort1(temp_storage.sort1).SortDescendingBlockedToStriped((DType(&)[1]) score, (int(&)[1]) idx);
break;
case 2:
BlockRadixSort2(temp_storage.sort2).SortDescendingBlockedToStriped((DType(&)[2]) score, (int(&)[2]) idx);
break;
case 3:
BlockRadixSort3(temp_storage.sort3).SortDescendingBlockedToStriped((DType(&)[3]) score, (int(&)[3]) idx);
break;
case 4:
BlockRadixSort4(temp_storage.sort4).SortDescendingBlockedToStriped((DType(&)[4]) score, (int(&)[4]) idx);
break;
case 5:
BlockRadixSort5(temp_storage.sort5).SortDescendingBlockedToStriped((DType(&)[5]) score, (int(&)[5]) idx);
break;
case 6:
BlockRadixSort6(temp_storage.sort6).SortDescendingBlockedToStriped((DType(&)[6]) score, (int(&)[6]) idx);
break;
case 7:
BlockRadixSort7(temp_storage.sort7).SortDescendingBlockedToStriped((DType(&)[7]) score, (int(&)[7]) idx);
break;
case 8:
BlockRadixSort8(temp_storage.sort8).SortDescendingBlockedToStriped((DType(&)[8]) score, (int(&)[8]) idx);
break;
default: assert(false);
}
__syncthreads();
int outBlockOffset = N * keepTopK;
int topkItems = (keepTopK + (Threads - 1)) / Threads;
for (int i = 0; i < topkItems; ++i)
{
int curI = i * blockDim.x + threadIdx.x;
if (curI < keepTopK)
{
BBox oB = {(BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f};
DType oS = 0.0f;
BoxType oL = -1;
if (curI < finalTopK && idx[i] >= 0 && float(score[i]) > MinValidScore)
{
oB = ((BBox*) inBbox)[blockOffset + inBboxRefIdx[blockOffset + idx[i]]];
oS = score[i];
oL = (BoxType) inLabel[blockOffset + idx[i]];
}
outDetections[(outBlockOffset + curI) * 6] = oB.y1;
outDetections[(outBlockOffset + curI) * 6 + 1] = oB.x1;
outDetections[(outBlockOffset + curI) * 6 + 2] = oB.y2;
outDetections[(outBlockOffset + curI) * 6 + 3] = oB.x2;
outDetections[(outBlockOffset + curI) * 6 + 4] = oL;
outDetections[(outBlockOffset + curI) * 6 + 5] = oS;
}
}
}
RefineDetectionWorkSpace::RefineDetectionWorkSpace(
const int batchSize, const int sampleCount, const RefineNMSParameters& param, const nvinfer1::DataType inType)
: argMaxScoreDims(sampleCount, 1)
, argMaxBboxDims(sampleCount, 4)
, argMaxLabelDims(sampleCount, 1)
, sortClassScoreDims(sampleCount, 1)
, sortClassLabelDims(sampleCount, 1)
, sortClassSampleIdxDims(sampleCount + 1, 1)
, sortClassPosDims(param.numClasses + 1, 1)
, sortNMSMarkDims(sampleCount, 1)
{
size_t sumSize = 0;
const nvinfer1::DataType type = nvinfer1::DataType::kFLOAT;
// resource
// arMaxScore : [N, samples] : m_Type
argMaxScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(argMaxScoreDims) * typeSize(type) * batchSize);
argMaxBboxOffset = sumSize;
// argMaxBbox : [N, samples, 4] : m_Type
sumSize += AlignMem(dimVolume(argMaxBboxDims) * typeSize(type) * batchSize);
argMaxLabelOffset = sumSize;
// argMaxLabel : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(argMaxLabelDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassScoreOffset = sumSize;
// sortClassScore : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassScoreDims) * typeSize(type) * batchSize);
sortClassLabelOffset = sumSize;
// sortClassLabel : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassLabelDims) * typeSize(type) * batchSize);
sortClassSampleIdxOffset = sumSize;
// sortClassSampleIdx : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortClassSampleIdxDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassValidCountOffset = sumSize;
// sortClassValidCount : [N, 1] : kINT32
sumSize += AlignMem(dimVolume(sortClassValidCountDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassPosOffset = sumSize;
// sortClassPos : [N, numClasses+1] : kINT32
sumSize += AlignMem(dimVolume(sortClassPosDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortNMSMarkOffset = sumSize;
// sortNMSMark : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortNMSMarkDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
totalSize = sumSize;
}
ProposalWorkSpace::ProposalWorkSpace(const int batchSize, const int inputCnt, const int sampleCount,
const RefineNMSParameters& param, const nvinfer1::DataType inType)
: preRefineScoreDims(inputCnt, 1)
, preRefineSortedScoreDims(inputCnt, 1)
, preRefineBboxDims(inputCnt, 4)
, argMaxScoreDims(sampleCount, 1)
, argMaxBboxDims(sampleCount, 4)
, argMaxLabelDims(sampleCount, 1)
, sortClassScoreDims(sampleCount, 1)
, sortClassLabelDims(sampleCount, 1)
, sortClassSampleIdxDims(sampleCount, 1)
, sortClassPosDims(param.numClasses + 1, 1)
, sortNMSMarkDims(sampleCount, 1)
{
size_t sumSize = 0;
const nvinfer1::DataType type = nvinfer1::DataType::kFLOAT;
// resource
// temp storage size for sorting scores
tempStorageOffset = sumSize;
sumSize += (1 << 23) * batchSize;
// preRefineScore : [N, inputcnt, 1] // extracted foreground score from inputs[0]
preRefineScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(preRefineScoreDims) * typeSize(type) * batchSize);
// preRefineSortedScore: [N, inputcnt, 1]
preRefineSortedScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(preRefineSortedScoreDims) * typeSize(type) * batchSize);
// preRefineBbox: [N, inputcnt, 4] // sorted bbox
preRefineBboxOffset = sumSize;
sumSize += AlignMem(dimVolume(preRefineBboxDims) * typeSize(type) * batchSize);
// arMaxScore : [N, samples] : m_Type
argMaxScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(argMaxScoreDims) * typeSize(type) * batchSize);
argMaxBboxOffset = sumSize;
// argMaxBbox : [N, samples, 4] : m_Type
sumSize += AlignMem(dimVolume(argMaxBboxDims) * typeSize(type) * batchSize);
argMaxLabelOffset = sumSize;
// argMaxLabel : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(argMaxLabelDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassScoreOffset = sumSize;
// sortClassScore : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassScoreDims) * typeSize(type) * batchSize);
sortClassLabelOffset = sumSize;
// sortClassLabel : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassLabelDims) * typeSize(type) * batchSize);
sortClassSampleIdxOffset = sumSize;
// sortClassSampleIdx : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortClassSampleIdxDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassValidCountOffset = sumSize;
// sortClassValidCount : [N, 1] : kINT32
sumSize += AlignMem(dimVolume(sortClassValidCountDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassPosOffset = sumSize;
// sortClassPos : [N, numClasses+1] : kINT32
sumSize += AlignMem(dimVolume(sortClassPosDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortNMSMarkOffset = sumSize;
// sortNMSMark : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortNMSMarkDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
totalSize = sumSize;
}
MultilevelProposeROIWorkSpace::MultilevelProposeROIWorkSpace(const int batchSize, const int inputCnt,
const int sampleCount, const RefineNMSParameters& param, const nvinfer1::DataType inType)
: preRefineSortedScoreDims(inputCnt, 1)
, preRefineBboxDims(inputCnt, 4)
, argMaxScoreDims(sampleCount, 1)
, argMaxBboxDims(sampleCount, 4)
, argMaxLabelDims(sampleCount, 1)
, sortClassScoreDims(sampleCount, 1)
, sortClassLabelDims(sampleCount, 1)
, sortClassSampleIdxDims(sampleCount + 1, 1)
, sortClassPosDims(param.numClasses + 1, 1)
, sortNMSMarkDims(sampleCount, 1)
{
size_t sumSize = 0;
const nvinfer1::DataType type = inType;
// resource
// temp storage size for sorting scores
tempStorageOffset = sumSize;
sumSize += (1 << 23) * batchSize;
// preRefineSortedScore: [N, inputcnt, 1]
preRefineSortedScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(preRefineSortedScoreDims) * typeSize(type) * batchSize);
// preRefineBbox: [N, inputcnt, 4] // sorted bbox
preRefineBboxOffset = sumSize;
sumSize += AlignMem(dimVolume(preRefineBboxDims) * typeSize(type) * batchSize);
// argMaxScore : [N, samples] : m_Type
argMaxScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(argMaxScoreDims) * typeSize(type) * batchSize);
argMaxBboxOffset = sumSize;
// argMaxBbox : [N, samples, 4] : m_Type
sumSize += AlignMem(dimVolume(argMaxBboxDims) * typeSize(type) * batchSize);
argMaxLabelOffset = sumSize;
// argMaxLabel : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(argMaxLabelDims) * typeSize(type) * batchSize);
sortClassScoreOffset = sumSize;
// sortClassScore : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassScoreDims) * typeSize(type) * batchSize);
sortClassLabelOffset = sumSize;
// sortClassLabel : [N, samples] : m_Type
sumSize += AlignMem(dimVolume(sortClassLabelDims) * typeSize(type) * batchSize);
sortClassSampleIdxOffset = sumSize;
// sortClassSampleIdx : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortClassSampleIdxDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassValidCountOffset = sumSize;
// sortClassValidCount : [N, 1] : kINT32
sumSize += AlignMem(dimVolume(sortClassValidCountDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortClassPosOffset = sumSize;
// sortClassPos : [N, numClasses+1] : kINT32
sumSize += AlignMem(dimVolume(sortClassPosDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
sortNMSMarkOffset = sumSize;
// sortNMSMark : [N, samples] : kINT32
sumSize += AlignMem(dimVolume(sortNMSMarkDims) * typeSize(nvinfer1::DataType::kINT32) * batchSize);
totalSize = sumSize;
}
ConcatTopKWorkSpace::ConcatTopKWorkSpace(
const int batchSize, const int concatCnt, const int topK, const nvinfer1::DataType inType)
: concatedScoreDims(concatCnt * topK, 1)
, concatedBBoxDims(concatCnt * topK, 4)
, sortedScoreDims(concatCnt * topK, 1)
, sortedBBoxDims(concatCnt * topK, 4)
{
size_t sumSize = 0;
// const nvinfer1::DataType type = nvinfer1::DataType::kFLOAT;
const nvinfer1::DataType type = inType;
// resource
// temp storage size for sorting scores
tempStorageOffset = sumSize;
sumSize += (1 << 23) * batchSize;
// concatedScoreOffset: [N, concatCnt*topK, 1]
concatedScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(concatedScoreDims) * typeSize(type) * batchSize);
// concatedBBoxOffset: [N, concatCnt*topK, 4]
concatedBBoxOffset = sumSize;
sumSize += AlignMem(dimVolume(concatedBBoxDims) * typeSize(type) * batchSize);
// sortedScoreOffset: [N, concatCnt * topK, 1]
sortedScoreOffset = sumSize;
sumSize += AlignMem(dimVolume(sortedScoreDims) * typeSize(type) * batchSize);
// sortedBBoxOffset: [N, concatCnt * topK, 4]
sortedBBoxOffset = sumSize;
sumSize += AlignMem(dimVolume(sortedBBoxDims) * typeSize(type) * batchSize);
totalSize = sumSize;
}
template <int Threads>
cudaError_t argMaxGroup(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int NClass,
const void* inScore, const void* inBbox, const void* validSamples, void* outScore, void* outLabel, void* outBbox)
{
int gridX = nAlignDown(dMIN(samples, 512 / N), 32);
gridX = dMAX(gridX, 1);
dim3 gridDim = {static_cast<unsigned int>(gridX), static_cast<unsigned int>(N), 1};
dim3 threads = {Threads, 1, 1};
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
argMaxGroup_kernel<float, float, Threads><<<gridDim, threads, 0, stream>>>(
samples, 0, NClass, inScore, inBbox, validSamples, outScore, outLabel, outBbox);
break;
case nvinfer1::DataType::kHALF: break;
case nvinfer1::DataType::kBF16:
case nvinfer1::DataType::kINT64: PLUGIN_FAIL("Unsupported data type");
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
template <int Threads>
cudaError_t argMaxWOBackground(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int NClass,
const void* inScore, const void* inBbox, const void* validSamples, void* outScore, void* outLabel, void* outBbox)
{
int gridX = nAlignDown(dMIN(samples, 512 / N), 32);
gridX = dMAX(gridX, 1);
dim3 gridDim = {static_cast<unsigned int>(gridX), static_cast<unsigned int>(N), 1};
dim3 threads = {Threads, 1, 1};
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
argMaxGroup_kernel<float, float, Threads><<<gridDim, threads, 0, stream>>>(
samples, 1, NClass, inScore, inBbox, validSamples, outScore, outLabel, outBbox);
break;
case nvinfer1::DataType::kHALF: break;
case nvinfer1::DataType::kBF16:
case nvinfer1::DataType::kINT64: PLUGIN_FAIL("Unsupported data type");
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
template <int Threads, int ItermPerThreads>
cudaError_t sortPerClass(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int NClass, int background,
float scoreThreshold, const void* inSampleValidCount, const void* inScorePtr, const void* inLabelPtr,
const void* inBboxPtr, void* outclassStartPosPtr, void* outScorePtr, void* outLabelPtr, void* outSampleIdxPtr,
void* outValidSampleCountPtr)
{
int blocks = N;
int threads = Threads;
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
sortPerClass_kernel<float, float, Threads, ItermPerThreads><<<blocks, threads, 0, stream>>>(samples, NClass,
background, scoreThreshold, inSampleValidCount, inScorePtr, inLabelPtr, inBboxPtr, outclassStartPosPtr,
outScorePtr, outLabelPtr, outSampleIdxPtr, outValidSampleCountPtr);
break;
case nvinfer1::DataType::kHALF:
sortPerClass_kernel_half<Threads, ItermPerThreads><<<blocks, threads, 0, stream>>>(samples, NClass,
background, scoreThreshold, inSampleValidCount, inScorePtr, inLabelPtr, inBboxPtr, outclassStartPosPtr,
outScorePtr, outLabelPtr, outSampleIdxPtr, outValidSampleCountPtr);
break;
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
};
template <int Threads>
cudaError_t PerClassNMS(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int NClass,
const float nmsThreshold, const void* validSampleCount,
// const void *inScore,
const void* inLabel, const void* inBbox, const void* inBboxRefIdx, const void* classStarts, void* outFlagSamples)
{
int blocks = N;
int threads = Threads;
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
PerClassNMS_kernel<float, float, Threads><<<blocks, threads, 0, stream>>>(samples, NClass, nmsThreshold,
validSampleCount, inLabel, inBbox, inBboxRefIdx, classStarts, outFlagSamples);
break;
case nvinfer1::DataType::kHALF:
PerClassNMS_half_kernel<Threads><<<blocks, threads, 0, stream>>>(samples, NClass, nmsThreshold,
validSampleCount, inLabel, inBbox, inBboxRefIdx, classStarts, outFlagSamples);
break;
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
template <int Threads>
cudaError_t KeepTopKGather(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int keepTopK,
const void* validSampleCountPtr, const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr,
const void* inBboxRefIdxPtr, const void* inFlagSamplesPtr, void* outDetections, int proposal)
{
int blocks = N;
int threads = Threads;
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
if (proposal)
{
TopKGatherProposal_kernel<float, float, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr,
outDetections);
}
else
{
TopKGather_kernel<float, float, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr,
outDetections);
}
break;
case nvinfer1::DataType::kHALF: break;
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
// TopKGather For TLT RPN Proposal
// gridDim.x : batch-N
// blockDim.x : Threads
// ItemsPerThreads : = divUp(samples, Threads)
// outDetectionCount : int [N], must be set 0 before kernel
#define MaxItemsPerThreads 8
template <typename DType, typename BoxType, int Threads = 256>
__global__ void TopKGatherBoxScore_kernel(int samples, int keepTopK, const void* validSampleCountPtr,
const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr, const void* inBboxRefIdxPtr,
const void* inFlagSamplesPtr, void* outScorePtr, void* outBboxPtr)
{
typedef cub::BlockRadixSort<DType, Threads, 1, int> BlockRadixSort1;
typedef cub::BlockRadixSort<DType, Threads, 2, int> BlockRadixSort2;
typedef cub::BlockRadixSort<DType, Threads, 3, int> BlockRadixSort3;
typedef cub::BlockRadixSort<DType, Threads, 4, int> BlockRadixSort4;
typedef cub::BlockRadixSort<DType, Threads, 5, int> BlockRadixSort5;
typedef cub::BlockRadixSort<DType, Threads, 6, int> BlockRadixSort6;
typedef cub::BlockRadixSort<DType, Threads, 7, int> BlockRadixSort7;
typedef cub::BlockRadixSort<DType, Threads, 8, int> BlockRadixSort8;
__shared__ union {
typename BlockRadixSort8::TempStorage sort8;
typename BlockRadixSort7::TempStorage sort7;
typename BlockRadixSort6::TempStorage sort6;
typename BlockRadixSort5::TempStorage sort5;
typename BlockRadixSort4::TempStorage sort4;
typename BlockRadixSort3::TempStorage sort3;
typename BlockRadixSort2::TempStorage sort2;
typename BlockRadixSort1::TempStorage sort1;
} temp_storage;
assert(MaxItemsPerThreads * Threads >= samples);
typedef BBoxT<BoxType> BBox;
const int* validSampleCount = static_cast<const int*>(validSampleCountPtr);
const DType* inScore = static_cast<const DType*>(inScorePtr);
const BBox* inBbox = static_cast<const BBox*>(inBboxPtr);
const int* inBboxRefIdx = static_cast<const int*>(inBboxRefIdxPtr);
const int* inFlagSamples = static_cast<const int*>(inFlagSamplesPtr);
BBox* outBbox = static_cast<BBox*>(outBboxPtr);
DType* outScore = static_cast<DType*>(outScorePtr);
int N = blockIdx.x;
int blockOffset = N * samples;
int validSamples = validSampleCount[N];
int finalTopK = dMIN(keepTopK, validSamples);
int idx[MaxItemsPerThreads];
DType score[MaxItemsPerThreads];
int totalItems = (validSamples + (blockDim.x - 1)) / blockDim.x;
for (int ite = 0; ite < totalItems; ++ite)
{
int curIdx = ite * blockDim.x + threadIdx.x;
if (curIdx < validSamples && inFlagSamples[blockOffset + curIdx])
{
idx[ite] = curIdx;
score[ite] = inScore[blockOffset + curIdx];
}
else
{
idx[ite] = -1;
score[ite] = 0.0f;
}
}
switch (totalItems)
{
case 0: break;
case 1:
BlockRadixSort1(temp_storage.sort1).SortDescendingBlockedToStriped((DType(&)[1]) score, (int(&)[1]) idx);
break;
case 2:
BlockRadixSort2(temp_storage.sort2).SortDescendingBlockedToStriped((DType(&)[2]) score, (int(&)[2]) idx);
break;
case 3:
BlockRadixSort3(temp_storage.sort3).SortDescendingBlockedToStriped((DType(&)[3]) score, (int(&)[3]) idx);
break;
case 4:
BlockRadixSort4(temp_storage.sort4).SortDescendingBlockedToStriped((DType(&)[4]) score, (int(&)[4]) idx);
break;
case 5:
BlockRadixSort5(temp_storage.sort5).SortDescendingBlockedToStriped((DType(&)[5]) score, (int(&)[5]) idx);
break;
case 6:
BlockRadixSort6(temp_storage.sort6).SortDescendingBlockedToStriped((DType(&)[6]) score, (int(&)[6]) idx);
break;
case 7:
BlockRadixSort7(temp_storage.sort7).SortDescendingBlockedToStriped((DType(&)[7]) score, (int(&)[7]) idx);
break;
case 8:
BlockRadixSort8(temp_storage.sort8).SortDescendingBlockedToStriped((DType(&)[8]) score, (int(&)[8]) idx);
break;
default: assert(false);
}
__syncthreads();
int outBlockOffset = N * keepTopK;
int topkItems = (keepTopK + (Threads - 1)) / Threads;
for (int i = 0; i < topkItems; ++i)
{
int curI = i * blockDim.x + threadIdx.x;
if (curI < keepTopK)
{
BBox oB = {(BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f, (BoxType) 0.0f};
DType oS = 0.0f;
if (curI < finalTopK && idx[i] >= 0)
{
oB = ((BBox*) inBbox)[blockOffset + inBboxRefIdx[blockOffset + idx[i]]];
oS = score[i];
}
((BBox*) outBbox)[outBlockOffset + curI] = oB;
outScore[outBlockOffset + curI] = oS;
}
}
}
template <int Threads>
cudaError_t KeepTopKGatherBoxScore(cudaStream_t stream, int N, nvinfer1::DataType dtype, int samples, int keepTopK,
const void* validSampleCountPtr, const void* inScorePtr, const void* inLabelPtr, const void* inBboxPtr,
const void* inBboxRefIdxPtr, const void* inFlagSamplesPtr, void* outScores, void* outDetections, int proposal)
{
int blocks = N;
int threads = Threads;
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
if (proposal)
{
TopKGatherBoxScore_kernel<float, float, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr, outScores,
outDetections);
}
else
{
TopKGather_kernel<float, float, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr,
outDetections);
}
break;
case nvinfer1::DataType::kHALF:
if (proposal)
{
TopKGatherBoxScore_kernel<__half, __half, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr, outScores,
outDetections);
}
else
{
TopKGather_kernel<__half, __half, Threads><<<blocks, threads, 0, stream>>>(samples, keepTopK,
validSampleCountPtr, inScorePtr, inLabelPtr, inBboxPtr, inBboxRefIdxPtr, inFlagSamplesPtr,
outDetections);
}
break;
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
cudaError_t RefineBatchClassNMS(cudaStream_t stream, int N, int samples, nvinfer1::DataType dtype,
const RefineNMSParameters& param, const RefineDetectionWorkSpace& refineOffset, void* workspace,
const void* inScores, const void* inDelta, const void* inCountValid, const void* inROI, void* outDetections)
{
int NClass = param.numClasses;
int8_t* wsPtr = static_cast<int8_t*>(workspace);
void* argMaxScorePtr = wsPtr + refineOffset.argMaxScoreOffset;
void* argMaxLabelPtr = wsPtr + refineOffset.argMaxLabelOffset;
void* argMaxBBoxPtr = wsPtr + refineOffset.argMaxBboxOffset;
void* sortClassScorePtr = wsPtr + refineOffset.sortClassScoreOffset;
void* sortClassLabelPtr = wsPtr + refineOffset.sortClassLabelOffset;
void* sortClassSampleIdxPtr = wsPtr + refineOffset.sortClassSampleIdxOffset;
void* sortClassValidCountPtr = wsPtr + refineOffset.sortClassValidCountOffset;
void* sortClassPosPtr = wsPtr + refineOffset.sortClassPosOffset;
void* sortNMSMarkPtr = wsPtr + refineOffset.sortNMSMarkOffset;
cudaError_t status = cudaSuccess;
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassValidCountPtr, 0, N * sizeof(int), stream));
if (NClass > 1)
{ // multiple classes
status = argMaxGroup<32>(stream, N, dtype, samples, NClass, inScores, inDelta, inCountValid, argMaxScorePtr,
argMaxLabelPtr, argMaxBBoxPtr); // argMaxBBoxPtr means delta of bboxes
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
}
else
{ // Only one class
argMaxScorePtr = const_cast<void*>(inScores);
argMaxBBoxPtr = const_cast<void*>(inDelta);
int threads = 512;
int blocks = (N * samples + threads - 1) / threads;
blocks = dMIN(blocks, 8);
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
resetMemValue_kernel<float><<<blocks, threads, 0, stream>>>(argMaxLabelPtr, N * samples, 0);
break;
}
case nvinfer1::DataType::kHALF: { break;
}
default: PLUGIN_ASSERT(false);
}
}
status = ApplyDelta2Bboxes(stream, N, samples, inROI, argMaxBBoxPtr, argMaxBBoxPtr);
PLUGIN_ASSERT(status == cudaSuccess);
if (samples <= 1024)
{
status = sortPerClass<256, 4>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 2048)
{
status = sortPerClass<256, 8>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 4096)
{
status = sortPerClass<256, 16>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else
{
PLUGIN_ASSERT(false && "unsupported sortPerClass");
return cudaErrorLaunchFailure;
}
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
status = PerClassNMS<256>(stream, N, dtype, samples, NClass, param.iouThreshold, sortClassValidCountPtr,
// sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortClassPosPtr, sortNMSMarkPtr);
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
status = KeepTopKGather<256>(stream, N, dtype, samples, param.keepTopK, sortClassValidCountPtr, sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortNMSMarkPtr, outDetections, 0);
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
return status;
}
cudaError_t DetectionPostProcess(cudaStream_t stream, int N, int samples, const float* regWeight,
const float inputHeight, const float inputWidth, nvinfer1::DataType dtype, const RefineNMSParameters& param,
const RefineDetectionWorkSpace& refineOffset, void* workspace, const void* inScores, const void* inDelta,
const void* inCountValid, const void* inROI, void* outDetections)
{
int NClass = param.numClasses;
int8_t* wsPtr = static_cast<int8_t*>(workspace);
void* argMaxScorePtr = wsPtr + refineOffset.argMaxScoreOffset;
void* argMaxLabelPtr = wsPtr + refineOffset.argMaxLabelOffset;
void* argMaxBBoxPtr = wsPtr + refineOffset.argMaxBboxOffset;
void* sortClassScorePtr = wsPtr + refineOffset.sortClassScoreOffset;
void* sortClassLabelPtr = wsPtr + refineOffset.sortClassLabelOffset;
void* sortClassSampleIdxPtr = wsPtr + refineOffset.sortClassSampleIdxOffset;
void* sortClassValidCountPtr = wsPtr + refineOffset.sortClassValidCountOffset;
void* sortClassPosPtr = wsPtr + refineOffset.sortClassPosOffset;
void* sortNMSMarkPtr = wsPtr + refineOffset.sortNMSMarkOffset;
cudaError_t status = cudaSuccess;
PLUGIN_CUASSERT(cudaMemsetAsync(argMaxScorePtr, 0, N * samples * sizeof(float), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(argMaxBBoxPtr, 0, N * samples * 4 * sizeof(float), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassValidCountPtr, 0, N * sizeof(int), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassPosPtr, 0, N * (NClass + 1) * sizeof(int), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassSampleIdxPtr, 0, N * (samples + 1) * sizeof(int), stream));
if (NClass > 1)
{ // multiple classes
status = argMaxWOBackground<32>(stream, N, dtype, samples, NClass, inScores, inDelta, inCountValid,
argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr); // argMaxBBoxPtr means delta of bboxes
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
}
else
{ // Only one class
argMaxScorePtr = const_cast<void*>(inScores);
argMaxBBoxPtr = const_cast<void*>(inDelta);
int threads = 512;
int blocks = (N * samples + threads - 1) / threads;
blocks = dMIN(blocks, 8);
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
resetMemValue_kernel<float><<<blocks, threads, 0, stream>>>(argMaxLabelPtr, N * samples, 0);
break;
}
case nvinfer1::DataType::kHALF: { break;
}
default: PLUGIN_ASSERT(false);
}
}
status = DecodeBBoxes(stream, N, samples, regWeight, inputHeight, inputWidth, inROI, argMaxBBoxPtr, argMaxBBoxPtr, dtype);
PLUGIN_ASSERT(status == cudaSuccess);
if (samples <= 1024)
{
status = sortPerClass<256, 4>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 2048)
{
status = sortPerClass<256, 8>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 4096)
{
status = sortPerClass<256, 16>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else
{
PLUGIN_ASSERT(false && "unsupported sortPerClass");
return cudaErrorLaunchFailure;
}
PLUGIN_ASSERT(status == cudaSuccess);
PLUGIN_CUASSERT(status);
status = PerClassNMS<256>(stream, N, dtype, samples, NClass, param.iouThreshold, sortClassValidCountPtr,
// sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortClassPosPtr, sortNMSMarkPtr);
PLUGIN_CUASSERT(status);
status = KeepTopKGather<256>(stream, N, dtype, samples, param.keepTopK, sortClassValidCountPtr, sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortNMSMarkPtr, outDetections, 0);
PLUGIN_CUASSERT(status);
return status;
}
struct BF_SCORE
{
float bg, fg;
};
// in_scores : [N, samples, 2]
// output_score : [N, samples, 1]
__global__ void extract_fg_kernel(int samples, const void* in_scores, void* output_score)
{
const BF_SCORE* in = static_cast<const BF_SCORE*>(in_scores);
float* out = static_cast<float*>(output_score);
int N = blockIdx.x;
int blockOffset = N * samples;
int totalItems = (samples + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < samples)
{
out[blockOffset + cur_id] = in[blockOffset + cur_id].fg;
}
}
}
__global__ void set_offset_kernel(int stride, int size, int* output)
{
// One block, because batch size shouldn't be too large.
for (int i = threadIdx.x; i < size; i += blockDim.x)
{
output[i] = i * stride;
}
}
template <typename Dtype>
__global__ void resample_kernel(int orig_size, int sample_size, const void* orig_score_ptr, const void* orig_bbox_ptr,
void* sampled_score_ptr, void* sampled_bbox_ptr)
{
const Dtype* in_score = static_cast<const Dtype*>(orig_score_ptr);
const BBoxT<Dtype>* in_bbox = static_cast<const BBoxT<Dtype>*>(orig_bbox_ptr);
Dtype* out_score = static_cast<Dtype*>(sampled_score_ptr);
BBoxT<Dtype>* out_bbox = static_cast<BBoxT<Dtype>*>(sampled_bbox_ptr);
int N = blockIdx.x;
int blockOffset_in = N * orig_size;
int blockOffset_out = N * sample_size;
int realSampleCnt = dMIN(sample_size, orig_size);
int totalItems = (realSampleCnt + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < realSampleCnt)
{
out_score[blockOffset_out + cur_id] = in_score[blockOffset_in + cur_id];
out_bbox[blockOffset_out + cur_id] = in_bbox[blockOffset_in + cur_id];
}
}
}
cudaError_t proposalRefineBatchClassNMS(cudaStream_t stream, int N, int inputCnt, int samples, nvinfer1::DataType dtype,
const RefineNMSParameters& param, const ProposalWorkSpace& proposalOffset, void* workspace,
const void* inScores, //[N, inputcnt, 2]
const void* inDelta, //[N, inputcnt, 4]
const void* inCountValid,
const void* inAnchors, //[N, inputcnt, 4]
void* outProposals)
{
int8_t* wsPtr = static_cast<int8_t*>(workspace);
void* tempStoragePtr = wsPtr + proposalOffset.tempStorageOffset;
void* preRefineScorePtr = wsPtr + proposalOffset.preRefineScoreOffset;
void* preRefineSortedScorePtr = wsPtr + proposalOffset.preRefineSortedScoreOffset;
void* preRefineBboxPtr = wsPtr + proposalOffset.preRefineBboxOffset;
void* argMaxScorePtr = wsPtr + proposalOffset.argMaxScoreOffset;
void* argMaxLabelPtr = wsPtr + proposalOffset.argMaxLabelOffset;
void* argMaxBBoxPtr = wsPtr + proposalOffset.argMaxBboxOffset;
void* sortClassScorePtr = wsPtr + proposalOffset.sortClassScoreOffset;
void* sortClassLabelPtr = wsPtr + proposalOffset.sortClassLabelOffset;
void* sortClassSampleIdxPtr = wsPtr + proposalOffset.sortClassSampleIdxOffset;
void* sortClassValidCountPtr = wsPtr + proposalOffset.sortClassValidCountOffset;
void* sortClassPosPtr = wsPtr + proposalOffset.sortClassPosOffset;
void* sortNMSMarkPtr = wsPtr + proposalOffset.sortNMSMarkOffset;
cudaError_t status = cudaSuccess;
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassValidCountPtr, 0, N * sizeof(int), stream));
// extract foreground score
extract_fg_kernel<<<N, dMIN(inputCnt, 1024), 0, stream>>>(inputCnt, inScores, preRefineScorePtr);
PLUGIN_CUASSERT(cudaGetLastError());
// Here, inDelta are converted to normalize coordinates based on anchors
status = ApplyDelta2Bboxes(stream, N, inputCnt, inAnchors, inDelta, const_cast<void*>(inDelta));
PLUGIN_CUASSERT(status);
// sort the score
// d_key_in: preRefineScorePtr [N, inputCnt, 1]
// d_key_out: preRefineSortedScorePtr
// d_values_in: inDelta [N, inputCnt, 4]
// d_values_out: preRefineBboxPtr
// num_items: inputCnt*N
// num_segments: N
// offsets: [0, inputCnt, inputCnt*2, ..., ]
int* offsets = static_cast<int*>(tempStoragePtr);
set_offset_kernel<<<1, 1024, 0, stream>>>(inputCnt, N + 1, offsets);
PLUGIN_ASSERT(cudaGetLastError() == cudaSuccess);
tempStoragePtr = static_cast<void*>(static_cast<int*>(tempStoragePtr) + (N + 1));
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending(NULL, temp_storage_bytes, (float*) preRefineScorePtr,
(float*) preRefineSortedScorePtr, (BBoxT<float>*) inDelta, (BBoxT<float>*) preRefineBboxPtr, N * inputCnt, N,
offsets, offsets + 1, 0, 8 * sizeof(float), stream);
PLUGIN_ASSERT((1 << 23) * (size_t) N > temp_storage_bytes);
cub::DeviceSegmentedRadixSort::SortPairsDescending(tempStoragePtr, temp_storage_bytes, (float*) preRefineScorePtr,
(float*) preRefineSortedScorePtr, (BBoxT<float>*) inDelta, (BBoxT<float>*) preRefineBboxPtr, N * inputCnt, N,
offsets, offsets + 1, 0, 8 * sizeof(float), stream);
int NClass = param.numClasses;
PLUGIN_ASSERT(NClass == 1);
if (NClass == 1)
{ // Only one class
resample_kernel<float><<<N, dMIN(samples, 1024), 0, stream>>>(
inputCnt, samples, preRefineSortedScorePtr, preRefineBboxPtr, argMaxScorePtr, argMaxBBoxPtr);
int threads = 512;
int blocks = (N * samples + threads - 1) / threads;
blocks = dMIN(blocks, 8);
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
resetMemValue_kernel<float><<<blocks, threads, 0, stream>>>(argMaxLabelPtr, N * samples, 0);
break;
}
case nvinfer1::DataType::kHALF: { break;
}
default: PLUGIN_ASSERT(false);
}
}
if (samples <= 1024)
{
status = sortPerClass<256, 4>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 2048)
{
status = sortPerClass<256, 8>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 4096)
{
status = sortPerClass<256, 16>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else
{
PLUGIN_ASSERT(false && "unsupported sortPerClass");
return cudaErrorLaunchFailure;
}
PLUGIN_CUASSERT(status);
status = PerClassNMS<256>(stream, N, dtype, samples, NClass, param.iouThreshold, sortClassValidCountPtr,
// sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortClassPosPtr, sortNMSMarkPtr);
PLUGIN_CUASSERT(status);
status = KeepTopKGather<256>(stream, N, dtype, samples, param.keepTopK, sortClassValidCountPtr, sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortNMSMarkPtr, outProposals, 1);
PLUGIN_CUASSERT(status);
return status;
}
template<typename Dtype>
void score_bbox_cub_sort(void* tempStorage,
const void* inScore,
void* sortedScore,
const void* inBBox,
void* sortedBBox,
int totalCnt,
int segCnt,
int* offsets,
cudaStream_t stream
)
{
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedRadixSort::SortPairsDescending(NULL, temp_storage_bytes, (Dtype*) inScore, (Dtype*) sortedScore,
(BBoxT<Dtype>*) inBBox, (BBoxT<Dtype>*) sortedBBox, totalCnt, segCnt, offsets, offsets + 1, 0,
8 * sizeof(Dtype), stream);
PLUGIN_CUASSERT(cudaGetLastError());
cub::DeviceSegmentedRadixSort::SortPairsDescending(tempStorage, temp_storage_bytes, (Dtype*) inScore,
(Dtype*) sortedScore, (BBoxT<Dtype>*) inBBox, (BBoxT<Dtype>*) sortedBBox, totalCnt, segCnt, offsets,
offsets + 1, 0, 8 * sizeof(Dtype), stream);
PLUGIN_CUASSERT(cudaGetLastError());
}
cudaError_t MultilevelPropose(cudaStream_t stream, int N, int inputCnt, int samples, const float* regWeight,
const float inputHeight, const float inputWidth, nvinfer1::DataType dtype, const RefineNMSParameters& param,
const MultilevelProposeROIWorkSpace& proposalOffset, void* workspace,
const void* inScore, //[N, inputcnt, 1]
const void* inDelta, //[N, inputcnt, 4]
void* inCountValid,
const void* inAnchors, //[N, inputcnt, 4]
void* outScore, void* outBbox)
{
int8_t* wsPtr = static_cast<int8_t*>(workspace);
void* tempStoragePtr = wsPtr + proposalOffset.tempStorageOffset;
void* preRefineSortedScorePtr = wsPtr + proposalOffset.preRefineSortedScoreOffset;
void* preRefineBboxPtr = wsPtr + proposalOffset.preRefineBboxOffset;
void* argMaxScorePtr = wsPtr + proposalOffset.argMaxScoreOffset;
void* argMaxLabelPtr = wsPtr + proposalOffset.argMaxLabelOffset;
void* argMaxBBoxPtr = wsPtr + proposalOffset.argMaxBboxOffset;
void* sortClassScorePtr = wsPtr + proposalOffset.sortClassScoreOffset;
void* sortClassLabelPtr = wsPtr + proposalOffset.sortClassLabelOffset;
void* sortClassSampleIdxPtr = wsPtr + proposalOffset.sortClassSampleIdxOffset;
void* sortClassValidCountPtr = wsPtr + proposalOffset.sortClassValidCountOffset;
void* sortClassPosPtr = wsPtr + proposalOffset.sortClassPosOffset;
void* sortNMSMarkPtr = wsPtr + proposalOffset.sortNMSMarkOffset;
cudaError_t status = cudaSuccess;
int NClass = param.numClasses;
PLUGIN_ASSERT(NClass == 1);
PLUGIN_CUASSERT(cudaMemsetAsync(argMaxScorePtr, 0, N * samples * sizeof(dtype), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(argMaxBBoxPtr, 0, N * samples * 4 * sizeof(dtype), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassValidCountPtr, 0, N * sizeof(int), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassPosPtr, 0, N * (NClass + 1) * sizeof(int), stream));
PLUGIN_CUASSERT(cudaMemsetAsync(sortClassSampleIdxPtr, 0, N * (samples + 1) * sizeof(int), stream));
PLUGIN_CUASSERT(cudaGetLastError());
// Here, inDelta are converted to normalize coordinates based on anchors
status = DecodeBBoxes(
stream, N, inputCnt, regWeight, inputHeight, inputWidth, inAnchors, inDelta, const_cast<void*>(inDelta), dtype);
PLUGIN_CUASSERT(cudaGetLastError());
// sort the score
// d_key_in: preRefineScorePtr [N, inputCnt, 1]
// d_key_out: preRefineSortedScorePtr
// d_values_in: inDelta [N, inputCnt, 4]
// d_values_out: preRefineBboxPtr
// num_items: inputCnt*N
// num_segments: N
// offsets: [0, inputCnt, inputCnt*2, ..., ]
int* offsets = static_cast<int*>(tempStoragePtr);
set_offset_kernel<<<1, 1024, 0, stream>>>(inputCnt, N + 1, offsets);
PLUGIN_CUASSERT(cudaGetLastError());
tempStoragePtr = static_cast<void*>(static_cast<int*>(tempStoragePtr) + (N + 1));
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
score_bbox_cub_sort<float>(tempStoragePtr, inScore, preRefineSortedScorePtr,
inDelta, preRefineBboxPtr, N * inputCnt, N,
offsets, stream);
break;
}
case nvinfer1::DataType::kHALF:
{
score_bbox_cub_sort<__half>(tempStoragePtr, inScore, preRefineSortedScorePtr,
inDelta, preRefineBboxPtr, N * inputCnt, N,
offsets, stream);
break;
}
default: PLUGIN_ASSERT(false);
}
if (NClass == 1)
{ // Only one class
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
resample_kernel<float><<<N, dMIN(samples, 1024), 0, stream>>>(
inputCnt, samples, preRefineSortedScorePtr, preRefineBboxPtr, argMaxScorePtr, argMaxBBoxPtr);
PLUGIN_CUASSERT(cudaGetLastError());
break;
}
case nvinfer1::DataType::kHALF:
{
resample_kernel<__half><<<N, dMIN(samples, 1024), 0, stream>>>(
inputCnt, samples, preRefineSortedScorePtr, preRefineBboxPtr, argMaxScorePtr, argMaxBBoxPtr);
PLUGIN_CUASSERT(cudaGetLastError());
break;
}
default: PLUGIN_ASSERT(false);
}
int threads = 512;
int blocks = (N * samples + threads - 1) / threads;
blocks = dMIN(blocks, 8);
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
resetMemValue_kernel<float><<<blocks, threads, 0, stream>>>(argMaxLabelPtr, N * samples, 0);
PLUGIN_CUASSERT(cudaGetLastError());
break;
}
case nvinfer1::DataType::kHALF:
{
resetMemValue_kernel<__half><<<blocks, threads, 0, stream>>>(argMaxLabelPtr, N * samples, 0);
PLUGIN_CUASSERT(cudaGetLastError());
break;
}
default: PLUGIN_ASSERT(false);
}
}
if (samples <= 1024)
{
status = sortPerClass<256, 4>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 2048)
{
status = sortPerClass<256, 8>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else if (samples <= 4096)
{
status = sortPerClass<256, 16>(stream, N, dtype, samples, NClass, param.backgroundLabelId, param.scoreThreshold,
inCountValid, argMaxScorePtr, argMaxLabelPtr, argMaxBBoxPtr, sortClassPosPtr, sortClassScorePtr,
sortClassLabelPtr, sortClassSampleIdxPtr, sortClassValidCountPtr);
}
else
{
PLUGIN_FAIL("Unsupported sortPerClass");
return cudaErrorLaunchFailure;
}
PLUGIN_CUASSERT(cudaGetLastError());
status = PerClassNMS<1024>(stream, N, dtype, samples, NClass, param.iouThreshold, sortClassValidCountPtr,
// sortClassScorePtr,
sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortClassPosPtr, sortNMSMarkPtr);
PLUGIN_CUASSERT(cudaGetLastError());
status = KeepTopKGatherBoxScore<512>(stream, N, dtype, samples, param.keepTopK, sortClassValidCountPtr,
sortClassScorePtr, sortClassLabelPtr, argMaxBBoxPtr, sortClassSampleIdxPtr, sortNMSMarkPtr, outScore, outBbox,
1);
PLUGIN_CUASSERT(cudaGetLastError());
return status;
}
struct BBOX
{
float y1, x1, y2, x2;
};
struct DELTA
{
float dy, dx, logdh, logdw;
};
struct DELTA_HALF
{
__half dy, dx, logdh, logdw;
};
__global__ void decode_bboxes_kernel(int samples, const void* anchors, const void* delta, const float* regWeight,
const float inputHeight, const float inputWidth, void* outputBbox, float bboxClipThresh)
{
const BBOX* anchors_in = static_cast<const BBOX*>(anchors);
const DELTA* delta_in = static_cast<const DELTA*>(delta);
BBOX* bbox_out = static_cast<BBOX*>(outputBbox);
int N = blockIdx.x;
int blockOffset = N * samples;
int totalItems = (samples + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < samples)
{
BBOX cur_anchor_yxyx = anchors_in[blockOffset + cur_id];
// convert yxyx -> cyxhw
// cy, cx, h, w
/*BBOX cur_anchor_cyxhw;*/
float cur_anchor_h = (cur_anchor_yxyx.y2 - cur_anchor_yxyx.y1 + 1.0);
float cur_anchor_w = (cur_anchor_yxyx.x2 - cur_anchor_yxyx.x1 + 1.0); // w
float cur_anchor_yc = cur_anchor_yxyx.y1 + cur_anchor_h * 0.5; // cy
float cur_anchor_xc = cur_anchor_yxyx.x1 + cur_anchor_w * 0.5; // cx
DELTA cur_delta = delta_in[blockOffset + cur_id];
// divided by regWeight
cur_delta.dy /= regWeight[0];
cur_delta.dx /= regWeight[1];
cur_delta.logdh /= regWeight[2];
cur_delta.logdw /= regWeight[3];
cur_delta.logdh = dMIN(cur_delta.logdh, bboxClipThresh);
cur_delta.logdw = dMIN(cur_delta.logdw, bboxClipThresh);
// apply delta
float decoded_box_yc = cur_anchor_yc + cur_delta.dy * cur_anchor_h;
float decoded_box_xc = cur_anchor_xc + cur_delta.dx * cur_anchor_w;
float decoded_box_h = expf(cur_delta.logdh) * cur_anchor_h;
float decoded_box_w = expf(cur_delta.logdw) * cur_anchor_w;
float decoded_box_ymin = decoded_box_yc - 0.5 * decoded_box_h;
float decoded_box_xmin = decoded_box_xc - 0.5 * decoded_box_w;
float decoded_box_ymax = decoded_box_ymin + decoded_box_h - 1.0;
float decoded_box_xmax = decoded_box_xmin + decoded_box_w - 1.0;
// clip bbox: a more precision clip method based on real window could be implemented
decoded_box_ymin = dMAX(dMIN(decoded_box_ymin, inputHeight - 1.0), 0.0);
decoded_box_xmin = dMAX(dMIN(decoded_box_xmin, inputWidth - 1.0), 0.0);
decoded_box_ymax = dMAX(dMIN(decoded_box_ymax, inputHeight - 1.0), 0.0);
decoded_box_xmax = dMAX(dMIN(decoded_box_xmax, inputWidth - 1.0), 0.0);
bbox_out[blockOffset + cur_id].y1 = decoded_box_ymin;
bbox_out[blockOffset + cur_id].x1 = decoded_box_xmin;
bbox_out[blockOffset + cur_id].y2 = decoded_box_ymax;
bbox_out[blockOffset + cur_id].x2 = decoded_box_xmax;
}
}
}
__global__ void decode_bboxes_kernel_half(int samples, const void* anchors, const void* delta, const float* regWeight,
const float inputHeight, const float inputWidth, void* outputBbox, float bboxClipThresh)
{
const BBoxT<float>* anchors_in = static_cast<const BBoxT<float>*>(anchors);
const DELTA_HALF* delta_in = static_cast<const DELTA_HALF*>(delta);
BBoxT<__half>* bbox_out = static_cast<BBoxT<__half>*>(outputBbox);
int N = blockIdx.x;
int blockOffset = N * samples;
int totalItems = (samples + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < samples)
{
BBoxT<float> cur_anchor_yxyx = anchors_in[blockOffset + cur_id];
// convert yxyx -> cyxhw
// cy, cx, h, w
float cur_anchor_h = (cur_anchor_yxyx.y2 - cur_anchor_yxyx.y1 + 1.0);
float cur_anchor_w = (cur_anchor_yxyx.x2 - cur_anchor_yxyx.x1 + 1.0); // w
float cur_anchor_yc = cur_anchor_yxyx.y1 + cur_anchor_h * 0.5; // cy
float cur_anchor_xc = cur_anchor_yxyx.x1 + cur_anchor_w * 0.5; // cx
DELTA_HALF cur_delta_half = delta_in[blockOffset + cur_id];
DELTA cur_delta;
cur_delta.dy = __half2float(cur_delta_half.dy);
cur_delta.dx = __half2float(cur_delta_half.dx);
cur_delta.logdh = __half2float(cur_delta_half.logdh);
cur_delta.logdw = __half2float(cur_delta_half.logdw);
// divided by regWeight
cur_delta.dy /= regWeight[0];
cur_delta.dx /= regWeight[1];
cur_delta.logdh /= regWeight[2];
cur_delta.logdw /= regWeight[3];
cur_delta.logdh = dMIN(cur_delta.logdh, bboxClipThresh);
cur_delta.logdw = dMIN(cur_delta.logdw, bboxClipThresh);
// apply delta
float decoded_box_yc = cur_anchor_yc + cur_delta.dy * cur_anchor_h;
float decoded_box_xc = cur_anchor_xc + cur_delta.dx * cur_anchor_w;
float decoded_box_h = expf(cur_delta.logdh) * cur_anchor_h;
float decoded_box_w = expf(cur_delta.logdw) * cur_anchor_w;
float decoded_box_ymin = decoded_box_yc - 0.5 * decoded_box_h;
float decoded_box_xmin = decoded_box_xc - 0.5 * decoded_box_w;
float decoded_box_ymax = decoded_box_ymin + decoded_box_h - 1.0;
float decoded_box_xmax = decoded_box_xmin + decoded_box_w - 1.0;
// clip bbox: a more precision clip method based on real window could be implemented
decoded_box_ymin = dMAX(dMIN(decoded_box_ymin, inputHeight - 1.0), 0.0);
decoded_box_xmin = dMAX(dMIN(decoded_box_xmin, inputWidth - 1.0), 0.0);
decoded_box_ymax = dMAX(dMIN(decoded_box_ymax, inputHeight - 1.0), 0.0);
decoded_box_xmax = dMAX(dMIN(decoded_box_xmax, inputWidth - 1.0), 0.0);
bbox_out[blockOffset + cur_id].y1 = __float2half(decoded_box_ymin);
bbox_out[blockOffset + cur_id].x1 = __float2half(decoded_box_xmin);
bbox_out[blockOffset + cur_id].y2 = __float2half(decoded_box_ymax);
bbox_out[blockOffset + cur_id].x2 = __float2half(decoded_box_xmax);
}
}
}
cudaError_t DecodeBBoxes(cudaStream_t stream, int N,
int samples, // number of anchors per image
const float* regWeight, const float inputHeight, const float inputWidth,
const void* anchors, // [N, anchors, (y1, x1, y2, x2)]
const void* delta, //[N, anchors, (dy, dx, log(dh), log(dw)])
void* outputBbox, //[N, anchors, (y1, x1, y2, x2)]
nvinfer1::DataType dtype
)
{
int blocks = N;
int threads = dMIN(samples, 1024);
// delta multiply bbox_std
// apply delta steps:
// cy = anchor_cy + dy*height
// cx = anchor_cx + dx*weight
// h = exp(dh)*anchor_h
// w = exp(dw)*anchor_w
// clip the bbox in absolute coordinates
float bboxClipThresh = log(1000.0f / 16.0f);
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
decode_bboxes_kernel<<<blocks, threads, 0, stream>>>(
samples, anchors, delta, regWeight, inputHeight, inputWidth, outputBbox, bboxClipThresh);
break;
}
case nvinfer1::DataType::kHALF:
{
decode_bboxes_kernel_half<<<blocks, threads, 0, stream>>>(
samples, anchors, delta, regWeight, inputHeight, inputWidth, outputBbox, bboxClipThresh);
break;
}
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
__global__ void apply_delta_kernel(int samples, const void* anchors, const void* delta, void* outputBbox)
{
const BBOX* anchors_in = static_cast<const BBOX*>(anchors);
const DELTA* delta_in = static_cast<const DELTA*>(delta);
BBOX* bbox_out = static_cast<BBOX*>(outputBbox);
int N = blockIdx.x;
int blockOffset = N * samples;
int totalItems = (samples + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < samples)
{
BBOX cur_anchor_yxyx = anchors_in[blockOffset + cur_id];
// convert yxyx -> cyxhw
// cy, cx, h, w
BBOX cur_anchor_cyxhw;
cur_anchor_cyxhw.y1 = (cur_anchor_yxyx.y1 + cur_anchor_yxyx.y2) / 2.f;
cur_anchor_cyxhw.x1 = (cur_anchor_yxyx.x1 + cur_anchor_yxyx.x2) / 2.f;
cur_anchor_cyxhw.y2 = (cur_anchor_yxyx.y2 - cur_anchor_yxyx.y1);
cur_anchor_cyxhw.x2 = (cur_anchor_yxyx.x2 - cur_anchor_yxyx.x1);
DELTA cur_delta = delta_in[blockOffset + cur_id];
// multiply std_dev
cur_delta.dy *= 0.1f;
cur_delta.dx *= 0.1f;
cur_delta.logdh *= 0.2f;
cur_delta.logdw *= 0.2f;
// apply delta
cur_anchor_cyxhw.y1 += cur_delta.dy * cur_anchor_cyxhw.y2;
cur_anchor_cyxhw.x1 += cur_delta.dx * cur_anchor_cyxhw.x2;
cur_anchor_cyxhw.y2 *= expf(cur_delta.logdh);
cur_anchor_cyxhw.x2 *= expf(cur_delta.logdw);
cur_anchor_yxyx.y1 = cur_anchor_cyxhw.y1 - 0.5f * cur_anchor_cyxhw.y2;
cur_anchor_yxyx.x1 = cur_anchor_cyxhw.x1 - 0.5f * cur_anchor_cyxhw.x2;
cur_anchor_yxyx.y2 = cur_anchor_yxyx.y1 + cur_anchor_cyxhw.y2;
cur_anchor_yxyx.x2 = cur_anchor_yxyx.x1 + cur_anchor_cyxhw.x2;
// clip bbox: a more precision clip method based on real window could be implemented
cur_anchor_yxyx.y1 = dMAX(dMIN(cur_anchor_yxyx.y1, 1.f), 0.f);
cur_anchor_yxyx.x1 = dMAX(dMIN(cur_anchor_yxyx.x1, 1.f), 0.f);
cur_anchor_yxyx.y2 = dMAX(dMIN(cur_anchor_yxyx.y2, 1.f), 0.f);
cur_anchor_yxyx.x2 = dMAX(dMIN(cur_anchor_yxyx.x2, 1.f), 0.f);
bbox_out[blockOffset + cur_id].y1 = cur_anchor_yxyx.y1;
bbox_out[blockOffset + cur_id].x1 = cur_anchor_yxyx.x1;
bbox_out[blockOffset + cur_id].y2 = cur_anchor_yxyx.y2;
bbox_out[blockOffset + cur_id].x2 = cur_anchor_yxyx.x2;
}
}
}
cudaError_t ApplyDelta2Bboxes(cudaStream_t stream, int N,
int samples, // number of anchors per image
const void* anchors, // [N, anchors, (y1, x1, y2, x2)]
const void* delta, //[N, anchors, (dy, dx, log(dh), log(dw)])
void* outputBbox //[N, anchors, (y1, x1, y2, x2)]
)
{
int blocks = N;
int threads = dMIN(samples, 1024);
// delta multiply bbox_std
// apply delta steps:
// cy = anchor_cy + dy*height
// cx = anchor_cx + dx*weight
// h = exp(dh)*anchor_h
// w = exp(dw)*anchor_w
// clip the bbox
apply_delta_kernel<<<blocks, threads, 0, stream>>>(samples, anchors, delta, outputBbox);
return cudaGetLastError();
}
template <typename Tfeat>
__device__ inline Tfeat interpolateBilinear(const Tfeat* src, xy_t srcDims, float y, float x)
{
const int y0 = static_cast<int>(y);
const float yAlpha = y - static_cast<float>(y0);
const int x0 = static_cast<int>(x);
const float xAlpha = x - static_cast<float>(x0);
assert(y0 < srcDims.y);
assert(x0 < srcDims.x);
const int y1 = (yAlpha == 0) ? y0 : y0 + 1; // ceil
const int x1 = (xAlpha == 0) ? x0 : x0 + 1; // ceil
assert(y1 < srcDims.y);
assert(x1 < srcDims.x);
const Tfeat src00 = src[(y0) *srcDims.x + (x0)];
const Tfeat src01 = src[(y0) *srcDims.x + (x1)];
const Tfeat src10 = src[(y1) *srcDims.x + (x0)];
const Tfeat src11 = src[(y1) *srcDims.x + (x1)];
const Tfeat src0 = src00 * (1.0F - xAlpha) + src01 * xAlpha;
const Tfeat src1 = src10 * (1.0F - xAlpha) + src11 * xAlpha;
return src0 * (1.0F - yAlpha) + src1 * yAlpha;
}
template <>
__device__ inline __half interpolateBilinear(const __half* src, xy_t srcDims, float y, float x)
{
const int y0 = static_cast<int>(y);
const float yAlpha = y - static_cast<float>(y0);
const int x0 = static_cast<int>(x);
const float xAlpha = x - static_cast<float>(x0);
assert(y0 < srcDims.y);
assert(x0 < srcDims.x);
const int y1 = (yAlpha == 0) ? y0 : y0 + 1; // ceil
const int x1 = (xAlpha == 0) ? x0 : x0 + 1; // ceil
assert(y1 < srcDims.y);
assert(x1 < srcDims.x);
const __half src00 = src[(y0) *srcDims.x + (x0)];
const __half src01 = src[(y0) *srcDims.x + (x1)];
const __half src10 = src[(y1) *srcDims.x + (x0)];
const __half src11 = src[(y1) *srcDims.x + (x1)];
const __half src0 = add_fb(mul_fb(src00, (1.0F - xAlpha)), mul_fb(src01, xAlpha));
const __half src1 = add_fb(mul_fb(src10, (1.0F - xAlpha)), mul_fb(src11, xAlpha));
return add_fb(mul_fb(src0, (1.0F - yAlpha)), mul_fb(src1, yAlpha));
}
template <typename Trois, typename Tfeat>
__global__ void roiAlign_kernel(xy_t const imageSize, int32_t const featureCount, int32_t const roiCount,
float const firstThreshold, int32_t const transformCoords, bool const absCoords, bool const swapCoords,
bool const plusOneCoords, int32_t const samplingRatio, Trois const* rois, Tfeat const* P2, xy_t const P2dims,
Tfeat const* P3, xy_t const P3dims, Tfeat const* P4, xy_t const P4dims, Tfeat const* P5, xy_t const P5dims,
Tfeat* pooled, xy_t const poolDims)
{
int32_t const batch = blockIdx.x;
int32_t const feature = blockIdx.y;
int32_t const roiIdx = blockIdx.z;
Trois const* roi = rois + 4 * (batch * roiCount + roiIdx);
float y1, x1, y2, x2, hw;
if (swapCoords)
{
y1 = min(roi[0], roi[2]);
x1 = min(roi[1], roi[3]);
y2 = max(roi[0], roi[2]);
x2 = max(roi[1], roi[3]);
}
else
{
x1 = min(roi[0], roi[2]);
y1 = min(roi[1], roi[3]);
x2 = max(roi[0], roi[2]);
y2 = max(roi[1], roi[3]);
}
if (absCoords)
{
y1 = max(0.F, min(static_cast<float>(imageSize.y), y1)) / imageSize.y;
x1 = max(0.F, min(static_cast<float>(imageSize.x), x1)) / imageSize.x;
y2 = max(0.F, min(static_cast<float>(imageSize.y), y2)) / imageSize.y;
x2 = max(0.F, min(static_cast<float>(imageSize.x), x2)) / imageSize.x;
}
else
{
y1 = max(0.F, min(1.F, y1));
x1 = max(0.F, min(1.F, x1));
y2 = max(0.F, min(1.F, y2));
x2 = max(0.F, min(1.F, x2));
}
if (plusOneCoords)
{
hw = (y2 - y1 + 1.F / imageSize.y) * (x2 - x1 + 1.F / imageSize.x);
}
else
{
hw = (y2 - y1) * (x2 - x1);
}
Tfeat const* src = P2;
xy_t srcDims = P2dims;
int32_t iP = 2;
float threshold = firstThreshold;
if (hw > threshold)
{
src = P3;
srcDims = P3dims;
++iP;
}
threshold *= 4;
if (hw > threshold)
{
src = P4;
srcDims = P4dims;
++iP;
}
threshold *= 4;
if (hw > threshold)
{
src = P5;
srcDims = P5dims;
++iP;
}
src += srcDims.x * srcDims.y * (batch * featureCount + feature);
Tfeat* dst = pooled + poolDims.x * poolDims.y * (batch * roiCount * featureCount + roiIdx * featureCount + feature);
float yStart, xStart, yEnd, xEnd, yDelta, xDelta;
float samplingOffset;
if (transformCoords == -1)
{
// Back-Compatibility with old PyramidROIAlign implementation.
samplingOffset = 0.F;
yStart = y1 * (srcDims.y - 1);
xStart = x1 * (srcDims.x - 1);
yEnd = y2 * (srcDims.y - 1);
xEnd = x2 * (srcDims.x - 1);
yDelta = (yEnd - yStart) / (poolDims.y - 1);
xDelta = (xEnd - xStart) / (poolDims.x - 1);
}
else
{
float inputOffset;
if (transformCoords == 0) // No Half Pixel
{
inputOffset = 0.F;
samplingOffset = 0.F;
}
if (transformCoords == 1) // Output Half Pixel
{
inputOffset = 0.F;
samplingOffset = 0.5F;
}
if (transformCoords == 2) // Half Pixel
{
inputOffset = 0.5F;
samplingOffset = 0.5F;
}
yStart = y1 * srcDims.y - inputOffset;
xStart = x1 * srcDims.x - inputOffset;
yEnd = y2 * srcDims.y - inputOffset;
xEnd = x2 * srcDims.x - inputOffset;
yDelta = (yEnd - yStart) / poolDims.y;
xDelta = (xEnd - xStart) / poolDims.x;
}
int32_t const samplingRatioX
= samplingRatio > 0 ? samplingRatio : max(1, static_cast<int32_t>(ceilf((xEnd - xStart) / poolDims.x)));
int32_t const samplingRatioY
= samplingRatio > 0 ? samplingRatio : max(1, static_cast<int32_t>(ceilf((yEnd - yStart) / poolDims.y)));
int32_t const samplingCount = samplingRatioX * samplingRatioY;
for (int32_t outIdx = threadIdx.x; outIdx < poolDims.x * poolDims.y; outIdx += blockDim.x)
{
int32_t xx = outIdx % poolDims.x;
int32_t yy = outIdx / poolDims.x;
Tfeat* out = dst + poolDims.x * yy + xx;
Tfeat result = 0;
for (int32_t iy = 0; iy < samplingRatioY; iy++)
{
float ySample = yStart + yDelta * yy;
ySample += yDelta * (iy + samplingOffset) / samplingRatioY;
ySample = min(max(ySample, 0.F), srcDims.y - 1.F);
for (int32_t ix = 0; ix < samplingRatioX; ix++)
{
float xSample = xStart + xDelta * xx;
xSample += xDelta * (ix + samplingOffset) / samplingRatioX;
xSample = min(max(xSample, 0.F), srcDims.x - 1.F);
result += interpolateBilinear(src, srcDims, ySample, xSample);
}
}
*out = result / samplingCount;
}
}
cudaError_t roiAlign(cudaStream_t const stream, int32_t const batchSize, xy_t const imageSize,
int32_t const featureCount, int32_t const roiCount, float const firstThreshold, int32_t const transformCoords,
bool const absCoords, bool const swapCoords, bool const plusOneCoords, int32_t const samplingRatio,
void const* rois, void const* const layers[], xy_t const* layerDims, void* const pooled, xy_t const poolDims)
{
dim3 const blocks(batchSize, featureCount, roiCount);
int32_t const threads(min(256, poolDims.x * poolDims.y));
roiAlign_kernel<<<blocks, threads, 0, stream>>>(imageSize, featureCount, roiCount, firstThreshold, transformCoords,
absCoords, swapCoords, plusOneCoords, samplingRatio, static_cast<float const*>(rois),
static_cast<float const*>(layers[0]), layerDims[0], static_cast<float const*>(layers[1]), layerDims[1],
static_cast<float const*>(layers[2]), layerDims[2], static_cast<float const*>(layers[3]), layerDims[3],
static_cast<float*>(pooled), poolDims);
return cudaGetLastError();
}
template <typename Trois, typename Tfeat>
__global__ void roiAlignHalfCenter_kernel(int featureCount, int roiCount,
float threshold, int inputHeight, int inputWidth, const void* rois_,
const void* const P2_, const xy_t P2dims, const void* const P3_, const xy_t P3dims, const void* const P4_, const xy_t P4dims,
const void* const P5_, const xy_t P5dims, const void* const P6_, const xy_t P6dims,
void* pooled_, const xy_t poolDims)
{
const Trois* rois = static_cast<const Trois*>(rois_);
const Tfeat* P2 = static_cast<const Tfeat*>(P2_);
const Tfeat* P3 = static_cast<const Tfeat*>(P3_);
const Tfeat* P4 = static_cast<const Tfeat*>(P4_);
const Tfeat* P5 = static_cast<const Tfeat*>(P5_);
const Tfeat* P6 = static_cast<const Tfeat*>(P6_);
Tfeat* pooled = static_cast<Tfeat* >(pooled_);
const int batch = blockIdx.x;
const int feature = blockIdx.y;
const int roiIdx = blockIdx.z;
const int total_item_cnt = poolDims.x * poolDims.y;
for (int itemIdx = threadIdx.x; itemIdx < total_item_cnt; itemIdx += blockDim.x)
{
const Trois* roi = rois + 4 * (batch * roiCount + roiIdx);
const float y1 = roi[0];
const float x1 = roi[1];
const float y2 = roi[2];
const float x2 = roi[3];
if (!(0 <= y1 && y1 <= inputHeight && 0 <= x1 && x1 <= inputWidth && 0 <= y2 && y2 <= inputHeight && 0 <= x2
&& x2 <= inputWidth && y1 < y2 && x1 < x2))
{
continue;
}
else
{
}
const float hw = (y2 - y1) * (x2 - x1);
const Tfeat* src = P2;
xy_t srcDims = P2dims;
int iP = 2;
float threshold_per_item = threshold;
if (hw > threshold_per_item)
{
src = P3;
srcDims = P3dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P4;
srcDims = P4dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P5;
srcDims = P5dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P6;
srcDims = P6dims;
++iP;
}
src += srcDims.x * srcDims.y * (batch * featureCount + feature);
Tfeat* dst
= pooled + poolDims.x * poolDims.y * (batch * roiCount * featureCount + roiIdx * featureCount + feature) + itemIdx;
float scale_to_level = 1.0f;
for (int i = 0; i < iP; i++)
{
scale_to_level *= 2.0f;
}
const float yStart = y1 / scale_to_level;
const float xStart = x1 / scale_to_level;
const float yEnd = y2 / scale_to_level;
const float xEnd = x2 / scale_to_level;
const float yDelta = (yEnd - yStart) / (poolDims.y);
const float xDelta = (xEnd - xStart) / (poolDims.x);
const int yy = itemIdx / poolDims.y;
const int xx = itemIdx % poolDims.x;
const float ySample = dMIN(dMAX(yStart + yDelta * (yy + 0.5), 0.0f), srcDims.y - 1.0f);
const float xSample = dMIN(dMAX(xStart + xDelta * (xx + 0.5), 0.0f), srcDims.x - 1.0f);
Tfeat result = interpolateBilinear<Tfeat>(src, srcDims, ySample, xSample);
*dst = result;
}
}
template <>
__global__ void roiAlignHalfCenter_kernel<__half, __half>(int featureCount, int roiCount,
float threshold, int inputHeight, int inputWidth, const void* rois_,
const void* const P2_, const xy_t P2dims, const void* const P3_, const xy_t P3dims, const void* const P4_, const xy_t P4dims,
const void* const P5_, const xy_t P5dims, const void* const P6_, const xy_t P6dims,
void* pooled_, const xy_t poolDims)
{
const __half* rois = static_cast<const __half*>(rois_);
const __half* P2 = static_cast<const __half*>(P2_);
const __half* P3 = static_cast<const __half*>(P3_);
const __half* P4 = static_cast<const __half*>(P4_);
const __half* P5 = static_cast<const __half*>(P5_);
const __half* P6 = static_cast<const __half*>(P6_);
__half* pooled = static_cast<__half* >(pooled_);
const int batch = blockIdx.x;
const int feature = blockIdx.y;
const int roiIdx = blockIdx.z;
const int total_item_cnt = poolDims.x * poolDims.y;
for (int itemIdx = threadIdx.x; itemIdx < total_item_cnt; itemIdx += blockDim.x)
{
const __half* roi = rois + 4 * (batch * roiCount + roiIdx);
const float y1 = __half2float(roi[0]);
const float x1 = __half2float(roi[1]);
const float y2 = __half2float(roi[2]);
const float x2 = __half2float(roi[3]);
if (!(0 <= y1 && y1 <= inputHeight && 0 <= x1 && x1 <= inputWidth && 0 <= y2 && y2 <= inputHeight && 0 <= x2
&& x2 <= inputWidth && y1 < y2 && x1 < x2))
{
continue;
}
else
{
}
const float hw = (y2 - y1) * (x2 - x1);
const __half* src = P2;
xy_t srcDims = P2dims;
int iP = 2;
float threshold_per_item = threshold;
if (hw > threshold_per_item)
{
src = P3;
srcDims = P3dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P4;
srcDims = P4dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P5;
srcDims = P5dims;
++iP;
}
threshold_per_item *= 4;
if (hw > threshold_per_item)
{
src = P6;
srcDims = P6dims;
++iP;
}
src += srcDims.x * srcDims.y * (batch * featureCount + feature);
__half* dst
= pooled + poolDims.x * poolDims.y * (batch * roiCount * featureCount + roiIdx * featureCount + feature) + itemIdx;
float scale_to_level = 1.0f;
for (int i = 0; i < iP; i++)
{
scale_to_level *= 2.0f;
}
const float yStart = y1 / scale_to_level;
const float xStart = x1 / scale_to_level;
const float yEnd = y2 / scale_to_level;
const float xEnd = x2 / scale_to_level;
const float yDelta = (yEnd - yStart) / (poolDims.y);
const float xDelta = (xEnd - xStart) / (poolDims.x);
const int yy = itemIdx / poolDims.y;
const int xx = itemIdx % poolDims.x;
const float ySample = dMIN(dMAX(yStart + yDelta * (yy + 0.5), 0.0f), srcDims.y - 1.0f);
const float xSample = dMIN(dMAX(xStart + xDelta * (xx + 0.5), 0.0f), srcDims.x - 1.0f);
__half result = interpolateBilinear<__half>(src, srcDims, ySample, xSample);
*dst = result;
}
}
cudaError_t roiAlignHalfCenter(cudaStream_t stream, int batchSize, int featureCount, int roiCount, float firstThreshold,
int inputHeight, int inputWidth, const void* rois, const void* const layers[], const xy_t* layerDims,
void* pooled, const xy_t poolDims, const nvinfer1::DataType dtype)
{
const dim3 blocks(batchSize, featureCount, roiCount);
const int threads(64);
switch (dtype){
case nvinfer1::DataType::kFLOAT:
{
roiAlignHalfCenter_kernel<float, float><<<blocks, threads, 0, stream>>>(featureCount, roiCount,
firstThreshold, inputHeight, inputWidth, rois, layers[0], layerDims[0], layers[1], layerDims[1],
layers[2], layerDims[2], layers[3], layerDims[3], layers[4], layerDims[4], pooled, poolDims);
break;
}
case nvinfer1::DataType::kHALF:
{
roiAlignHalfCenter_kernel<__half, __half><<<blocks, threads, 0, stream>>>(featureCount, roiCount,
firstThreshold, inputHeight, inputWidth, rois, layers[0], layerDims[0], layers[1], layerDims[1],
layers[2], layerDims[2], layers[3], layerDims[3], layers[4], layerDims[4], pooled, poolDims);
break;
}
default: PLUGIN_ASSERT(false);
}
return cudaGetLastError();
}
__global__ void resize_nearest_kernel_2d(int nbatch, float scale, int2 osize, float const* idata, int istride,
int ibatchstride, float* odata, int ostride, int obatchstride)
{
int x0 = threadIdx.x + blockIdx.x * blockDim.x;
int y0 = threadIdx.y + blockIdx.y * blockDim.y;
int z0 = blockIdx.z;
for (int batch = z0; batch < nbatch; batch += gridDim.z)
{
for (int oy = y0; oy < osize.y; oy += blockDim.y * gridDim.y)
{
for (int ox = x0; ox < osize.x; ox += blockDim.x * gridDim.x)
{
int ix = int(ox / scale);
int iy = int(oy / scale);
odata[batch * obatchstride + oy * ostride + ox] = idata[batch * ibatchstride + iy * istride + ix];
}
}
}
}
void resizeNearest(dim3 grid, dim3 block, cudaStream_t stream, int nbatch, float scale, int2 osize, float const* idata,
int istride, int ibatchstride, float* odata, int ostride, int obatchstride)
{
resize_nearest_kernel_2d<<<grid, block, 0, stream>>>(
nbatch, scale, osize, idata, istride, ibatchstride, odata, ostride, obatchstride);
}
struct BOX
{
float y1, x1, y2, x2;
};
struct DETECTION
{
float y1, x1, y2, x2, class_id, score;
};
__global__ void specialslice_kernel(int samples, const void* idata, void* odata)
{
int N = blockIdx.x;
int blockOffset = N * samples;
int totalItems = (samples + (blockDim.x - 1)) / blockDim.x;
const DETECTION* in_detections = static_cast<const DETECTION*>(idata);
BOX* out_bboxes = static_cast<BOX*>(odata);
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < samples)
{
out_bboxes[blockOffset + cur_id].y1 = in_detections[blockOffset + cur_id].y1;
out_bboxes[blockOffset + cur_id].x1 = in_detections[blockOffset + cur_id].x1;
out_bboxes[blockOffset + cur_id].y2 = in_detections[blockOffset + cur_id].y2;
out_bboxes[blockOffset + cur_id].x2 = in_detections[blockOffset + cur_id].x2;
}
}
}
void specialSlice(cudaStream_t stream, int batch_size, int boxes_cnt, const void* idata, void* odata)
{
int blocks = batch_size;
int threads = dMIN(boxes_cnt, 2048);
specialslice_kernel<<<blocks, threads, 0, stream>>>(boxes_cnt, idata, odata);
}
template <typename Dtype>
__global__ void concatenate(int featureCnt, int sampleCnt, const void* const* inScores, const void* const* inBBox,
void* outScore, void* outBBox)
{
int N = blockIdx.x;
int outBlockOffset = N * sampleCnt * featureCnt;
int inBlockOffset = N * sampleCnt;
int itemsPerThread = (sampleCnt + blockDim.x - 1) / blockDim.x;
Dtype* outScorePtr = static_cast<Dtype*>(outScore);
BBoxT<Dtype>* outBBoxPtr = static_cast<BBoxT<Dtype>*>(outBBox);
for (int fId = 0; fId < featureCnt; fId++)
{
const Dtype* fInScorePtr = static_cast<const Dtype*>(inScores[fId]);
const BBoxT<Dtype>* fInBBoxPtr = static_cast<const BBoxT<Dtype>*>(inBBox[fId]);
int featureOffset = fId * sampleCnt;
for (int i = 0; i < itemsPerThread; i++)
{
int curId = i * blockDim.x + threadIdx.x;
if (curId < sampleCnt)
{
outScorePtr[outBlockOffset + featureOffset + curId] = fInScorePtr[inBlockOffset + curId];
outBBoxPtr[outBlockOffset + featureOffset + curId] = fInBBoxPtr[inBlockOffset + curId];
}
}
}
}
template <typename Dtype>
__global__ void resampleBBox_kernel(int orig_size, int sample_size, const void* orig_bbox_ptr, void* sampled_bbox_ptr)
{
const BBoxT<Dtype>* in_bbox = static_cast<const BBoxT<Dtype>*>(orig_bbox_ptr);
BBoxT<Dtype>* out_bbox = static_cast<BBoxT<Dtype>*>(sampled_bbox_ptr);
int N = blockIdx.x;
int blockOffset_in = N * orig_size;
int blockOffset_out = N * sample_size;
int totalItems = (sample_size + (blockDim.x - 1)) / blockDim.x;
for (int i = 0; i < totalItems; i++)
{
int cur_id = i * blockDim.x + threadIdx.x;
if (cur_id < sample_size)
{
out_bbox[blockOffset_out + cur_id] = in_bbox[blockOffset_in + cur_id];
}
}
}
cudaError_t ConcatTopK(cudaStream_t stream, int N, int featureCnt, int topK, nvinfer1::DataType dtype, void* workspace,
const ConcatTopKWorkSpace& spaceOffset, void** inScores, void** inBBox, void* outProposals)
{
// Prepare Offset
int8_t* wsPtr = static_cast<int8_t*>(workspace);
void* tempStoragePtr = wsPtr + spaceOffset.tempStorageOffset;
void* concatedScorePtr = wsPtr + spaceOffset.concatedScoreOffset;
void* concatedBBoxPtr = wsPtr + spaceOffset.concatedBBoxOffset;
void* sortedScorePtr = wsPtr + spaceOffset.sortedScoreOffset;
void* sortedBBoxPtr = wsPtr + spaceOffset.sortedBBoxOffset;
int blocks = N; // batch_size
int threads = dMIN(topK, 2048);
// Concat Scores and inBBox
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
concatenate<float>
<<<blocks, threads, 0, stream>>>(featureCnt, topK, inScores, inBBox, concatedScorePtr, concatedBBoxPtr);
PLUGIN_CUASSERT(cudaGetLastError());
break;
case nvinfer1::DataType::kHALF:
concatenate<__half>
<<<blocks, threads, 0, stream>>>(featureCnt, topK, inScores, inBBox, concatedScorePtr, concatedBBoxPtr);
PLUGIN_CUASSERT(cudaGetLastError());
break;
default: PLUGIN_ASSERT(false);
}
// Sort and sample topK
int itemCnt = topK * featureCnt;
int* offsets = static_cast<int*>(tempStoragePtr);
set_offset_kernel<<<1, 1024, 0, stream>>>(itemCnt, N + 1, offsets);
PLUGIN_ASSERT(cudaGetLastError() == cudaSuccess);
tempStoragePtr = static_cast<void*>(static_cast<int*>(tempStoragePtr) + (N + 1));
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
{
score_bbox_cub_sort<float>(tempStoragePtr, concatedScorePtr, sortedScorePtr,
concatedBBoxPtr, sortedBBoxPtr, N * itemCnt, N,
offsets, stream);
break;
}
case nvinfer1::DataType::kHALF:
{
score_bbox_cub_sort<__half>(tempStoragePtr, concatedScorePtr, sortedScorePtr,
concatedBBoxPtr, sortedBBoxPtr, N * itemCnt, N,
offsets, stream);
break;
}
default: PLUGIN_ASSERT(false);
}
// Sample
switch (dtype)
{
case nvinfer1::DataType::kFLOAT:
resampleBBox_kernel<float><<<N, dMIN(topK, 1024), 0, stream>>>(itemCnt, topK, sortedBBoxPtr, outProposals);
PLUGIN_CUASSERT(cudaGetLastError());
break;
case nvinfer1::DataType::kHALF:
resampleBBox_kernel<__half><<<N, dMIN(topK, 1024), 0, stream>>>(itemCnt, topK, sortedBBoxPtr, outProposals);
PLUGIN_CUASSERT(cudaGetLastError());
break;
default: PLUGIN_ASSERT(false);
}
PLUGIN_ASSERT(cudaGetLastError() == cudaSuccess);
return cudaGetLastError();
}
Reference in New Issue View Git Blame Copy Permalink