wehub-resource-sync
262 lines
14 KiB
C++
262 lines
14 KiB
C++
/*
|
|
* 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
|