chore: import upstream snapshot with attribution

plugin/common/kernels/CMakeLists.txt

@@ -0,0 +1,49 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+#
+
+add_plugin_source(
+    allClassNMS.cu
+    bboxDeltas2Proposals.cu
+    common.cu
+    cropAndResizeKernel.cu
+    decodeBbox3DKernels.cu
+    extractFgScores.cu
+    gatherTopDetections.cu
+    generateAnchors.cu
+    gridAnchorLayer.cu
+    kernel.h
+    lReLU.cu
+    maskRCNNKernels.cu
+    maskRCNNKernels.h
+    nmsLayer.cu
+    normalizeLayer.cu
+    permuteData.cu
+    pillarScatterKernels.cu
+    priorBoxLayer.cu
+    proposalKernel.cu
+    proposalsForward.cu
+    reducedMathPlugin.h
+    regionForward.cu
+    reorgForward.cu
+    roiPooling.cu
+    rproiInferenceFused.cu
+    saturate.h
+    sortScoresPerClass.cu
+    sortScoresPerImage.cu
+    topkLastDim.cu
+    voxelGeneratorKernels.cu
+)

plugin/common/kernels/allClassNMS.cu

@@ -0,0 +1,394 @@
+/*
+ * 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_fp16.h>
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T_BBOX>
+__device__ float bboxSize(const Bbox<T_BBOX>& bbox, const bool normalized)
+{
+    if (float(bbox.xmax) < float(bbox.xmin) || float(bbox.ymax) < float(bbox.ymin))
+    {
+        // If bbox is invalid (e.g. xmax < xmin or ymax < ymin), return 0.
+        return 0;
+    }
+    else
+    {
+        float width = float(bbox.xmax) - float(bbox.xmin);
+        float height = float(bbox.ymax) - float(bbox.ymin);
+        if (normalized)
+        {
+            return width * height;
+        }
+        else
+        {
+            // If bbox is not within range [0, 1].
+            return (width + 1.f) * (height + 1.f);
+        }
+    }
+}
+
+template <typename T_BBOX>
+__device__ void intersectBbox(
+    const Bbox<T_BBOX>& bbox1,
+    const Bbox<T_BBOX>& bbox2,
+    Bbox<T_BBOX>* intersect_bbox)
+{
+    if (bbox2.xmin > bbox1.xmax || bbox2.xmax < bbox1.xmin || bbox2.ymin > bbox1.ymax || bbox2.ymax < bbox1.ymin)
+    {
+        // Return [0, 0, 0, 0] if there is no intersection.
+        intersect_bbox->xmin = T_BBOX(0);
+        intersect_bbox->ymin = T_BBOX(0);
+        intersect_bbox->xmax = T_BBOX(0);
+        intersect_bbox->ymax = T_BBOX(0);
+    }
+    else
+    {
+        intersect_bbox->xmin = max(bbox1.xmin, bbox2.xmin);
+        intersect_bbox->ymin = max(bbox1.ymin, bbox2.ymin);
+        intersect_bbox->xmax = min(bbox1.xmax, bbox2.xmax);
+        intersect_bbox->ymax = min(bbox1.ymax, bbox2.ymax);
+    }
+}
+
+
+template <>
+__device__ void intersectBbox<__half>(
+    const Bbox<__half>& bbox1,
+    const Bbox<__half>& bbox2,
+    Bbox<__half>* intersect_bbox)
+{
+    if (float(bbox2.xmin) > float(bbox1.xmax)
+        || float(bbox2.xmax) < float(bbox1.xmin)
+        || float(bbox2.ymin) > float(bbox1.ymax)
+        || float(bbox2.ymax) < float(bbox1.ymin))
+    {
+        // Return [0, 0, 0, 0] if there is no intersection.
+        intersect_bbox->xmin = __half(0);
+        intersect_bbox->ymin = __half(0);
+        intersect_bbox->xmax = __half(0);
+        intersect_bbox->ymax = __half(0);
+    }
+    else
+    {
+        intersect_bbox->xmin = max(float(bbox1.xmin), float(bbox2.xmin));
+        intersect_bbox->ymin = max(float(bbox1.ymin), float(bbox2.ymin));
+        intersect_bbox->xmax = min(float(bbox1.xmax), float(bbox2.xmax));
+        intersect_bbox->ymax = min(float(bbox1.ymax), float(bbox2.ymax));
+    }
+}
+
+
+template <typename T_BBOX>
+__device__ Bbox<T_BBOX> getDiagonalMinMaxSortedBox(const Bbox<T_BBOX>& bbox1)
+{
+    Bbox<T_BBOX> result;
+    result.xmin = min(bbox1.xmin, bbox1.xmax);
+    result.xmax = max(bbox1.xmin, bbox1.xmax);
+
+    result.ymin = min(bbox1.ymin, bbox1.ymax);
+    result.ymax = max(bbox1.ymin, bbox1.ymax);
+    return result;
+}
+
+template <>
+__device__ Bbox<__half> getDiagonalMinMaxSortedBox(const Bbox<__half>& bbox1)
+{
+    Bbox<__half> result;
+    result.xmin = min(float(bbox1.xmin), float(bbox1.xmax));
+    result.xmax = max(float(bbox1.xmin), float(bbox1.xmax));
+
+    result.ymin = min(float(bbox1.ymin), float(bbox1.ymax));
+    result.ymax = max(float(bbox1.ymin), float(bbox1.ymax));
+    return result;
+}
+
+template <typename T_BBOX>
+__device__ float jaccardOverlap(
+    const Bbox<T_BBOX>& bbox1, const Bbox<T_BBOX>& bbox2, const bool normalized, const bool caffeSemantics)
+{
+    Bbox<T_BBOX> intersect_bbox;
+
+    Bbox<T_BBOX> localbbox1 = getDiagonalMinMaxSortedBox(bbox1);
+    Bbox<T_BBOX> localbbox2 = getDiagonalMinMaxSortedBox(bbox2);
+
+    intersectBbox(localbbox1, localbbox2, &intersect_bbox);
+
+    float intersect_width, intersect_height;
+    // Only when using Caffe semantics, IOU calculation adds "1" to width and height if bbox is not normalized.
+    // https://github.com/weiliu89/caffe/blob/ssd/src/caffe/util/bbox_util.cpp#L92-L97
+    if (normalized || !caffeSemantics)
+    {
+        intersect_width = float(intersect_bbox.xmax) - float(intersect_bbox.xmin);
+        intersect_height = float(intersect_bbox.ymax) - float(intersect_bbox.ymin);
+    }
+    else
+    {
+        intersect_width = float(intersect_bbox.xmax) - float(intersect_bbox.xmin) + float(T_BBOX(1));
+        intersect_height = float(intersect_bbox.ymax) - float(intersect_bbox.ymin) + float(T_BBOX(1));
+    }
+    if (intersect_width > 0 && intersect_height > 0)
+    {
+        float intersect_size = intersect_width * intersect_height;
+        float bbox1_size = bboxSize(localbbox1, normalized);
+        float bbox2_size = bboxSize(localbbox2, normalized);
+        return intersect_size / (bbox1_size + bbox2_size - intersect_size);
+    }
+    else
+    {
+        return 0.;
+    }
+}
+
+template <typename T_BBOX>
+__device__ void emptyBboxInfo(
+    BboxInfo<T_BBOX>* bbox_info)
+{
+    bbox_info->conf_score = T_BBOX(0);
+    bbox_info->label = -2; // -1 is used for all labels when shared_location is ture
+    bbox_info->bbox_idx = -1;
+    bbox_info->kept = false;
+}
+/********** new NMS for only score and index array **********/
+
+template <typename T_SCORE, typename T_BBOX, int TSIZE>
+__global__ void allClassNMS_kernel(const int num, const int num_classes, const int num_preds_per_class, const int top_k,
+    const float nms_threshold, const bool share_location, const bool isNormalized,
+    T_BBOX* bbox_data, // bbox_data should be float to preserve location information
+    T_SCORE* beforeNMS_scores, int* beforeNMS_index_array, T_SCORE* afterNMS_scores, int* afterNMS_index_array,
+    bool flipXY, const float score_shift, bool caffeSemantics)
+{
+    //__shared__ bool kept_bboxinfo_flag[CAFFE_CUDA_NUM_THREADS * TSIZE];
+    extern __shared__ bool kept_bboxinfo_flag[];
+
+    for (int i = 0; i < num; i++)
+    {
+        int32_t const offset = i * num_classes * num_preds_per_class + blockIdx.x * num_preds_per_class;
+        // Should not write data beyond [offset, top_k).
+        int32_t const max_idx = offset + top_k;
+        // Should not read beyond [offset, num_preds_per_class).
+        int32_t const max_read_idx = offset + min(top_k, num_preds_per_class);
+        int32_t const bbox_idx_offset = i * num_preds_per_class * (share_location ? 1 : num_classes);
+
+        // local thread data
+        int loc_bboxIndex[TSIZE];
+        Bbox<T_BBOX> loc_bbox[TSIZE];
+
+        // initialize Bbox, Bboxinfo, kept_bboxinfo_flag
+        // Eliminate shared memory RAW hazard
+        __syncthreads();
+#pragma unroll
+        for (int t = 0; t < TSIZE; t++)
+        {
+            const int cur_idx = threadIdx.x + blockDim.x * t;
+            const int item_idx = offset + cur_idx;
+            // Init all output data
+            if (item_idx < max_idx)
+            {
+                // Do not access data if it exceeds read boundary
+                if (item_idx < max_read_idx)
+                {
+                    loc_bboxIndex[t] = beforeNMS_index_array[item_idx];
+                }
+                else
+                {
+                    loc_bboxIndex[t] = -1;
+                }
+
+                if (loc_bboxIndex[t] != -1)
+                {
+                    const int bbox_data_idx = share_location ? (loc_bboxIndex[t] % num_preds_per_class + bbox_idx_offset) : loc_bboxIndex[t];
+
+                    loc_bbox[t].xmin = flipXY ? bbox_data[bbox_data_idx * 4 + 1] : bbox_data[bbox_data_idx * 4 + 0];
+                    loc_bbox[t].ymin = flipXY ? bbox_data[bbox_data_idx * 4 + 0] : bbox_data[bbox_data_idx * 4 + 1];
+                    loc_bbox[t].xmax = flipXY ? bbox_data[bbox_data_idx * 4 + 3] : bbox_data[bbox_data_idx * 4 + 2];
+                    loc_bbox[t].ymax = flipXY ? bbox_data[bbox_data_idx * 4 + 2] : bbox_data[bbox_data_idx * 4 + 3];
+                    kept_bboxinfo_flag[cur_idx] = true;
+                }
+                else
+                {
+                    kept_bboxinfo_flag[cur_idx] = false;
+                }
+            }
+            else
+            {
+                kept_bboxinfo_flag[cur_idx] = false;
+            }
+        }
+
+        // filter out overlapped boxes with lower scores
+        int ref_item_idx = offset;
+
+        int32_t ref_bbox_idx = -1;
+        if (ref_item_idx < max_read_idx)
+        {
+            ref_bbox_idx = share_location
+                ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset)
+                : beforeNMS_index_array[ref_item_idx];
+        }
+        while ((ref_bbox_idx != -1) && ref_item_idx < max_read_idx)
+        {
+            Bbox<T_BBOX> ref_bbox;
+            ref_bbox.xmin = flipXY ? bbox_data[ref_bbox_idx * 4 + 1] : bbox_data[ref_bbox_idx * 4 + 0];
+            ref_bbox.ymin = flipXY ? bbox_data[ref_bbox_idx * 4 + 0] : bbox_data[ref_bbox_idx * 4 + 1];
+            ref_bbox.xmax = flipXY ? bbox_data[ref_bbox_idx * 4 + 3] : bbox_data[ref_bbox_idx * 4 + 2];
+            ref_bbox.ymax = flipXY ? bbox_data[ref_bbox_idx * 4 + 2] : bbox_data[ref_bbox_idx * 4 + 3];
+
+            // Eliminate shared memory RAW hazard
+            __syncthreads();
+
+            for (int t = 0; t < TSIZE; t++)
+            {
+                const int cur_idx = threadIdx.x + blockDim.x * t;
+                const int item_idx = offset + cur_idx;
+
+                if ((kept_bboxinfo_flag[cur_idx]) && (item_idx > ref_item_idx))
+                {
+                    if (jaccardOverlap(ref_bbox, loc_bbox[t], isNormalized, caffeSemantics) > nms_threshold)
+                    {
+                        kept_bboxinfo_flag[cur_idx] = false;
+                    }
+                }
+            }
+            __syncthreads();
+
+            do
+            {
+                ref_item_idx++;
+            } while (ref_item_idx < max_read_idx && !kept_bboxinfo_flag[ref_item_idx - offset]);
+
+            // Move to next valid point
+            if (ref_item_idx < max_read_idx)
+            {
+                ref_bbox_idx = share_location
+                    ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset)
+                    : beforeNMS_index_array[ref_item_idx];
+            }
+        }
+
+        // store data
+        for (int t = 0; t < TSIZE; t++)
+        {
+            const int cur_idx = threadIdx.x + blockDim.x * t;
+            const int read_item_idx = offset + cur_idx;
+            const int write_item_idx = (i * num_classes * top_k + blockIdx.x * top_k) + cur_idx;
+            /*
+             * If not not keeping the bbox
+             * Set the score to 0
+             * Set the bounding box index to -1
+             */
+            if (read_item_idx < max_idx)
+            {
+                afterNMS_scores[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? T_SCORE(beforeNMS_scores[read_item_idx]) : T_SCORE(score_shift);
+                afterNMS_index_array[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? loc_bboxIndex[t] : -1;
+            }
+        }
+    }
+}
+
+template <typename T_SCORE, typename T_BBOX>
+pluginStatus_t allClassNMS_gpu(cudaStream_t stream, const int num, const int num_classes, const int num_preds_per_class,
+    const int top_k, const float nms_threshold, const bool share_location, const bool isNormalized, void* bbox_data,
+    void* beforeNMS_scores, void* beforeNMS_index_array, void* afterNMS_scores, void* afterNMS_index_array, bool flipXY,
+    const float score_shift, bool caffeSemantics)
+{
+#define P(tsize) allClassNMS_kernel<T_SCORE, T_BBOX, (tsize)>
+
+    void (*kernel[8])(const int, const int, const int, const int, const float, const bool, const bool, T_BBOX*,
+        T_SCORE*, int*, T_SCORE*, int*, bool, const float, bool)
+        = {
+            P(1),
+            P(2),
+            P(3),
+            P(4),
+            P(5),
+            P(6),
+            P(7),
+            P(8),
+        };
+
+    const int BS = 512;
+    const int GS = num_classes;
+    const int t_size = (top_k + BS - 1) / BS;
+
+    kernel[t_size - 1]<<<GS, BS, BS * t_size * sizeof(bool), stream>>>(num, num_classes, num_preds_per_class, top_k,
+        nms_threshold, share_location, isNormalized, (T_BBOX*) bbox_data, (T_SCORE*) beforeNMS_scores,
+        (int*) beforeNMS_index_array, (T_SCORE*) afterNMS_scores, (int*) afterNMS_index_array, flipXY, score_shift,
+        caffeSemantics);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// allClassNMS LAUNCH CONFIG
+typedef pluginStatus_t (*nmsFunc)(cudaStream_t, const int, const int, const int, const int, const float, const bool,
+    const bool, void*, void*, void*, void*, void*, bool, const float, bool);
+
+struct nmsLaunchConfigSSD
+{
+    DataType t_score;
+    DataType t_bbox;
+    nmsFunc function;
+
+    nmsLaunchConfigSSD(DataType t_score, DataType t_bbox)
+        : t_score(t_score)
+        , t_bbox(t_bbox)
+        , function(nullptr)
+    {
+    }
+    nmsLaunchConfigSSD(DataType t_score, DataType t_bbox, nmsFunc function)
+        : t_score(t_score)
+        , t_bbox(t_bbox)
+        , function(function)
+    {
+    }
+    bool operator==(nmsLaunchConfigSSD const& other) const
+    {
+        return t_score == other.t_score && t_bbox == other.t_bbox;
+    }
+};
+
+static std::array<nmsLaunchConfigSSD, 2> nmsSsdLCOptions = {
+    nmsLaunchConfigSSD(DataType::kFLOAT, DataType::kFLOAT, allClassNMS_gpu<float, float>),
+    nmsLaunchConfigSSD(DataType::kHALF, DataType::kHALF, allClassNMS_gpu<__half, __half>)
+};
+
+pluginStatus_t allClassNMS(cudaStream_t stream, const int num, const int num_classes, const int num_preds_per_class,
+    const int top_k, const float nms_threshold, const bool share_location, const bool isNormalized,
+    const DataType DT_SCORE, const DataType DT_BBOX, void* bbox_data, void* beforeNMS_scores,
+    void* beforeNMS_index_array, void* afterNMS_scores, void* afterNMS_index_array, bool flipXY,
+    const float score_shift, bool caffeSemantics)
+{
+    nmsLaunchConfigSSD lc = nmsLaunchConfigSSD(DT_SCORE, DT_BBOX);
+    for (unsigned i = 0; i < nmsSsdLCOptions.size(); ++i)
+    {
+        if (lc == nmsSsdLCOptions[i])
+        {
+            DEBUG_PRINTF("all class nms kernel %d\n", i);
+            return nmsSsdLCOptions[i].function(stream, num, num_classes, num_preds_per_class, top_k, nms_threshold,
+                share_location, isNormalized, bbox_data, beforeNMS_scores, beforeNMS_index_array, afterNMS_scores,
+                afterNMS_index_array, flipXY, score_shift, caffeSemantics);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/bboxDeltas2Proposals.cu

@@ -0,0 +1,290 @@
+/*
+ * 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 <algorithm>
+#include <array>
+#include <math.h>
+#include <stdio.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+using std::max;
+using std::min;
+
+// BBD2P KERNEL
+template <typename T_DELTAS,
+          DLayout_t L_DELTAS,
+          typename TV_PROPOSALS,
+          DLayout_t L_PROPOSALS,
+          typename T_FGSCORES,
+          DLayout_t L_FGSCORES>
+__global__ void bboxDeltas2Proposals_kernel(
+    int N,
+    int A,
+    int H,
+    int W,
+    const float* __restrict__ anchors,
+    const float* __restrict__ imInfo,
+    int featureStride,
+    float minSize,
+    const T_DELTAS* __restrict__ deltas,
+    TV_PROPOSALS* __restrict__ proposals,
+    T_FGSCORES* __restrict__ scores)
+{
+    int tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    if (tid < N * A * H * W)
+    { // TODO this can be a loop.
+        // Find out the index of bounding box for the output
+        int cnt = tid;
+        // width index
+        int w = cnt % W;
+        cnt = cnt / W;
+        // height index
+        int h = cnt % H;
+        cnt = cnt / H;
+        // anchor box index
+        int a = cnt % A;
+        cnt = cnt / A;
+        // batch index
+        int n = cnt;
+        int hw = h * W + w;
+
+        // Get the height and width of the original input image
+        float imHeight = imInfo[3 * n];
+        float imWidth = imInfo[3 * n + 1];
+
+        // Point to the right anchor box
+        float4 anchor = ((float4*) anchors)[a];
+        // Get anchor box coordinates
+        float a_ctr_x = anchor.x;
+        float a_ctr_y = anchor.y;
+        float a_w = anchor.z;
+        float a_h = anchor.w;
+
+        // NCHW format
+        // Find out the starting position of the bounding box in the input (predicted bounding box offsets)
+        int id = ((tid - hw) * 4) + hw;
+        T_DELTAS dx;
+        T_DELTAS dy;
+        T_DELTAS dw;
+        T_DELTAS dh;
+        if (L_DELTAS == NCHW)
+        {
+            // The offsets between adjacent coordinates on linear memory is H * W
+            dx = deltas[id];
+            dy = deltas[id + 1 * H * W];
+            dw = deltas[id + 2 * H * W];
+            dh = deltas[id + 3 * H * W];
+        }
+        // NC4HW format
+        else if (L_DELTAS == NC4HW)
+        {
+            dx = deltas[tid * 4 + 0];
+            dy = deltas[tid * 4 + 1];
+            dw = deltas[tid * 4 + 2];
+            dh = deltas[tid * 4 + 3];
+        }
+        /*
+         * Calculate the coordinates of decoded bounding box on the original input image scale
+         * Only works if param.minBoxSize == param.featureStride
+         */
+        float ctr_x = a_ctr_x + w * featureStride;
+        float ctr_y = a_ctr_y + h * featureStride;
+        // float ctr_x = (w + 0.5) * featureStride;
+        // float ctr_y = (h + 0.5) * featureStride;
+
+        /*
+         * Decode the predicted bounding box
+         * The decoded bounding boxes has coordinates of [x_topleft, y_topleft, x_bottomright, y_bottomright]
+         * The units are in pixels
+         */
+        ctr_x = ctr_x + dx * a_w;
+        ctr_y = ctr_y + dy * a_h;
+        float b_w = __expf(dw) * a_w;
+        float b_h = __expf(dh) * a_h;
+        float bx = ctr_x - (b_w / 2);
+        float by = ctr_y - (b_h / 2);
+        float bz = ctr_x + (b_w / 2);
+        float bw = ctr_y + (b_h / 2);
+
+        TV_PROPOSALS bbox;
+        // Make sure that the decoded bouding box go outside of the original input image
+        bbox.x = fminf(fmaxf(bx, 0.0f), imWidth - 1.0f);
+        bbox.y = fminf(fmaxf(by, 0.0f), imHeight - 1.0f);
+        bbox.z = fminf(fmaxf(bz, 0.0f), imWidth - 1.0f);
+        bbox.w = fminf(fmaxf(bw, 0.0f), imHeight - 1.0f);
+
+        // Put the decoded bounding box information to the outputs
+        if (L_PROPOSALS == NC4HW)
+        {
+            proposals[tid] = bbox;
+        }
+
+        int ininf = 0xff800000;
+        float ninf = *(float*) &ininf;
+        // minBoxSize at the original input image scale
+        float scaledMinSize = minSize * imInfo[3 * n + 2];
+        // Set the objectness score to -inf if the predicted bounding box has edgth length less than the minimal box size expected.
+        if (bbox.z - bbox.x + 1 < scaledMinSize || bbox.w - bbox.y + 1 < scaledMinSize)
+        {
+            if (L_FGSCORES == NCHW)
+                scores[tid] = ninf;
+        }
+    }
+}
+
+// BBD2P KERNEL LAUNCHER
+template <typename T_DELTAS,
+          DLayout_t L_DELTAS,
+          typename TV_PROPOSALS,
+          DLayout_t L_PROPOSALS,
+          typename T_FGSCORES,
+          DLayout_t L_FGSCORES>
+pluginStatus_t bboxDeltas2Proposals_gpu(cudaStream_t stream,
+                                       int N,
+                                       int A,
+                                       int H,
+                                       int W,
+                                       const float* imInfo,
+                                       int featureStride,
+                                       float minBoxSize,
+                                       const float* anchors,
+                                       const void* deltas,
+                                       void* propos,
+                                       void* scores)
+{
+    const int BS = 32;
+    const int GS = ((N * A * H * W) + BS - 1) / BS;
+
+    bboxDeltas2Proposals_kernel<T_DELTAS, L_DELTAS, TV_PROPOSALS, L_PROPOSALS, T_FGSCORES, L_FGSCORES><<<GS, BS, 0, stream>>>(N, A, H, W,
+                                                                                                                              anchors,
+                                                                                                                              imInfo,
+                                                                                                                              featureStride,
+                                                                                                                              minBoxSize,
+                                                                                                                              (T_DELTAS*) deltas,
+                                                                                                                              (TV_PROPOSALS*) propos,
+                                                                                                                              (T_FGSCORES*) scores);
+
+    DEBUG_PRINTF("&&&& [bboxD2P] POST LAUNCH\n");
+    DEBUG_PRINTF("&&&& [bboxD2P] PROPOS %u\n", hash(propos, N * A * H * W * 4 * sizeof(float)));
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// BBD2P LAUNCH CONFIG {{{
+typedef pluginStatus_t (*bd2pFun)(cudaStream_t,
+                                 int,
+                                 int,
+                                 int,
+                                 int,
+                                 const float*,
+                                 int,
+                                 float,
+                                 const float*,
+                                 const void*,
+                                 void*,
+                                 void*);
+
+struct bd2pLaunchConfig
+{
+    DataType t_deltas;
+    DLayout_t l_deltas;
+    DataType t_proposals;
+    DLayout_t l_proposals;
+    DataType t_scores;
+    DLayout_t l_scores;
+    bd2pFun function;
+
+    bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals,
+        DataType t_scores, DLayout_t l_scores)
+        : t_deltas(t_deltas)
+        , l_deltas(l_deltas)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_scores(t_scores)
+        , l_scores(l_scores)
+        , function(nullptr)
+    {
+    }
+
+    bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals, DataType t_scores, DLayout_t l_scores, bd2pFun function)
+        : t_deltas(t_deltas)
+        , l_deltas(l_deltas)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_scores(t_scores)
+        , l_scores(l_scores)
+        , function(function)
+    {
+    }
+
+    bool operator==(bd2pLaunchConfig const& other) const
+    {
+        return t_deltas == other.t_deltas && l_deltas == other.l_deltas && t_proposals == other.t_proposals && l_proposals == other.l_proposals && t_scores == other.t_scores && l_scores == other.l_scores;
+    }
+};
+
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+static std::array<bd2pLaunchConfig, 2> bd2pLCOptions = {
+    bd2pLaunchConfig(FLOAT32, NC4HW, FLOAT32, NC4HW, FLOAT32, NCHW, bboxDeltas2Proposals_gpu<float, NC4HW, float4, NC4HW, float, NCHW>),
+    bd2pLaunchConfig(FLOAT32, NCHW, FLOAT32, NC4HW, FLOAT32, NCHW, bboxDeltas2Proposals_gpu<float, NCHW, float4, NC4HW, float, NCHW>)};
+
+// BBD2P
+pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
+                                   const int N,
+                                   const int A,
+                                   const int H,
+                                   const int W,
+                                   const int featureStride,
+                                   const float minBoxSize,
+                                   const float* imInfo,
+                                   const float* anchors,
+                                   const DataType t_deltas,
+                                   const DLayout_t l_deltas,
+                                   const void* deltas,
+                                   const DataType t_proposals,
+                                   const DLayout_t l_proposals,
+                                   void* proposals,
+                                   const DataType t_scores,
+                                   const DLayout_t l_scores,
+                                   void* scores)
+{
+    bd2pLaunchConfig lc = bd2pLaunchConfig(t_deltas, l_deltas, t_proposals, l_proposals, t_scores, l_scores);
+    for (unsigned i = 0; i < bd2pLCOptions.size(); i++)
+    {
+        if (lc == bd2pLCOptions[i])
+        {
+            DEBUG_PRINTF("BBD2P kernel %d\n", i);
+            return bd2pLCOptions[i].function(stream,
+                                          N, A, H, W,
+                                          imInfo,
+                                          featureStride,
+                                          minBoxSize,
+                                          anchors,
+                                          deltas,
+                                          proposals,
+                                          scores);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/common.cu

@@ -0,0 +1,175 @@
+/*
+ * 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.h"
+#include <cub/cub.cuh>
+#include <stdint.h>
+
+using namespace nvinfer1::pluginInternal;
+
+#define CUDA_MEM_ALIGN 256
+namespace nvinfer1
+{
+namespace plugin
+{
+// HASH
+unsigned int hash(const void* array_, size_t size)
+{
+    // Apply hashing only when debugging RPN codes.
+    if (DEBUG_ENABLE)
+    {
+        const char* array_const;
+        char* array;
+        PLUGIN_CHECK_CUDA(cudaMallocHost((void**) &array, size));
+        PLUGIN_CHECK_CUDA(cudaMemcpy(array, array_, size, cudaMemcpyDeviceToHost));
+        array_const = array;
+        unsigned int hash = 45599;
+        for (size_t i = 0; i < size; i++)
+        {
+            unsigned int value = array_const[i];
+            hash = hash * 1487 + value;
+            hash = hash * 317;
+            hash = hash % 105359;
+        }
+        return hash;
+    }
+    else
+    {
+        return 0;
+    }
+}
+
+// ALIGNPTR
+int8_t* alignPtr(int8_t* ptr, uintptr_t to)
+{
+    uintptr_t addr = (uintptr_t) ptr;
+    if (addr % to)
+    {
+        addr += to - addr % to;
+    }
+    return (int8_t*) addr;
+}
+
+// NEXTWORKSPACEPTR
+int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize)
+{
+    uintptr_t addr = (uintptr_t) ptr;
+    addr += previousWorkspaceSize;
+    return alignPtr((int8_t*) addr, CUDA_MEM_ALIGN);
+}
+
+// CALCULATE TOTAL WORKSPACE SIZE
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int count)
+{
+    size_t total = 0;
+    for (int i = 0; i < count; i++)
+    {
+        total += workspaces[i];
+        if (workspaces[i] % CUDA_MEM_ALIGN)
+        {
+            total += CUDA_MEM_ALIGN - (workspaces[i] % CUDA_MEM_ALIGN);
+        }
+    }
+    return total;
+}
+
+using nvinfer1::DataType;
+
+// DATA TYPE SIZE
+size_t dataTypeSize(const DataType dtype)
+{
+    switch (dtype)
+    {
+    case DataType::kINT8: return sizeof(char);
+    case DataType::kHALF: return sizeof(short);
+    case DataType::kFLOAT: return sizeof(float);
+    case DataType::kBF16:
+    case DataType::kINT64: PLUGIN_FAIL("Unsupported data type");
+    default: return 0;
+    }
+}
+
+// CUB
+/*
+size_t cubSortFloatIntPairsWorkspaceSize(int num_items, int num_segments)
+{
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+    (int *)NULL, temp_storage_bytes,
+    (const float *)NULL, (float *)NULL,
+    (const int *)NULL, (int *)NULL,
+    num_items,     // # items
+    num_segments,  // # segments
+    (const int *)NULL, (const int *)NULL);
+    return temp_storage_bytes;
+}
+
+size_t cubSortFloatBboxInfoPairsWorkspaceSize(int num_items, int num_segments)
+{
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+    (int *)NULL, temp_storage_bytes,
+    (const float *)NULL, (float *)NULL,
+    (const BboxInfo<float> *)NULL, (BboxInfo<float> *)NULL,
+    num_items,     // # items
+    num_segments,  // # segments
+    (const int *)NULL, (const int *)NULL);
+    return temp_storage_bytes;
+}
+*/
+
+template <unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta) __global__
+    void setUniformOffsets_kernel(const int num_segments, const int offset, int* d_offsets)
+{
+    const int idx = blockIdx.x * nthds_per_cta + threadIdx.x;
+    if (idx <= num_segments)
+        d_offsets[idx] = idx * offset;
+}
+
+void setUniformOffsets(cudaStream_t stream, const int num_segments, const int offset, int* d_offsets)
+{
+    const int BS = 32;
+    const int GS = (num_segments + 1 + BS - 1) / BS;
+    setUniformOffsets_kernel<BS><<<GS, BS, 0, stream>>>(num_segments, offset, d_offsets);
+}
+
+const char* cublasGetErrorString(cublasStatus_t error)
+{
+    switch (error)
+    {
+    case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS";
+    case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED";
+    case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
+    case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
+    case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
+    case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
+    case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
+    case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
+#if CUDA_VERSION >= 6000
+    case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
+#endif
+#if CUDA_VERSION >= 6050
+    case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR";
+#endif
+    }
+    return "Unknown cublas status";
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/cropAndResizeKernel.cu

@@ -0,0 +1,146 @@
+/*
+ * 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 <cuda.h>
+#include <cuda_runtime.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T>
+__global__ void cropAndResizeKernel(const int nthreads, const T* image_ptr, const float* boxes_ptr,
+                                    int num_boxes, int batch, int image_height, int image_width,
+                                    int crop_height, int crop_width, int depth,
+                                    float extrapolation_value, float* crops_ptr)
+{
+    for (int out_idx = threadIdx.x + blockIdx.x * blockDim.x ; out_idx < nthreads;
+            out_idx += blockDim.x * gridDim.x)
+    {
+        int idx =  out_idx;
+        const int x = idx % crop_width;
+        idx /= crop_width;
+        const int y = idx % crop_height;
+        idx /= crop_height;
+        const int d = idx % depth;
+        const int b = idx / depth;
+        const float y1 = boxes_ptr[b * 4];
+        const float x1 = boxes_ptr[b * 4 + 1];
+        const float y2 = boxes_ptr[b * 4 + 2];
+        const float x2 = boxes_ptr[b * 4 + 3];
+        //each image has num_boxes of boxes, so we simply divide to get the box index.
+        const int b_in = b / num_boxes;
+
+        if (b_in < 0 || b_in >= batch)
+        {
+            continue;
+        }
+
+        const float height_scale =
+            (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1)
+            : 0;
+        const float width_scale =
+            (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) : 0;
+        const float in_y = (crop_height > 1)
+                            ? y1 * (image_height - 1) + y * height_scale
+                            : 0.5 * (y1 + y2) * (image_height - 1);
+
+        if (in_y < 0 || in_y > image_height - 1)
+        {
+            crops_ptr[out_idx] = extrapolation_value;
+            continue;
+        }
+
+        const float in_x = (crop_width > 1)
+                           ? x1 * (image_width - 1) + x * width_scale
+                           : 0.5 * (x1 + x2) * (image_width - 1);
+
+        if (in_x < 0 || in_x > image_width - 1)
+        {
+            crops_ptr[out_idx] = extrapolation_value;
+            continue;
+        }
+
+        const int top_y_index = floorf(in_y);
+        const int bottom_y_index = ceilf(in_y);
+        const float y_lerp = in_y - top_y_index;
+        const int left_x_index = floorf(in_x);
+        const int right_x_index = ceilf(in_x);
+        const float x_lerp = in_x - left_x_index;
+        const float top_left(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               top_y_index) *
+                                               image_width +
+                                               left_x_index]));
+        const float top_right(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                                top_y_index) *
+                                                image_width +
+                                                right_x_index]));
+        const float bottom_left(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               bottom_y_index) *
+                                               image_width +
+                                               left_x_index]));
+        const float bottom_right(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               bottom_y_index) *
+                                               image_width +
+                                               right_x_index]));
+        const float top = top_left + (top_right - top_left) * x_lerp;
+        const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp;
+        crops_ptr[out_idx] = top + (bottom - top) * y_lerp;
+    }
+}
+
+
+
+
+
+int cropAndResizeInference(
+    cudaStream_t stream,
+    int n,
+    const void* image,
+    const void* rois,
+    int batch_size,
+    int input_height,
+    int input_width,
+    int num_boxes,
+    int crop_height,
+    int crop_width,
+    int depth,
+    void* output)
+{
+    int output_volume = batch_size * num_boxes * crop_height * crop_width * depth;
+    int block_size = 1024;
+    int grid_size = (output_volume + block_size - 1 ) / block_size;
+    cropAndResizeKernel<float> <<< grid_size, block_size, 0, stream>>>(output_volume,
+            static_cast<const float*>(image),
+            static_cast<const float*>(rois),
+            num_boxes,
+            batch_size,
+            input_height,
+            input_width,
+            crop_height,
+            crop_width,
+            depth,
+            0.0f,
+            static_cast<float*>(output));
+    return 0;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/decodeBbox3DKernels.cu

@@ -0,0 +1,188 @@
+/*
+ * 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 <iostream>
+#include <cuda_runtime_api.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define checkCudaErrors(status_)                                  \
+{                                                                 \
+  auto const status = status_;                                    \
+  if (status != 0)                                                \
+  {                                                               \
+    std::cout << "Cuda failure: " << cudaGetErrorString(status)   \
+              << " at line " << __LINE__                          \
+              << " in file " << __FILE__                          \
+              << " error status: " << status                      \
+              << std::endl;                                       \
+              exit(EXIT_FAILURE);                                           \
+    }                                                             \
+}
+
+#ifndef M_PI
+    #define M_PI 3.14159265358979323846
+#endif
+
+__device__ float sigmoid(const float x) { return 1.0f / (1.0f + expf(-x)); }
+
+__global__ void postprocess_kernal(const float *cls_input,
+                                        float const* box_input,
+                                        const float *dir_cls_input,
+                                        float *anchors,
+                                        float *anchors_bottom_height,
+                                        float *bndbox_output,
+                                        int *object_counter,
+                                        const float min_x_range,
+                                        const float max_x_range,
+                                        const float min_y_range,
+                                        const float max_y_range,
+                                        const int feature_x_size,
+                                        const int feature_y_size,
+                                        const int num_anchors,
+                                        const int num_classes,
+                                        const int num_box_values,
+                                        const float score_thresh,
+                                        const float dir_offset,
+                                        const float dir_limit_offset,
+                                        const int num_dir_bins)
+{
+  int max_box_num = feature_x_size * feature_y_size * num_anchors;
+  int loc_index =blockIdx.x;
+  int batch_idx = blockIdx.x / (feature_x_size * feature_y_size);
+  int loc_index_in_frame = blockIdx.x % (feature_x_size * feature_y_size);
+  int ith_anchor = threadIdx.x;
+  if (ith_anchor >= num_anchors)
+  {
+      return;
+  }
+  int col = loc_index_in_frame % feature_x_size;
+  int row = loc_index_in_frame / feature_x_size;
+  float x_offset = min_x_range + col * (max_x_range - min_x_range) / (feature_x_size - 1);
+  float y_offset = min_y_range + row * (max_y_range - min_y_range) / (feature_y_size - 1);
+  int cls_offset = loc_index * num_classes * num_anchors + ith_anchor * num_classes;
+  float dev_cls[2] = {-1, 0};
+  const float *scores = cls_input + cls_offset;
+  float max_score = sigmoid(scores[0]);
+  int cls_id = 0;
+  for (int i = 1; i < num_classes; i++) {
+    float cls_score = sigmoid(scores[i]);
+    if (cls_score > max_score) {
+      max_score = cls_score;
+      cls_id = i;
+    }
+  }
+  dev_cls[0] = static_cast<float>(cls_id);
+  dev_cls[1] = max_score;
+  if (dev_cls[1] >= score_thresh)
+  {
+    int box_offset = loc_index * num_anchors * num_box_values + ith_anchor * num_box_values;
+    int dir_cls_offset = loc_index * num_anchors * 2 + ith_anchor * 2;
+    float *anchor_ptr = anchors + ith_anchor * 4;
+    float z_offset = anchor_ptr[2] / 2 + anchors_bottom_height[ith_anchor / 2];
+    float anchor[7] = {x_offset, y_offset, z_offset, anchor_ptr[0], anchor_ptr[1], anchor_ptr[2], anchor_ptr[3]};
+    float const* box_encodings = box_input + box_offset;
+    float xa = anchor[0];
+    float ya = anchor[1];
+    float za = anchor[2];
+    float dxa = anchor[3];
+    float dya = anchor[4];
+    float dza = anchor[5];
+    float ra = anchor[6];
+    float diagonal = sqrtf(dxa * dxa + dya * dya);
+    float be0 = box_encodings[0] * diagonal + xa;
+    float be1 = box_encodings[1] * diagonal + ya;
+    float be2 = box_encodings[2] * dza + za;
+    float be3 = expf(box_encodings[3]) * dxa;
+    float be4 = expf(box_encodings[4]) * dya;
+    float be5 = expf(box_encodings[5]) * dza;
+    float be6 = box_encodings[6] + ra;
+    float yaw;
+    int dir_label = dir_cls_input[dir_cls_offset] > dir_cls_input[dir_cls_offset + 1] ? 0 : 1;
+    float period = 2.0f * float(M_PI) / num_dir_bins;
+    float val = be6 - dir_offset;
+    float dir_rot = val - floor(val / period + dir_limit_offset) * period;
+    yaw = dir_rot + dir_offset + period * dir_label;
+    int resCount = atomicAdd(object_counter + batch_idx, 1);
+    float *data = bndbox_output + (batch_idx * max_box_num + resCount) * 9;
+    data[0] = be0;
+    data[1] = be1;
+    data[2] = be2;
+    data[3] = be3;
+    data[4] = be4;
+    data[5] = be5;
+    data[6] = yaw;
+    data[7] = dev_cls[0];
+    data[8] = dev_cls[1];
+  }
+}
+
+
+void  decodeBbox3DLaunch(
+  const int batch_size,
+  const float *cls_input,
+  const float *box_input,
+  const float *dir_cls_input,
+  float *anchors,
+  float *anchors_bottom_height,
+  float *bndbox_output,
+  int *object_counter,
+  const float min_x_range,
+  const float max_x_range,
+  const float min_y_range,
+  const float max_y_range,
+  const int feature_x_size,
+  const int feature_y_size,
+  const int num_anchors,
+  const int num_classes,
+  const int num_box_values,
+  const float score_thresh,
+  const float dir_offset,
+  const float dir_limit_offset,
+  const int num_dir_bins,
+  cudaStream_t stream)
+{
+  int bev_size = batch_size * feature_x_size * feature_y_size;
+  dim3 threads (num_anchors);
+  dim3 blocks (bev_size);
+  postprocess_kernal<<<blocks, threads, 0, stream>>>
+                (cls_input,
+                 box_input,
+                 dir_cls_input,
+                 anchors,
+                 anchors_bottom_height,
+                 bndbox_output,
+                 object_counter,
+                 min_x_range,
+                 max_x_range,
+                 min_y_range,
+                 max_y_range,
+                 feature_x_size,
+                 feature_y_size,
+                 num_anchors,
+                 num_classes,
+                 num_box_values,
+                 score_thresh,
+                 dir_offset,
+                 dir_limit_offset,
+                 num_dir_bins);
+  checkCudaErrors(cudaGetLastError());
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/extractFgScores.cu

@@ -0,0 +1,81 @@
+/*
+ * 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"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T>
+pluginStatus_t extractFgScores_gpu(cudaStream_t stream, int N, int A, int H, int W, const void* scores, void* fgScores)
+{
+    // Copy all the objectness scores for one batch
+    size_t size = A * H * W * sizeof(T);
+    for (int n = 0; n < N; n++)
+    {
+        // Find out the starting pointer of the objectness scores in the input
+        size_t offset_ld = (n * 2 + 1) * A * H * W;
+        // Find out the starting pointer of the objectness scores in the output
+        size_t offset_st = n * A * H * W;
+        CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size, cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+template <typename T>
+pluginStatus_t extractFgScores_cpu(int N,
+                                  int A,
+                                  int H,
+                                  int W,
+                                  const void* scores,
+                                  void* fgScores)
+{
+    size_t size = A * H * W * sizeof(T);
+    for (int n = 0; n < N; n++)
+    {
+        size_t offset_ld = (n * 2 + 1) * A * H * W;
+        size_t offset_st = n * A * H * W;
+        memcpy(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size);
+    }
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t extractFgScores(cudaStream_t stream,
+                              const int N,
+                              const int A,
+                              const int H,
+                              const int W,
+                              const DataType t_scores,
+                              const DLayout_t l_scores,
+                              const void* scores,
+                              const DataType t_fgScores,
+                              const DLayout_t l_fgScores,
+                              void* fgScores)
+{
+    if (l_fgScores != NCHW || l_scores != NCHW)
+        return STATUS_BAD_PARAM;
+
+    if (t_fgScores != DataType::kFLOAT)
+        return STATUS_BAD_PARAM;
+
+    if (t_scores != DataType::kFLOAT)
+        return STATUS_BAD_PARAM;
+
+    return extractFgScores_gpu<float>(stream, N, A, H, W, scores, fgScores);
+}
+} // namespace plugin
+} // namespace nvinfer1

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

plugin/common/kernels/generateAnchors.cu

@@ -0,0 +1,102 @@
+/*
+ * 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 <cstdio>
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t generateAnchors_cpu(
+    int numRatios, float* ratios, int numScales, float* scales, int baseSize, float* anchors)
+{
+#ifdef DEBUG
+    DEBUG_PRINTF("Generating Anchors with:\n");
+    DEBUG_PRINTF("Scales:");
+    for (int s = 0; s < numScales; ++s)
+    {
+        DEBUG_PRINTF("%f\t", scales[s]);
+    }
+    DEBUG_PRINTF("\n");
+    DEBUG_PRINTF("Ratios:");
+    for (int r = 0; r < numRatios; ++r)
+    {
+        DEBUG_PRINTF("%f\t", ratios[r]);
+    }
+    DEBUG_PRINTF("\n");
+#endif
+
+    if ((numScales <= 0) || (numRatios <= 0) || (baseSize <= 0))
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    // Generate parameters for numRatios * numScales general anchor boxes
+    for (int r = 0; r < numRatios; ++r)
+    {
+        for (int s = 0; s < numScales; ++s)
+        {
+            int id = r * numScales + s;
+            float scale = scales[s];
+            float ratio = ratios[r];
+            float bs = baseSize;
+            float ws = round(sqrt((float) (bs * bs) / ratio));
+            float hs = round(ws * ratio);
+            // Width: bs / sqrt(ratio) * scale
+            // Height: bs * sqrt(ratio) * scale
+            ws *= scale;
+            hs *= scale;
+
+            // x_anchor_ctr
+            /*
+             * This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
+             * Because the center of anchor box in the original input raw image scale will not be dependent on this.
+             */
+            anchors[id * 4] = (bs - 1) / 2;
+            // y_anchor_ctr
+            /*
+             * This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
+             * Because the center of anchor box in the original input raw image scale will not be dependent on this.
+             */
+            anchors[id * 4 + 1] = (bs - 1) / 2;
+            // w_anchor
+            anchors[id * 4 + 2] = ws;
+            // h_anchor
+            anchors[id * 4 + 3] = hs;
+        }
+    }
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t generateAnchors(cudaStream_t stream,
+                              int numRatios,
+                              float* ratios,
+                              int numScales,
+                              float* scales,
+                              int baseSize,
+                              float* anchors)
+{
+    // Each anchor box has 4 parameters
+    int ac = numRatios * numScales * 4;
+    float* anchors_cpu;
+    CSC(cudaMallocHost((void**) &anchors_cpu, sizeof(float) * ac), STATUS_FAILURE);
+    pluginStatus_t status = generateAnchors_cpu(numRatios, ratios, numScales, scales, baseSize, anchors_cpu);
+    CSC(cudaMemcpyAsync(anchors, anchors_cpu, sizeof(float) * ac, cudaMemcpyHostToDevice, stream), STATUS_FAILURE);
+    CSC(cudaFreeHost(anchors_cpu), STATUS_FAILURE);
+    return status;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/gridAnchorLayer.cu

@@ -0,0 +1,132 @@
+/*
+ * 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 "reducedMathPlugin.h"
+#include <iostream>
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+using nvinfer1::plugin::ReducedDivisor;
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void gridAnchorKernel(const GridAnchorParameters param,
+    const int numAspectRatios, ReducedDivisor divObj, const float* widths, const float* heights, float* outputData)
+{
+    // output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
+    const int dim = param.H * param.W * numAspectRatios;
+    /*
+     * Parameters used to calculate the bounding box coordinates back to input image scale
+     * Normally we calculate the anchorStride = image_input_size (in pixel) / feature_map_size
+     * Here we do not use image_input_size for the moment
+     * Instead we use 1.0
+     * The coordinates calculated are scaled by the input image size.
+     * Most of the coordinates will be in a range of [0, 1], except for the bounding box coordinates going outside of
+     * the image Every coordinate will go back to the pixel coordinates in the input image if being multiplied by
+     * image_input_size.
+     */
+    float anchorStrideH = (1.0F / param.H);
+    float anchorStrideW = (1.0F / param.W);
+    float anchorOffsetH = 0.5F * anchorStrideH;
+    float anchorOffsetW = 0.5F * anchorStrideW;
+
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= dim)
+        return;
+    int arId, currIndex;
+    divObj.divmod(tid, currIndex, arId);
+
+    const int w = currIndex % param.W;
+    const int h = currIndex / param.W;
+
+    // Center coordinates
+    float yC = h * anchorStrideH + anchorOffsetH;
+    float xC = w * anchorStrideW + anchorOffsetW;
+
+    // x_min, y_min
+    float xMin = xC - 0.5 * widths[arId];
+    float yMin = yC - 0.5 * heights[arId];
+
+    // x_max, y_max
+    float xMax = xC + 0.5 * widths[arId];
+    float yMax = yC + 0.5 * heights[arId];
+
+    outputData[tid * 4] = xMin;
+    outputData[tid * 4 + 1] = yMin;
+    outputData[tid * 4 + 2] = xMax;
+    outputData[tid * 4 + 3] = yMax;
+
+    // Remember to move the output cursor
+    float* output = outputData + dim * 4;
+
+    // Simply copying the variance
+    output[tid * 4] = param.variance[0];
+    output[tid * 4 + 1] = param.variance[1];
+    output[tid * 4 + 2] = param.variance[2];
+    output[tid * 4 + 3] = param.variance[3];
+}
+
+pluginStatus_t anchorGridInference(cudaStream_t stream, const GridAnchorParameters param, const int numAspectRatios,
+    const void* widths, const void* heights, void* outputData)
+{
+    const int dim = param.H * param.W * numAspectRatios;
+    ReducedDivisor divObj(numAspectRatios);
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t anchorGridInference(cudaStream_t stream, const GridAnchorParameters param, const int numAspectRatios,
+    const void* widths, const void* heights, void* outputData)
+{
+    const int dim = param.H * param.W * numAspectRatios;
+    ReducedDivisor divObj(numAspectRatios);
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/kernel.h

@@ -0,0 +1,261 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
+ */
+
+#ifndef TRT_KERNEL_H
+#define TRT_KERNEL_H
+
+#include "common/plugin.h"
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cuda_fp16.h>
+
+#define DEBUG_ENABLE 0
+namespace nvinfer1
+{
+namespace plugin
+{
+typedef enum
+{
+    NCHW = 0,
+    NC4HW = 1,
+    NC32HW = 2
+} DLayout_t;
+#ifndef TRT_RPNLAYER_H
+
+pluginStatus_t allClassNMS(cudaStream_t stream, int32_t num, int32_t num_classes, int32_t num_preds_per_class,
+    int32_t top_k, float nms_threshold, bool share_location, bool isNormalized, nvinfer1::DataType DT_SCORE,
+    nvinfer1::DataType DT_BBOX, void* bbox_data, void* beforeNMS_scores, void* beforeNMS_index_array,
+    void* afterNMS_scores, void* afterNMS_index_array, bool flipXY, float const score_shift, bool caffeSemantics);
+
+pluginStatus_t nmsInference(cudaStream_t stream, int32_t N, int32_t boxesSize, int32_t scoresSize, bool shareLocation,
+    int32_t backgroundLabelId, int32_t numPredsPerClass, int32_t numClasses, int32_t topK, int32_t keepTopK,
+    float scoreThreshold, float iouThreshold, nvinfer1::DataType DT_BBOX, void const* locData,
+    nvinfer1::DataType DT_SCORE, void const* confData, void* keepCount, void* nmsedBoxes, void* nmsedScores,
+    void* nmsedClasses, void* workspace, bool isNormalized = true, bool confSigmoid = false, bool clipBoxes = true,
+    int32_t scoreBits = 16, bool caffeSemantics = true);
+
+pluginStatus_t gatherTopDetections(cudaStream_t stream, bool shareLocation, int32_t numImages, int32_t numPredsPerClass,
+    int32_t numClasses, int32_t topK, int32_t keepTopK, nvinfer1::DataType DT_BBOX, nvinfer1::DataType DT_SCORE,
+    void const* indices, void const* scores, void const* bboxData, void* keepCount, void* topDetections,
+    float const scoreShift);
+
+size_t detectionForwardBBoxDataSize(int32_t N, int32_t C1, nvinfer1::DataType DT_BBOX);
+
+size_t detectionForwardBBoxPermuteSize(bool shareLocation, int32_t N, int32_t C1, nvinfer1::DataType DT_BBOX);
+
+size_t sortScoresPerClassWorkspaceSize(
+    int32_t num, int32_t num_classes, int32_t num_preds_per_class, nvinfer1::DataType DT_CONF);
+
+size_t sortScoresPerImageWorkspaceSize(int32_t num_images, int32_t num_items_per_image, nvinfer1::DataType DT_SCORE);
+
+pluginStatus_t sortScoresPerImage(cudaStream_t stream, int32_t num_images, int32_t num_items_per_image,
+    nvinfer1::DataType DT_SCORE, void* unsorted_scores, void* unsorted_bbox_indices, void* sorted_scores,
+    void* sorted_bbox_indices, void* workspace, int32_t score_bits);
+
+pluginStatus_t sortScoresPerClass(cudaStream_t stream, int32_t num, int32_t num_classes, int32_t num_preds_per_class,
+    int32_t background_label_id, float confidence_threshold, nvinfer1::DataType DT_SCORE, void* conf_scores_gpu,
+    void* index_array_gpu, void* workspace, int32_t const score_bits, float const score_shift);
+
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int32_t count);
+
+char const* cublasGetErrorString(nvinfer1::pluginInternal::cublasStatus_t error);
+
+pluginStatus_t permuteData(cudaStream_t stream, int32_t nthreads, int32_t num_classes, int32_t num_data,
+    int32_t num_dim, nvinfer1::DataType DT_DATA, bool confSigmoid, void const* data, void* new_data);
+
+size_t detectionForwardPreNMSSize(int32_t N, int32_t C2);
+
+size_t detectionForwardPostNMSSize(int32_t N, int32_t numClasses, int32_t topK);
+
+size_t normalizePluginWorkspaceSize(bool acrossSpatial, int32_t C, int32_t H, int32_t W);
+
+pluginStatus_t normalizeInference(cudaStream_t stream, nvinfer1::pluginInternal::cublasHandle_t handle,
+    bool acrossSpatial, bool channelShared, int32_t N, int32_t C, int32_t H, int32_t W, float eps, void const* scale,
+    void const* inputData, void* outputData, void* workspace);
+
+pluginStatus_t scatterNDInference(cudaStream_t stream, int32_t* outputDims, int32_t nOutputDims, int32_t sliceRank,
+    int32_t nRows, int32_t rowSize, int32_t CopySize, int32_t sizeOfElementInBytes, void const* index,
+    void const* updates, void const* data, void* output, void* workspace);
+
+pluginStatus_t priorBoxInference(cudaStream_t stream, nvinfer1::plugin::PriorBoxParameters param, int32_t H, int32_t W,
+    int32_t numPriors, int32_t numAspectRatios, void const* minSize, void const* maxSize, void const* aspectRatios,
+    void* outputData);
+
+pluginStatus_t lReLUInference(cudaStream_t stream, int32_t n, float negativeSlope, void const* input, void* output);
+
+pluginStatus_t reorgInference(cudaStream_t stream, int32_t batch, int32_t C, int32_t H, int32_t W, int32_t stride,
+    void const* input, void* output);
+
+pluginStatus_t anchorGridInference(cudaStream_t stream, nvinfer1::plugin::GridAnchorParameters param,
+    int32_t numAspectRatios, void const* aspectRatios, void const* scales, void* outputData);
+
+pluginStatus_t regionInference(cudaStream_t stream, int32_t batch, int32_t C, int32_t H, int32_t W, int32_t num,
+    int32_t coords, int32_t classes, bool hasSoftmaxTree, nvinfer1::plugin::softmaxTree const* smTree,
+    void const* input, void* output);
+
+// GENERATE ANCHORS
+// For now it takes host pointers - ratios and scales but
+// in GPU MODE anchors should be device pointer
+pluginStatus_t generateAnchors(cudaStream_t stream,
+    int32_t numRatios, // number of ratios
+    float* ratios,     // ratio array
+    int32_t numScales, // number of scales
+    float* scales,     // scale array
+    int32_t baseSize,  // size of the base anchor (baseSize x baseSize)
+    float* anchors);   // output anchors (numRatios x numScales)
+
+// BBD2P
+pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
+    int32_t N,                     // batch size
+    int32_t A,                     // number of anchors
+    int32_t H,                     // last feature map H
+    int32_t W,                     // last feature map W
+    int32_t featureStride,         // feature stride
+    float minBoxSize,              // minimum allowed box size before scaling
+    float const* imInfo,           // image info (nrows, ncols, image scale)
+    float const* anchors,          // input anchors
+    nvinfer1::DataType tDeltas,    // type of input deltas
+    DLayout_t lDeltas,             // layout of input deltas
+    void const* deltas,            // input deltas
+    nvinfer1::DataType tProposals, // type of output proposals
+    DLayout_t lProposals,          // layout of output proposals
+    void* proposals,               // output proposals
+    nvinfer1::DataType tScores,    // type of output scores
+    DLayout_t lScores,             // layout of output scores
+    void* scores);                 // output scores (the score associated with too small box will be set to -inf)
+
+// NMS
+pluginStatus_t nms(cudaStream_t stream,
+    int32_t N,                     // batch size
+    int32_t R,                     // number of ROIs (region of interest) per image
+    int32_t preNmsTop,             // number of proposals before applying NMS
+    int32_t nmsMaxOut,             // number of remaining proposals after applying NMS
+    float iouThreshold,            // IoU threshold
+    nvinfer1::DataType tFgScores,  // type of foreground scores
+    DLayout_t lFgScores,           // layout of foreground scores
+    void* fgScores,                // foreground scores
+    nvinfer1::DataType tProposals, // type of proposals
+    DLayout_t lProposals,          // layout of proposals
+    void const* proposals,         // proposals
+    void* workspace,               // workspace
+    nvinfer1::DataType tRois,      // type of ROIs
+    void* rois);                   // ROIs
+
+// WORKSPACE SIZES
+size_t proposalsForwardNMSWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+size_t proposalsForwardBboxWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W);
+
+size_t proposalForwardFgScoresWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W);
+
+size_t proposalsInferenceWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+size_t RPROIInferenceFusedWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+// PROPOSALS INFERENCE
+pluginStatus_t proposalsInference(cudaStream_t stream, int32_t N, int32_t A, int32_t H, int32_t W,
+    int32_t featureStride, int32_t preNmsTop, int32_t nmsMaxOut, float iouThreshold, float minBoxSize,
+    float const* imInfo, float const* anchors, nvinfer1::DataType tScores, DLayout_t lScores, void const* scores,
+    nvinfer1::DataType tDeltas, DLayout_t lDeltas, void const* deltas, void* workspace, nvinfer1::DataType tRois,
+    void* rois);
+
+// EXTRACT FG SCORES
+pluginStatus_t extractFgScores(cudaStream_t stream, int32_t N, int32_t A, int32_t H, int32_t W,
+    nvinfer1::DataType tScores, DLayout_t lScores, void const* scores, nvinfer1::DataType tFgScores,
+    DLayout_t lFgScores, void* fgScores);
+
+// ROI INFERENCE
+pluginStatus_t roiInference(cudaStream_t stream,
+    int32_t const R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    int32_t const N,        // Batch size
+    int32_t const C,        // Channels
+    int32_t const H,        // Input feature map H
+    int32_t const W,        // Input feature map W
+    int32_t const poolingH, // Output feature map H
+    int32_t const poolingW, // Output feature map W
+    float const spatialScale, nvinfer1::DataType const tRois, void const* rois, nvinfer1::DataType const tFeatureMap,
+    DLayout_t const lFeatureMap, void const* featureMap, nvinfer1::DataType const tTop, DLayout_t const lTop, void* top,
+    size_t deviceSmemSize);
+
+// ROI FORWARD
+pluginStatus_t roiForward(cudaStream_t stream,
+    int32_t R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    int32_t N,        // Batch size
+    int32_t C,        // Channels
+    int32_t H,        // Input feature map H
+    int32_t W,        // Input feature map W
+    int32_t poolingH, // Output feature map H
+    int32_t poolingW, // Output feature map W
+    float spatialScale, nvinfer1::DataType tRois, void const* rois, nvinfer1::DataType tFeatureMap,
+    DLayout_t lFeatureMap, void const* featureMap, nvinfer1::DataType tTop, DLayout_t lTop, void* top, int32_t* maxIds);
+
+// RP ROI Fused INFERENCE
+pluginStatus_t RPROIInferenceFused(cudaStream_t stream, int32_t N, int32_t A, int32_t C, int32_t H, int32_t W,
+    int32_t poolingH, int32_t poolingW, int32_t featureStride, int32_t preNmsTop, int32_t nmsMaxOut, float iouThreshold,
+    float minBoxSize, float spatialScale, float const* imInfo, float const* anchors, nvinfer1::DataType tScores,
+    DLayout_t lScores, void const* scores, nvinfer1::DataType tDeltas, DLayout_t lDeltas, void const* deltas,
+    nvinfer1::DataType tFeatureMap, DLayout_t lFeatureMap, void const* featureMap, void* workspace,
+    nvinfer1::DataType tRois, void* rois, nvinfer1::DataType tTop, DLayout_t lTop, void* top, size_t deviceSmemSize);
+
+// GENERATE ANCHORS CPU
+pluginStatus_t generateAnchors_cpu(
+    int32_t numRatios, float* ratios, int32_t numScales, float* scales, int32_t baseSize, float* anchors);
+
+int32_t cropAndResizeInference(cudaStream_t stream, int32_t n, void const* image, void const* rois, int32_t batch_size,
+    int32_t input_height, int32_t input_width, int32_t num_boxes, int32_t crop_height, int32_t crop_width,
+    int32_t depth, void* output);
+
+int32_t proposalInference_gpu(cudaStream_t stream, void const* rpn_prob, void const* rpn_regr, int32_t batch_size,
+    int32_t input_height, int32_t input_width, int32_t rpn_height, int32_t rpn_width, int32_t MAX_BOX_NUM,
+    int32_t RPN_PRE_NMS_TOP_N, float* ANCHOR_SIZES, int32_t anc_size_num, float* ANCHOR_RATIOS, int32_t anc_ratio_num,
+    float rpn_std_scaling, int32_t rpn_stride, float bbox_min_size, float nms_iou_threshold, void* workspace,
+    void* output);
+
+size_t _get_workspace_size(
+    int32_t N, int32_t anc_size_num, int32_t anc_ratio_num, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+void decodeBbox3DLaunch(int32_t const batch_size, float const* cls_input, float const* box_input,
+    float const* dir_cls_input, float* anchors, float* anchors_bottom_height, float* bndbox_output,
+    int32_t* object_counter, float const min_x_range, float const max_x_range, float const min_y_range,
+    float const max_y_range, int32_t const feature_x_size, int32_t const feature_y_size, int32_t const num_anchors,
+    int32_t const num_classes, int32_t const num_box_values, float const score_thresh, float const dir_offset,
+    float const dir_limit_offset, int32_t const num_dir_bins, cudaStream_t stream = 0);
+
+template <typename Element>
+int32_t pillarScatterKernelLaunch(int32_t batch_size, int32_t max_pillar_num, int32_t num_features,
+    Element const* pillar_features_data, uint32_t const* coords_data, uint32_t const* params_data, uint32_t featureX,
+    uint32_t featureY, Element* spatial_feature_data, cudaStream_t stream);
+
+void generateVoxels_launch(int32_t batch_size, int32_t max_num_points, float* points, uint32_t* points_size,
+    float min_x_range, float max_x_range, float min_y_range, float max_y_range, float min_z_range, float max_z_range,
+    float pillar_x_size, float pillar_y_size, float pillar_z_size, int32_t grid_y_size, int32_t grid_x_size,
+    int32_t num_point_values, int32_t max_points_per_voxel, uint32_t* mask, float* voxels, cudaStream_t stream);
+
+void generateBaseFeatures_launch(int32_t batch_size, uint32_t* mask, float* voxels, int32_t grid_y_size,
+    int32_t grid_x_size, uint32_t* pillar_num, int32_t max_pillar_num, int32_t max_points_per_voxel,
+    int32_t num_point_values, float* voxel_features, uint32_t* voxel_num_points, uint32_t* coords, cudaStream_t stream);
+
+int32_t generateFeatures_launch(int32_t batch_size, int32_t dense_pillar_num, float* voxel_features,
+    uint32_t* voxel_num_points, uint32_t* coords, uint32_t* params, float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z, uint32_t voxel_features_size, uint32_t max_points,
+    uint32_t max_voxels, uint32_t num_point_values, float* features, cudaStream_t stream);
+
+#endif // TRT_RPNLAYER_H
+} // namespace plugin
+} // namespace nvinfer1
+#endif

