chore: import upstream snapshot with attribution

plugin/common/kernels/gatherTopDetections.cu

@@ -0,0 +1,230 @@
+/*
+ * 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/kernels/kernel.h"
+#include "common/plugin.h"
+#include <cuda_fp16.h>
+#include <array>
+
+
+namespace nvinfer1
+{
+namespace plugin
+{
+inline __device__ __half minus_fb(const __half& a, const __half& b)
+{
+#if __CUDA_ARCH__ >= 530
+    return a - b;
+#else
+    return __float2half(__half2float(a) - __half2float(b));
+#endif
+}
+
+inline __device__ float minus_fb(const float & a, const float & b) {
+    return a - b;
+}
+
+template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void gatherTopDetections_kernel(
+        const bool shareLocation,
+        const int numImages,
+        const int numPredsPerClass,
+        const int numClasses,
+        const int topK,
+        const int keepTopK,
+        const int* indices,
+        const T_SCORE* scores,
+        const T_BBOX* bboxData,
+        int* keepCount,
+        T_BBOX* topDetections,
+        const T_SCORE score_shift)
+{
+    if (keepTopK > topK)
+        return;
+    for (int i = blockIdx.x * nthds_per_cta + threadIdx.x;
+         i < numImages * keepTopK;
+         i += gridDim.x * nthds_per_cta)
+    {
+        const int imgId = i / keepTopK;
+        const int detId = i % keepTopK;
+        const int offset = imgId * numClasses * topK;
+        const int index = indices[offset + detId];
+        const T_SCORE score = scores[offset + detId];
+        /*
+         * It is also likely that there is "bad bounding boxes" in the keepTopK bounding boxes.
+         * We set the bounding boxes parameters as the parameters shown below.
+         * These data will only show up at the end of keepTopK bounding boxes since the bounding boxes were sorted previously.
+         * It is also not going to affect the count of valid bounding boxes (keepCount).
+         * These data will probably never be used (because we have keepCount).
+         */
+        if (index == -1)
+        {
+            topDetections[i * 7] = imgId;  // image id
+            topDetections[i * 7 + 1] = -1; // label
+            topDetections[i * 7 + 2] = 0;  // confidence score
+            // score==0 will not pass the VisualizeBBox check
+            topDetections[i * 7 + 3] = 0;   // bbox xmin
+            topDetections[i * 7 + 4] = 0;   // bbox ymin
+            topDetections[i * 7 + 5] = 0;   // bbox xmax
+            topDetections[i * 7 + 6] = 0;   // bbox ymax
+        }
+        else
+        {
+            const int bboxOffset = imgId * (shareLocation ? numPredsPerClass : (numClasses * numPredsPerClass));
+            const int bboxId = ((shareLocation ? (index % numPredsPerClass)
+                        : index % (numClasses * numPredsPerClass)) + bboxOffset) * 4;
+            topDetections[i * 7] = imgId;                                                            // image id
+            topDetections[i * 7 + 1] = (index % (numClasses * numPredsPerClass)) / numPredsPerClass; // label
+            topDetections[i * 7 + 2] = score;                                                        // confidence score
+            // subtract 1.0 score shift we added in sortScorePerClass
+            topDetections[i * 7 + 2] = minus_fb(topDetections[i * 7 + 2], score_shift);
+            const T_BBOX xMin = bboxData[bboxId];
+            const T_BBOX yMin = bboxData[bboxId + 1];
+            const T_BBOX xMax = bboxData[bboxId + 2];
+            const T_BBOX yMax = bboxData[bboxId + 3];
+            // clipped bbox xmin
+            topDetections[i * 7 + 3] = __saturatef(xMin);
+            // clipped bbox ymin
+            topDetections[i * 7 + 4] = __saturatef(yMin);
+            // clipped bbox xmax
+            topDetections[i * 7 + 5] = __saturatef(xMax);
+            // clipped bbox ymax
+            topDetections[i * 7 + 6] = __saturatef(yMax);
+            // Atomic add to increase the count of valid keepTopK bounding boxes
+            // Without having to do manual sync.
+            atomicAdd(&keepCount[i / keepTopK], 1);
+        }
+    }
+}
+
+template <typename T_BBOX, typename T_SCORE>
+pluginStatus_t gatherTopDetections_gpu(
+    cudaStream_t stream,
+    const bool shareLocation,
+    const int numImages,
+    const int numPredsPerClass,
+    const int numClasses,
+    const int topK,
+    const int keepTopK,
+    const void* indices,
+    const void* scores,
+    const void* bboxData,
+    void* keepCount,
+    void* topDetections,
+    const float score_shift
+)
+{
+    CSC(cudaMemsetAsync(keepCount, 0, numImages * sizeof(int), stream), STATUS_FAILURE);
+    const int BS = 32;
+    const int GS = 32;
+    gatherTopDetections_kernel<T_BBOX, T_SCORE, BS><<<GS, BS, 0, stream>>>(shareLocation, numImages, numPredsPerClass,
+                                                                           numClasses, topK, keepTopK,
+                                                                           (int*) indices, (T_SCORE*) scores, (T_BBOX*) bboxData,
+                                                                           (int*) keepCount, (T_BBOX*) topDetections,
+                                                                           T_SCORE(score_shift));
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// gatherTopDetections LAUNCH CONFIG
+typedef pluginStatus_t (*gtdFunc)(cudaStream_t,
+                               const bool,
+                               const int,
+                               const int,
+                               const int,
+                               const int,
+                               const int,
+                               const void*,
+                               const void*,
+                               const void*,
+                               void*,
+                               void*,
+                               const float);
+struct gtdLaunchConfig
+{
+    DataType t_bbox;
+    DataType t_score;
+    gtdFunc function;
+
+    gtdLaunchConfig(DataType t_bbox, DataType t_score)
+        : t_bbox(t_bbox)
+        , t_score(t_score)
+        , function(nullptr)
+    {
+    }
+    gtdLaunchConfig(DataType t_bbox, DataType t_score, gtdFunc function)
+        : t_bbox(t_bbox)
+        , t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(gtdLaunchConfig const& other) const
+    {
+        return t_bbox == other.t_bbox && t_score == other.t_score;
+    }
+};
+
+using nvinfer1::DataType;
+
+static std::array<gtdLaunchConfig, 2> gtdLCOptions = {
+    gtdLaunchConfig(DataType::kFLOAT, DataType::kFLOAT, gatherTopDetections_gpu<float, float>),
+    gtdLaunchConfig(DataType::kHALF, DataType::kHALF, gatherTopDetections_gpu<__half, __half>)
+};
+
+pluginStatus_t gatherTopDetections(
+    cudaStream_t stream,
+    const bool shareLocation,
+    const int numImages,
+    const int numPredsPerClass,
+    const int numClasses,
+    const int topK,
+    const int keepTopK,
+    const DataType DT_BBOX,
+    const DataType DT_SCORE,
+    const void* indices,
+    const void* scores,
+    const void* bboxData,
+    void* keepCount,
+    void* topDetections,
+    const float score_shift)
+{
+    gtdLaunchConfig lc = gtdLaunchConfig(DT_BBOX, DT_SCORE);
+    for (unsigned i = 0; i < gtdLCOptions.size(); ++i)
+    {
+        if (lc == gtdLCOptions[i])
+        {
+            DEBUG_PRINTF("gatherTopDetections kernel %d\n", i);
+            return gtdLCOptions[i].function(stream,
+                                          shareLocation,
+                                          numImages,
+                                          numPredsPerClass,
+                                          numClasses,
+                                          topK,
+                                          keepTopK,
+                                          indices,
+                                          scores,
+                                          bboxData,
+                                          keepCount,
+                                          topDetections,
+                                          score_shift);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1