TRT 9.0 EA - OSS (#930)

Signed-off-by: Samurdhi Karunaratne <[email protected]>

Author: samurdhikaru

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "third_party/onnx"]
 	path = third_party/onnx
 	url = https://github.com/onnx/onnx.git
-	branch = v1.13.1
+	branch = v1.14.0

CMakeLists.txt

@@ -27,9 +27,9 @@ add_definitions("-DSOURCE_LENGTH=${SOURCE_LENGTH}")
 #--------------------------------------------------
 # Version information
 #--------------------------------------------------
-set(ONNX2TRT_MAJOR 8)
-set(ONNX2TRT_MINOR 6)
-set(ONNX2TRT_PATCH 1)
+set(ONNX2TRT_MAJOR 9)
+set(ONNX2TRT_MINOR 0)
+set(ONNX2TRT_PATCH 0)
 set(ONNX2TRT_VERSION "${ONNX2TRT_MAJOR}.${ONNX2TRT_MINOR}.${ONNX2TRT_PATCH}" CACHE STRING "ONNX2TRT version")
 
 #--------------------------------------------------
@@ -49,6 +49,7 @@ set(IMPORTER_SOURCES
   RNNHelpers.cpp
   OnnxAttrs.cpp
   ConditionalHelpers.cpp
+  bfloat16.cpp
 )
 
 if (BUILD_ONNXIFI)

ImporterContext.cpp

@@ -65,12 +65,6 @@ void ImporterContext::registerTensor(TensorOrWeights tensor, std::string const&
         }
         else if (tensor.is_weights())
         {
-            auto const& weights = tensor.weights();
-            if (tensor.weights().type == ::ONNX_NAMESPACE::TensorProto::INT64)
-            {
-                tensor = ShapedWeights{::ONNX_NAMESPACE::TensorProto::INT32,
-                    convertINT64(reinterpret_cast<int64_t*>(weights.values), weights.shape, this), weights.shape};
-            }
             // It may be possible for nested subgraphs to have different values for the same initializer.
             // For multiple name scopes - use unique name to keep track of weights.
             if (!mBaseNameScopeStack.empty())
@@ -118,7 +112,14 @@ void ImporterContext::registerLayer(nvinfer1::ILayer* layer, std::string const&
         std::string const& uniqueName = generateUniqueName(mLayerNames, name);
 
         auto* ctx = this; // To enable logging.
-        LOG_VERBOSE("Registering layer: " << uniqueName << " for ONNX node: " << basename);
+        if (node != nullptr)
+        {
+            LOG_VERBOSE("Registering layer: " << uniqueName << " for ONNX node: " << basename);
+        }
+        else
+        {
+            LOG_VERBOSE("Registering layer: " << uniqueName << " required by ONNX-TRT");
+        }
 
         layer->setName(uniqueName.c_str());
         if (layer->getType() == nvinfer1::LayerType::kCONSTANT)
@@ -133,7 +134,7 @@ void ImporterContext::registerLayer(nvinfer1::ILayer* layer, std::string const&
     }
     if (node != nullptr)
     {
-        processMetadata(*node, layer);
+        processMetadata(this, *node, layer);
     }
 }
 

ImporterContext.hpp

@@ -114,6 +114,15 @@ class ImporterContext final : public IImporterContext
     //!
     std::vector<StringMap<std::pair<bool, TensorOrWeights>>> mBaseNameScopeStack;
 
+    //! Map holding FunctionProtos
+    StringMap<::ONNX_NAMESPACE::FunctionProto> mLocalFunctions;
+
+    //! Vector to hold current local function names
+    std::vector<std::string> mLocalFunctionStack;
+
+    //! Vector to hold expected graph outputs
+    std::vector<::ONNX_NAMESPACE::ValueInfoProto> mGraphOutputNames;
+
 public:
     ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
         : mNetwork(network)
@@ -123,6 +132,7 @@ class ImporterContext final : public IImporterContext
     }
     nvinfer1::INetworkDefinition* network() override
     {
+        assert(mNetwork != nullptr);
         return mNetwork;
     }
     StringMap<TensorOrWeights>& tensors() override
@@ -322,6 +332,18 @@ class ImporterContext final : public IImporterContext
     {
         mConvertDoubleOutOfBoundsLogged = logged;
     }
+    StringMap<::ONNX_NAMESPACE::FunctionProto>& localFunctions() override
+    {
+        return mLocalFunctions;
+    }
+    std::vector<std::string>& localFunctionStack() override
+    {
+        return mLocalFunctionStack;
+    }
+    std::vector<::ONNX_NAMESPACE::ValueInfoProto>& getGraphOutputNames() override
+    {
+        return mGraphOutputNames;
+    }
 
 private:
     std::string const& generateUniqueName(std::set<std::string>& namesSet, const std::string& basename)

LoopHelpers.cpp

@@ -8,13 +8,17 @@
 namespace onnx2trt
 {
 
-nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int32_t initial)
+nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int64_t initial)
 {
-    nvinfer1::ITensor* initialTensor = addConstantScalar(ctx, initial, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0);
-    nvinfer1::ITensor* one = addConstantScalar(ctx, 1, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0);
+    nvinfer1::ITensor* initialTensor
+        = addConstantScalar(ctx, initial, ::ONNX_NAMESPACE::TensorProto::INT64, nvinfer1::Dims{1, 1})->getOutput(0);
+    nvinfer1::ITensor* one
+        = addConstantScalar(ctx, 1, ::ONNX_NAMESPACE::TensorProto::INT64, nvinfer1::Dims{1, 1})->getOutput(0);
 
     auto counter = loop->addRecurrence(*initialTensor);
-    nvinfer1::ITensor* addOne = ctx->network()->addElementWise(*counter->getOutput(0), *one, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
+    nvinfer1::ITensor* addOne = ctx->network()
+                                    ->addElementWise(*counter->getOutput(0), *one, nvinfer1::ElementWiseOperation::kSUM)
+                                    ->getOutput(0);
     counter->setInput(1, *addOne);
     return counter->getOutput(0);
 }

LoopHelpers.hpp

@@ -11,6 +11,6 @@
 namespace onnx2trt
 {
 
-nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int32_t initial = 0);
+nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int64_t initial = 0);
 
 } // namespace onnx2trt