44
55#include " ConditionalHelpers.hpp"
66#include " ModelImporter.hpp"
7- #include " onnx2trt_utils .hpp"
7+ #include " importerUtils .hpp"
88#include " toposort.hpp"
99
1010namespace onnx2trt
@@ -34,7 +34,7 @@ SubgraphPortsMap::const_iterator findLayer(const SubgraphPortsMap& inputs, const
3434
3535// Add an ConditionalInputLayer between `layer` and its inputs.
3636// I.e. input[inIdx] -> layer ==> input[inIdx] -> ConditionalInputLayer -> layer.
37- Status addConditionalInputLayer (IImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
37+ Status addConditionalInputLayer (ImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
3838 nvinfer1::ILayer& layer, int32_t inIdx)
3939{
4040 auto input = layer.getInput (inIdx);
@@ -54,27 +54,28 @@ Status addConditionalInputLayer(IImporterContext* ctx, nvinfer1::IIfConditional*
5454 nvinfer1::IIfConditionalInputLayer* inputLayer = nullptr ;
5555 if (it == inputsMap.end ())
5656 {
57- inputLayer = conditional->addInput (*input);
57+ inputLayer = N_CHECK ( conditional->addInput (*input) );
5858 inputsMap[name] = inputLayer;
5959 const std::string inputLayerName (name);
6060 ctx->registerLayer (inputLayer, inputLayerName + " _InputLayer" , nullptr );
6161 // Note: Since multiple conditionals may use the same external tensor, check unique names for output tensors of
6262 // IfConditionalInputLayers to avoid tensor name duplication.
6363 ctx->registerTensor (
64- TensorOrWeights{inputLayer->getOutput (0 )}, inputLayerName + " _InputLayer_output" , /* checkUniqueName*/ true );
64+ TensorOrWeights{N_CHECK ( inputLayer->getOutput (0 ) )}, inputLayerName + " _InputLayer_output" , /* checkUniqueName*/ true );
6565 }
6666 else
6767 {
6868 // An InputLayer may in the inputsMap if it has several consumers.
6969 inputLayer = it->second ;
7070 }
71- layer.setInput (inIdx, *(inputLayer->getOutput (0 )));
71+ auto ifOutput = N_CHECK (inputLayer->getOutput (0 ));
72+ layer.setInput (inIdx, *ifOutput);
7273 return Status::success ();
7374};
7475
7576// Take a snapshot of the network before and after parsing the subgraph and return a list
7677// of newly added network layers.
77- Status importSubgraph (IImporterContext * ctx, ::ONNX_NAMESPACE::GraphProto const & subgraph,
78+ Status importSubgraph (ImporterContext * ctx, ::ONNX_NAMESPACE::GraphProto const & subgraph,
7879 std::vector<nvinfer1::ILayer*>& newLayers, std::vector<TensorOrWeights>& subgraphTensors)
7980{
8081 auto net = ctx->network ();
@@ -84,7 +85,7 @@ Status importSubgraph(IImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const&
8485 NameScope nameScope (*ctx);
8586
8687 std::vector<Status> errors{};
87- CHECK (onnx2trt::parseGraph (ctx, subgraph, errors));
88+ CHECK_STATUS (onnx2trt::parseGraph (ctx, subgraph, errors));
8889
8990 for (int32_t i = 0 ; i < subgraph.output_size (); ++i)
9091 {
@@ -101,7 +102,7 @@ Status importSubgraph(IImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const&
101102}
102103
103104// Add an IConditionalInputLayer to `layer`'s inputs, if they don't already exist.
104- Status addConditionalInputIfNeeded (IImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
105+ Status addConditionalInputIfNeeded (ImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
105106 nvinfer1::ILayer& layer, SubgraphPortsMap subgraphInputsMap)
106107{
107108 // Return all of the layer's inputs that are external to the subgraph that
@@ -128,7 +129,7 @@ Status addConditionalInputIfNeeded(IImporterContext* ctx, nvinfer1::IIfCondition
128129}
129130
130131// Add IConditionalInputLayers to `layer`'s inputs.
131- Status addIfInputLayers (IImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
132+ Status addIfInputLayers (ImporterContext * ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
132133 const std::vector<nvinfer1::ILayer*>& newLayers)
133134{
134135 // Find all of the tensors entering the subgraph.
0 commit comments