wehub-resource-sync
726 lines
30 KiB
Plaintext
726 lines
30 KiB
Plaintext
/*
|
|
* 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 "common/bboxUtils.h"
|
|
#include "cub/cub.cuh"
|
|
#include "cuda_runtime_api.h"
|
|
|
|
#include "efficientNMSInference.cuh"
|
|
#include "efficientNMSInference.h"
|
|
|
|
#define NMS_TILES 5
|
|
|
|
using namespace nvinfer1;
|
|
using namespace nvinfer1::plugin;
|
|
|
|
template <typename T>
|
|
__device__ float IOU(EfficientNMSParameters param, BoxCorner<T> box1, BoxCorner<T> box2)
|
|
{
|
|
// Regardless of the selected box coding, IOU is always performed in BoxCorner coding.
|
|
// The boxes are copied so that they can be reordered without affecting the originals.
|
|
BoxCorner<T> b1 = box1;
|
|
BoxCorner<T> b2 = box2;
|
|
b1.reorder();
|
|
b2.reorder();
|
|
float intersectArea = BoxCorner<T>::intersect(b1, b2).area();
|
|
if (intersectArea <= 0.f)
|
|
{
|
|
return 0.f;
|
|
}
|
|
float unionArea = b1.area() + b2.area() - intersectArea;
|
|
if (unionArea <= 0.f)
|
|
{
|
|
return 0.f;
|
|
}
|
|
return intersectArea / unionArea;
|
|
}
|
|
|
|
template <typename T, typename Tb>
|
|
__device__ BoxCorner<T> DecodeBoxes(EfficientNMSParameters param, int boxIdx, int anchorIdx,
|
|
const Tb* __restrict__ boxesInput, const Tb* __restrict__ anchorsInput)
|
|
{
|
|
// The inputs will be in the selected coding format, as well as the decoding function. But the decoded box
|
|
// will always be returned as BoxCorner.
|
|
Tb box = boxesInput[boxIdx];
|
|
if (!param.boxDecoder)
|
|
{
|
|
return BoxCorner<T>(box);
|
|
}
|
|
Tb anchor = anchorsInput[anchorIdx];
|
|
box.reorder();
|
|
anchor.reorder();
|
|
return BoxCorner<T>(box.decode(anchor));
|
|
}
|
|
|
|
template <typename T, typename Tb>
|
|
__device__ void MapNMSData(EfficientNMSParameters param, int idx, int imageIdx, const Tb* __restrict__ boxesInput,
|
|
const Tb* __restrict__ anchorsInput, const int* __restrict__ topClassData, const int* __restrict__ topAnchorsData,
|
|
const int* __restrict__ topNumData, const T* __restrict__ sortedScoresData, const int* __restrict__ sortedIndexData,
|
|
T& scoreMap, int& classMap, BoxCorner<T>& boxMap, int& boxIdxMap)
|
|
{
|
|
// idx: Holds the NMS box index, within the current batch.
|
|
// idxSort: Holds the batched NMS box index, which indexes the (filtered, but sorted) score buffer.
|
|
// scoreMap: Holds the score that corresponds to the indexed box being processed by NMS.
|
|
if (idx >= topNumData[imageIdx])
|
|
{
|
|
return;
|
|
}
|
|
int idxSort = imageIdx * param.numScoreElements + idx;
|
|
scoreMap = sortedScoresData[idxSort];
|
|
|
|
// idxMap: Holds the re-mapped index, which indexes the (filtered, but unsorted) buffers.
|
|
// classMap: Holds the class that corresponds to the idx'th sorted score being processed by NMS.
|
|
// anchorMap: Holds the anchor that corresponds to the idx'th sorted score being processed by NMS.
|
|
int idxMap = imageIdx * param.numScoreElements + sortedIndexData[idxSort];
|
|
classMap = topClassData[idxMap];
|
|
int anchorMap = topAnchorsData[idxMap];
|
|
|
|
// boxIdxMap: Holds the re-re-mapped index, which indexes the (unfiltered, and unsorted) boxes input buffer.
|
|
boxIdxMap = -1;
|
|
if (param.shareLocation) // Shape of boxesInput: [batchSize, numAnchors, 1, 4]
|
|
{
|
|
boxIdxMap = imageIdx * param.numAnchors + anchorMap;
|
|
}
|
|
else // Shape of boxesInput: [batchSize, numAnchors, numClasses, 4]
|
|
{
|
|
int batchOffset = imageIdx * param.numAnchors * param.numClasses;
|
|
int anchorOffset = anchorMap * param.numClasses;
|
|
boxIdxMap = batchOffset + anchorOffset + classMap;
|
|
}
|
|
// anchorIdxMap: Holds the re-re-mapped index, which indexes the (unfiltered, and unsorted) anchors input buffer.
|
|
int anchorIdxMap = -1;
|
|
if (param.shareAnchors) // Shape of anchorsInput: [1, numAnchors, 4]
|
|
{
|
|
anchorIdxMap = anchorMap;
|
|
}
|
|
else // Shape of anchorsInput: [batchSize, numAnchors, 4]
|
|
{
|
|
anchorIdxMap = imageIdx * param.numAnchors + anchorMap;
|
|
}
|
|
// boxMap: Holds the box that corresponds to the idx'th sorted score being processed by NMS.
|
|
boxMap = DecodeBoxes<T, Tb>(param, boxIdxMap, anchorIdxMap, boxesInput, anchorsInput);
|
|
}
|
|
|
|
template <typename T>
|
|
__device__ void WriteNMSResult(EfficientNMSParameters param, int* __restrict__ numDetectionsOutput,
|
|
T* __restrict__ nmsScoresOutput, int* __restrict__ nmsClassesOutput, BoxCorner<T>* __restrict__ nmsBoxesOutput,
|
|
T threadScore, int threadClass, BoxCorner<T> threadBox, int imageIdx, unsigned int resultsCounter)
|
|
{
|
|
int outputIdx = imageIdx * param.numOutputBoxes + resultsCounter - 1;
|
|
if (param.scoreSigmoid)
|
|
{
|
|
nmsScoresOutput[outputIdx] = sigmoid_mp(threadScore);
|
|
}
|
|
else if (param.scoreBits > 0)
|
|
{
|
|
nmsScoresOutput[outputIdx] = add_mp(threadScore, (T) -1);
|
|
}
|
|
else
|
|
{
|
|
nmsScoresOutput[outputIdx] = threadScore;
|
|
}
|
|
nmsClassesOutput[outputIdx] = threadClass;
|
|
if (param.clipBoxes)
|
|
{
|
|
nmsBoxesOutput[outputIdx] = threadBox.clip((T) 0, (T) 1);
|
|
}
|
|
else
|
|
{
|
|
nmsBoxesOutput[outputIdx] = threadBox;
|
|
}
|
|
numDetectionsOutput[imageIdx] = resultsCounter;
|
|
}
|
|
|
|
__device__ void WriteONNXResult(EfficientNMSParameters param, int* outputIndexData, int* __restrict__ nmsIndicesOutput,
|
|
int imageIdx, int threadClass, int boxIdxMap)
|
|
{
|
|
int index = boxIdxMap % param.numAnchors;
|
|
int idx = atomicAdd((unsigned int*) &outputIndexData[0], 1);
|
|
nmsIndicesOutput[idx * 3 + 0] = imageIdx;
|
|
nmsIndicesOutput[idx * 3 + 1] = threadClass;
|
|
nmsIndicesOutput[idx * 3 + 2] = index;
|
|
}
|
|
|
|
__global__ void PadONNXResult(EfficientNMSParameters param, int* outputIndexData, int* __restrict__ nmsIndicesOutput)
|
|
{
|
|
if (threadIdx.x > 0)
|
|
{
|
|
return;
|
|
}
|
|
int pidx = outputIndexData[0] - 1;
|
|
if (pidx < 0)
|
|
{
|
|
return;
|
|
}
|
|
for (int idx = pidx + 1; idx < param.batchSize * param.numOutputBoxes; idx++)
|
|
{
|
|
nmsIndicesOutput[idx * 3 + 0] = nmsIndicesOutput[pidx * 3 + 0];
|
|
nmsIndicesOutput[idx * 3 + 1] = nmsIndicesOutput[pidx * 3 + 1];
|
|
nmsIndicesOutput[idx * 3 + 2] = nmsIndicesOutput[pidx * 3 + 2];
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Tb>
|
|
__global__ void EfficientNMS(EfficientNMSParameters param, const int* topNumData, int* outputIndexData,
|
|
int* outputClassData, const int* sortedIndexData, const T* __restrict__ sortedScoresData,
|
|
const int* __restrict__ topClassData, const int* __restrict__ topAnchorsData, const Tb* __restrict__ boxesInput,
|
|
const Tb* __restrict__ anchorsInput, int* __restrict__ numDetectionsOutput, T* __restrict__ nmsScoresOutput,
|
|
int* __restrict__ nmsClassesOutput, int* __restrict__ nmsIndicesOutput, BoxCorner<T>* __restrict__ nmsBoxesOutput)
|
|
{
|
|
unsigned int thread = threadIdx.x;
|
|
unsigned int imageIdx = blockIdx.y;
|
|
unsigned int tileSize = blockDim.x;
|
|
if (imageIdx >= param.batchSize)
|
|
{
|
|
return;
|
|
}
|
|
|
|
int numSelectedBoxes = min(topNumData[imageIdx], param.numSelectedBoxes);
|
|
int numTiles = (numSelectedBoxes + tileSize - 1) / tileSize;
|
|
if (thread >= numSelectedBoxes)
|
|
{
|
|
return;
|
|
}
|
|
|
|
__shared__ int blockState;
|
|
__shared__ unsigned int resultsCounter;
|
|
if (thread == 0)
|
|
{
|
|
blockState = 0;
|
|
resultsCounter = 0;
|
|
}
|
|
|
|
int threadState[NMS_TILES];
|
|
unsigned int boxIdx[NMS_TILES];
|
|
T threadScore[NMS_TILES];
|
|
int threadClass[NMS_TILES];
|
|
BoxCorner<T> threadBox[NMS_TILES];
|
|
int boxIdxMap[NMS_TILES];
|
|
for (int tile = 0; tile < numTiles; tile++)
|
|
{
|
|
threadState[tile] = 0;
|
|
boxIdx[tile] = thread + tile * blockDim.x;
|
|
MapNMSData<T, Tb>(param, boxIdx[tile], imageIdx, boxesInput, anchorsInput, topClassData, topAnchorsData,
|
|
topNumData, sortedScoresData, sortedIndexData, threadScore[tile], threadClass[tile], threadBox[tile],
|
|
boxIdxMap[tile]);
|
|
}
|
|
|
|
// Iterate through all boxes to NMS against.
|
|
for (int i = 0; i < numSelectedBoxes; i++)
|
|
{
|
|
int tile = i / tileSize;
|
|
|
|
if (boxIdx[tile] == i)
|
|
{
|
|
// Iteration lead thread, figure out what the other threads should do,
|
|
// this will be signaled via the blockState shared variable.
|
|
if (threadState[tile] == -1)
|
|
{
|
|
// Thread already dead, this box was already dropped in a previous iteration,
|
|
// because it had a large IOU overlap with another lead thread previously, so
|
|
// it would never be kept anyway, therefore it can safely be skip all IOU operations
|
|
// in this iteration.
|
|
blockState = -1; // -1 => Signal all threads to skip iteration
|
|
}
|
|
else if (threadState[tile] == 0)
|
|
{
|
|
// As this box will be kept, this is a good place to find what index in the results buffer it
|
|
// should have, as this allows to perform an early loop exit if there are enough results.
|
|
if (resultsCounter >= param.numOutputBoxes)
|
|
{
|
|
blockState = -2; // -2 => Signal all threads to do an early loop exit.
|
|
}
|
|
else
|
|
{
|
|
// Thread is still alive, because it has not had a large enough IOU overlap with
|
|
// any other kept box previously. Therefore, this box will be kept for sure. However,
|
|
// we need to check against all other subsequent boxes from this position onward,
|
|
// to see how those other boxes will behave in future iterations.
|
|
blockState = 1; // +1 => Signal all (higher index) threads to calculate IOU against this box
|
|
threadState[tile] = 1; // +1 => Mark this box's thread to be kept and written out to results
|
|
|
|
// If the numOutputBoxesPerClass check is enabled, write the result only if the limit for this
|
|
// class on this image has not been reached yet. Other than (possibly) skipping the write, this
|
|
// won't affect anything else in the NMS threading.
|
|
bool write = true;
|
|
if (param.numOutputBoxesPerClass >= 0)
|
|
{
|
|
int classCounterIdx = imageIdx * param.numClasses + threadClass[tile];
|
|
write = (outputClassData[classCounterIdx] < param.numOutputBoxesPerClass);
|
|
outputClassData[classCounterIdx]++;
|
|
}
|
|
if (write)
|
|
{
|
|
// This branch is visited by one thread per iteration, so it's safe to do non-atomic increments.
|
|
resultsCounter++;
|
|
if (param.outputONNXIndices)
|
|
{
|
|
WriteONNXResult(
|
|
param, outputIndexData, nmsIndicesOutput, imageIdx, threadClass[tile], boxIdxMap[tile]);
|
|
}
|
|
else
|
|
{
|
|
WriteNMSResult<T>(param, numDetectionsOutput, nmsScoresOutput, nmsClassesOutput,
|
|
nmsBoxesOutput, threadScore[tile], threadClass[tile], threadBox[tile], imageIdx,
|
|
resultsCounter);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
// This state should never be reached, but just in case...
|
|
blockState = 0; // 0 => Signal all threads to not do any updates, nothing happens.
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (blockState == -2)
|
|
{
|
|
// This is the signal to exit from the loop.
|
|
return;
|
|
}
|
|
|
|
if (blockState == -1)
|
|
{
|
|
// This is the signal for all threads to just skip this iteration, as no IOU's need to be checked.
|
|
continue;
|
|
}
|
|
|
|
// Grab a box and class to test the current box against. The test box corresponds to iteration i,
|
|
// therefore it will have a lower index than the current thread box, and will therefore have a higher score
|
|
// than the current box because it's located "before" in the sorted score list.
|
|
T testScore;
|
|
int testClass;
|
|
BoxCorner<T> testBox;
|
|
int testBoxIdxMap;
|
|
MapNMSData<T, Tb>(param, i, imageIdx, boxesInput, anchorsInput, topClassData, topAnchorsData, topNumData,
|
|
sortedScoresData, sortedIndexData, testScore, testClass, testBox, testBoxIdxMap);
|
|
|
|
for (int tile = 0; tile < numTiles; tile++)
|
|
{
|
|
bool ignoreClass = true;
|
|
if (!param.classAgnostic)
|
|
{
|
|
ignoreClass = threadClass[tile] == testClass;
|
|
}
|
|
|
|
// IOU
|
|
if (boxIdx[tile] > i && // Make sure two different boxes are being tested, and that it's a higher index;
|
|
boxIdx[tile] < numSelectedBoxes && // Make sure the box is within numSelectedBoxes;
|
|
blockState == 1 && // Signal that allows IOU checks to be performed;
|
|
threadState[tile] == 0 && // Make sure this box hasn't been either dropped or kept already;
|
|
ignoreClass && // Compare only boxes of matching classes when classAgnostic is false;
|
|
lte_mp(threadScore[tile], testScore) && // Make sure the sorting order of scores is as expected;
|
|
IOU<T>(param, threadBox[tile], testBox) >= param.iouThreshold) // And... IOU overlap.
|
|
{
|
|
// Current box overlaps with the box tested in this iteration, this box will be skipped.
|
|
threadState[tile] = -1; // -1 => Mark this box's thread to be dropped.
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
cudaError_t EfficientNMSLauncher(EfficientNMSParameters& param, int* topNumData, int* outputIndexData,
|
|
int* outputClassData, int* sortedIndexData, T* sortedScoresData, int* topClassData, int* topAnchorsData,
|
|
const void* boxesInput, const void* anchorsInput, int* numDetectionsOutput, T* nmsScoresOutput,
|
|
int* nmsClassesOutput, int* nmsIndicesOutput, void* nmsBoxesOutput, cudaStream_t stream)
|
|
{
|
|
unsigned int tileSize = param.numSelectedBoxes / NMS_TILES;
|
|
if (param.numSelectedBoxes <= 512)
|
|
{
|
|
tileSize = 512;
|
|
}
|
|
if (param.numSelectedBoxes <= 256)
|
|
{
|
|
tileSize = 256;
|
|
}
|
|
|
|
const dim3 blockSize = {tileSize, 1, 1};
|
|
const dim3 gridSize = {1, (unsigned int) param.batchSize, 1};
|
|
|
|
if (param.boxCoding == 0)
|
|
{
|
|
EfficientNMS<T, BoxCorner<T>><<<gridSize, blockSize, 0, stream>>>(param, topNumData, outputIndexData,
|
|
outputClassData, sortedIndexData, sortedScoresData, topClassData, topAnchorsData,
|
|
(BoxCorner<T>*) boxesInput, (BoxCorner<T>*) anchorsInput, numDetectionsOutput, nmsScoresOutput,
|
|
nmsClassesOutput, nmsIndicesOutput, (BoxCorner<T>*) nmsBoxesOutput);
|
|
}
|
|
else if (param.boxCoding == 1)
|
|
{
|
|
// Note that nmsBoxesOutput is always coded as BoxCorner<T>, regardless of the input coding type.
|
|
EfficientNMS<T, BoxCenterSize<T>><<<gridSize, blockSize, 0, stream>>>(param, topNumData, outputIndexData,
|
|
outputClassData, sortedIndexData, sortedScoresData, topClassData, topAnchorsData,
|
|
(BoxCenterSize<T>*) boxesInput, (BoxCenterSize<T>*) anchorsInput, numDetectionsOutput, nmsScoresOutput,
|
|
nmsClassesOutput, nmsIndicesOutput, (BoxCorner<T>*) nmsBoxesOutput);
|
|
}
|
|
|
|
if (param.outputONNXIndices)
|
|
{
|
|
PadONNXResult<<<1, 1, 0, stream>>>(param, outputIndexData, nmsIndicesOutput);
|
|
}
|
|
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
__global__ void EfficientNMSFilterSegments(EfficientNMSParameters param, const int* __restrict__ topNumData,
|
|
int* __restrict__ topOffsetsStartData, int* __restrict__ topOffsetsEndData)
|
|
{
|
|
int imageIdx = threadIdx.x;
|
|
if (imageIdx > param.batchSize)
|
|
{
|
|
return;
|
|
}
|
|
topOffsetsStartData[imageIdx] = imageIdx * param.numScoreElements;
|
|
topOffsetsEndData[imageIdx] = imageIdx * param.numScoreElements + topNumData[imageIdx];
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void EfficientNMSFilter(EfficientNMSParameters param, const T* __restrict__ scoresInput,
|
|
int* __restrict__ topNumData, int* __restrict__ topIndexData, int* __restrict__ topAnchorsData,
|
|
T* __restrict__ topScoresData, int* __restrict__ topClassData)
|
|
{
|
|
int elementIdx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
int imageIdx = blockDim.y * blockIdx.y + threadIdx.y;
|
|
|
|
// Boundary Conditions
|
|
if (elementIdx >= param.numScoreElements || imageIdx >= param.batchSize)
|
|
{
|
|
return;
|
|
}
|
|
|
|
// Shape of scoresInput: [batchSize, numAnchors, numClasses]
|
|
int scoresInputIdx = imageIdx * param.numScoreElements + elementIdx;
|
|
|
|
// For each class, check its corresponding score if it crosses the threshold, and if so select this anchor,
|
|
// and keep track of the maximum score and the corresponding (argmax) class id
|
|
T score = scoresInput[scoresInputIdx];
|
|
if (gte_mp(score, (T) param.scoreThreshold))
|
|
{
|
|
// Unpack the class and anchor index from the element index
|
|
int classIdx = elementIdx % param.numClasses;
|
|
int anchorIdx = elementIdx / param.numClasses;
|
|
|
|
// If this is a background class, ignore it.
|
|
if (classIdx == param.backgroundClass)
|
|
{
|
|
return;
|
|
}
|
|
|
|
// Use an atomic to find an open slot where to write the selected anchor data.
|
|
if (topNumData[imageIdx] >= param.numScoreElements)
|
|
{
|
|
return;
|
|
}
|
|
int selectedIdx = atomicAdd((unsigned int*) &topNumData[imageIdx], 1);
|
|
if (selectedIdx >= param.numScoreElements)
|
|
{
|
|
topNumData[imageIdx] = param.numScoreElements;
|
|
return;
|
|
}
|
|
|
|
// Shape of topScoresData / topClassData: [batchSize, numScoreElements]
|
|
int topIdx = imageIdx * param.numScoreElements + selectedIdx;
|
|
|
|
if (param.scoreBits > 0)
|
|
{
|
|
score = add_mp(score, (T) 1);
|
|
if (gt_mp(score, (T) (2.f - 1.f / 1024.f)))
|
|
{
|
|
// Ensure the incremented score fits in the mantissa without changing the exponent
|
|
score = (2.f - 1.f / 1024.f);
|
|
}
|
|
}
|
|
|
|
topIndexData[topIdx] = selectedIdx;
|
|
topAnchorsData[topIdx] = anchorIdx;
|
|
topScoresData[topIdx] = score;
|
|
topClassData[topIdx] = classIdx;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void EfficientNMSDenseIndex(EfficientNMSParameters param, int* __restrict__ topNumData,
|
|
int* __restrict__ topIndexData, int* __restrict__ topAnchorsData, int* __restrict__ topOffsetsStartData,
|
|
int* __restrict__ topOffsetsEndData, T* __restrict__ topScoresData, int* __restrict__ topClassData)
|
|
{
|
|
int elementIdx = blockDim.x * blockIdx.x + threadIdx.x;
|
|
int imageIdx = blockDim.y * blockIdx.y + threadIdx.y;
|
|
|
|
if (elementIdx >= param.numScoreElements || imageIdx >= param.batchSize)
|
|
{
|
|
return;
|
|
}
|
|
|
|
int dataIdx = imageIdx * param.numScoreElements + elementIdx;
|
|
int anchorIdx = elementIdx / param.numClasses;
|
|
int classIdx = elementIdx % param.numClasses;
|
|
if (param.scoreBits > 0)
|
|
{
|
|
T score = topScoresData[dataIdx];
|
|
if (lt_mp(score, (T) param.scoreThreshold))
|
|
{
|
|
score = (T) 1;
|
|
}
|
|
else if (classIdx == param.backgroundClass)
|
|
{
|
|
score = (T) 1;
|
|
}
|
|
else
|
|
{
|
|
score = add_mp(score, (T) 1);
|
|
if (gt_mp(score, (T) (2.f - 1.f / 1024.f)))
|
|
{
|
|
// Ensure the incremented score fits in the mantissa without changing the exponent
|
|
score = (2.f - 1.f / 1024.f);
|
|
}
|
|
}
|
|
topScoresData[dataIdx] = score;
|
|
}
|
|
else
|
|
{
|
|
T score = topScoresData[dataIdx];
|
|
if (lt_mp(score, (T) param.scoreThreshold))
|
|
{
|
|
topScoresData[dataIdx] = -(1 << 15);
|
|
}
|
|
else if (classIdx == param.backgroundClass)
|
|
{
|
|
topScoresData[dataIdx] = -(1 << 15);
|
|
}
|
|
}
|
|
|
|
topIndexData[dataIdx] = elementIdx;
|
|
topAnchorsData[dataIdx] = anchorIdx;
|
|
topClassData[dataIdx] = classIdx;
|
|
|
|
if (elementIdx == 0)
|
|
{
|
|
// Saturate counters
|
|
topNumData[imageIdx] = param.numScoreElements;
|
|
topOffsetsStartData[imageIdx] = imageIdx * param.numScoreElements;
|
|
topOffsetsEndData[imageIdx] = (imageIdx + 1) * param.numScoreElements;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
cudaError_t EfficientNMSFilterLauncher(EfficientNMSParameters& param, const T* scoresInput, int* topNumData,
|
|
int* topIndexData, int* topAnchorsData, int* topOffsetsStartData, int* topOffsetsEndData, T* topScoresData,
|
|
int* topClassData, cudaStream_t stream)
|
|
{
|
|
const unsigned int elementsPerBlock = 512;
|
|
const unsigned int imagesPerBlock = 1;
|
|
const unsigned int elementBlocks = (param.numScoreElements + elementsPerBlock - 1) / elementsPerBlock;
|
|
const unsigned int imageBlocks = (param.batchSize + imagesPerBlock - 1) / imagesPerBlock;
|
|
const dim3 blockSize = {elementsPerBlock, imagesPerBlock, 1};
|
|
const dim3 gridSize = {elementBlocks, imageBlocks, 1};
|
|
|
|
float kernelSelectThreshold = 0.007f;
|
|
if (param.scoreSigmoid)
|
|
{
|
|
// Inverse Sigmoid
|
|
if (param.scoreThreshold <= 0.f)
|
|
{
|
|
param.scoreThreshold = -(1 << 15);
|
|
}
|
|
else
|
|
{
|
|
param.scoreThreshold = logf(param.scoreThreshold / (1.f - param.scoreThreshold));
|
|
}
|
|
kernelSelectThreshold = logf(kernelSelectThreshold / (1.f - kernelSelectThreshold));
|
|
// Disable Score Bits Optimization
|
|
param.scoreBits = -1;
|
|
}
|
|
|
|
if (param.scoreThreshold < kernelSelectThreshold)
|
|
{
|
|
// A full copy of the buffer is necessary because sorting will scramble the input data otherwise.
|
|
PLUGIN_CHECK_CUDA(cudaMemcpyAsync(topScoresData, scoresInput,
|
|
param.batchSize * param.numScoreElements * sizeof(T), cudaMemcpyDeviceToDevice, stream));
|
|
|
|
EfficientNMSDenseIndex<T><<<gridSize, blockSize, 0, stream>>>(param, topNumData, topIndexData, topAnchorsData,
|
|
topOffsetsStartData, topOffsetsEndData, topScoresData, topClassData);
|
|
}
|
|
else
|
|
{
|
|
EfficientNMSFilter<T><<<gridSize, blockSize, 0, stream>>>(
|
|
param, scoresInput, topNumData, topIndexData, topAnchorsData, topScoresData, topClassData);
|
|
|
|
EfficientNMSFilterSegments<<<1, param.batchSize, 0, stream>>>(
|
|
param, topNumData, topOffsetsStartData, topOffsetsEndData);
|
|
}
|
|
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
template <typename T>
|
|
size_t EfficientNMSSortWorkspaceSize(int batchSize, int numScoreElements)
|
|
{
|
|
size_t sortedWorkspaceSize = 0;
|
|
cub::DoubleBuffer<T> keysDB(nullptr, nullptr);
|
|
cub::DoubleBuffer<int> valuesDB(nullptr, nullptr);
|
|
cub::DeviceSegmentedRadixSort::SortPairsDescending(nullptr, sortedWorkspaceSize, keysDB, valuesDB,
|
|
numScoreElements, batchSize, (const int*) nullptr, (const int*) nullptr);
|
|
return sortedWorkspaceSize;
|
|
}
|
|
|
|
size_t EfficientNMSWorkspaceSize(int batchSize, int numScoreElements, int numClasses, DataType datatype)
|
|
{
|
|
size_t total = 0;
|
|
const size_t align = 256;
|
|
// Counters
|
|
// 3 for Filtering
|
|
// 1 for Output Indexing
|
|
// C for Max per Class Limiting
|
|
size_t size = (3 + 1 + numClasses) * batchSize * sizeof(int);
|
|
total += size + (size % align ? align - (size % align) : 0);
|
|
// Int Buffers
|
|
for (int i = 0; i < 4; i++)
|
|
{
|
|
size = batchSize * numScoreElements * sizeof(int);
|
|
total += size + (size % align ? align - (size % align) : 0);
|
|
}
|
|
// Float Buffers
|
|
for (int i = 0; i < 2; i++)
|
|
{
|
|
size = batchSize * numScoreElements * dataTypeSize(datatype);
|
|
total += size + (size % align ? align - (size % align) : 0);
|
|
}
|
|
// Sort Workspace
|
|
if (datatype == DataType::kHALF)
|
|
{
|
|
size = EfficientNMSSortWorkspaceSize<__half>(batchSize, numScoreElements);
|
|
total += size + (size % align ? align - (size % align) : 0);
|
|
}
|
|
else if (datatype == DataType::kFLOAT)
|
|
{
|
|
size = EfficientNMSSortWorkspaceSize<float>(batchSize, numScoreElements);
|
|
total += size + (size % align ? align - (size % align) : 0);
|
|
}
|
|
|
|
return total;
|
|
}
|
|
|
|
template <typename T>
|
|
T* EfficientNMSWorkspace(void* workspace, size_t& offset, size_t elements)
|
|
{
|
|
T* buffer = (T*) ((size_t) workspace + offset);
|
|
size_t align = 256;
|
|
size_t size = elements * sizeof(T);
|
|
size_t sizeAligned = size + (size % align ? align - (size % align) : 0);
|
|
offset += sizeAligned;
|
|
return buffer;
|
|
}
|
|
|
|
template <typename T>
|
|
pluginStatus_t EfficientNMSDispatch(EfficientNMSParameters param, const void* boxesInput, const void* scoresInput,
|
|
const void* anchorsInput, void* numDetectionsOutput, void* nmsBoxesOutput, void* nmsScoresOutput,
|
|
void* nmsClassesOutput, void* nmsIndicesOutput, void* workspace, cudaStream_t stream)
|
|
{
|
|
// Clear Outputs (not all elements will get overwritten by the kernels, so safer to clear everything out)
|
|
if (param.outputONNXIndices)
|
|
{
|
|
CSC(cudaMemsetAsync(nmsIndicesOutput, 0xFF, param.batchSize * param.numOutputBoxes * 3 * sizeof(int), stream), STATUS_FAILURE);
|
|
}
|
|
else
|
|
{
|
|
CSC(cudaMemsetAsync(numDetectionsOutput, 0x00, param.batchSize * sizeof(int), stream), STATUS_FAILURE);
|
|
CSC(cudaMemsetAsync(nmsScoresOutput, 0x00, param.batchSize * param.numOutputBoxes * sizeof(T), stream), STATUS_FAILURE);
|
|
CSC(cudaMemsetAsync(nmsBoxesOutput, 0x00, param.batchSize * param.numOutputBoxes * 4 * sizeof(T), stream), STATUS_FAILURE);
|
|
CSC(cudaMemsetAsync(nmsClassesOutput, 0x00, param.batchSize * param.numOutputBoxes * sizeof(int), stream), STATUS_FAILURE);
|
|
}
|
|
|
|
// Empty Inputs
|
|
if (param.numScoreElements < 1)
|
|
{
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
// Counters Workspace
|
|
size_t workspaceOffset = 0;
|
|
int countersTotalSize = (3 + 1 + param.numClasses) * param.batchSize;
|
|
int* topNumData = EfficientNMSWorkspace<int>(workspace, workspaceOffset, countersTotalSize);
|
|
int* topOffsetsStartData = topNumData + param.batchSize;
|
|
int* topOffsetsEndData = topNumData + 2 * param.batchSize;
|
|
int* outputIndexData = topNumData + 3 * param.batchSize;
|
|
int* outputClassData = topNumData + 4 * param.batchSize;
|
|
CSC(cudaMemsetAsync(topNumData, 0x00, countersTotalSize * sizeof(int), stream), STATUS_FAILURE);
|
|
cudaError_t status = cudaGetLastError();
|
|
CSC(status, STATUS_FAILURE);
|
|
|
|
// Other Buffers Workspace
|
|
int* topIndexData
|
|
= EfficientNMSWorkspace<int>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
int* topClassData
|
|
= EfficientNMSWorkspace<int>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
int* topAnchorsData
|
|
= EfficientNMSWorkspace<int>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
int* sortedIndexData
|
|
= EfficientNMSWorkspace<int>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
T* topScoresData = EfficientNMSWorkspace<T>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
T* sortedScoresData
|
|
= EfficientNMSWorkspace<T>(workspace, workspaceOffset, param.batchSize * param.numScoreElements);
|
|
size_t sortedWorkspaceSize = EfficientNMSSortWorkspaceSize<T>(param.batchSize, param.numScoreElements);
|
|
char* sortedWorkspaceData = EfficientNMSWorkspace<char>(workspace, workspaceOffset, sortedWorkspaceSize);
|
|
cub::DoubleBuffer<T> scoresDB(topScoresData, sortedScoresData);
|
|
cub::DoubleBuffer<int> indexDB(topIndexData, sortedIndexData);
|
|
|
|
// Kernels
|
|
status = EfficientNMSFilterLauncher<T>(param, (T*) scoresInput, topNumData, topIndexData, topAnchorsData,
|
|
topOffsetsStartData, topOffsetsEndData, topScoresData, topClassData, stream);
|
|
CSC(status, STATUS_FAILURE);
|
|
|
|
status = cub::DeviceSegmentedRadixSort::SortPairsDescending(sortedWorkspaceData, sortedWorkspaceSize, scoresDB,
|
|
indexDB, param.batchSize * param.numScoreElements, param.batchSize, topOffsetsStartData, topOffsetsEndData,
|
|
param.scoreBits > 0 ? (10 - param.scoreBits) : 0, param.scoreBits > 0 ? 10 : sizeof(T) * 8, stream);
|
|
CSC(status, STATUS_FAILURE);
|
|
|
|
status = EfficientNMSLauncher<T>(param, topNumData, outputIndexData, outputClassData, indexDB.Current(),
|
|
scoresDB.Current(), topClassData, topAnchorsData, boxesInput, anchorsInput, (int*) numDetectionsOutput,
|
|
(T*) nmsScoresOutput, (int*) nmsClassesOutput, (int*) nmsIndicesOutput, nmsBoxesOutput, stream);
|
|
CSC(status, STATUS_FAILURE);
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
pluginStatus_t EfficientNMSInference(EfficientNMSParameters param, const void* boxesInput, const void* scoresInput,
|
|
const void* anchorsInput, void* numDetectionsOutput, void* nmsBoxesOutput, void* nmsScoresOutput,
|
|
void* nmsClassesOutput, void* nmsIndicesOutput, void* workspace, cudaStream_t stream)
|
|
{
|
|
if (param.datatype == DataType::kFLOAT)
|
|
{
|
|
param.scoreBits = -1;
|
|
return EfficientNMSDispatch<float>(param, boxesInput, scoresInput, anchorsInput, numDetectionsOutput,
|
|
nmsBoxesOutput, nmsScoresOutput, nmsClassesOutput, nmsIndicesOutput, workspace, stream);
|
|
}
|
|
else if (param.datatype == DataType::kHALF)
|
|
{
|
|
if (param.scoreBits <= 0 || param.scoreBits > 10)
|
|
{
|
|
param.scoreBits = -1;
|
|
}
|
|
return EfficientNMSDispatch<__half>(param, boxesInput, scoresInput, anchorsInput, numDetectionsOutput,
|
|
nmsBoxesOutput, nmsScoresOutput, nmsClassesOutput, nmsIndicesOutput, workspace, stream);
|
|
}
|
|
else
|
|
{
|
|
return STATUS_NOT_SUPPORTED;
|
|
}
|
|
}
|