ONNX-TensorRT 9.1 Update

Signed-off-by: Kevin Chen <[email protected]>

CMakeLists.txt

@@ -28,7 +28,7 @@ add_definitions("-DSOURCE_LENGTH=${SOURCE_LENGTH}")
 # Version information
 #--------------------------------------------------
 set(ONNX2TRT_MAJOR 9)
-set(ONNX2TRT_MINOR 0)
+set(ONNX2TRT_MINOR 1)
 set(ONNX2TRT_PATCH 0)
 set(ONNX2TRT_VERSION "${ONNX2TRT_MAJOR}.${ONNX2TRT_MINOR}.${ONNX2TRT_PATCH}" CACHE STRING "ONNX2TRT version")
 
@@ -40,14 +40,15 @@ set(IMPORTER_SOURCES
   NvOnnxParser.cpp
   ModelImporter.cpp
   builtin_op_importers.cpp
-  onnx2trt_utils.cpp
-  onnxErrorRecorder.cpp
+  builtin_op_static_checkers.cpp
   ImporterContext.cpp
+  onnx2trt_utils.cpp
   ShapedWeights.cpp
   ShapeTensor.cpp
   LoopHelpers.cpp
   RNNHelpers.cpp
   OnnxAttrs.cpp
+  onnxErrorRecorder.cpp
   ConditionalHelpers.cpp
   bfloat16.cpp
 )
@@ -116,11 +117,18 @@ set_target_properties(nvonnxparser PROPERTIES
   SOVERSION ${ONNX2TRT_MAJOR}
   LINK_DEPENDS ${PARSER_LINKER_SCRIPT}
   LINK_FLAGS "-Wl,--version-script=${PARSER_LINKER_SCRIPT}"
+  ARCHIVE_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
+  LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
+  RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
 )
 add_library(nvonnxparser_static STATIC ${IMPORTER_SOURCES})
 target_include_directories(nvonnxparser_static PUBLIC ${ONNX_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${CUDA_INCLUDE_DIR})
 target_link_libraries(nvonnxparser_static PUBLIC onnx_proto ${PROTOBUF_LIBRARY} ${TENSORRT_LIBRARY})
-
+set_target_properties(nvonnxparser_static PROPERTIES
+  ARCHIVE_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
+  LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
+  RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
+)
 # --------------------------------
 # Onnxifi library
 # --------------------------------

ConditionalHelpers.cpp

@@ -75,20 +75,21 @@ Status addConditionalInputLayer(IImporterContext* ctx, nvinfer1::IIfConditional*
 // Take a snapshot of the network before and after parsing the subgraph and return a list
 // of newly added network layers.
 Status importSubgraph(IImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const& subgraph,
-    std::vector<nvinfer1::ILayer*>& newLayers, StringMap<TensorOrWeights>& subgraphTensors)
+    std::vector<nvinfer1::ILayer*>& newLayers, std::vector<TensorOrWeights>& subgraphTensors)
 {
     auto net = ctx->network();
     int32_t beforeSubgraph = net->getNbLayers();
 
     // Establish scope for names local to the subgraph.
     NameScope nameScope(*ctx);
 
-    CHECK(onnx2trt::parseGraph(ctx, subgraph));
+    std::vector<Status> errors{};
+    CHECK(onnx2trt::parseGraph(ctx, subgraph, errors));
 
     for (int32_t i = 0; i < subgraph.output_size(); ++i)
     {
         std::string name = subgraph.output(i).name();
-        subgraphTensors.emplace(std::make_pair(name, ctx->tensors().at(name)));
+        subgraphTensors.push_back(ctx->tensors().at(name));
     }
 
     for (int32_t i = beforeSubgraph; i < net->getNbLayers(); i++)
@@ -145,81 +146,6 @@ Status addIfInputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditi
     return Status::success();
 }
 
-// Add an IConditionalOutputLayer to `layer`'s outputs.
-Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditional,
-    ::ONNX_NAMESPACE::GraphProto const& thenGraph, std::vector<nvinfer1::ILayer*> const& thenLayers,
-    StringMap<TensorOrWeights> const& thenSubgraphTensors, ::ONNX_NAMESPACE::GraphProto const& elseGraph,
-    std::vector<nvinfer1::ILayer*> const& elseLayers, StringMap<TensorOrWeights> const& elseSubgraphTensors,
-    std::vector<TensorOrWeights>& graphOutputs)
-{
-    // Reported outputs are outputs that the ONNX model reports as subgraph outputs.  This list is
-    // not sufficient because it may produce names that are not fully compatible with TensorRT's naming.
-    // We use this list to help find the subgraph (SG) output tensors.
-    auto getReportedOutputs
-        = [&ctx](const ::ONNX_NAMESPACE::GraphProto& body, std::vector<std::string>& reportedOutputs) {
-              // Assuming that the subgraph was imported already, we can iterate on its output tensors.
-              const auto nbOutputs = body.output_size();
-              for (auto i = 0; i < nbOutputs; i++)
-              {
-                  reportedOutputs.emplace_back(body.output(i).name());
-              }
-          };
-
-    std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>> thenOutputs;
-    std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>> elseOutputs;
-
-    std::vector<std::string> thenReportedOutputs;
-    getReportedOutputs(thenGraph, thenReportedOutputs);
-    getSubgraphOutputs(thenLayers, thenOutputs, thenReportedOutputs);
-    std::vector<std::string> elseReportedOutputs;
-    getReportedOutputs(elseGraph, elseReportedOutputs);
-    getSubgraphOutputs(elseLayers, elseOutputs, elseReportedOutputs);
-
-    // Retrieve the output tensors of a subgraph (tensors exiting the subgraph).
-    auto getSubgraphOutputTensors
-        = [](IImporterContext* ctx, std::vector<nvinfer1::ITensor*>& sgOutputs, SubgraphPortsMap& subgraphOutputs,
-              ::ONNX_NAMESPACE::GraphProto const& subgraph, std::vector<nvinfer1::ILayer*> subgraphLayers,
-              StringMap<TensorOrWeights> const& subgraphTensors) {
-              for (auto const& pair : subgraphOutputs)
-              {
-                  sgOutputs.push_back(pair.first);
-              }
-
-              if (sgOutputs.empty())
-              {
-                  // No new layers, so we can't deduce the outputs and have to use what ONNX tells us.
-                  const int32_t nbOutputs = subgraph.output_size();
-                  for (int32_t outIdx = 0; outIdx < nbOutputs; outIdx++)
-                  {
-                      const auto thenName = subgraph.output(outIdx).name();
-                      TensorOrWeights tw = subgraphTensors.at(thenName);
-                      auto* thenTensor = &convertToTensor(tw, ctx);
-                      sgOutputs.push_back(thenTensor);
-                  }
-              }
-          };
-
-    std::vector<nvinfer1::ITensor*> thenOutputTensors;
-    getSubgraphOutputTensors(ctx, thenOutputTensors, thenOutputs, thenGraph, thenLayers, thenSubgraphTensors);
-
-    std::vector<nvinfer1::ITensor*> elseSGOutputTensors;
-    getSubgraphOutputTensors(ctx, elseSGOutputTensors, elseOutputs, elseGraph, elseLayers, elseSubgraphTensors);
-
-    ASSERT(thenOutputTensors.size() == elseSGOutputTensors.size()
-            && "The then/else branches of an If operator must have the same number of outputs.",
-        ErrorCode::kINVALID_NODE);
-
-    // Add an ConditionalOutputLayer with one output and two inputs
-    // (one from the thenGraph and another from the elseGraph).
-    for (size_t i = 0; i < elseSGOutputTensors.size(); i++)
-    {
-        auto* outputLayer = conditional->addOutput(*thenOutputTensors[i], *elseSGOutputTensors[i]);
-        ctx->registerLayer(outputLayer, std::string(conditional->getName()) + "_OutputLayer", nullptr);
-        graphOutputs.emplace_back(outputLayer->getOutput(0));
-    }
-    return Status::success();
-}
-
 // Given a subgraph, find all of its external inputs/outputs (tensors entering/exiting the subgraph).
 Status getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
     std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalOutputs, bool extractOutputs,