plugin/common/kernels/lReLU.cu

@@ -0,0 +1,49 @@
+/*
+ * 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"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__
+    void pReLUKernel(const int n, const float negativeSlope, const float* input, float* output)
+{
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x; i < n; i += gridDim.x * nthdsPerCTA)
+    {
+        output[i] = input[i] > 0 ? input[i] : input[i] * negativeSlope;
+    }
+}
+
+pluginStatus_t lReLUGPU(cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
+{
+    const int BS = 512;
+    const int GS = (n + BS - 1) / BS;
+    pReLUKernel<BS><<<GS, BS, 0, stream>>>(n, negativeSlope,
+                                           (const float*) input,
+                                           (float*) output);
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t lReLUInference(
+    cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
+{
+    return lReLUGPU(stream, n, negativeSlope, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/maskRCNNKernels.h

@@ -0,0 +1,276 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 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.
+ */
+
+#ifndef TRT_MASKRCNN_UTILS_H
+#define TRT_MASKRCNN_UTILS_H
+
+#include "NvInfer.h"
+#include "common/plugin.h"
+
+inline size_t nAlignUp(size_t x, size_t align)
+{
+    size_t mask = align - 1;
+    PLUGIN_ASSERT((align & mask) == 0); // power of 2
+    return (x + mask) & (~mask);
+}
+
+inline size_t nAlignDown(size_t x, size_t align)
+{
+    size_t mask = align - 1;
+    PLUGIN_ASSERT((align & mask) == 0); // power of 2
+    return (x) & (~mask);
+}
+
+inline size_t dimVolume(const nvinfer1::Dims& dims)
+{
+    size_t volume = 1;
+    for (int32_t i = 0; i < dims.nbDims; ++i)
+        volume *= dims.d[i];
+
+    return volume;
+}
+
+inline size_t typeSize(const nvinfer1::DataType type)
+{
+    switch (type)
+    {
+    case nvinfer1::DataType::kFLOAT: return sizeof(float);
+    case nvinfer1::DataType::kBF16: return sizeof(uint16_t);
+    case nvinfer1::DataType::kHALF: return sizeof(uint16_t);
+    case nvinfer1::DataType::kINT8: return sizeof(uint8_t);
+    case nvinfer1::DataType::kINT32: return sizeof(int32_t);
+    case nvinfer1::DataType::kINT64: return sizeof(int64_t);
+    case nvinfer1::DataType::kBOOL: return sizeof(bool);
+    case nvinfer1::DataType::kUINT8: return sizeof(uint8_t);
+    case nvinfer1::DataType::kFP8:
+    case nvinfer1::DataType::kINT4:
+    case nvinfer1::DataType::kFP4:
+    case nvinfer1::DataType::kE8M0: PLUGIN_FAIL("Unsupported data type");
+    }
+    return 0;
+}
+
+#define AlignMem(x) nAlignUp(x, 256)
+
+struct RefineNMSParameters
+{
+    int32_t backgroundLabelId, numClasses, keepTopK;
+    float scoreThreshold, iouThreshold;
+};
+
+struct RefineDetectionWorkSpace
+{
+    RefineDetectionWorkSpace(const int32_t batchSize, const int32_t sampleCount, const RefineNMSParameters& param,
+        const nvinfer1::DataType type);
+
+    RefineDetectionWorkSpace() = default;
+
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct ProposalWorkSpace
+{
+    ProposalWorkSpace(const int32_t batchSize, const int32_t inputCnt, const int32_t sampleCount,
+        const RefineNMSParameters& param, const nvinfer1::DataType type);
+
+    ProposalWorkSpace() = default;
+
+    nvinfer1::DimsHW preRefineScoreDims;
+    nvinfer1::DimsHW preRefineSortedScoreDims;
+    nvinfer1::DimsHW preRefineBboxDims;
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t tempStorageOffset = 0;
+    size_t preRefineScoreOffset = 0;
+    size_t preRefineSortedScoreOffset = 0;
+    size_t preRefineBboxOffset = 0;
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct MultilevelProposeROIWorkSpace
+{
+    MultilevelProposeROIWorkSpace(const int32_t batchSize, const int32_t inputCnt, const int32_t sampleCount,
+        const RefineNMSParameters& param, const nvinfer1::DataType type);
+
+    MultilevelProposeROIWorkSpace() = default;
+
+    nvinfer1::DimsHW preRefineSortedScoreDims;
+    nvinfer1::DimsHW preRefineBboxDims;
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t tempStorageOffset = 0;
+    size_t preRefineSortedScoreOffset = 0;
+    size_t preRefineBboxOffset = 0;
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct ConcatTopKWorkSpace
+{
+    ConcatTopKWorkSpace(
+        const int32_t batchSize, const int32_t concatCnt, const int32_t topK, const nvinfer1::DataType inType);
+
+    ConcatTopKWorkSpace() = default;
+
+    nvinfer1::DimsHW concatedScoreDims;
+    nvinfer1::DimsHW concatedBBoxDims;
+    nvinfer1::DimsHW sortedScoreDims;
+    nvinfer1::DimsHW sortedBBoxDims;
+
+    size_t tempStorageOffset = 0;
+    size_t concatedScoreOffset = 0;
+    size_t concatedBBoxOffset = 0;
+    size_t sortedScoreOffset = 0;
+    size_t sortedBBoxOffset = 0;
+    size_t totalSize = 0;
+};
+
+cudaError_t RefineBatchClassNMS(cudaStream_t stream, int32_t N, int32_t 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);
+
+cudaError_t DetectionPostProcess(cudaStream_t stream, int32_t N, int32_t 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);
+
+cudaError_t proposalRefineBatchClassNMS(cudaStream_t stream, int32_t N,
+    int32_t inputCnt, // candidate anchors
+    int32_t samples,  // preNMS_topK
+    nvinfer1::DataType dtype, const RefineNMSParameters& param, const ProposalWorkSpace& proposalOffset,
+    void* workspace, const void* inScores, const void* inDelta, const void* inCountValid, const void* inAnchors,
+    void* outProposals);
+
+// inScores: [N, anchorsCnt, 1]
+// inDelta: [N, anchorsCnt, 4]
+// outScores: [N, topK, 1]
+// outBbox: [N, topK, 4]
+cudaError_t MultilevelPropose(cudaStream_t stream, int32_t N,
+    int32_t inputCnt, // candidate anchors number among feature map
+    int32_t samples,  // pre nms cnt
+    const float* regWeight, const float inputHeight, const float inputWidth, nvinfer1::DataType dtype,
+    const RefineNMSParameters& param, const MultilevelProposeROIWorkSpace& proposalOffset, void* workspace,
+    const void* inScore, const void* inDelta, void* inCountValid, const void* inAnchors, void* outScores,
+    void* outBbox);
+
+// inScores: [N, topK, 1] * featureCnt
+// inBboxes: [N, topK, 4] * featureCnt
+// outProposals: [N, topK, 4]
+cudaError_t ConcatTopK(cudaStream_t stream, int32_t N, int32_t featureCnt, int32_t topK, nvinfer1::DataType dtype,
+    void* workspace, const ConcatTopKWorkSpace& spaceOffset, void** inScores, void** inBBox, void* outProposals);
+
+cudaError_t DecodeBBoxes(cudaStream_t stream, int32_t N,
+    int32_t 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, nvinfer1::DataType dtype);
+
+cudaError_t ApplyDelta2Bboxes(cudaStream_t stream, int32_t N,
+    int32_t 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);
+
+struct xy_t
+{
+    int32_t y;
+    int32_t x;
+
+    xy_t()
+        : y(0)
+        , x(0)
+    {
+    }
+    xy_t(int32_t y_, int32_t x_)
+        : y(y_)
+        , x(x_)
+    {
+    }
+};
+// PYRAMID ROIALIGN
+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* pooled, xy_t const poolDims);
+
+cudaError_t roiAlignHalfCenter(cudaStream_t stream, int32_t batchSize, int32_t featureCount, int32_t roiCount,
+    float firstThreshold,
+
+    int32_t inputHeight, int32_t inputWidth, const void* rois, const void* const layers[], const xy_t* layerDims,
+
+    void* pooled, const xy_t poolDims, const nvinfer1::DataType dtype);
+
+// RESIZE NEAREST
+void resizeNearest(dim3 grid, dim3 block, cudaStream_t stream, int32_t nbatch, float scale, int2 osize,
+    float const* idata, int32_t istride, int32_t ibatchstride, float* odata, int32_t ostride, int32_t obatchstride);
+// SPECIAL SLICE
+void specialSlice(cudaStream_t stream, int32_t batch_size, int32_t boxes_cnt, const void* idata, void* odata);
+
+#endif // TRT_MASKRCNN_UTILS_H

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

plugin/common/kernels/normalizeLayer.cu

@@ -0,0 +1,294 @@
+/*
+ * 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 "common/cublasWrapper.h"
+
+using namespace nvinfer1::pluginInternal;
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define CUBLAS_CHECK(condition)                                                                                        \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cublasStatus_t status = condition;                                                                             \
+        if (status != CUBLAS_STATUS_SUCCESS)                                                                           \
+        {                                                                                                              \
+            printf("%s %d CUBLAS FAIL %s\n", __FILE__, __LINE__, cublasGetErrorString(status));                        \
+        }                                                                                                              \
+    } while (0)
+
+size_t normalizePluginWorkspaceSize(bool acrossSpatial, int C, int H, int W)
+{
+    if (acrossSpatial)
+        return sizeof(float) * C * H * W;
+    else
+        return (size_t) 0;
+}
+
+template <unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void normalizeNotAcrossSpatialKernel(
+        const bool channelShared,
+        const int N,
+        const int C,
+        const int H,
+        const int W,
+        const float eps,
+        const float* scale,
+        float* inputData,
+        float* outputData)
+{
+    const int dim = C * H * W;
+    const int spatialDim = H * W;
+    const int tile = 32;
+    const int numTile = (spatialDim + tile - 1) / tile;
+    for (int n = blockIdx.x; n < N * numTile; n += gridDim.x)
+    {
+        float* input = inputData + (n / numTile) * dim;
+        float* output = outputData + (n / numTile) * dim;
+        __shared__ float sum[tile];
+        float localsum = 0.0F;
+        for (int i = threadIdx.x; i < tile; i += nthds_per_cta)
+        {
+            sum[i] = 0.0F;
+        }
+        __syncthreads();
+        for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+        {
+            int row = i / tile;
+            int col = (n % numTile) * tile + i % tile;
+            float data = 0.0F;
+            if (col < spatialDim)
+                data = input[row * spatialDim + col];
+            localsum += data * data;
+        }
+        atomicAdd(&sum[threadIdx.x & 31], localsum);
+        __syncthreads();
+        for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+        {
+            int row = i / tile;
+            int col = (n % numTile) * tile + i % tile;
+            if (col < spatialDim)
+            {
+                int offset = row * spatialDim + col;
+                output[offset] = input[offset] / sqrt(sum[threadIdx.x & 31] + eps);
+            }
+        }
+        if (channelShared)
+        {
+            for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+            {
+                int row = i / tile;
+                int col = (n % numTile) * tile + i % tile;
+                if (col < spatialDim)
+                    output[row * spatialDim + col] *= scale[0];
+            }
+        }
+        else
+        {
+            for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+            {
+                int row = i / tile;
+                int col = (n % numTile) * tile + i % tile;
+                if (col < spatialDim)
+                    output[row * spatialDim + col] *= scale[row];
+            }
+        }
+    }
+}
+
+pluginStatus_t normalizeNotAcrossSpatialGpu(
+    cudaStream_t stream,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData)
+{
+    const int BS = 128;
+    const int GS = 256;
+    // assumes warp size == 32
+    PLUGIN_ASSERT(BS % 32 == 0);
+    normalizeNotAcrossSpatialKernel<BS><<<GS, BS, 0, stream>>>(
+        channelShared, N, C, H, W, eps, (const float*) scale, (float*) inputData, (float*) outputData);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+__global__ void squareKernel(
+    const int n,
+    const float* x,
+    float* y)
+{
+    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
+         i < n; i += gridDim.x * blockDim.x)
+    {
+        y[i] = x[i] * x[i];
+    }
+}
+
+__global__ void scalChannelKernel(
+    const int n,
+    const int spatialDim,
+    const float* inputData,
+    const float* scale,
+    float* outputData)
+{
+    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
+         i < n; i += gridDim.x * blockDim.x)
+    {
+        // scale factors are indepedent across different channels
+        // scale[i / spatialDim]: find the right scale factor for specific channels
+        outputData[i] = inputData[i] / scale[i / spatialDim];
+    }
+}
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t normalizeInference(
+    cudaStream_t stream,
+    cublasHandle_t handle,
+    const bool acrossSpatial,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData,
+    void* workspace)
+{
+    CublasWrapper& mCublasWrapper = getCublasWrapper();
+    const int dim = C * H * W;
+    // Normalization is conducted for each sample from the batch indepdently
+    if (acrossSpatial)
+    {
+        float* input = (float*) const_cast<void*>(inputData);
+        float* output = (float*) outputData;
+        float* buffer = (float*) workspace;
+        for (int n = 0; n < N; ++n)
+        {
+            // Take the square of each element in the input
+            squareKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, input, buffer);
+            float normsqr = 0.0F;
+            // Sum up all the squared elements
+            CUBLAS_CHECK(mCublasWrapper.cublasSasum(handle, dim, buffer, 1, &normsqr));
+            // Make a copy of the input to the output
+            CUBLAS_CHECK(mCublasWrapper.cublasScopy(handle, dim, input, 1, output, 1));
+            // Calculate the inverse of the square root of the sum
+            // Use eps to prevent being divided by zero
+            normsqr = 1 / sqrt(normsqr + eps);
+            // Scale all the outputs by normsqr
+            CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, &normsqr, output, 1));
+            // If channel shared is true, scale all the outputs
+            if (channelShared)
+            {
+                CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, (float*) scale, output, 1));
+            }
+            // Use different scale factors for different channels
+            else
+            {
+                // scale the output according to channels
+                scalChannelKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, H * W, output, (float*) scale, output);
+            }
+            // Move cursors
+            input += dim;
+            output += dim;
+        }
+        return STATUS_SUCCESS;
+    }
+    // Normalization ignoring the batch
+    else
+    {
+        return normalizeNotAcrossSpatialGpu(stream, channelShared, N, C, H, W, eps, scale, inputData, outputData);
+    }
+}
+} // namespace plugin
+} // namespace nvinfer1
+
+pluginStatus_t normalizeInference(
+    cudaStream_t stream,
+    cublasHandle_t handle,
+    const bool acrossSpatial,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData,
+    void* workspace)
+{
+    const int dim = C * H * W;
+    // Normalization is conducted for each sample from the batch indepdently
+    if (acrossSpatial)
+    {
+        float* input = (float*) const_cast<void*>(inputData);
+        float* output = (float*) outputData;
+        float* buffer = (float*) workspace;
+        CublasWrapper& mCublasWrapper = getCublasWrapper();
+        for (int n = 0; n < N; ++n)
+        {
+            // Take the square of each element in the input
+            squareKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, input, buffer);
+            float normsqr = 0.0F;
+            // Sum up all the squared elements
+            CUBLAS_CHECK(mCublasWrapper.cublasSasum(handle, dim, buffer, 1, &normsqr));
+            // Make a copy of the input to the output
+            CUBLAS_CHECK(mCublasWrapper.cublasScopy(handle, dim, input, 1, output, 1));
+            // Calculate the inverse of the square root of the sum
+            // Use eps to prevent being divided by zero
+            normsqr = 1 / sqrt(normsqr + eps);
+            // Scale all the outputs by normsqr
+            CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, &normsqr, output, 1));
+            // If channel shared is true, scale all the outputs
+            if (channelShared)
+            {
+                CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, (float*) scale, output, 1));
+            }
+            // Use different scale factors for different channels
+            else
+            {
+                // scale the output according to channels
+                scalChannelKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, H * W, output, (float*) scale, output);
+            }
+            // Move cursors
+            input += dim;
+            output += dim;
+        }
+        return STATUS_SUCCESS;
+    }
+    // Normalization ignoring the batch
+    else
+    {
+        return normalizeNotAcrossSpatialGpu(stream, channelShared, N, C, H, W, eps, scale, inputData, outputData);
+    }
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/permuteData.cu

@@ -0,0 +1,114 @@
+/*
+ * 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 <array>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+template <typename Dtype, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta) __global__ void permuteData_kernel(const int nthreads, const int num_classes,
+    const int num_data, const int num_dim, bool confSigmoid, const Dtype* data, Dtype* new_data)
+{
+    // data format: [batch_size, num_data, num_classes, num_dim]
+    for (int index = blockIdx.x * nthds_per_cta + threadIdx.x; index < nthreads; index += nthds_per_cta * gridDim.x)
+    {
+        const int i = index % num_dim;
+        const int c = (index / num_dim) % num_classes;
+        const int d = (index / num_dim / num_classes) % num_data;
+        const int n = index / num_dim / num_classes / num_data;
+        const int new_index = ((n * num_classes + c) * num_data + d) * num_dim + i;
+        float result = data[index];
+        if (confSigmoid)
+            result = exp(result) / (1 + exp(result));
+
+        new_data[new_index] = result;
+    }
+    // new data format: [batch_size, num_classes, num_data, num_dim]
+}
+
+template <typename Dtype>
+pluginStatus_t permuteData_gpu(
+    cudaStream_t stream,
+    const int nthreads,
+    const int num_classes,
+    const int num_data,
+    const int num_dim,
+    bool confSigmoid,
+    const void* data,
+    void* new_data)
+{
+    const int BS = 512;
+    const int GS = (nthreads + BS - 1) / BS;
+    permuteData_kernel<Dtype, BS><<<GS, BS, 0, stream>>>(nthreads, num_classes, num_data, num_dim, confSigmoid,
+                                                         (const Dtype*) data, (Dtype*) new_data);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// permuteData LAUNCH CONFIG
+typedef pluginStatus_t (*pdFunc)(cudaStream_t, const int, const int, const int, const int, bool, const void*, void*);
+
+struct pdLaunchConfig
+{
+    DataType t_data;
+    pdFunc function;
+
+    pdLaunchConfig(DataType t_data)
+        : t_data(t_data)
+        , function(nullptr)
+    {
+    }
+    pdLaunchConfig(DataType t_data, pdFunc function)
+        : t_data(t_data)
+        , function(function)
+    {
+    }
+    bool operator==(pdLaunchConfig const& other) const
+    {
+        return t_data == other.t_data;
+    }
+};
+
+static std::array<pdLaunchConfig, 2> pdLCOptions = {
+    pdLaunchConfig(DataType::kFLOAT, permuteData_gpu<float>), pdLaunchConfig(DataType::kHALF, permuteData_gpu<__half>)};
+
+pluginStatus_t permuteData(cudaStream_t stream, const int nthreads, const int num_classes, const int num_data,
+    const int num_dim, const DataType DT_DATA, bool confSigmoid, const void* data, void* new_data)
+{
+    pdLaunchConfig lc = pdLaunchConfig(DT_DATA);
+    for (unsigned i = 0; i < pdLCOptions.size(); ++i)
+    {
+        if (lc == pdLCOptions[i])
+        {
+            DEBUG_PRINTF("permuteData kernel %d\n", i);
+            return pdLCOptions[i].function(stream,
+                                         nthreads,
+                                         num_classes,
+                                         num_data,
+                                         num_dim,
+                                         confSigmoid,
+                                         data,
+                                         new_data);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/pillarScatterKernels.cu

@@ -0,0 +1,114 @@
+/*
+ * 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 <iostream>
+#include <cuda_runtime_api.h>
+#include <cuda_fp16.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+const int PILLARS_PER_BLOCK = 64;
+const int PILLAR_FEATURE_SIZE = 64;
+
+template <typename Element>
+__global__ void scatterBEV_kernel(const Element *pillar_features_data,
+          const unsigned int *coords_data, const unsigned int *params_data,
+          unsigned int featureX, unsigned int featureY,
+          Element *spatial_feature_data)
+{
+    int pillar_idx = blockIdx.x * PILLARS_PER_BLOCK + threadIdx.x;
+    int valid_pillars_inBlock = PILLARS_PER_BLOCK;
+    const int num_pillars = params_data[0];
+    int valid_blocks = (num_pillars+PILLARS_PER_BLOCK-1)/PILLARS_PER_BLOCK;
+    if(blockIdx.x >= valid_blocks) return;
+    if(blockIdx.x == (valid_blocks-1)) {
+      valid_pillars_inBlock = num_pillars % PILLARS_PER_BLOCK;
+    }
+    valid_pillars_inBlock = (valid_pillars_inBlock==0) ? PILLARS_PER_BLOCK : valid_pillars_inBlock;
+    __shared__ Element pillarSM[PILLARS_PER_BLOCK][PILLAR_FEATURE_SIZE]; //pillar*64
+    for (int i = 0; i < valid_pillars_inBlock; i++)
+    {
+      pillarSM[i][threadIdx.x] = pillar_features_data[ (blockIdx.x * PILLARS_PER_BLOCK +i)*PILLAR_FEATURE_SIZE + threadIdx.x];
+    }
+    __syncthreads();
+    if(pillar_idx >= num_pillars) return;
+    int4 coord = ((const int4 *)coords_data)[pillar_idx];
+    int x = coord.w;
+    int y = coord.z;
+    for (int i = 0; i < PILLAR_FEATURE_SIZE; i++)
+    {
+      spatial_feature_data[i*featureY*featureX + y*featureX + x] = pillarSM[threadIdx.x][i];
+    }
+}
+
+template <typename Element>
+int pillarScatterKernelLaunch(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const Element *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  Element *spatial_feature_data,
+  cudaStream_t stream)
+{
+    dim3 blocks( (featureX*featureY+PILLARS_PER_BLOCK-1)/PILLARS_PER_BLOCK);
+    dim3 threads(PILLARS_PER_BLOCK);
+    for (int b = 0; b < batch_size; b++) {
+      scatterBEV_kernel<Element><<<blocks, threads, 0, stream>>>
+        (pillar_features_data + b*max_pillar_num*num_features,
+         coords_data + b*max_pillar_num*4,
+         params_data + b,
+         featureX,
+         featureY,
+         spatial_feature_data + b*num_features*featureX*featureY
+        );
+      auto err = cudaGetLastError();
+      if (cudaSuccess != err) {
+          fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
+          return -1;
+      }
+    }
+    return 0;
+}
+
+template int pillarScatterKernelLaunch<half>(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const half *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  half *spatial_feature_data,
+  cudaStream_t stream);
+
+template int pillarScatterKernelLaunch<float>(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const float *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  float *spatial_feature_data,
+  cudaStream_t stream);
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/priorBoxLayer.cu

@@ -0,0 +1,202 @@
+/*
+ * 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 "NvInferPluginUtils.h"
+#include "common/kernels/kernel.h"
+#include "reducedMathPlugin.h"
+#include <iostream>
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+using nvinfer1::plugin::ReducedDivisor;
+
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void priorBoxKernel(PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const float* minSize, const float* maxSize,
+    const float* aspectRatios, float* outputData)
+{
+    // output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
+    const int dim = H * W * numPriors;
+    const bool haveMaxSize = param.numMaxSize > 0;
+    const int dimAR = (haveMaxSize ? 1 : 0) + numAspectRatios;
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+         i < dim; i += gridDim.x * nthdsPerCTA)
+    {
+        const int w = (i / numPriors) % W;
+        const int h = (i / numPriors) / W;
+        // Usually param.offset == 0.5
+        // Calucate the center of prior box at the input image scale
+        const float centerX = (w + param.offset) * param.stepW;
+        const float centerY = (h + param.offset) * param.stepH;
+        // Minimum size index
+        const int minSizeId = (i / dimAR) % param.numMinSize;
+        // Aspect ratio index
+        const int arId = i % dimAR;
+        // Generate square pior box of aspect ratio of 1.0, edge length of minSize[minSizeId]
+        if (arId == 0)
+        {
+            const float boxW = minSize[minSizeId];
+            const float boxH = boxW;
+            float x, y, z, w;
+            // Calculate [x_topleft, y_topleft, x_bottomright, y_bottomright]
+            // Coordinates were scaled to [0, 1] against the width or height of the original input image
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            // If we decided to clip the prior box make sure all the bounding box are inside the original input image
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            // Copy the bounding box coordinates to output
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+        // If have maxSize
+        // Generate square pior box for aspect ratio of 1.0, edge length of sqrt(minSize[minSizeId] * maxSize[minSizeId])
+        // Described in SSD paper page 6
+        else if (haveMaxSize && arId == 1)
+        {
+            const float boxW = sqrt(minSize[minSizeId] * maxSize[minSizeId]);
+            const float boxH = boxW;
+            float x, y, z, w;
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+        // Generate other bouding boxes with aspect ratios of not one.
+        else
+        {
+            const int arOffset = haveMaxSize ? arId - 1 : arId; // skip aspectRatios[0] which is 1
+            const float boxW = minSize[minSizeId] * sqrt(aspectRatios[arOffset]);
+            const float boxH = minSize[minSizeId] / sqrt(aspectRatios[arOffset]);
+            float x, y, z, w;
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+    }
+    // Simply copy variance to from the parameter to output
+    float* output = outputData + dim * 4;
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+         i < dim; i += gridDim.x * nthdsPerCTA)
+    {
+        float x, y, z, w;
+        x = param.variance[0];
+        y = param.variance[1];
+        z = param.variance[2];
+        w = param.variance[3];
+        output[i * 4] = x;
+        output[i * 4 + 1] = y;
+        output[i * 4 + 2] = z;
+        output[i * 4 + 3] = w;
+    }
+}
+
+pluginStatus_t priorBoxGpu(
+    cudaStream_t stream,
+    const PriorBoxParameters param,
+    const int H,
+    const int W,
+    const int numPriors,
+    const int numAspectRatios,
+    const void* minSize,
+    const void* maxSize,
+    const void* aspectRatios,
+    void* outputData)
+{
+    const int dim = H * W * numPriors;
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        priorBoxKernel<BS><<<GS, BS, 0, stream>>>(param, H, W, numPriors, numAspectRatios,
+                                                  (const float*) minSize, (const float*) maxSize,
+                                                  (const float*) aspectRatios, (float*) outputData);
+        CSC(cudaGetLastError(), STATUS_FAILURE);
+        return STATUS_SUCCESS;
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        priorBoxKernel<BS><<<GS, BS, 0, stream>>>(param, H, W, numPriors, numAspectRatios,
+                                                  (const float*) minSize, (const float*) maxSize,
+                                                  (const float*) aspectRatios, (float*) outputData);
+        CSC(cudaGetLastError(), STATUS_FAILURE);
+        return STATUS_SUCCESS;
+    }
+}
+
+pluginStatus_t priorBoxInference(cudaStream_t stream, const PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const void* minSize, const void* maxSize, const void* aspectRatios,
+    void* outputData)
+{
+    PLUGIN_ASSERT(param.numMaxSize >= 0);
+    if (param.numMaxSize)
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios, minSize, maxSize, aspectRatios, outputData);
+    else
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios,
+                           minSize, nullptr, aspectRatios, outputData);
+}
+
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t priorBoxInference(cudaStream_t stream, const PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const void* minSize, const void* maxSize, const void* aspectRatios,
+    void* outputData)
+{
+    PLUGIN_ASSERT(param.numMaxSize >= 0);
+    if (param.numMaxSize)
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios, minSize, maxSize, aspectRatios, outputData);
+    else
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios,
+                           minSize, nullptr, aspectRatios, outputData);
+}
+} // namespace nvinfer1
+} // namespace plugin

plugin/common/kernels/proposalKernel.cu

@@ -0,0 +1,693 @@
+/*
+ * 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 "NvInfer.h"
+#include "common/plugin.h"
+#include <algorithm>
+#include <assert.h>
+#include <cub/cub.cuh>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <functional>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <vector>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define PLUGIN_CHECK_CUDA(call)                                                                                        \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cudaError_t status = call;                                                                                     \
+        if (status != cudaSuccess)                                                                                     \
+        {                                                                                                              \
+            return status;                                                                                             \
+        }                                                                                                              \
+    } while (0)
+template <typename TFloat>
+struct Bbox
+{
+    TFloat x1, y1, x2, y2;
+};
+
+typedef nvinfer1::DataType DType_t;
+
+typedef enum
+{
+    NCHW = 0,
+    NC4HW = 1
+} DLayout_t;
+
+typedef pluginStatus_t frcnnStatus_t;
+
+#define DEBUG_RPN_ENABLE 0
+
+#define FRCNN_ASSERT_PARAM(exp)                                                         \
+    do                                                                                  \
+    {                                                                                   \
+        if (!(exp))                                                                     \
+        {                                                                               \
+            DEBUG_FPRINTF(stderr, "Bad param - " #exp ", %s:%d\n", __FILE__, __LINE__); \
+            return STATUS_BAD_PARAM;                                                    \
+        }                                                                               \
+    } while (0)
+
+#define DEBUG_FPRINTF(...)        \
+    do                            \
+    {                             \
+        if (DEBUG_RPN_ENABLE)     \
+        {                         \
+            fprintf(__VA_ARGS__); \
+        }                         \
+    } while (0)
+
+#define CUDA_MEM_ALIGN 256
+
+unsigned int hash(const void* array_, size_t size);
+int8_t* alignPtr(int8_t* ptr, uintptr_t to);
+__global__ void setOffset(int stride, int size, int* output);
+frcnnStatus_t nms(cudaStream_t stream,
+    const int N,
+    const int R,
+    const int preNmsTop,
+    const int nmsMaxOut,
+    const float iouThreshold,
+    const DType_t t_fgScores,
+    const DLayout_t l_fgScores,
+    void* fgScores,
+    const DType_t t_proposals,
+    const DLayout_t l_proposals,
+    const void* proposals,
+    void* workspace,
+    const DType_t t_rois,
+    void* rois);
+int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize);
+
+
+template <typename TFloat>
+__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
+{
+    TFloat left = max(a.x1, b.x1), right = min(a.x2, b.x2);
+    TFloat top = max(a.y1, b.y1), bottom = min(a.y2, b.y2);
+    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.x2 - a.x1 + (TFloat) 1) * (a.y2 - a.y1 + (TFloat) 1);
+    TFloat Sb = (b.x2 - b.x1 + (TFloat) 1) * (b.y2 - b.y1 + (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.x1;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.y1;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.x2;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.y2;
+            }
+        }
+
+        __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].x1;
+                    filtered[cnt * 4 + 1] = t[i].y1;
+                    filtered[cnt * 4 + 2] = t[i].x2;
+                    filtered[cnt * 4 + 3] = t[i].y2;
+                    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>
+frcnnStatus_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)
+    };
+    FRCNN_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;
+}
+// }}}
+
+// NMS GPU {{{
+template <typename T_SCORES, typename T_ROIS>
+frcnnStatus_t nmsGpu(cudaStream_t stream, const int N, const int R, const int preNmsTop, const int nmsMaxOut,
+    const float iouThreshold, void* fgScores, const void* proposals, void* workspace, void* rois)
+{
+    // CUB's bug workaround:
+    // To work properly for large batch size CUB segmented sort needs ridiculous
+    // workspace alignment.
+    constexpr uintptr_t kALIGNMENT = 1 << 20;
+
+    int8_t* vworkspace = alignPtr((int8_t*) workspace, kALIGNMENT);
+    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)));
+    frcnnStatus_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, kALIGNMENT);
+    // 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, kALIGNMENT);
+    Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
+    vworkspace = (int8_t*) (proposalsOut + N * R);
+    vworkspace = alignPtr(vworkspace, kALIGNMENT);
+    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;
+}
+// }}}
+
+typedef frcnnStatus_t (*nmsFun)(cudaStream_t,
+                                const int,   // N
+                                const int,   // R
+                                const int,   // preNmsTop
+                                const int,   // nmsMaxOut
+                                const float, // iouThreshold
+                                void*,       // fgScores
+                                const void*, // proposals,
+                                void*,       // workspace,
+                                void*);      // rois
+
+struct nmsLaunchConfig
+{
+    DType_t t_fgScores;
+    DLayout_t l_fgScores;
+    DType_t t_proposals;
+    DLayout_t l_proposals;
+    DType_t t_rois;
+    nmsFun function;
+
+    nmsLaunchConfig(DType_t t_fgScores,
+                    DLayout_t l_fgScores,
+                    DType_t t_proposals,
+                    DLayout_t l_proposals,
+                    DType_t 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(
+        DType_t t_fgScores, DLayout_t l_fgScores, DType_t t_proposals, DLayout_t l_proposals, DType_t t_rois)
+        : t_fgScores(t_fgScores)
+        , l_fgScores(l_fgScores)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_rois(t_rois)
+        , function(nullptr)
+    {
+    }
+
+    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);
+    }
+};
+
+static std::vector<nmsLaunchConfig> nmsLCVec;
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+
+__global__ void _inverse_transform_gpu(const float* RPN_prob, const float* RPN_regr, int N,
+                                       int INPUT_H, int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE,
+                                       float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size,
+                                       float* fg_scores, float* proposal_out)
+{
+    int nthreads = N * RPN_H * RPN_W * anc_size_num * anc_ratio_num;
+    int num_ancs = anc_size_num * anc_ratio_num;
+
+    for (int out_idx = threadIdx.x + blockDim.x * blockIdx.x; out_idx < nthreads;
+            out_idx += blockDim.x * gridDim.x)
+    {
+        //input RPN_regr: (N, A4, H, W), thread: (N, A, H, W)
+        int idx = out_idx;
+        int w = idx % RPN_W;
+        idx /= RPN_W;
+        int h = idx % RPN_H;
+        idx /= RPN_H;
+        int a = idx % num_ancs;
+        int n = idx / num_ancs;
+        // normalize by RPN_STD_SCALING
+        int ptr_1 = ((((n * num_ancs) + a) * 4) * RPN_H + h) * RPN_W + w;
+        int ptr_2 = ((((n * num_ancs) + a) * 4 + 1) * RPN_H + h) * RPN_W + w;
+        int ptr_3 = ((((n * num_ancs) + a) * 4 + 2) * RPN_H + h) * RPN_W + w;
+        int ptr_4 = ((((n * num_ancs) + a) * 4 + 3) * RPN_H + h) * RPN_W + w;
+        float tx = RPN_regr[ptr_1] / RPN_STD_SCALING;
+        float ty = RPN_regr[ptr_2] / RPN_STD_SCALING;
+        float tw = RPN_regr[ptr_3] / RPN_STD_SCALING;
+        float th = RPN_regr[ptr_4] / RPN_STD_SCALING;
+        // do inverse transform
+        int ar = a % anc_ratio_num;
+        int as = a / anc_ratio_num;
+        float anchor_w = ANCHOR_SIZES[as] * ANCHOR_RATIOS[ar];
+        float anchor_h = ANCHOR_SIZES[as] / ANCHOR_RATIOS[ar];
+        float anchor_cx = (w + 0.5f) * RPN_STRIDE;
+        float anchor_cy = (h + 0.5f) * RPN_STRIDE;
+        float cx1 = anchor_cx + anchor_w * tx;
+        float cy1 = anchor_cy + anchor_h * ty;
+        float w1 = __expf(tw) * anchor_w;
+        float h1 = __expf(th) * anchor_h;
+        tx = cx1 - w1 / 2.0f;
+        ty = cy1 - h1 / 2.0f;
+        tw = w1;
+        th = h1;
+        tw += tx;
+        th += ty;
+        // clip to min
+        tx = (tx >= 0.0f) ? tx : 0.0f;
+        ty = (ty >= 0.0f) ? ty : 0.0f;
+        tw = (tw >= 0.0f) ? tw : 0.0f;
+        th = (th >= 0.0f) ? th : 0.0f;
+        //clip to max
+        tx = (tx <= INPUT_W - 1.0f) ? tx : (INPUT_W - 1.0f);
+        ty = (ty <= INPUT_H - 1.0f) ? ty : (INPUT_H - 1.0f);
+        tw = (tw <= INPUT_W - 1.0f) ? tw : (INPUT_W - 1.0f);
+        th = (th <= INPUT_H - 1.0f) ? th : (INPUT_H - 1.0f);
+        // filter out small boxes by setting the confidence to -inf
+        int ininf = 0xff800000;
+        float ninf = *(float*) &ininf;
+
+        if (tw - tx <= bbox_min_size || th - ty <= bbox_min_size)
+        {
+            fg_scores[out_idx] = ninf;
+        }
+
+        // copy to proposal_out, output shape: (N, A, H, W, 4)
+        proposal_out[out_idx * 4] = tx;
+        proposal_out[out_idx * 4 + 1] = ty;
+        proposal_out[out_idx * 4 + 2] = tw;
+        proposal_out[out_idx * 4 + 3] = th;
+    }
+}
+
+
+
+cudaError_t _inverse_transform_wrapper(const float* RPN_prob, const float* RPN_regr, int N, int INPUT_H,
+                                int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE,  float* ANCHOR_SIZES,
+                                int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size, float* fg_scores,
+                                float* proposal_out, cudaStream_t stream)
+{
+    const int block_size = 1024;
+    const int grid_size = (N * anc_size_num * anc_ratio_num * RPN_H * RPN_W + block_size - 1) /
+                          (block_size);
+    _inverse_transform_gpu <<< grid_size, block_size, 0, stream>>> (RPN_prob, RPN_regr, N, INPUT_H,
+            INPUT_W, RPN_H, RPN_W, RPN_STD_SCALING, RPN_STRIDE, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS,
+            anc_ratio_num, bbox_min_size, fg_scores, proposal_out);
+
+    return cudaGetLastError();
+}
+
+size_t _proposalsForwardNMSWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return N * A * H * W * 5 * 5 * sizeof(float) + (1 << 22);
+}
+
+size_t _proposalsForwardBboxWorkspaceSize(int N, int A, int H, int W)
+{
+    return N * A * H * W * 4 * sizeof(float);
+}
+
+
+size_t _proposalForwardFgScoresWorkspaceSize(int N, int A, int H, int W)
+{
+    return N * A * H * W * sizeof(float);
+}
+
+
+size_t anchors_buf_size(int anc_size_num, int anc_ratio_num)
+{
+    return (anc_size_num + anc_ratio_num) * sizeof(float);
+}
+
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int count);
+
+size_t _get_workspace_size(int N,
+                           int anc_size_num,
+                           int anc_ratio_num,
+                           int H,
+                           int W,
+                           int nmsMaxOut)
+{
+    size_t wss[4];
+    int A = anc_size_num * anc_ratio_num;
+    wss[0] = _proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+    wss[1] = _proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    wss[2] = _proposalForwardFgScoresWorkspaceSize(N, A, H, W);
+    wss[3] = anchors_buf_size(anc_size_num, anc_ratio_num);
+    return calculateTotalWorkspaceSize(wss, 4);
+}
+
+
+
+template <typename T>
+frcnnStatus_t extractFgScores_gpu(cudaStream_t stream,
+                                  int N,
+                                  int A,
+                                  int H,
+                                  int W,
+                                  const void* scores,
+                                  void* fgScores)
+{
+    //TODO custom kernel for this
+    size_t size = A * H * W * sizeof(T);
+
+    for (int n = 0; n < N; n++)
+    {
+        size_t offset_ld = n * A * H * W;
+        size_t offset_st = n * A * H * W;
+        CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size,
+                            cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+cudaError_t _copy_anchors_to_gpu(cudaStream_t stream, float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS,
+    int anc_ratio_num, void* anchor_size_buf)
+{
+    PLUGIN_CHECK_CUDA(cudaMemcpyAsync(anchor_size_buf, static_cast<void*>(ANCHOR_SIZES), sizeof(float) * anc_size_num,
+        cudaMemcpyHostToDevice, stream));
+    PLUGIN_CHECK_CUDA(cudaMemcpyAsync(static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num),
+        static_cast<void*>(ANCHOR_RATIOS), sizeof(float) * anc_ratio_num, cudaMemcpyHostToDevice, stream));
+
+    return cudaSuccess;
+}
+
+
+__global__ void _normalize_rois_kernel(float* roi_after_nms, int nthreads, int width, int height)
+{
+    for(int i = threadIdx.x + blockDim.x * blockIdx.x; i < nthreads; i += blockDim.x * gridDim.x)
+    {
+        float x1 = roi_after_nms[i * 4];
+        float y1 = roi_after_nms[i * 4 + 1];
+        float x2 = roi_after_nms[i * 4 + 2];
+        float y2 = roi_after_nms[i * 4 + 3];
+        roi_after_nms[i * 4] = y1 / (height - 1.0f);
+        roi_after_nms[i * 4 + 1] = x1 / (width - 1.0f);
+        roi_after_nms[i * 4 + 2] = y2 / (height - 1.0f);
+        roi_after_nms[i * 4 + 3] = x2 / (width - 1.0f);
+    }
+}
+
+
+
+cudaError_t _normalize_rois(float* roi_after_nms, int n, int max_box_num, int input_width,
+                     int input_height, cudaStream_t stream)
+{
+    const int block_size = 1024;
+    const int grid_size = (n * max_box_num + block_size - 1) / block_size;
+    _normalize_rois_kernel <<< grid_size, block_size, 0, stream>>>(roi_after_nms, n * max_box_num,
+            input_width, input_height);
+
+    return cudaGetLastError();
+}
+
+
+int proposalInference_gpu(
+    cudaStream_t stream,
+    const void* rpn_prob,
+    const void* rpn_regr,
+    int batch_size,
+    int input_height,
+    int input_width,
+    int rpn_height,
+    int rpn_width,
+    int MAX_BOX_NUM,
+    int RPN_PRE_NMS_TOP_N,
+    float* ANCHOR_SIZES,
+    int anc_size_num,
+    float* ANCHOR_RATIOS,
+    int anc_ratio_num,
+    float rpn_std_scaling,
+    int rpn_stride,
+    float bbox_min_size,
+    float nms_iou_threshold,
+    void * workspace,
+    void* output)
+{
+    size_t nmsWorkspaceSize = _proposalsForwardNMSWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
+                              rpn_height, rpn_width, MAX_BOX_NUM);
+    void* nmsWorkspace = workspace;
+    size_t proposalsSize = _proposalsForwardBboxWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
+                           rpn_height, rpn_width);
+    const DType_t t_proposals = nvinfer1::DataType::kFLOAT;
+    const DLayout_t l_proposals = NC4HW;
+    void* proposals = nextWorkspacePtr((int8_t*) nmsWorkspace, nmsWorkspaceSize);
+    void* fg_scores = nextWorkspacePtr((int8_t*) proposals, proposalsSize);
+    size_t fg_scores_size = _proposalForwardFgScoresWorkspaceSize(batch_size,
+                            anc_size_num * anc_ratio_num, rpn_height, rpn_width);
+    void* anchor_size_buf = nextWorkspacePtr((int8_t*) fg_scores, fg_scores_size);
+    void* anchor_ratio_buf = static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num);
+    frcnnStatus_t status;
+    PLUGIN_CHECK_CUDA(
+        _copy_anchors_to_gpu(stream, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS, anc_ratio_num, anchor_size_buf));
+
+    status = extractFgScores_gpu<float>(
+        stream, batch_size, anc_size_num * anc_ratio_num, rpn_height, rpn_width, rpn_prob, fg_scores);
+    PLUGIN_ASSERT(status == STATUS_SUCCESS);
+    PLUGIN_CHECK_CUDA(
+        _inverse_transform_wrapper(static_cast<const float*>(rpn_prob), static_cast<const float*>(rpn_regr), batch_size,
+            input_height, input_width, rpn_height, rpn_width, rpn_std_scaling, rpn_stride,
+            static_cast<float*>(anchor_size_buf), anc_size_num, static_cast<float*>(anchor_ratio_buf), anc_ratio_num,
+            bbox_min_size, static_cast<float*>(fg_scores), static_cast<float*>(proposals), stream));
+
+    status = nms(stream, batch_size, anc_size_num * anc_ratio_num * rpn_height * rpn_width, RPN_PRE_NMS_TOP_N,
+        MAX_BOX_NUM, nms_iou_threshold, nvinfer1::DataType::kFLOAT, NCHW, fg_scores, t_proposals, l_proposals,
+        proposals, workspace, nvinfer1::DataType::kFLOAT, output);
+
+    PLUGIN_ASSERT(status == STATUS_SUCCESS);
+
+    PLUGIN_CHECK_CUDA(
+        _normalize_rois(static_cast<float*>(output), batch_size, MAX_BOX_NUM, input_width, input_height, stream));
+
+    return STATUS_SUCCESS;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/proposalsForward.cu

@@ -0,0 +1,161 @@
+/*
+ * 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"
+
+using namespace nvinfer1;
+namespace nvinfer1
+{
+namespace plugin
+{
+// PROPOSALS INFERENCE
+pluginStatus_t proposalsInference(cudaStream_t stream, const int N, const int A, const int H, const int W,
+    const int featureStride, const int preNmsTop, const int nmsMaxOut, const float iouThreshold, const float minBoxSize,
+    const float* imInfo, const float* anchors, const DataType t_scores, const DLayout_t l_scores, const void* scores,
+    const DataType t_deltas, const DLayout_t l_deltas, const void* deltas, void* workspace, const DataType t_rois,
+    void* rois)
+{
+    /*
+     * N: batch size
+     * A: number of anchor boxes per grid cell on feature map
+     * H: height of feature map
+     * W: width of feature map
+     */
+    if (imInfo == NULL || anchors == NULL || scores == NULL || deltas == NULL || workspace == NULL || rois == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    DEBUG_PRINTF("&&&& IM INFO %u\n", hash(imInfo, N * 3 * sizeof(float)));
+    // anchors: anchor boxes
+    /*
+     * The following line of code looks somewhat incorrect because it sounds like we always have 9 fixed anchor boxes.
+     * The "corrected" implementation should be
+     * DEBUG_PRINTF("&&&& ANCHORS %u\n", A * 4 * sizeof(float)));
+     */
+    DEBUG_PRINTF("&&&& ANCHORS %u\n", hash(anchors, 9 * 4 * sizeof(float)));
+    // scores: objectness of each predicted bounding boxes
+    // 2: softmax, instead of sigmoid, was used for binary objectness classifcation in Faster R-CNN
+    DEBUG_PRINTF("&&&& SCORES  %u\n", hash(scores, N * A * 2 * H * W * sizeof(float)));
+    // deltas: predicted bounding box offsets
+    DEBUG_PRINTF("&&&& DELTAS  %u\n", hash(deltas, N * A * 4 * H * W * sizeof(float)));
+
+    size_t nmsWorkspaceSize = proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+
+    void* nmsWorkspace = workspace;
+
+    size_t proposalsSize = proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    const DataType t_proposals = nvinfer1::DataType::kFLOAT;
+    const DLayout_t l_proposals = NC4HW;
+    void* proposals = nextWorkspacePtr((int8_t*) nmsWorkspace, nmsWorkspaceSize);
+
+    const DataType t_fgScores = t_scores;
+    const DLayout_t l_fgScores = NCHW;
+    void* fgScores = nextWorkspacePtr((int8_t*) proposals, proposalsSize);
+
+    pluginStatus_t status;
+
+    /*
+     * Only the second probability value of the objectness (probability of being a object) from the scores will be extracted.
+     * Because the first probability (probability of not being a object) value is redundant.
+     */
+    status = extractFgScores(stream,
+                             N, A, H, W,
+                             t_scores, l_scores, scores,
+                             t_fgScores, l_fgScores, fgScores);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& FG SCORES %u\n", hash((void*) fgScores, N * A * H * W * sizeof(float)));
+    DEBUG_PRINTF("&&&& DELTAS %u\n", hash((void*) proposals, N * A * H * W * 4 * sizeof(float)));
+
+    /*
+     * Decode predicted bounding boxes.
+     * Decoded predicted bounding boxes were at the raw input image scale.
+     */
+    status = bboxDeltas2Proposals(stream,
+                                  N, A, H, W,
+                                  featureStride,
+                                  minBoxSize,
+                                  imInfo,
+                                  anchors,
+                                  t_deltas, l_deltas, deltas,
+                                  t_proposals, l_proposals, proposals,
+                                  t_fgScores, l_fgScores, fgScores);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& PROPOSALS %u\n", hash((void*) proposals, N * A * H * W * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& FG SCORES %u\n", hash((void*) fgScores, N * A * H * W * sizeof(float)));
+
+    /*
+     * Non maximum suppression using objectness scores to get the most representative bounding boxes (ROIs).
+     * The rois were at the feature map scale.
+     */
+    status = nms(stream,
+                 N,
+                 A * H * W,
+                 preNmsTop,
+                 nmsMaxOut,
+                 iouThreshold,
+                 t_fgScores, l_fgScores, fgScores,
+                 t_proposals, l_proposals, proposals,
+                 nmsWorkspace,
+                 t_rois, rois);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& ROIS %u\n", hash((void*) rois, N * nmsMaxOut * 4 * sizeof(float)));
+    return STATUS_SUCCESS;
+}
+
+// WORKSPACE SIZES
+size_t proposalsForwardNMSWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return N * A * H * W * 5 * 5 * sizeof(float) + (1 << 22);
+}
+
+size_t proposalsForwardBboxWorkspaceSize(int N,
+                                         int A,
+                                         int H,
+                                         int W)
+{
+    return N * A * H * W * 4 * sizeof(float);
+}
+size_t proposalForwardFgScoresWorkspaceSize(int N,
+                                            int A,
+                                            int H,
+                                            int W)
+{
+    return N * A * H * W * sizeof(float);
+}
+
+size_t proposalsInferenceWorkspaceSize(int N,
+                                       int A,
+                                       int H,
+                                       int W,
+                                       int nmsMaxOut)
+{
+    size_t wss[3];
+    wss[0] = proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+    wss[1] = proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    wss[2] = proposalForwardFgScoresWorkspaceSize(N, A, H, W);
+    return calculateTotalWorkspaceSize(wss, 3);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/reducedMathPlugin.h

@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+
+#ifndef _REDUCED_MATH_PLUGIN_H
+#define _REDUCED_MATH_PLUGIN_H
+#include <cstdint>
+// Dynamically strength-reduced div and mod
+//
+// Ideas taken from Sean Baxter's MGPU library.
+// These classes provide for reduced complexity division and modulus
+// on integers, for the case where the same divisor or modulus will
+// be used repeatedly.
+
+namespace nvinfer1
+{
+namespace plugin
+{
+namespace detail
+{
+
+void findDivisor(int32_t denom, uint32_t& mul_coeff, uint32_t& shift_coeff);
+
+__host__ __device__ __forceinline__ uint32_t umulhi(uint32_t x, uint32_t y)
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 100
+    return __umulhi(x, y);
+#else
+    uint64_t z = (uint64_t) x * (uint64_t) y;
+    return (uint32_t) (z >> 32);
+#endif
+}
+
+// This is a weird implementation that returns div_up(0,1)=0 but
+// div_up(0,2)=1 (wrong) -- just do not use it with a=0.
+__host__ __device__ inline int32_t div_up(int32_t a, int32_t b)
+{
+    return (a - 1) / b + 1;
+}
+
+} // end namespace detail
+
+class ReducedDivisor
+{
+public:
+    ReducedDivisor() {}
+    __host__ __forceinline__ ReducedDivisor(int32_t _y)
+        : y(_y)
+    {
+        detail::findDivisor(y, mul_coeff, shift_coeff);
+    }
+    __host__ __device__ __forceinline__ ReducedDivisor(uint32_t _mul_coeff, uint32_t _shift_coeff, int32_t _y)
+        : mul_coeff(_mul_coeff)
+        , shift_coeff(_shift_coeff)
+        , y(_y)
+    {
+    }
+    __host__ __device__ __forceinline__ int32_t div(int32_t x) const
+    {
+        // if dividing by 1, then findDivisor wouldn't have worked because
+        // mul_coeff would have had to be 2^32, which can't be represented,
+        // so we have to special case that one.
+        return (y != 1) ? detail::umulhi((uint32_t) x, mul_coeff) >> shift_coeff : x;
+    }
+    __host__ __device__ __forceinline__ int32_t mod(int32_t x) const
+    {
+        return x - (div(x) * y);
+    }
+    __host__ __device__ __forceinline__ void divmod(int32_t x, int32_t& q, int32_t& mod) const
+    {
+        q = div(x);
+        mod = x - (q * y);
+    }
+    __host__ __device__ __forceinline__ int32_t get() const
+    {
+        return y;
+    }
+    inline __host__ void get_mul_shift(uint32_t& mul, uint32_t& shift)
+    {
+        mul = mul_coeff;
+        shift = shift_coeff;
+    }
+
+protected:
+    uint32_t mul_coeff{};
+    uint32_t shift_coeff{};
+    int32_t y{};
+};
+
+} // namespace plugin
+
+} // namespace nvinfer1
+#endif /*_REDUCED_MATH_PLUGIN_H*/

