chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,391 @@
+/*
+ * 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 "reorgPlugin.h"
+
+#include <memory>
+#include <string_view>
+
+namespace nvinfer1::plugin
+{
+static char const* const kREORG_PLUGIN_STATIC_VERSION{"1"};
+static char const* const kREORG_PLUGIN_DYNAMIC_VERSION{"2"};
+static char const* const kREORG_PLUGIN_NAME{"Reorg_TRT"};
+
+template <class TBaseClass>
+Reorg<TBaseClass>::Reorg(int32_t strideValue)
+    : stride(strideValue)
+{
+}
+
+template <class TBaseClass>
+int32_t Reorg<TBaseClass>::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+template <class TBaseClass>
+int32_t Reorg<TBaseClass>::initialize() noexcept
+{
+    return STATUS_SUCCESS;
+}
+
+template <class TBaseClass>
+void Reorg<TBaseClass>::terminate() noexcept
+{
+}
+
+template <class TBaseClass>
+char const* Reorg<TBaseClass>::getPluginType() const noexcept
+{
+    return kREORG_PLUGIN_NAME;
+}
+
+template <class TBaseClass>
+void Reorg<TBaseClass>::destroy() noexcept
+{
+    delete this;
+}
+
+// Set plugin namespace
+template <class TBaseClass>
+void Reorg<TBaseClass>::setPluginNamespace(char const* pluginNamespace) noexcept
+{
+    mPluginNamespace = pluginNamespace;
+}
+
+template <class TBaseClass>
+char const* Reorg<TBaseClass>::getPluginNamespace() const noexcept
+{
+    return mPluginNamespace.c_str();
+}
+
+// Return the DataType of the plugin output at the requested index
+template <class TBaseClass>
+DataType Reorg<TBaseClass>::getOutputDataType(
+    int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    // Only 1 input and 1 output from the plugin layer
+    PLUGIN_ASSERT(index == 0);
+
+    // Only DataType::kFLOAT is acceptable by the plugin layer
+    return DataType::kFLOAT;
+}
+
+// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
+template <class TBaseClass>
+void Reorg<TBaseClass>::attachToContext(
+    cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
+{
+}
+
+// Detach the plugin object from its execution context.
+template <class TBaseClass>
+void Reorg<TBaseClass>::detachFromContext() noexcept
+{
+}
+
+ReorgDynamic::ReorgDynamic(int32_t stride)
+    : Reorg<IPluginV2DynamicExt>(stride)
+{
+}
+
+ReorgDynamic::ReorgDynamic(void const* buffer, size_t length)
+{
+    char const* d = reinterpret_cast<char const*>(buffer);
+    char const* a = d;
+    stride = read<int32_t>(d);
+    PLUGIN_VALIDATE(d == a + length);
+}
+
+char const* ReorgDynamic::getPluginVersion() const noexcept
+{
+    return kREORG_PLUGIN_DYNAMIC_VERSION;
+}
+
+size_t ReorgDynamic::getSerializationSize() const noexcept
+{
+    // stride
+    return sizeof(int32_t);
+}
+
+size_t ReorgDynamic::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
+    PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    return 0;
+}
+
+void ReorgDynamic::serialize(void* buffer) const noexcept
+{
+    char *d = reinterpret_cast<char*>(buffer), *a = d;
+    write(d, stride);
+    PLUGIN_ASSERT(d == a + getSerializationSize());
+}
+
+DimsExprs ReorgDynamic::getOutputDimensions(
+    int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
+{
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(outputIndex == 0);
+    DimsExprs output{3, {}};
+    auto const* strideExpr = exprBuilder.constant(stride);
+    auto const* strideSquareExpr = exprBuilder.constant(stride * stride);
+    output.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[0], *strideSquareExpr);
+    output.d[1] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[0].d[1], *strideExpr);
+    output.d[2] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[0].d[2], *strideExpr);
+    return output;
+}
+
+bool ReorgDynamic::supportsFormatCombination(
+    int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    PLUGIN_ASSERT(pos >= 0 && pos <= 1);
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 1);
+    return (inOut[pos].type == DataType::kFLOAT && inOut[pos].format == PluginFormat::kLINEAR);
+}
+
+void ReorgDynamic::configurePlugin(
+    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 1);
+    PLUGIN_ASSERT(in->desc.type == DataType::kFLOAT);
+    PLUGIN_ASSERT(out->desc.type == DataType::kFLOAT);
+    PLUGIN_ASSERT(in->desc.format == PluginFormat::kLINEAR);
+    PLUGIN_ASSERT(out->desc.format == PluginFormat::kLINEAR);
+    PLUGIN_ASSERT(stride > 0);
+
+    int32_t H = in->desc.dims.d[2];
+    int32_t W = in->desc.dims.d[3];
+    PLUGIN_ASSERT(H % stride == 0);
+    PLUGIN_ASSERT(W % stride == 0);
+}
+
+int32_t ReorgDynamic::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    void const* inputData = inputs[0];
+    void* outputData = outputs[0];
+    int32_t const N = inputDesc[0].dims.d[0];
+    int32_t const C = inputDesc[0].dims.d[1];
+    int32_t const H = inputDesc[0].dims.d[2];
+    int32_t const W = inputDesc[0].dims.d[3];
+    pluginStatus_t status = reorgInference(stream, N, C, H, W, stride, inputData, outputData);
+    return status;
+}
+
+IPluginV2DynamicExt* ReorgDynamic::clone() const noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<ReorgDynamic>(stride);
+        plugin->setPluginNamespace(mPluginNamespace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+ReorgStatic::ReorgStatic(int32_t stride)
+    : Reorg<IPluginV2Ext>(stride)
+{
+}
+
+ReorgStatic::ReorgStatic(int32_t C, int32_t H, int32_t W, int32_t stride)
+    : Reorg<IPluginV2Ext>(stride)
+    , C(C)
+    , H(H)
+    , W(W)
+{
+}
+
+ReorgStatic::ReorgStatic(void const* buffer, size_t length)
+{
+    char const* d = reinterpret_cast<char const*>(buffer);
+    char const* a = d;
+    C = read<int32_t>(d);
+    H = read<int32_t>(d);
+    W = read<int32_t>(d);
+    stride = read<int32_t>(d);
+    PLUGIN_VALIDATE(d == a + length);
+}
+
+char const* ReorgStatic::getPluginVersion() const noexcept
+{
+    return kREORG_PLUGIN_STATIC_VERSION;
+}
+
+size_t ReorgStatic::getWorkspaceSize(int32_t maxBatchSize) const noexcept
+{
+    return 0;
+}
+
+size_t ReorgStatic::getSerializationSize() const noexcept
+{
+    // C, H, W, stride
+    return sizeof(int32_t) * 4;
+}
+
+void ReorgStatic::serialize(void* buffer) const noexcept
+{
+    char *d = reinterpret_cast<char*>(buffer), *a = d;
+    write(d, C);
+    write(d, H);
+    write(d, W);
+    write(d, stride);
+    PLUGIN_ASSERT(d == a + getSerializationSize());
+}
+
+Dims ReorgStatic::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
+{
+    PLUGIN_ASSERT(nbInputDims == 1);
+    PLUGIN_ASSERT(index == 0);
+    return Dims3(inputs[0].d[0] * stride * stride, inputs[0].d[1] / stride, inputs[0].d[2] / stride);
+}
+
+int32_t ReorgStatic::enqueue(
+    int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    void const* inputData = inputs[0];
+    void* outputData = outputs[0];
+    pluginStatus_t status = reorgInference(stream, batchSize, C, H, W, stride, inputData, outputData);
+    return status;
+}
+
+bool ReorgStatic::supportsFormat(DataType type, PluginFormat format) const noexcept
+{
+    return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
+}
+
+IPluginV2Ext* ReorgStatic::clone() const noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<ReorgStatic>(C, H, W, stride);
+        plugin->setPluginNamespace(mPluginNamespace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+// Configure the layer with input and output data types.
+void ReorgStatic::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(*inputTypes == DataType::kFLOAT && floatFormat == PluginFormat::kLINEAR);
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 1);
+    PLUGIN_ASSERT(stride > 0);
+    C = inputDims[0].d[0];
+    H = inputDims[0].d[1];
+    W = inputDims[0].d[2];
+    PLUGIN_ASSERT(H % stride == 0);
+    PLUGIN_ASSERT(W % stride == 0);
+}
+
+template <class TPluginClass>
+ReorgPluginCreator<TPluginClass>::ReorgPluginCreator()
+{
+    mPluginAttributes.clear();
+    mPluginAttributes.emplace_back(PluginField("stride", nullptr, PluginFieldType::kINT32, 1));
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+template <class TPluginClass>
+char const* ReorgPluginCreator<TPluginClass>::getPluginName() const noexcept
+{
+    return kREORG_PLUGIN_NAME;
+}
+
+template <class TPluginClass>
+char const* ReorgPluginCreator<TPluginClass>::getPluginVersion() const noexcept
+{
+    if (std::is_same_v<TPluginClass, ReorgStatic>)
+    {
+        return kREORG_PLUGIN_STATIC_VERSION;
+    }
+    else if (std::is_same_v<TPluginClass, ReorgDynamic>)
+    {
+        return kREORG_PLUGIN_DYNAMIC_VERSION;
+    }
+    return "";
+}
+
+template <class TPluginClass>
+PluginFieldCollection const* ReorgPluginCreator<TPluginClass>::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+template <class TPluginClass>
+IPluginV2Ext* ReorgPluginCreator<TPluginClass>::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
+{
+    try
+    {
+        using namespace std::string_view_literals;
+        PluginField const* fields = fc->fields;
+        PLUGIN_VALIDATE(fc->nbFields == 1);
+        PLUGIN_VALIDATE(fields[0].type == PluginFieldType::kINT32);
+        PLUGIN_VALIDATE(fields[0].name == "stride"sv);
+        stride = static_cast<int32_t>(*(static_cast<int32_t const*>(fields[0].data)));
+
+        PLUGIN_VALIDATE(stride > 0);
+
+        auto obj = std::make_unique<TPluginClass>(stride);
+        obj->setPluginNamespace(mNamespace.c_str());
+        return obj.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+template <class TPluginClass>
+IPluginV2Ext* ReorgPluginCreator<TPluginClass>::deserializePlugin(
+    char const* name, void const* serialData, size_t serialLength) noexcept
+{
+    try
+    {
+        // This object will be deleted when the network is destroyed, which will
+        // call ReorgPlugin::destroy()
+        auto obj = std::make_unique<TPluginClass>(serialData, serialLength);
+        obj->setPluginNamespace(mNamespace.c_str());
+        return obj.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+template class ReorgPluginCreator<ReorgStatic>;
+template class ReorgPluginCreator<ReorgDynamic>;
+
+} // namespace nvinfer1::plugin