ConditionalHelpers.hpp

@@ -34,19 +34,12 @@ Status getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
 // Take a snapshot of the network before and after parsing the subgraph and return a list
 // of newly added network layers.
 Status importSubgraph(IImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const& subgraph,
-    std::vector<nvinfer1::ILayer*>& newLayers, StringMap<TensorOrWeights>& subgraphTensors);
+    std::vector<nvinfer1::ILayer*>& newLayers, std::vector<TensorOrWeights>& subgraphTensors);
 
 using InputsMap = std::unordered_map<std::string, nvinfer1::IIfConditionalInputLayer*>;
 
 // Add IIfConditionalInputLayers to the inputs of the subgraph indicated by `subgraph`.
 onnx2trt::Status addIfInputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
     const std::vector<nvinfer1::ILayer*>& newLayers);
 
-// Add IIfConditionalOutputLayers to the outputs of the subgraph indicated by `subgraph`.
-onnx2trt::Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditional,
-    ::ONNX_NAMESPACE::GraphProto const& thenGraph, std::vector<nvinfer1::ILayer*> const& thenLayers,
-    StringMap<TensorOrWeights> const& thenSubgraphTensors, ::ONNX_NAMESPACE::GraphProto const& elseGraph,
-    std::vector<nvinfer1::ILayer*> const& elseLayers, StringMap<TensorOrWeights> const& elseSubgraphTensors,
-    std::vector<TensorOrWeights>& graphOutputs);
-
 } // namespace onnx2trt

ImporterContext.hpp

@@ -7,7 +7,6 @@
 #include "onnx2trt.hpp"
 #include "onnx2trt_utils.hpp"
 #include "onnxErrorRecorder.hpp"
-#include "onnx/common/stl_backports.h"
 #include <list>
 #include <string>
 #include <unordered_map>
@@ -117,8 +116,8 @@ class ImporterContext final : public IImporterContext
     //! Map holding FunctionProtos
     StringMap<::ONNX_NAMESPACE::FunctionProto> mLocalFunctions;
 
-    //! Vector to hold current local function names
-    std::vector<std::string> mLocalFunctionStack;
+    //! Vector to hold current local function names and attributes
+    std::vector<std::pair<std::string, StringMap<::ONNX_NAMESPACE::AttributeProto const*>>> mLocalFunctionStack;
 
     //! Vector to hold expected graph outputs
     std::vector<::ONNX_NAMESPACE::ValueInfoProto> mGraphOutputNames;
@@ -127,7 +126,7 @@ class ImporterContext final : public IImporterContext
     ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
         : mNetwork(network)
         , mLogger(logger)
-        , mErrorWrapper(ONNX_NAMESPACE::make_unique<ErrorRecorderWrapper>(mNetwork, logger))
+        , mErrorWrapper(std::make_unique<ErrorRecorderWrapper>(mNetwork, logger))
     {
     }
     nvinfer1::INetworkDefinition* network() override
@@ -336,7 +335,7 @@ class ImporterContext final : public IImporterContext
     {
         return mLocalFunctions;
     }
-    std::vector<std::string>& localFunctionStack() override
+    std::vector<std::pair<std::string, StringMap<::ONNX_NAMESPACE::AttributeProto const*>>>& localFunctionStack() override
     {
         return mLocalFunctionStack;
     }