wehub-resource-sync
442 lines
16 KiB
Plaintext
442 lines
16 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 "common/kernels/kernel.h"
|
|
#include "cuda_runtime_api.h"
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cub/cub.cuh>
|
|
#include <functional>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
|
|
using namespace nvinfer1;
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
// CUB's bug workaround:
|
|
// To work properly for large batch size CUB segmented sort needs ridiculous
|
|
// workspace alignment.
|
|
const uintptr_t ALIGNMENT = 1 << 20;
|
|
|
|
// IOU
|
|
template <typename TFloat>
|
|
__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
|
|
{
|
|
TFloat left = max(a.xmin, b.xmin), right = min(a.xmax, b.xmax);
|
|
TFloat top = max(a.ymin, b.ymin), bottom = min(a.ymax, b.ymax);
|
|
TFloat width = max((TFloat)(right - left + (TFloat) 1.0), (TFloat) 0.0);
|
|
TFloat height = max((TFloat)(bottom - top + (TFloat) 1.0), (TFloat) 0.0);
|
|
TFloat interS = width * height;
|
|
TFloat Sa = (a.xmax - a.xmin + (TFloat) 1) * (a.ymax - a.ymin + (TFloat) 1);
|
|
TFloat Sb = (b.xmax - b.xmin + (TFloat) 1) * (b.ymax - b.ymin + (TFloat) 1);
|
|
return (float) interS / (float) (Sa + Sb - interS);
|
|
}
|
|
|
|
// NMS KERNEL FOR SMALL BATCH SIZE
|
|
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
|
|
__global__ __launch_bounds__(DIM) void nmsKernel1(const int propSize,
|
|
Bbox<T_PROPOSALS> const* __restrict__ preNmsProposals,
|
|
T_ROIS* __restrict__ afterNmsProposals,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
__shared__ bool kept_boxes[TSIZE * DIM];
|
|
int kept = 0;
|
|
int batch_offset = blockIdx.x * propSize;
|
|
int max_box_idx = batch_offset + preNmsTopN;
|
|
int batch_offset_out = blockIdx.x * afterNmsTopN;
|
|
|
|
int flag_idx[TSIZE];
|
|
int boxes_idx[TSIZE];
|
|
Bbox<T_PROPOSALS> cur_boxes[TSIZE];
|
|
|
|
// initialize kept_boxes
|
|
#pragma unroll
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
boxes_idx[i] = threadIdx.x + batch_offset + DIM * i;
|
|
flag_idx[i] = threadIdx.x + DIM * i;
|
|
|
|
if (boxes_idx[i] < max_box_idx)
|
|
{
|
|
cur_boxes[i] = preNmsProposals[boxes_idx[i]];
|
|
kept_boxes[flag_idx[i]] = true;
|
|
}
|
|
else
|
|
{
|
|
kept_boxes[flag_idx[i]] = false;
|
|
boxes_idx[i] = -1.0f;
|
|
flag_idx[i] = -1.0f;
|
|
}
|
|
}
|
|
|
|
int ref_box_idx = 0 + batch_offset;
|
|
|
|
// remove the overlapped boxes
|
|
while ((kept < afterNmsTopN) && (ref_box_idx < max_box_idx))
|
|
{
|
|
Bbox<T_PROPOSALS> ref_box;
|
|
ref_box = preNmsProposals[ref_box_idx];
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
if (boxes_idx[i] > ref_box_idx)
|
|
{
|
|
if (IoU(ref_box, cur_boxes[i]) > nmsThres)
|
|
{
|
|
kept_boxes[flag_idx[i]] = false;
|
|
}
|
|
}
|
|
else if (boxes_idx[i] == ref_box_idx)
|
|
{
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 0] = ref_box.xmin;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.ymin;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.xmax;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.ymax;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
do
|
|
{
|
|
ref_box_idx++;
|
|
} while (!kept_boxes[ref_box_idx - batch_offset] && ref_box_idx < max_box_idx);
|
|
|
|
kept++;
|
|
}
|
|
}
|
|
|
|
// NMS KERNEL FOR LARGE BATCH SIZE
|
|
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
|
|
__global__ __launch_bounds__(DIM) void nmsKernel2(const int propSize,
|
|
Bbox<T_PROPOSALS> const* __restrict__ proposals,
|
|
T_ROIS* __restrict__ filtered,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
Bbox<T_PROPOSALS> const* cProposals = proposals + blockIdx.x * propSize;
|
|
|
|
Bbox<T_PROPOSALS> t[TSIZE];
|
|
uint64_t del = 0;
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
if (i < TSIZE - 1 || i * DIM + threadIdx.x < preNmsTopN)
|
|
{
|
|
t[i] = cProposals[i * DIM + threadIdx.x];
|
|
}
|
|
}
|
|
|
|
__shared__ Bbox<T_PROPOSALS> last;
|
|
__shared__ bool kept;
|
|
__shared__ int foundBatch;
|
|
if (threadIdx.x == 0)
|
|
foundBatch = 0;
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
for (int j = 0; j < DIM; j++)
|
|
{
|
|
int offset = i * DIM;
|
|
int index = offset + j;
|
|
if (index >= preNmsTopN)
|
|
break;
|
|
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == j)
|
|
{
|
|
kept = 0 == (del & ((uint64_t) 1 << i));
|
|
last = t[i];
|
|
|
|
if (kept)
|
|
{
|
|
int cnt = blockIdx.x * afterNmsTopN + foundBatch;
|
|
filtered[cnt * 4 + 0] = t[i].xmin;
|
|
filtered[cnt * 4 + 1] = t[i].ymin;
|
|
filtered[cnt * 4 + 2] = t[i].xmax;
|
|
filtered[cnt * 4 + 3] = t[i].ymax;
|
|
foundBatch++;
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (foundBatch == afterNmsTopN)
|
|
{
|
|
return;
|
|
}
|
|
|
|
if (kept)
|
|
{
|
|
Bbox<T_PROPOSALS> test = last;
|
|
|
|
for (int k = 0; k < TSIZE; k++)
|
|
{
|
|
if (index < k * DIM + threadIdx.x
|
|
&& IoU<T_PROPOSALS>(test, t[k]) > nmsThres)
|
|
{
|
|
del |= (uint64_t) 1 << k;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// NMS LAUNCH
|
|
template <typename T_PROPOSALS, DLayout_t L_PROPOSALS, typename T_ROIS>
|
|
pluginStatus_t nmsLaunch(cudaStream_t stream,
|
|
const int batch,
|
|
const int propSize,
|
|
void* proposals,
|
|
void* filtered,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
const int blockSize = 1024;
|
|
|
|
#define P1(tsize) nmsKernel1<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
|
|
#define P2(tsize) nmsKernel2<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
|
|
|
|
void (*kernel[64])(int, Bbox<T_PROPOSALS> const*, T_ROIS*, int, float, int) = {
|
|
P1(1), P1(2), P1(3), P1(4), P1(5), P1(6), P1(7), P1(8), P1(9), P1(10), P1(11), P1(12), P2(13), P2(14), P2(15), P2(16),
|
|
P2(17), P2(18), P2(19), P2(20), P2(21), P2(22), P2(23), P2(24), P2(25), P2(26), P2(27), P2(28), P2(29), P2(30), P2(31), P2(32),
|
|
P2(33), P2(34), P2(35), P2(36), P2(37), P2(38), P2(39), P2(40), P2(41), P2(42), P2(43), P2(44), P2(45), P2(46), P2(47), P2(48),
|
|
P2(49), P2(50), P2(51), P2(52), P2(53), P2(54), P2(55), P2(56), P2(57), P2(58), P2(59), P2(60), P2(61), P2(62), P2(63), P2(64)};
|
|
|
|
ASSERT_PARAM(preNmsTopN < 64 * blockSize);
|
|
|
|
CSC(cudaMemsetAsync(filtered, 0, batch * afterNmsTopN * 4 * sizeof(T_ROIS), stream), STATUS_FAILURE);
|
|
|
|
kernel[(preNmsTopN + blockSize - 1) / blockSize - 1]<<<batch, blockSize, 0, stream>>>(propSize,
|
|
(Bbox<T_PROPOSALS>*) proposals,
|
|
(T_ROIS*) filtered,
|
|
preNmsTopN,
|
|
nmsThres,
|
|
afterNmsTopN);
|
|
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
// SET OFFSET
|
|
// Works for up to 2Gi elements (cub's limitation)!
|
|
__global__ void setOffset(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;
|
|
}
|
|
}
|
|
|
|
// NMS GPU
|
|
template <typename T_SCORES, typename T_ROIS>
|
|
pluginStatus_t nmsGpu(cudaStream_t stream,
|
|
const int N,
|
|
const int R,
|
|
const int preNmsTop,
|
|
const int nmsMaxOut,
|
|
const float iouThreshold,
|
|
//const float minBoxSize,
|
|
//const float * imInfo,
|
|
void* fgScores,
|
|
const void* proposals,
|
|
void* workspace,
|
|
void* rois)
|
|
{
|
|
int8_t* vworkspace = alignPtr((int8_t*) workspace, ALIGNMENT);
|
|
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
|
|
|
|
pluginStatus_t error;
|
|
|
|
DEBUG_PRINTF("&&&& [NMS] DISCARD\n");
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
|
|
|
|
// Generate offsets
|
|
int* offsets = (int*) vworkspace;
|
|
setOffset<<<1, 1024, 0, stream>>>(R, N + 1, offsets);
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
|
|
vworkspace = vworkspace + N + 1;
|
|
vworkspace = alignPtr(vworkspace, ALIGNMENT);
|
|
|
|
// Sort (batched)
|
|
std::size_t tempStorageBytes = 0;
|
|
|
|
cub::DeviceSegmentedRadixSort::SortPairsDescending(
|
|
NULL, tempStorageBytes,
|
|
(T_SCORES*) fgScores, (T_SCORES*) fgScores,
|
|
(Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposals,
|
|
N * R, N,
|
|
offsets, offsets + 1, 0, 8 * sizeof(T_SCORES), stream);
|
|
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
|
|
T_SCORES* scoresOut = (T_SCORES*) vworkspace;
|
|
vworkspace = (int8_t*) (scoresOut + N * R);
|
|
vworkspace = alignPtr(vworkspace, ALIGNMENT);
|
|
Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
|
|
vworkspace = (int8_t*) (proposalsOut + N * R);
|
|
vworkspace = alignPtr(vworkspace, ALIGNMENT);
|
|
|
|
cub::DeviceSegmentedRadixSort::SortPairsDescending(
|
|
vworkspace, tempStorageBytes,
|
|
(T_SCORES*) fgScores, (T_SCORES*) scoresOut,
|
|
(Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposalsOut,
|
|
N * R, N,
|
|
offsets, offsets + 1,
|
|
0, 8 * sizeof(T_SCORES), stream);
|
|
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
|
|
DEBUG_PRINTF("&&&& [NMS] POST CUB\n");
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposalsOut, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(scoresOut, N * R * sizeof(float)));
|
|
|
|
error = nmsLaunch<T_ROIS, NC4HW, T_ROIS>(stream,
|
|
N,
|
|
R,
|
|
proposalsOut,
|
|
rois,
|
|
preNmsTop,
|
|
iouThreshold,
|
|
nmsMaxOut);
|
|
|
|
DEBUG_PRINTF("&&&& [NMS] POST LAUNCH\n");
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(rois, N * nmsMaxOut * 4 * sizeof(float)));
|
|
|
|
if (error != STATUS_SUCCESS)
|
|
{
|
|
return error;
|
|
}
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
// NMS LAUNCH CONFIG
|
|
typedef pluginStatus_t (*nmsFun)(cudaStream_t,
|
|
const int, // N
|
|
const int, // R
|
|
const int, // preNmsTop
|
|
const int, // nmsMaxOut
|
|
const float, // iouThreshold
|
|
//const float, // minBoxSize
|
|
//const float *, // imInfo
|
|
void*, // fgScores
|
|
const void*, // proposals,
|
|
void*, // workspace,
|
|
void*); // rois
|
|
|
|
struct nmsLaunchConfig
|
|
{
|
|
DataType t_fgScores;
|
|
DLayout_t l_fgScores;
|
|
DataType t_proposals;
|
|
DLayout_t l_proposals;
|
|
DataType t_rois;
|
|
nmsFun function;
|
|
|
|
nmsLaunchConfig(DataType t_fgScores,
|
|
DLayout_t l_fgScores,
|
|
DataType t_proposals,
|
|
DLayout_t l_proposals,
|
|
DataType t_rois,
|
|
nmsFun function)
|
|
: t_fgScores(t_fgScores)
|
|
, l_fgScores(l_fgScores)
|
|
, t_proposals(t_proposals)
|
|
, l_proposals(l_proposals)
|
|
, t_rois(t_rois)
|
|
, function(function)
|
|
{
|
|
}
|
|
|
|
nmsLaunchConfig(DataType t_fgScores,
|
|
DLayout_t l_fgScores,
|
|
DataType t_proposals,
|
|
DLayout_t l_proposals,
|
|
DataType t_rois)
|
|
: t_fgScores(t_fgScores)
|
|
, l_fgScores(l_fgScores)
|
|
, t_proposals(t_proposals)
|
|
, l_proposals(l_proposals)
|
|
, t_rois(t_rois)
|
|
{
|
|
}
|
|
|
|
bool operator==(nmsLaunchConfig const& other) const
|
|
{
|
|
return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores) && (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals) && (t_rois == other.t_rois);
|
|
}
|
|
};
|
|
|
|
#define FLOAT32 nvinfer1::DataType::kFLOAT
|
|
static std::array<nmsLaunchConfig, 1> nmsLCOptions = {
|
|
nmsLaunchConfig(FLOAT32, NCHW, FLOAT32, NC4HW, FLOAT32, nmsGpu<float, float>)};
|
|
|
|
// NMS
|
|
pluginStatus_t nms(cudaStream_t stream,
|
|
const int N,
|
|
const int R,
|
|
const int preNmsTop,
|
|
const int nmsMaxOut,
|
|
const float iouThreshold,
|
|
const DataType t_fgScores,
|
|
const DLayout_t l_fgScores,
|
|
void* fgScores,
|
|
const DataType t_proposals,
|
|
const DLayout_t l_proposals,
|
|
const void* proposals,
|
|
void* workspace,
|
|
const DataType t_rois,
|
|
void* rois)
|
|
{
|
|
|
|
nmsLaunchConfig lc(t_fgScores, l_fgScores, t_proposals, l_proposals, t_rois);
|
|
for (unsigned i = 0; i < nmsLCOptions.size(); i++)
|
|
{
|
|
if (nmsLCOptions[i] == lc)
|
|
{
|
|
DEBUG_PRINTF("NMS KERNEL %d\n", i);
|
|
return nmsLCOptions[i].function(stream,
|
|
N, R,
|
|
preNmsTop,
|
|
nmsMaxOut,
|
|
iouThreshold,
|
|
fgScores,
|
|
proposals,
|
|
workspace,
|
|
rois);
|
|
}
|
|
}
|
|
return STATUS_BAD_PARAM;
|
|
}
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|