plugin/common/kernels/regionForward.cu

@@ -0,0 +1,140 @@
+/*
+ * 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"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void softmaxKernel(const float* input, const int n, const int batch,
+    const int batchOffset, const int groups, const int groupOffset, const int stride, const float temp, float* output)
+{
+    int id = blockIdx.x * nthdsPerCTA + threadIdx.x;
+    if (id < batch * groups)
+    {
+        int b = id / groups;
+        int g = id % groups;
+        float sum = 0.;
+        // Initialze largest to be the smallest float number.
+        float largest = -3.402823466e+38;
+        int offset = b * batchOffset + g * groupOffset;
+        // Find the largest digits before softmax
+        for (int i = 0; i < n; ++i)
+        {
+            float val = input[i * stride + offset];
+            largest = (val > largest) ? val : largest;
+        }
+        // Softmax for a group of candidate classes
+        // Calculate exponentials
+        for (int i = 0; i < n; ++i)
+        {
+            /*
+             * Here we used a trick to prevent numeric overflow
+             * xm = max{x_1, x_2, ..., x_n}
+             * e^{x_1} / (e^{x_1} + e^{x_2} + e^{x_n}) = e^{x_1 - xm} / (e^{x_1 - xm} + e^{x_2 - xm} + e^{x_n - xm})
+             */
+            float e = exp(input[i * stride + offset] / temp - largest / temp);
+            sum += e;
+            output[i * stride + offset] = e;
+        }
+        // Normalize
+        for (int i = 0; i < n; ++i)
+            output[i * stride + offset] /= sum;
+    }
+}
+
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA)
+    __global__ void activateKernel(float* data,
+                                   const int range)
+{
+    int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+    // Sigmoid function
+    if (i < range)
+        data[i] = 1. / (1. + exp(-data[i]));
+}
+
+pluginStatus_t regionGPU(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int num,
+    const int coords,
+    const int classes,
+    const bool hasSoftmaxTree,
+    const nvinfer1::plugin::softmaxTree* smTree,
+    const float* input,
+    float* output)
+{
+    const int BS = 512;
+    const int GS1 = (2 * H * W + BS - 1) / BS;
+    const int GS2 = (H * W + BS - 1) / BS;
+    // Applying sigmoid activations
+    for (int b = 0; b < batch; ++b)
+    {
+        for (int n = 0; n < num; ++n)
+        {
+            // Apply sigmoid activation for the encoded center coordinates t_x, and t_y
+            int index = b * C * H * W + n * H * W * (coords + classes + 1);
+            activateKernel<BS><<<GS1, BS, 0, stream>>>(output + index, 2 * H * W);
+            /*
+             * Apply sigmoid for the encoded objectness t_o
+             * + 4 * H * W because we want to skip the first four coordinates t_x, t_y, t_w, t_h
+             * Chaning 4 * H * W to + coords * H * W will not make this function more general
+             * Since we already assumed the information layout in the channels of the input tensor
+             */
+            index = b * C * H * W + n * H * W * (coords + classes + 1) + 4 * H * W;
+            activateKernel<BS><<<GS2, BS, 0, stream>>>(output + index, H * W);
+        }
+    }
+    const int GS3 = (batch * num * H * W + BS - 1) / BS;
+    // Applying softmax activations
+    if (hasSoftmaxTree)
+    {
+        // Softmax for hierarchical classification
+        // The first 5 elements are t_x, t_y, t_w, t_h, t_o which we don't need to apply softmax activation
+        int count = 5;
+        // Only groups and groupSize information is useful for this plugin
+        // Applying softmax activation sequentially for each group of candidate classes
+        for (int i = 0; i < smTree->groups; ++i)
+        {
+            int groupSize = smTree->groupSize[i];
+            softmaxKernel<BS><<<GS3, BS, 0, stream>>>(input + count * H * W, groupSize, batch * num, (C * H * W / num), H * W, 1, H * W, 1., output + count * H * W);
+            count += groupSize;
+        }
+    }
+    else
+    {
+        // Softmax for non-hierarchical classificiation
+        softmaxKernel<BS><<<GS3, BS, 0, stream>>>(input + 5 * H * W, classes, batch * num, (C * H * W / num), H * W, 1, H * W, 1., output + 5 * H * W);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t regionInference(cudaStream_t stream, const int batch, const int C, const int H, const int W,
+    const int num, const int coords, const int classes, const bool hasSoftmaxTree,
+    const nvinfer1::plugin::softmaxTree* smTree, const void* input, void* output)
+{
+    PLUGIN_CHECK(cudaMemcpyAsync(output, input, batch * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice, stream));
+    return regionGPU(
+        stream, batch, C, H, W, num, coords, classes, hasSoftmaxTree, smTree, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/reorgForward.cu

@@ -0,0 +1,104 @@
+/*
+ * 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 "reducedMathPlugin.h"
+
+using namespace nvinfer1::plugin; // for ReducedDivisor
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA)
+    __global__ void reorgKernel(
+        const float* input, // input tensor of shape (batch, C, H, W)
+        const int volume,   // note that volumes of input and output tensors are the same
+        ReducedDivisor batch,
+        ReducedDivisor C,
+        ReducedDivisor H,
+        ReducedDivisor W,
+        ReducedDivisor C_out,
+        ReducedDivisor stride,
+        float* output) // output tensor of shape (batch, C * stride * stride, H / stride, W / stride)
+{
+    /*
+     * Reference
+     * https://github.com/pjreddie/darknet/blob/f6d861736038da22c9eb0739dca84003c5a5e275/src/blas_kernels.cu#L370
+     * https://github.com/pjreddie/darknet/blob/f6d861736038da22c9eb0739dca84003c5a5e275/src/blas.c#L9
+     */
+
+    // outIndex is row-major position of input coordinates
+    for (int outIndex = blockIdx.x * nthdsPerCTA + threadIdx.x; outIndex < volume; outIndex += nthdsPerCTA)
+    {
+        int i = outIndex;
+
+        // calculate output coordinates from outIndex
+        int outW, outH, outC;
+        W.divmod(i, i, outW);
+        H.divmod(i, i, outH);
+        C.divmod(i, i, outC);
+        int outN = i;
+
+        // calculate input coordinates based on output coordinates
+        // offset is [0, 1, ..., stride * stride - 1] = posH * stride + posW
+        int offset, inC, posH, posW;
+        C_out.divmod(outC, offset, inC);
+        stride.divmod(offset, posH, posW);
+        int inH = outH * stride.get() + posH;
+        int inW = outW * stride.get() + posW;
+        int inN = outN;
+
+        // inIndex is row-major position of input coordinates
+        int inIndex = inW + W.get() * stride.get() * (inH + H.get() * stride.get() * (inC + C_out.get() * inN));
+
+        output[outIndex] = input[inIndex];
+    }
+}
+
+pluginStatus_t reorgGPU(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int stride,
+    const float* input,
+    float* output)
+{
+    const int BS = 512;                    // number of threads in one block
+    const int volume = batch * C * H * W;  // size of input tensor
+    const int GS = (volume + BS - 1) / BS; // number of blocks to launch, calculated so global number of threads is >= volume
+
+    ReducedDivisor C_out(C / (stride * stride));
+    reorgKernel<BS><<<GS, BS, 0, stream>>>(input, volume, ReducedDivisor(batch), ReducedDivisor(C), ReducedDivisor(H), ReducedDivisor(W), C_out, ReducedDivisor(stride), output);
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t reorgInference(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int stride,
+    const void* input,
+    void* output)
+{
+    return reorgGPU(stream, batch, C, H, W, stride, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/roiPooling.cu

@@ -0,0 +1,353 @@
+/*
+ * 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 <algorithm>
+#include <array>
+#include <assert.h>
+#include <cfloat>
+#include <cstdio>
+#include <math.h>
+#include <stdio.h>
+namespace nvinfer1
+{
+namespace plugin
+{
+// This macro is to control shared memory usage. If set to 1, kernel loads the whole feature map
+// into shared memory for reuse; If set to 0, kernel loads data from global memory directly.
+// Roi pooling performance is data dependent. You can test which value is better to your data.
+// If all bboxes are very small, 0 is recommended, otherwise, shared memory will load many unused
+// data; If bboxes have many overlaps, 1 is recommended to avoid duplicate loads.
+// 1 requires larger shared memory size. It may fail if it is larger than CUDA allowed per-block
+// shared memory size upper bound. Then you have to use 0.
+#define ROIPOOLING_FEATURE_MAP_USE_SHMEM 1
+
+template <typename T>
+__device__ T getMax();
+
+template <>
+__device__ __forceinline__ int8_t getMax<int8_t>()
+{
+    return INT8_MAX;
+}
+
+template <>
+__device__ __forceinline__ float getMax<float>()
+{
+    return FLT_MAX;
+}
+
+// ROI POOLING FORWARD KERNEL
+template <typename DATA_T, typename ROI_T, bool INFER_ONLY, bool FM_IN_SMEM>
+__global__ void ROIPoolingForwardKernelAligned(int32_t ROICount, const ROI_T* rois,
+    int32_t N, // feature map size
+    int32_t C, // feature map size
+    int32_t H, // feature map size
+    int32_t W, // feature map size
+    const DATA_T* featureMap, const int32_t poolingH, const int32_t poolingW, const float spatialScale, DATA_T* top,
+    int32_t* maxIds, int32_t fmapStep)
+{
+    extern __shared__ float smem[];
+    DATA_T* feature_shr = (DATA_T*) &smem[0];
+    int* rois_shr = nullptr;
+    if (FM_IN_SMEM)
+    {
+        rois_shr = (int*) &feature_shr[H * W];
+    }
+    else
+    {
+        rois_shr = (int*) &feature_shr[0];
+        feature_shr = nullptr;
+    }
+
+    const int batch = blockIdx.x / C;
+    const int channel = blockIdx.x % C;
+
+    // load ROIs to shared memory
+    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
+    {
+        int offset = j << 2;
+        float4 roi = reinterpret_cast<float4*>(const_cast<float*>(rois))[batch * ROICount + j];
+        // spatialScale = 1.0 / featureStride
+        // Convert the coordinates to feature map scale
+        rois_shr[offset] = round(roi.x * spatialScale);     //roi_start_w
+        rois_shr[offset + 1] = round(roi.y * spatialScale); //roi_start_h
+        rois_shr[offset + 2] = round(roi.z * spatialScale) - round(roi.x * spatialScale); //roi_length_w
+        rois_shr[offset + 3] = round(roi.w * spatialScale) - round(roi.y * spatialScale); // roi_length_h
+    }
+
+    // NC/xHW
+    int fmapOffset = blockIdx.x / fmapStep * H * W * fmapStep + blockIdx.x % fmapStep;
+
+    // Assumes #CTAs is just enough to cover all channels of all blocks
+    const DATA_T* bottom_data_offset = featureMap + fmapOffset;
+    if (FM_IN_SMEM)
+    {
+        // load the current channel to the shared memory
+        for (int j = threadIdx.x; j < H * W; j += blockDim.x)
+        {
+            feature_shr[j] = bottom_data_offset[j * fmapStep];
+        }
+    }
+    __syncthreads();
+
+    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
+    {
+        const int offset = j << 2;
+        // Force malformed ROIs to be 1x1
+        int roi_start_w = rois_shr[offset];
+        int roi_start_h = rois_shr[offset + 1];
+        int roi_width = max(rois_shr[offset + 2] + 1, 1);
+        int roi_height = max(rois_shr[offset + 3] + 1, 1);
+        float bin_size_h = static_cast<float>(roi_height) / static_cast<float>(poolingH);
+        float bin_size_w = static_cast<float>(roi_width) / static_cast<float>(poolingW);
+
+        for (int ph = 0; ph < poolingH; ++ph)
+        {
+            for (int pw = 0; pw < poolingW; ++pw)
+            {
+                int hstart = static_cast<int>(floor(static_cast<float>(ph) * bin_size_h));
+                int wstart = static_cast<int>(floor(static_cast<float>(pw) * bin_size_w));
+                int hend = static_cast<int>(ceil(static_cast<float>(ph + 1) * bin_size_h));
+                int wend = static_cast<int>(ceil(static_cast<float>(pw + 1) * bin_size_w));
+
+                // Add roi offsets and clip to input boundaries
+                // In fact, clipping should be done in the RPN, but just in case...
+                hstart = min(max(hstart + roi_start_h, 0), H);
+                hend = min(max(hend + roi_start_h, 0), H);
+                wstart = min(max(wstart + roi_start_w, 0), W);
+                wend = min(max(wend + roi_start_w, 0), W);
+                bool is_empty = (hend <= hstart) || (wend <= wstart);
+
+                // Define an empty pooling region to be zero
+                DATA_T maxval = is_empty ? 0 : -getMax<DATA_T>();
+                int maxId = -1;
+                DATA_T data = 0;
+                for (int h = hstart; h < hend; ++h)
+                {
+                    for (int w = wstart; w < wend; ++w)
+                    {
+                        int index = h * W + w;
+                        if (FM_IN_SMEM)
+                        {
+                            data = feature_shr[index];
+                        }
+                        else
+                        {
+                            data = bottom_data_offset[index * fmapStep];
+                        }
+                        if (data > maxval)
+                        {
+                            maxval = data;
+                            maxId = index;
+                        }
+                    }
+                }
+                top[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxval;
+                if (!INFER_ONLY)
+                {
+                    maxIds[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxId;
+                }
+            } //for:pw
+        }     //for:ph
+    }         // for:j
+}
+
+template <typename DATA_T, DLayout_t DATA_L, typename ROI_T, bool INFER_ONLY>
+pluginStatus_t ROIPoolingForwardKernelAlignedLauncher(cudaStream_t stream,
+    const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    const int N,        // Batch size
+    const int C,        // Channels
+    const int H,        // Input feature map H
+    const int W,        // Input feature map W
+    const int poolingH, // Output feature map H
+    const int poolingW, // Output feature map W
+    const float spatialScale, const void* rois, const void* featureMap, void* top, int* maxIds, size_t deviceSmemSize)
+{
+    size_t roiShmemSize = (R / N) * 4 * sizeof(ROI_T);
+
+#if ROIPOOLING_FEATURE_MAP_USE_SHMEM
+    size_t shmemSize = H * W * sizeof(DATA_T) + roiShmemSize;
+    const bool fmap_in_shmem = true;
+#else
+    size_t shmemSize = roiShmemSize;
+    const bool fmap_in_shmem = false;
+#endif
+
+    if (shmemSize > deviceSmemSize)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    // in the aligned version of ROI Pooling R should always be a multiple of N
+    PLUGIN_ASSERT(R % N == 0);
+
+    // NC/xHW
+    int32_t fmapStep = 1;
+    switch(DATA_L)
+    {
+    case NCHW: fmapStep = 1; break;
+    case NC4HW:
+        fmapStep = 4;
+        PLUGIN_ASSERT((N * C) % 4 == 0);
+        break;
+    case NC32HW:
+        fmapStep = 32;
+        PLUGIN_ASSERT((N * C) % 32 == 0);
+        break;
+    default: PLUGIN_ASSERT(false);
+    }
+
+    if (shmemSize > 48 * 1024)
+    {
+        PLUGIN_CHECK(cudaFuncSetAttribute(&ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, true>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, shmemSize));
+    }
+    ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, fmap_in_shmem><<<N * C, 256, shmemSize, stream>>>(R / N,
+        (const ROI_T*) rois,
+        N, // feature map size
+        C, // feature map size
+        H, // feature map size
+        W, // feature map size
+        (const DATA_T*) featureMap, poolingH, poolingW, spatialScale, (DATA_T*) top, maxIds, fmapStep);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// ROI POOLING LAUNCH CONFIG
+
+typedef pluginStatus_t (*roiFwd)(cudaStream_t,
+                                const int,   //R, // TOTAL number of rois -> ~nmsMaxOut * N
+                                const int,   //N, // Batch size
+                                const int,   //C, // Channels
+                                const int,   //H, // Input feature map H
+                                const int,   //W, // Input feature map W
+                                const int,   //poolingH, // Output feature map H
+                                const int,   //poolingW, // Output feature map W
+                                const float, //spatialScale,
+                                const void*, //rois,
+                                const void*, //featureMap,
+                                void*,       //top
+                                int*,        //maxIds
+                                size_t);     //device shmem size
+
+// struct
+struct roiFwdLaunchConfig
+{
+    DataType t_rois;
+    DataType t_featureMap;
+    DLayout_t l_featureMap;
+    DataType t_top;
+    DLayout_t l_top;
+    bool inferOnly;
+    roiFwd function;
+
+    roiFwdLaunchConfig(
+        DataType t_rois, DataType t_featureMap, DLayout_t l_featureMap, DataType t_top, DLayout_t l_top, bool inferOnly)
+        : t_rois(t_rois)
+        , t_featureMap(t_featureMap)
+        , l_featureMap(l_featureMap)
+        , t_top(t_top)
+        , l_top(l_top)
+        , inferOnly(inferOnly)
+        , function(nullptr)
+    {
+    }
+
+    roiFwdLaunchConfig(DataType t_rois,
+                       DataType t_featureMap,
+                       DLayout_t l_featureMap,
+                       DataType t_top,
+                       DLayout_t l_top,
+                       bool inferOnly,
+                       roiFwd function)
+        : t_rois(t_rois)
+        , t_featureMap(t_featureMap)
+        , l_featureMap(l_featureMap)
+        , t_top(t_top)
+        , l_top(l_top)
+        , inferOnly(inferOnly)
+        , function(function)
+    {
+    }
+
+    bool operator==(roiFwdLaunchConfig const& other) const
+    {
+        return (t_rois == other.t_rois)
+            && (t_featureMap == other.t_featureMap)
+            && (l_featureMap == other.l_featureMap)
+            && (t_top == other.t_top)
+            && (l_top == other.l_top)
+            && (inferOnly == other.inferOnly);
+    }
+};
+
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+#define INT8 nvinfer1::DataType::kINT8
+static std::array<roiFwdLaunchConfig, 6> roiFwdLCOptions = {
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, false, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, false>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NCHW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NCHW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NC4HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC4HW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NC32HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC32HW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NC4HW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NC4HW, float, true>)};
+
+// ROI INFERENCE
+pluginStatus_t roiInference(cudaStream_t stream,
+                           const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
+                           const int N,        // Batch size
+                           const int C,        // Channels
+                           const int H,        // Input feature map H
+                           const int W,        // Input feature map W
+                           const int poolingH, // Output feature map H
+                           const int poolingW, // Output feature map W
+                           const float spatialScale,
+                           const nvinfer1::DataType t_rois,
+                           const void* rois,
+                           const nvinfer1::DataType t_featureMap,
+                           const DLayout_t l_featureMap,
+                           const void* featureMap,
+                           const nvinfer1::DataType t_top,
+                           const DLayout_t l_top,
+                           void* top,
+                           size_t deviceSmemSize)
+{
+    if (featureMap == NULL || rois == NULL || top == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    DEBUG_PRINTF("&&&& ROIS %u\n", hash(rois, R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& FMAP %u\n", hash(featureMap, N * C * H * W * sizeof(float)));
+
+    roiFwdLaunchConfig rflc = roiFwdLaunchConfig(t_rois, t_featureMap, l_featureMap, t_top, l_top, true);
+    ASSERT_PARAM(R > 0);
+
+    for (unsigned i = 0; i < roiFwdLCOptions.size(); i++)
+    {
+        if (rflc == roiFwdLCOptions[i])
+        {
+            DEBUG_PRINTF("$$$$ ROI KERNEL %d\n", i);
+            return roiFwdLCOptions[i].function(stream,
+                                                R, N, C, H, W, poolingH, poolingW,
+                                                spatialScale, rois, featureMap, top, NULL, deviceSmemSize);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/rproiInferenceFused.cu

@@ -0,0 +1,97 @@
+/*
+ * 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"
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t RPROIInferenceFused(cudaStream_t stream, const int N, const int A, const int C, const int H, const int W,
+    const int poolingH, const int poolingW, const int featureStride, const int preNmsTop, const int nmsMaxOut,
+    const float iouThreshold, const float minBoxSize, const float spatialScale, const float* imInfo,
+    const float* anchors, const DataType t_scores, const DLayout_t l_scores, const void* scores,
+    const DataType t_deltas, const DLayout_t l_deltas, const void* deltas, const DataType t_featureMap,
+    const DLayout_t l_featureMap, const void* featureMap, void* workspaces, const DataType t_rois, void* rois,
+    const DataType t_top, const DLayout_t l_top, void* top, size_t deviceSmemSize)
+{
+    if (imInfo == NULL || anchors == NULL || scores == NULL || deltas == NULL || featureMap == NULL
+        || workspaces == NULL || rois == NULL || top == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    pluginStatus_t status;
+    // Region proposal inference
+    // Getting the region of interests (ROIs) bounding box coordinates from region proposals using non maximum suppression (NMS)
+    status = proposalsInference(stream,
+                                N,
+                                A,
+                                H,
+                                W,
+                                featureStride,
+                                preNmsTop,
+                                nmsMaxOut,
+                                iouThreshold,
+                                minBoxSize,
+                                imInfo,
+                                anchors,
+                                t_scores,
+                                l_scores,
+                                scores,
+                                t_deltas,
+                                l_deltas,
+                                deltas,
+                                workspaces,
+                                t_rois,
+                                rois);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+    // ROI inference
+    // ROI pooling for ROIs
+    status = roiInference(stream,
+                          N * nmsMaxOut, // TOTAL number of rois -> ~nmsMaxOut * N
+                          N,             // Batch size
+                          C,             // Channels
+                          H,             // Input feature map H
+                          W,             // Input feature map W
+                          poolingH,      // Output feature map H
+                          poolingW,      // Output feature map W
+                          spatialScale,
+                          t_rois,
+                          rois,
+                          t_featureMap,
+                          l_featureMap,
+                          featureMap,
+                          t_top,
+                          l_top,
+                          top,
+                          deviceSmemSize);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    return STATUS_SUCCESS;
+}
+
+size_t RPROIInferenceFusedWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return proposalsInferenceWorkspaceSize(N, A, H, W, nmsMaxOut);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/saturate.h

@@ -0,0 +1,41 @@
+/*
+ * 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.
+ */
+#ifndef TRT_SATURATE_H
+#define TRT_SATURATE_H
+
+#include <array>
+#include <cuda_fp16.h>
+
+template <typename T_BBOX>
+__device__ T_BBOX saturate(T_BBOX v)
+{
+    return max(min(v, T_BBOX(1)), T_BBOX(0));
+}
+
+template <>
+inline __device__ __half saturate(__half v)
+{
+#if __CUDA_ARCH__ >= 800
+    return __hmax(__hmin(v, __half(1)), __half(0));
+#elif __CUDA_ARCH__ >= 530
+    return __hge(v, __half(1)) ? __half(1) : (__hle(v, __half(0)) ? __half(0) : v);
+#else
+    return max(min(v, float(1)), float(0));
+#endif
+}
+
+#endif // TRT_SATURATE_H

plugin/common/kernels/sortScoresPerClass.cu

@@ -0,0 +1,297 @@
+/*
+ * 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/cub_helper.h"
+#include "common/kernels/kernel.h"
+#include "cub/cub.cuh"
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+
+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
+}
+
+// overload for float
+inline __device__ float add_fb(const float & a, const float & b) {
+    return a + b;
+}
+
+inline __device__ bool ge_fb(const __half & a, const __half & b) {
+#if __CUDA_ARCH__ >= 530
+    return a >= b;
+#else
+    return __half2float(a) >= __half2float(b);
+#endif
+}
+
+template <typename T_SCORE, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void prepareSortData(
+        const int num,
+        const int num_classes,
+        const int num_preds_per_class,
+        const int background_label_id,
+        const float confidence_threshold,
+        T_SCORE* conf_scores_gpu,
+        T_SCORE* temp_scores,
+        T_SCORE score_shift,
+        int* temp_idx,
+        int* d_offsets)
+{
+    // Prepare scores data for sort
+    const int cur_idx = blockIdx.x * nthds_per_cta + threadIdx.x;
+    const int numPredsPerBatch = num_classes * num_preds_per_class;
+    T_SCORE clip_val = T_SCORE(float(score_shift) + 1.f - 1.f / 1024.f);
+    if (cur_idx < numPredsPerBatch)
+    {
+        const int class_idx = cur_idx / num_preds_per_class;
+        for (int i = 0; i < num; i++)
+        {
+            const int targetIdx = i * numPredsPerBatch + cur_idx;
+            const T_SCORE score = conf_scores_gpu[targetIdx];
+
+            // "Clear" background labeled score and index
+            // Because we do not care about background
+            if (class_idx == background_label_id)
+            {
+                // Set scores to 0
+                // Set label = -1
+                // add shift of 1.0 to normalize the score values
+                // to the range [1, 2). 
+                // add a constant shift to scores will not change the sort
+                // result, but will help reduce the computation because
+                // we only need to sort the mantissa part of the floating-point
+                // numbers
+                temp_scores[targetIdx] = score_shift;
+                temp_idx[targetIdx] = -1;
+                conf_scores_gpu[targetIdx] = score_shift;
+            }
+            // "Clear" scores lower than threshold
+            else
+            {
+                if (float(score) > confidence_threshold)
+                {
+                    // add shift of 1.0 to normalize the score values
+                    // to the range [1, 2). 
+                    // add a constant shift to scores will not change the sort
+                    // result, but will help reduce the computation because
+                    // we only need to sort the mantissa part of the floating-point
+                    // numbers
+                    temp_scores[targetIdx] = add_fb(score, score_shift);
+                    if (float(score_shift) > 0.f && (ge_fb(temp_scores[targetIdx], clip_val)))
+                        temp_scores[targetIdx] = clip_val;
+                    temp_idx[targetIdx] = cur_idx + i * numPredsPerBatch;
+                }
+                else
+                {
+                    // Set scores to 0
+                    // Set label = -1
+                    // add shift of 1.0 to normalize the score values
+                    // to the range [1, 2). 
+                    // add a constant shift to scores will not change the sort
+                    // result, but will help reduce the computation because
+                    // we only need to sort the mantissa part of the floating-point
+                    // numbers
+                    temp_scores[targetIdx] = score_shift;
+                    temp_idx[targetIdx] = -1;
+                    conf_scores_gpu[targetIdx] = score_shift;
+                    // TODO: HERE writing memory too many times
+                }
+            }
+
+            if ((cur_idx % num_preds_per_class) == 0)
+            {
+                const int offset_ct = i * num_classes + cur_idx / num_preds_per_class;
+                d_offsets[offset_ct] = offset_ct * num_preds_per_class;
+                // set the last element in d_offset
+                if (blockIdx.x == 0 && threadIdx.x == 0)
+                    d_offsets[num * num_classes] = num * numPredsPerBatch;
+            }
+        }
+    }
+}
+
+template <typename T_SCORE>
+pluginStatus_t sortScoresPerClass_gpu(
+    cudaStream_t stream,
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const int background_label_id,
+    const float confidence_threshold,
+    void* conf_scores_gpu,
+    void* index_array_gpu,
+    void* workspace,
+    const int score_bits,
+    const float score_shift
+)
+{
+    const int num_segments = num * num_classes;
+    void* temp_scores = workspace;
+    const int arrayLen = num * num_classes * num_preds_per_class;
+    void* temp_idx = nextWorkspacePtr((int8_t*) temp_scores, arrayLen * sizeof(T_SCORE));
+    void* d_offsets = nextWorkspacePtr((int8_t*) temp_idx, arrayLen * sizeof(int));
+    size_t cubOffsetSize = (num_segments + 1) * sizeof(int);
+    void* cubWorkspace = nextWorkspacePtr((int8_t*) d_offsets, cubOffsetSize);
+
+    const int BS = 512;
+    const int GS = (num_classes * num_preds_per_class + BS - 1) / BS;
+    // prepare the score, index, and offsets for CUB radix sort
+    // also normalize the scores to the range [1, 2)
+    // so we only need to sort the mantissa of floating-point numbers
+    // since their sign bit and exponential bits are identical
+    // we will subtract the 1.0 shift in gatherTopDetections()
+    prepareSortData<T_SCORE, BS><<<GS, BS, 0, stream>>>(num, num_classes, num_preds_per_class,
+                                                        background_label_id, confidence_threshold,
+                                                        (T_SCORE*) conf_scores_gpu,
+                                                        (T_SCORE*) temp_scores,
+                                                        T_SCORE(score_shift),
+                                                        (int*) temp_idx,
+                                                        (int*) d_offsets);
+
+    size_t temp_storage_bytes = cubSortPairsWorkspaceSize<T_SCORE, int>(arrayLen, num_segments);
+    size_t begin_bit = 0;
+    size_t end_bit = sizeof(T_SCORE) * 8;
+    if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10)
+    {
+        // only sort score_bits in 10 mantissa bits.
+        end_bit = 10;
+        begin_bit = end_bit - score_bits;
+    }
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        cubWorkspace, temp_storage_bytes,
+        (const T_SCORE*) (temp_scores), (T_SCORE*) (conf_scores_gpu),
+        (const int*) (temp_idx), (int*) (index_array_gpu),
+        arrayLen, num_segments,
+        (const int*) d_offsets, (const int*) d_offsets + 1,
+        begin_bit, end_bit,
+        stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// sortScoresPerClass LAUNCH CONFIG
+typedef pluginStatus_t (*sspcFunc)(cudaStream_t,
+                                const int,
+                                const int,
+                                const int,
+                                const int,
+                                const float,
+                                void*,
+                                void*,
+                                void*,
+                                const int,
+                                const float);
+struct sspcLaunchConfig
+{
+    DataType t_score;
+    sspcFunc function;
+
+    sspcLaunchConfig(DataType t_score)
+        : t_score(t_score)
+    {
+    }
+    sspcLaunchConfig(DataType t_score, sspcFunc function)
+        : t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(sspcLaunchConfig const& other) const
+    {
+        return t_score == other.t_score;
+    }
+};
+
+static std::array<sspcLaunchConfig, 2> sspcLCOptions = {
+    sspcLaunchConfig(DataType::kFLOAT, sortScoresPerClass_gpu<float>),
+    sspcLaunchConfig(DataType::kHALF, sortScoresPerClass_gpu<__half>)
+};
+
+pluginStatus_t sortScoresPerClass(
+    cudaStream_t stream,
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const int background_label_id,
+    const float confidence_threshold,
+    const DataType DT_SCORE,
+    void* conf_scores_gpu,
+    void* index_array_gpu,
+    void* workspace,
+    const int score_bits,
+    const float score_shift
+)
+{
+    sspcLaunchConfig lc = sspcLaunchConfig(DT_SCORE);
+    for (unsigned i = 0; i < sspcLCOptions.size(); ++i)
+    {
+        if (lc == sspcLCOptions[i])
+        {
+            DEBUG_PRINTF("sortScoresPerClass kernel %d\n", i);
+            return sspcLCOptions[i].function(stream,
+                                           num,
+                                           num_classes,
+                                           num_preds_per_class,
+                                           background_label_id,
+                                           confidence_threshold,
+                                           conf_scores_gpu,
+                                           index_array_gpu,
+                                           workspace,
+                                           score_bits,
+                                           score_shift);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+
+size_t sortScoresPerClassWorkspaceSize(
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const DataType DT_CONF)
+{
+    size_t wss[4];
+    const int arrayLen = num * num_classes * num_preds_per_class;
+    wss[0] = arrayLen * dataTypeSize(DT_CONF);      // temp scores
+    wss[1] = arrayLen * sizeof(int);                // temp indices
+    wss[2] = (num * num_classes + 1) * sizeof(int); // offsets
+    if (DT_CONF == DataType::kFLOAT)
+    {
+        wss[3] = cubSortPairsWorkspaceSize<float, int>(arrayLen, num * num_classes); // cub workspace
+    }
+    else if (DT_CONF == DataType::kHALF)
+    {
+        wss[3] = cubSortPairsWorkspaceSize<__half, int>(arrayLen, num * num_classes); // cub workspace
+    }
+    else
+    {
+        printf("SCORE type not supported\n");
+        return (size_t) -1;
+    }
+
+    return calculateTotalWorkspaceSize(wss, 4);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/sortScoresPerImage.cu

@@ -0,0 +1,151 @@
+/*
+ * 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/cub_helper.h"
+#include "common/kernels/kernel.h"
+#include "cub/cub.cuh"
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+
+template <typename T_SCORE>
+pluginStatus_t sortScoresPerImage_gpu(cudaStream_t stream, const int num_images, const int num_items_per_image,
+    void* unsorted_scores, void* unsorted_bbox_indices, void* sorted_scores, void* sorted_bbox_indices, void* workspace,
+    int score_bits)
+{
+    void* d_offsets = workspace;
+    void* cubWorkspace = nextWorkspacePtr((int8_t*) d_offsets, (num_images + 1) * sizeof(int));
+
+    setUniformOffsets(stream, num_images, num_items_per_image, (int*) d_offsets);
+
+    const int arrayLen = num_images * num_items_per_image;
+    size_t temp_storage_bytes = cubSortPairsWorkspaceSize<T_SCORE, int>(arrayLen, num_images);
+    size_t begin_bit = 0;
+    size_t end_bit = sizeof(T_SCORE) * 8;
+    if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10)
+    {
+        end_bit = 10;
+        begin_bit = end_bit - score_bits;
+    }
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        cubWorkspace, temp_storage_bytes,
+        (const T_SCORE*) (unsorted_scores), (T_SCORE*) (sorted_scores),
+        (const int*) (unsorted_bbox_indices), (int*) (sorted_bbox_indices),
+        arrayLen, num_images,
+        (const int*) d_offsets, (const int*) d_offsets + 1,
+        begin_bit, end_bit,
+        stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// sortScoresPerImage LAUNCH CONFIG
+typedef pluginStatus_t (*sspiFunc)(cudaStream_t,
+                                const int,
+                                const int,
+                                void*,
+                                void*,
+                                void*,
+                                void*,
+                                void*,
+                                int);
+struct sspiLaunchConfig
+{
+    DataType t_score;
+    sspiFunc function;
+
+    sspiLaunchConfig(DataType t_score)
+        : t_score(t_score)
+    {
+    }
+    sspiLaunchConfig(DataType t_score, sspiFunc function)
+        : t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(sspiLaunchConfig const& other) const
+    {
+        return t_score == other.t_score;
+    }
+};
+
+static std::array<sspiLaunchConfig, 2> sspiLCOptions = {
+    sspiLaunchConfig(DataType::kFLOAT, sortScoresPerImage_gpu<float>),
+    sspiLaunchConfig(DataType::kHALF, sortScoresPerImage_gpu<__half>),
+};
+
+pluginStatus_t sortScoresPerImage(
+    cudaStream_t stream,
+    const int num_images,
+    const int num_items_per_image,
+    const DataType DT_SCORE,
+    void* unsorted_scores,
+    void* unsorted_bbox_indices,
+    void* sorted_scores,
+    void* sorted_bbox_indices,
+    void* workspace,
+    int score_bits
+)
+{
+    sspiLaunchConfig lc = sspiLaunchConfig(DT_SCORE);
+    for (unsigned i = 0; i < sspiLCOptions.size(); ++i)
+    {
+        if (lc == sspiLCOptions[i])
+        {
+            DEBUG_PRINTF("sortScoresPerImage kernel %d\n", i);
+            return sspiLCOptions[i].function(stream,
+                                           num_images,
+                                           num_items_per_image,
+                                           unsorted_scores,
+                                           unsorted_bbox_indices,
+                                           sorted_scores,
+                                           sorted_bbox_indices,
+                                           workspace,
+                                           score_bits);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+
+size_t sortScoresPerImageWorkspaceSize(
+    const int num_images,
+    const int num_items_per_image,
+    const DataType DT_SCORE)
+{
+    const int arrayLen = num_images * num_items_per_image;
+    size_t wss[2];
+    wss[0] = (num_images + 1) * sizeof(int); // offsets
+    if (DT_SCORE == DataType::kFLOAT)
+    {
+        wss[1] = cubSortPairsWorkspaceSize<float, int>(arrayLen, num_images); // cub workspace
+    }
+    else if (DT_SCORE == DataType::kHALF)
+    {
+        wss[1] = cubSortPairsWorkspaceSize<__half, int>(arrayLen, num_images); // cub workspace
+    }
+    else
+    {
+        printf("SCORE type not supported.\n");
+        return (size_t) -1;
+    }
+
+    return calculateTotalWorkspaceSize(wss, 2);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/voxelGeneratorKernels.cu

@@ -0,0 +1,412 @@
+/*
+ * 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 <iostream>
+#include <cuda_runtime_api.h>
+namespace nvinfer1
+{
+namespace plugin
+{
+__global__ void generateVoxels_kernel(
+        int max_num_points,
+        float *points, unsigned int* points_size,
+        float min_x_range, float max_x_range,
+        float min_y_range, float max_y_range,
+        float min_z_range, float max_z_range,
+        float pillar_x_size, float pillar_y_size, float pillar_z_size,
+        int grid_y_size, int grid_x_size, int num_point_values,
+        int max_points_per_voxel,
+        unsigned int *mask, float *voxels)
+{
+  int point_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int batch_idx = point_idx / max_num_points;
+  int point_idx_in_frame = point_idx % max_num_points;
+  if(point_idx_in_frame >= points_size[batch_idx]) return;
+  float px = points[num_point_values * point_idx];
+  float py = points[num_point_values * point_idx + 1];
+  float pz = points[num_point_values * point_idx + 2];
+  float pw = points[num_point_values * point_idx + 3];
+  float pt;
+  if (num_point_values == 5) {
+    pt = points[num_point_values * point_idx + 4];
+  }
+  if(px<min_x_range||px>=max_x_range
+    || py<min_y_range||py>=max_y_range
+    || pz<min_z_range||pz>=max_z_range) return;
+  int voxel_idx = floorf((px - min_x_range)/pillar_x_size);
+  int voxel_idy = floorf((py - min_y_range)/pillar_y_size);
+  unsigned int voxel_index = (batch_idx * grid_y_size + voxel_idy) * grid_x_size + voxel_idx;
+  unsigned int point_id = atomicAdd(&(mask[voxel_index]), 1);
+  if(point_id >= max_points_per_voxel) return;
+  float *address = voxels + (voxel_index*max_points_per_voxel + point_id)*num_point_values;
+  atomicExch(address+0, px);
+  atomicExch(address+1, py);
+  atomicExch(address+2, pz);
+  atomicExch(address+3, pw);
+  if (num_point_values == 5) {
+    atomicExch(address+4, pt);
+  }
+}
+
+__global__ void generateBaseFeatures_kernel(
+        int batch_size,
+        unsigned int *mask, float *voxels,
+        int grid_y_size, int grid_x_size,
+        unsigned int *pillar_num,
+        int max_pillar_num,
+        int max_points_per_voxel,
+        int num_point_values,
+        float *voxel_features,
+        unsigned int *voxel_num_points,
+        unsigned int *coords)
+{
+  int voxel_id = blockIdx.x * blockDim.x + threadIdx.x;
+  int voxel_idx = voxel_id % grid_x_size;
+  int voxel_idy = (voxel_id / grid_x_size) % grid_y_size;
+  int batch_id = voxel_id / (grid_y_size * grid_x_size);
+  if (batch_id >= batch_size) return;
+  unsigned int count = mask[voxel_id];
+  if( !(count>0) ) return;
+  count = count<max_points_per_voxel?count:max_points_per_voxel;
+  int current_pillarId = 0;
+  current_pillarId = atomicAdd(pillar_num + batch_id, 1);
+  voxel_num_points[batch_id * grid_y_size * grid_x_size + current_pillarId] = count;
+  int4 coord = {0, 0, voxel_idy, voxel_idx};
+  ((int4*)coords)[batch_id * max_pillar_num + current_pillarId] = coord;
+  for (int i=0; i<count; i++){
+    int inIndex = voxel_id*max_points_per_voxel + i;
+    int outIndex = (batch_id * grid_x_size * grid_y_size + current_pillarId)*max_points_per_voxel + i;
+    if (num_point_values == 4) {
+      ((float4*)voxel_features)[outIndex] = ((float4*)voxels)[inIndex];
+    }
+    else if (num_point_values == 5) {
+      for(int k=0; k<5;k++)
+          voxel_features[5 * outIndex + k] = voxels[5 * inIndex + k];
+    }
+  }
+}
+
+void generateVoxels_launch(
+        int batch_size, int max_num_points,
+        float *points, unsigned int* points_size,
+        float min_x_range, float max_x_range,
+        float min_y_range, float max_y_range,
+        float min_z_range, float max_z_range,
+        float pillar_x_size, float pillar_y_size, float pillar_z_size,
+        int grid_y_size, int grid_x_size, int num_point_values,
+        int max_points_per_voxel,
+        unsigned int *mask, float *voxels,
+        cudaStream_t stream)
+{
+  int threadNum = 256;
+  dim3 blocks((batch_size * max_num_points + threadNum - 1) / threadNum);
+  dim3 threads(threadNum);
+  generateVoxels_kernel<<<blocks, threads, 0, stream>>>
+      (max_num_points,
+        points, points_size,
+        min_x_range, max_x_range,
+        min_y_range, max_y_range,
+        min_z_range, max_z_range,
+        pillar_x_size, pillar_y_size, pillar_z_size,
+        grid_y_size, grid_x_size, num_point_values,
+        max_points_per_voxel,
+        mask, voxels);
+}
+
+void generateBaseFeatures_launch(
+        int batch_size,
+        unsigned int *mask, float *voxels,
+        int grid_y_size, int grid_x_size,
+        unsigned int *pillar_num,
+        int max_pillar_num,
+        int max_points_per_voxel,
+        int num_point_values,
+        float *voxel_features,
+        unsigned int *voxel_num_points,
+        unsigned int *coords,
+        cudaStream_t stream)
+{
+  int blockSize = 1024;
+  dim3 threads(blockSize);
+  dim3 blocks((batch_size * grid_y_size * grid_x_size + blockSize - 1) / blockSize);
+  generateBaseFeatures_kernel<<<blocks, threads, 0, stream>>>
+      (
+        batch_size,
+        mask, voxels, grid_y_size, grid_x_size,
+        pillar_num,
+        max_pillar_num,
+        max_points_per_voxel,
+        num_point_values,
+        voxel_features,
+        voxel_num_points,
+        coords
+      );
+}
+
+__global__ void generateFeatures_kernel(
+    int batch_size,
+    int dense_pillar_num,
+    float* voxel_features,
+    unsigned int* voxel_num_points,
+    unsigned int* coords, unsigned int *params,
+    float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z,
+    unsigned int voxel_features_size, unsigned int max_points,
+    unsigned int max_voxels,
+    float* features)
+{
+    int warp_size = max_points;
+    int pillar_idx = blockIdx.x * 4 + threadIdx.x/warp_size;
+    int point_idx = threadIdx.x % warp_size;
+    // In case the actual number of points is less than warp_size
+    // E.g., warp_size=32, max_points=20
+    if (point_idx >= max_points) return;
+    int batch_idx = pillar_idx / max_voxels;
+    if (batch_idx >= batch_size) return;
+    int pillar_idx_in_frame = pillar_idx % max_voxels;
+    int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
+    int pillar_idx_inBlock = threadIdx.x/warp_size;
+    // Limit number of voxels to max_voxels
+    unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
+    // Update max_voxel to actual number
+    if (pillar_idx_in_frame == 0 && point_idx == 0) {
+      params[batch_idx] = num_pillars;
+    }
+    if (pillar_idx_in_frame >= num_pillars) return;
+
+    //load src
+    __shared__ float pillarSM[4][64][5]; // up to 64 points per pillar
+    __shared__ float4 pillarSumSM[4]; //4*4
+    __shared__ int4 cordsSM[4]; //4*4
+    __shared__ int pointsNumSM[4]; //4
+    __shared__ float pillarOutSM[4][64][11]; // up to 11 features per point
+
+    if (point_idx == 0) {
+      pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
+      cordsSM[pillar_idx_inBlock] = ((int4*)coords)[dense_pillar_idx];
+      pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
+    }
+    for(int k=0; k<5; k++) {
+      pillarSM[pillar_idx_inBlock][point_idx][k] = voxel_features[5 * (dense_pillar_idx*max_points + point_idx) + k];
+    }
+    __syncthreads();
+    //calculate sm
+    if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x),  pillarSM[pillar_idx_inBlock][point_idx][0]);
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y),  pillarSM[pillar_idx_inBlock][point_idx][1]);
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z),  pillarSM[pillar_idx_inBlock][point_idx][2]);
+    }
+    __syncthreads();
+    //feature-mean
+    float4 mean;
+    float validPoints = pointsNumSM[pillar_idx_inBlock];
+    mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
+    mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
+    mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
+    mean.x  = pillarSM[pillar_idx_inBlock][point_idx][0] - mean.x;
+    mean.y  = pillarSM[pillar_idx_inBlock][point_idx][1] - mean.y;
+    mean.z  = pillarSM[pillar_idx_inBlock][point_idx][2] - mean.z;
+    //calculate offset
+    float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
+    float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
+    float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
+    //feature-offset
+    float4 center;
+    center.x  = pillarSM[pillar_idx_inBlock][point_idx][0] - x_offset;
+    center.y  = pillarSM[pillar_idx_inBlock][point_idx][1] - y_offset;
+    center.z  = pillarSM[pillar_idx_inBlock][point_idx][2] - z_offset;
+    //store output
+    if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+      for(int k=0; k<5; k++)
+        pillarOutSM[pillar_idx_inBlock][point_idx][k] = pillarSM[pillar_idx_inBlock][point_idx][k];
+      pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.x;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 1] = mean.y;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 2] = mean.z;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 3] = center.x;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 4] = center.y;
+      if (5 + 5 < voxel_features_size)
+        pillarOutSM[pillar_idx_inBlock][point_idx][warp_size + 5] = center.z;
+    } else {
+      for (int k = 0; k < voxel_features_size; k++)
+        pillarOutSM[pillar_idx_inBlock][point_idx][k] = 0;
+    }
+    __syncthreads();
+    for(int i = 0; i < voxel_features_size; i ++) {
+      int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
+      int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
+      features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
+    }
+
+}
+
+__global__ void generateFeatures_kernel_4x(
+  int batch_size,
+  int dense_pillar_num,
+  float* voxel_features,
+  unsigned int* voxel_num_points, unsigned int* coords,
+  unsigned int *params,
+  float voxel_x, float voxel_y, float voxel_z,
+  float range_min_x, float range_min_y, float range_min_z,
+  unsigned int voxel_features_size, unsigned int max_points,
+  unsigned int max_voxels,
+  float* features)
+{
+  int warp_size = max_points;
+  int pillar_idx = blockIdx.x * 4 + threadIdx.x / warp_size;
+  int point_idx = threadIdx.x % warp_size;
+  // In case the actual number of points is less than warp_size
+  // E.g., warp_size=32, max_points=20
+  if (point_idx >= max_points) return;
+  int batch_idx = pillar_idx / max_voxels;
+  if (batch_idx >= batch_size) return;
+  int pillar_idx_in_frame = pillar_idx % max_voxels;
+  int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
+  int pillar_idx_inBlock = threadIdx.x / warp_size;
+  // Limit number of voxels to max_voxels
+  unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
+  // Update max_voxel to actual number
+  if (pillar_idx_in_frame == 0 && point_idx == 0) {
+    params[batch_idx] = num_pillars;
+  }
+  if (pillar_idx_in_frame >= num_pillars) return;
+  //load src
+  __shared__ float4 pillarSM[4][64]; // up to 64 points per pillar
+  __shared__ float4 pillarSumSM[4]; //4*4
+  __shared__ int4 cordsSM[4]; //4*4
+  __shared__ int pointsNumSM[4]; //4
+  __shared__ float pillarOutSM[4][64][11]; // up to 11 output features per point
+
+  if (point_idx == 0) {
+    pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
+    cordsSM[pillar_idx_inBlock] = ((int4*)coords)[pillar_idx];
+    pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
+  }
+  pillarSM[pillar_idx_inBlock][point_idx] = ((float4*)voxel_features)[dense_pillar_idx*max_points + point_idx];
+  __syncthreads();
+  //calculate sm
+  if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x),  pillarSM[pillar_idx_inBlock][point_idx].x);
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y),  pillarSM[pillar_idx_inBlock][point_idx].y);
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z),  pillarSM[pillar_idx_inBlock][point_idx].z);
+  }
+  __syncthreads();
+  //feature-mean
+  float4 mean;
+  float validPoints = pointsNumSM[pillar_idx_inBlock];
+  mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
+  mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
+  mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
+  mean.x  = pillarSM[pillar_idx_inBlock][point_idx].x - mean.x;
+  mean.y  = pillarSM[pillar_idx_inBlock][point_idx].y - mean.y;
+  mean.z  = pillarSM[pillar_idx_inBlock][point_idx].z - mean.z;
+  //calculate offset
+  float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
+  float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
+  float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
+  //feature-offset
+  float4 center;
+  center.x  = pillarSM[pillar_idx_inBlock][point_idx].x - x_offset;
+  center.y  = pillarSM[pillar_idx_inBlock][point_idx].y - y_offset;
+  center.z  = pillarSM[pillar_idx_inBlock][point_idx].z - z_offset;
+  //store output
+  if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+    pillarOutSM[pillar_idx_inBlock][point_idx][0] = pillarSM[pillar_idx_inBlock][point_idx].x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][1] = pillarSM[pillar_idx_inBlock][point_idx].y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][2] = pillarSM[pillar_idx_inBlock][point_idx].z;
+    pillarOutSM[pillar_idx_inBlock][point_idx][3] = pillarSM[pillar_idx_inBlock][point_idx].w;
+    pillarOutSM[pillar_idx_inBlock][point_idx][4] = mean.x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][6] = mean.z;
+    pillarOutSM[pillar_idx_inBlock][point_idx][7] = center.x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][8] = center.y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][9] = center.z;
+
+  } else {
+    pillarOutSM[pillar_idx_inBlock][point_idx][0] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][1] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][2] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][3] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][4] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][5] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][6] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][7] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][8] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][9] = 0;
+  }
+  __syncthreads();
+  for(int i = 0; i < voxel_features_size; i ++) {
+    int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
+    int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
+    features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
+  }
+
+}
+
+int generateFeatures_launch(
+    int batch_size,
+    int dense_pillar_num,
+    float* voxel_features,
+    unsigned int* voxel_num_points,
+    unsigned int* coords,
+    unsigned int *params,
+    float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z,
+    unsigned int voxel_features_size, unsigned int max_points,
+    unsigned int max_voxels, unsigned int num_point_values,
+    float* features,
+    cudaStream_t stream)
+{
+    unsigned int warp_size = max_points;
+    dim3 blocks((batch_size * max_voxels + 3) / 4);
+    dim3 threads(4*warp_size);
+    if (num_point_values == 4) {
+      generateFeatures_kernel_4x<<<blocks, threads, 0, stream>>>
+      (batch_size,
+      dense_pillar_num,
+      voxel_features,
+      voxel_num_points,
+      coords,
+      params,
+      voxel_x, voxel_y, voxel_z,
+      range_min_x, range_min_y, range_min_z,
+      voxel_features_size, max_points,
+      max_voxels,
+      features);
+    }
+    else {
+      generateFeatures_kernel<<<blocks, threads, 0, stream>>>
+      (batch_size,
+      dense_pillar_num,
+      voxel_features,
+      voxel_num_points,
+      coords,
+      params,
+      voxel_x, voxel_y, voxel_z,
+      range_min_x, range_min_y, range_min_z,
+      voxel_features_size, max_points,
+      max_voxels,
+      features);
+    }
+    auto err = cudaGetLastError();
+    if (cudaSuccess != err) {
+        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
+        exit(-1);
+    }
+    return err;
+}
+} // namespace plugin
+} // namespace nvinfer1