chore: import upstream snapshot with attribution

plugin/pyramidROIAlignPlugin/pyramidROIAlignPlugin.cpp

@@ -0,0 +1,434 @@
+/*
+ * 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.
+ */
+#include "pyramidROIAlignPlugin.h"
+#include "common/plugin.h"
+#include <cuda_runtime_api.h>
+#include <math.h>
+#include <memory>
+#include <string_view>
+
+using namespace nvinfer1;
+using namespace plugin;
+using nvinfer1::plugin::PyramidROIAlign;
+using nvinfer1::plugin::PyramidROIAlignPluginCreator;
+
+namespace
+{
+char const* const kPYRAMIDROIALGIN_PLUGIN_VERSION{"1"};
+char const* const kPYRAMIDROIALGIN_PLUGIN_NAME{"PyramidROIAlign_TRT"};
+} // namespace
+
+PyramidROIAlignPluginCreator::PyramidROIAlignPluginCreator()
+{
+    mPluginAttributes.clear();
+    mPluginAttributes.emplace_back(PluginField("fpn_scale", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("pooled_size", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("image_size", nullptr, PluginFieldType::kINT32, 2));
+    mPluginAttributes.emplace_back(PluginField("roi_coords_absolute", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("roi_coords_swap", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("roi_coords_plusone", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("roi_coords_transform", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("sampling_ratio", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("legacy", nullptr, PluginFieldType::kINT32, 1));
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* PyramidROIAlignPluginCreator::getPluginName() const noexcept
+{
+    return kPYRAMIDROIALGIN_PLUGIN_NAME;
+}
+
+char const* PyramidROIAlignPluginCreator::getPluginVersion() const noexcept
+{
+    return kPYRAMIDROIALGIN_PLUGIN_VERSION;
+}
+
+PluginFieldCollection const* PyramidROIAlignPluginCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+IPluginV2Ext* PyramidROIAlignPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
+{
+    try
+    {
+        // Default values for the plugin creator, these will be used when the corresponding
+        // plugin field is not passed, allowing to have defaults for "optional" ONNX attributes.
+        int32_t pooledSize = 7;
+        int32_t transformCoords = 2;
+        bool absCoords = true;
+        bool swapCoords = false;
+        bool plusOneCoords = false;
+        bool legacy = false;
+        int32_t samplingRatio = 0;
+        xy_t imageSize = {dimToInt32(MaskRCNNConfig::IMAGE_SHAPE.d[1]), dimToInt32(MaskRCNNConfig::IMAGE_SHAPE.d[2])};
+        int32_t fpnScale = 224;
+
+        using namespace std::string_view_literals;
+        PluginField const* fields = fc->fields;
+        for (int32_t i = 0; i < fc->nbFields; ++i)
+        {
+            std::string_view const attrName = fields[i].name;
+            if (attrName == "fpn_scale"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                fpnScale = *(static_cast<int32_t const*>(fields[i].data));
+                PLUGIN_VALIDATE(fpnScale >= 1);
+            }
+            if (attrName == "pooled_size"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                pooledSize = *(static_cast<int32_t const*>(fields[i].data));
+                PLUGIN_VALIDATE(pooledSize >= 1);
+            }
+            if (attrName == "image_size"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                PLUGIN_VALIDATE(fields[i].length == 2);
+                auto const dims = static_cast<int32_t const*>(fields[i].data);
+                imageSize.y = dims[0];
+                imageSize.x = dims[1];
+                PLUGIN_VALIDATE(imageSize.y >= 1);
+                PLUGIN_VALIDATE(imageSize.x >= 1);
+            }
+            if (attrName == "roi_coords_absolute"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                absCoords = *(static_cast<int32_t const*>(fields[i].data));
+            }
+            if (attrName == "roi_coords_swap"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                swapCoords = *(static_cast<int32_t const*>(fields[i].data));
+            }
+            if (attrName == "roi_coords_plusone"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                plusOneCoords = *(static_cast<int32_t const*>(fields[i].data));
+            }
+            if (attrName == "roi_coords_transform"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                transformCoords = *(static_cast<int32_t const*>(fields[i].data));
+            }
+            if (attrName == "sampling_ratio"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                samplingRatio = *(static_cast<int32_t const*>(fields[i].data));
+                PLUGIN_VALIDATE(samplingRatio >= 0);
+            }
+            if (attrName == "legacy"sv)
+            {
+                PLUGIN_ASSERT(fields[i].type == PluginFieldType::kINT32);
+                legacy = *(static_cast<int32_t const*>(fields[i].data));
+            }
+        }
+        return new PyramidROIAlign(pooledSize, transformCoords, absCoords, swapCoords, plusOneCoords, samplingRatio,
+            legacy, imageSize, fpnScale);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+IPluginV2Ext* PyramidROIAlignPluginCreator::deserializePlugin(
+    char const* name, void const* data, size_t length) noexcept
+{
+    try
+    {
+        return new PyramidROIAlign(data, length);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+PyramidROIAlign::PyramidROIAlign(int32_t pooledSize, int32_t transformCoords, bool absCoords, bool swapCoords,
+    bool plusOneCoords, int32_t samplingRatio, bool legacy, xy_t imageSize, int32_t fpnScale)
+    : mPooledSize({pooledSize, pooledSize})
+    , mImageSize(imageSize)
+    , mFPNScale(fpnScale)
+    , mTransformCoords(transformCoords)
+    , mAbsCoords(absCoords)
+    , mSwapCoords(swapCoords)
+    , mPlusOneCoords(plusOneCoords)
+    , mSamplingRatio(samplingRatio)
+    , mIsLegacy(legacy)
+{
+    PLUGIN_VALIDATE(pooledSize >= 1);
+    PLUGIN_VALIDATE(samplingRatio >= 0);
+    PLUGIN_VALIDATE(fpnScale >= 1);
+}
+
+int32_t PyramidROIAlign::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+int32_t PyramidROIAlign::initialize() noexcept
+{
+    return 0;
+}
+
+void PyramidROIAlign::terminate() noexcept {}
+
+void PyramidROIAlign::destroy() noexcept
+{
+    delete this;
+}
+
+size_t PyramidROIAlign::getWorkspaceSize(int32_t) const noexcept
+{
+    return 0;
+}
+
+bool PyramidROIAlign::supportsFormat(DataType type, PluginFormat format) const noexcept
+{
+    return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
+}
+
+char const* PyramidROIAlign::getPluginType() const noexcept
+{
+    return kPYRAMIDROIALGIN_PLUGIN_NAME;
+}
+
+char const* PyramidROIAlign::getPluginVersion() const noexcept
+{
+    return kPYRAMIDROIALGIN_PLUGIN_VERSION;
+}
+
+IPluginV2Ext* PyramidROIAlign::clone() const noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<PyramidROIAlign>(*this);
+        plugin->setPluginNamespace(mNameSpace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+void PyramidROIAlign::setPluginNamespace(char const* libNamespace) noexcept
+{
+    mNameSpace = libNamespace;
+}
+
+char const* PyramidROIAlign::getPluginNamespace() const noexcept
+{
+    return mNameSpace.c_str();
+}
+
+void PyramidROIAlign::check_valid_inputs(nvinfer1::Dims const* inputs, int32_t nbInputDims)
+{
+    // to be compatible with tensorflow node's input:
+    // roi: [N, anchors, 4],
+    // feature_map list(4 maps): p2, p3, p4, p5
+    PLUGIN_ASSERT(nbInputDims == 1 + mFeatureMapCount);
+
+    nvinfer1::Dims rois = inputs[0];
+    PLUGIN_ASSERT(rois.nbDims == 2);
+    PLUGIN_ASSERT(rois.d[1] == 4);
+
+    for (int32_t i = 1; i < nbInputDims; ++i)
+    {
+        nvinfer1::Dims dims = inputs[i];
+
+        // CHW with the same #C
+        PLUGIN_ASSERT(dims.nbDims == 3 && dims.d[0] == inputs[i].d[0]);
+    }
+}
+
+Dims PyramidROIAlign::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
+{
+
+    check_valid_inputs(inputs, nbInputDims);
+    PLUGIN_ASSERT(index == 0);
+
+    nvinfer1::Dims result{};
+    result.nbDims = 4;
+
+    // mROICount
+    result.d[0] = inputs[0].d[0];
+    // mFeatureLength
+    result.d[1] = inputs[1].d[0];
+    // height
+    result.d[2] = mPooledSize.y;
+    // width
+    result.d[3] = mPooledSize.x;
+
+    return result;
+}
+
+int32_t PyramidROIAlign::enqueue(
+    int32_t batch_size, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    void* const pooled = outputs[0];
+    cudaError_t status;
+
+    // Support legacy UFF mode
+    if (mIsLegacy)
+    {
+        // Legacy values
+        mTransformCoords = -1;
+        mPlusOneCoords = 0;
+        mSwapCoords = true;
+        mAbsCoords = false;
+        mSamplingRatio = 1;
+        float const firstThreshold
+            = (224 * 224 * 2.F / (MaskRCNNConfig::IMAGE_SHAPE.d[1] * MaskRCNNConfig::IMAGE_SHAPE.d[2])) / (4.F * 4.F);
+        status = roiAlign(stream, batch_size, mImageSize, mFeatureLength, mROICount, firstThreshold, mTransformCoords,
+            mAbsCoords, mSwapCoords, mPlusOneCoords, mSamplingRatio, inputs[0], &inputs[1], mFeatureSpatialSize, pooled,
+            mPooledSize);
+    }
+    else
+    {
+        // As per FPN paper equation 1 (https://arxiv.org/pdf/1612.03144.pdf)
+        // the default 224 FPN scale corresponds to the canonical ImageNet size
+        // used to define the ROI scale threshold that samples from P4. Because the
+        // plugin works with normalized ROI coordinates, the FPN scale must be normalized
+        // by the input image size.
+        float const scale = static_cast<float>(mFPNScale);
+        float const normScale = sqrtf(scale * scale / (mImageSize.y * mImageSize.x));
+        // Furthermore, the roiAlign kernel expects a first threshold instead. This is
+        // the *area* of an ROI but for one level down, i.e. at the P2->P3 transition.
+        float const firstThreshold = normScale * normScale / 4.F;
+        status = roiAlign(stream, batch_size, mImageSize, mFeatureLength, mROICount, firstThreshold, mTransformCoords,
+            mAbsCoords, mSwapCoords, mPlusOneCoords, mSamplingRatio, inputs[0], &inputs[1], mFeatureSpatialSize, pooled,
+            mPooledSize);
+    }
+    return status;
+}
+
+size_t PyramidROIAlign::getSerializationSize() const noexcept
+{
+    return sizeof(int32_t) * 2 // mPooledSize
+        + sizeof(int32_t) * 2  // mImageSize
+        + sizeof(int32_t)      // mFeatureLength
+        + sizeof(int32_t)      // mROICount
+        + sizeof(int32_t)      // mFPNScale
+        + sizeof(int32_t)      // mTransformCoords
+        + sizeof(bool)         // mAbsCoords
+        + sizeof(bool)         // mSwapCoords
+        + sizeof(bool)         // mPlusOneCoords
+        + sizeof(int32_t)      // mSamplingRatio
+        + sizeof(bool)         // mIsLegacy
+        + sizeof(int32_t) * 8; // mFeatureSpatialSize
+}
+
+void PyramidROIAlign::serialize(void* buffer) const noexcept
+{
+    char *d = reinterpret_cast<char*>(buffer), *a = d;
+    write(d, mPooledSize.y);
+    write(d, mPooledSize.x);
+    write(d, mImageSize.y);
+    write(d, mImageSize.x);
+    write(d, mFeatureLength);
+    write(d, mROICount);
+    write(d, mFPNScale);
+    write(d, mTransformCoords);
+    write(d, mAbsCoords);
+    write(d, mSwapCoords);
+    write(d, mPlusOneCoords);
+    write(d, mSamplingRatio);
+    write(d, mIsLegacy);
+    write(d, mFeatureSpatialSize[0].y);
+    write(d, mFeatureSpatialSize[0].x);
+    write(d, mFeatureSpatialSize[1].y);
+    write(d, mFeatureSpatialSize[1].x);
+    write(d, mFeatureSpatialSize[2].y);
+    write(d, mFeatureSpatialSize[2].x);
+    write(d, mFeatureSpatialSize[3].y);
+    write(d, mFeatureSpatialSize[3].x);
+    PLUGIN_ASSERT(d == a + getSerializationSize());
+}
+
+PyramidROIAlign::PyramidROIAlign(void const* data, size_t length)
+{
+    deserialize(static_cast<int8_t const*>(data), length);
+}
+void PyramidROIAlign::deserialize(int8_t const* data, size_t length)
+{
+    auto const* d{data};
+    mPooledSize = {read<int32_t>(d), read<int32_t>(d)};
+    mImageSize = {read<int32_t>(d), read<int32_t>(d)};
+    mFeatureLength = read<int32_t>(d);
+    mROICount = read<int32_t>(d);
+    mFPNScale = read<int32_t>(d);
+    mTransformCoords = read<int32_t>(d);
+    mAbsCoords = read<bool>(d);
+    mSwapCoords = read<bool>(d);
+    mPlusOneCoords = read<bool>(d);
+    mSamplingRatio = read<int32_t>(d);
+    mIsLegacy = read<bool>(d);
+    mFeatureSpatialSize[0].y = read<int32_t>(d);
+    mFeatureSpatialSize[0].x = read<int32_t>(d);
+    mFeatureSpatialSize[1].y = read<int32_t>(d);
+    mFeatureSpatialSize[1].x = read<int32_t>(d);
+    mFeatureSpatialSize[2].y = read<int32_t>(d);
+    mFeatureSpatialSize[2].x = read<int32_t>(d);
+    mFeatureSpatialSize[3].y = read<int32_t>(d);
+    mFeatureSpatialSize[3].x = read<int32_t>(d);
+
+    PLUGIN_VALIDATE(d == data + length);
+}
+
+// Return the DataType of the plugin output at the requested index
+DataType PyramidROIAlign::getOutputDataType(
+    int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    // Only DataType::kFLOAT is acceptable by the plugin layer
+    return DataType::kFLOAT;
+}
+
+// Configure the layer with input and output data types.
+void PyramidROIAlign::configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims,
+    int32_t nbOutputs, DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
+    bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept
+{
+    PLUGIN_ASSERT(supportsFormat(inputTypes[0], floatFormat));
+    check_valid_inputs(inputDims, nbInputs);
+
+    PLUGIN_ASSERT(nbOutputs == 1);
+    PLUGIN_ASSERT(nbInputs == 1 + mFeatureMapCount);
+
+    mROICount = inputDims[0].d[0];
+    mFeatureLength = inputDims[1].d[0];
+
+    for (size_t layer = 0; layer < mFeatureMapCount; ++layer)
+    {
+        mFeatureSpatialSize[layer] = {dimToInt32(inputDims[layer + 1].d[1]), dimToInt32(inputDims[layer + 1].d[2])};
+    }
+}
+
+// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
+void PyramidROIAlign::attachToContext(
+    cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
+{
+}
+
+// Detach the plugin object from its execution context.
+void PyramidROIAlign::detachFromContext() noexcept {}