chore: import upstream snapshot with attribution

plugin/common/kernels/nmsLayer.cu

@@ -0,0 +1,441 @@
+/*
+ * 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