Skip to content

Commit

Permalink
ONNX-TensorRT 8.6-EA release
Browse files Browse the repository at this point in the history 
Signed-off-by: Kevin Chen <[email protected]>
  • Loading branch information
@kevinch-nv
kevinch-nv committed Mar 14, 2023
1 parent fd119fe commit 6872a94
Show file tree
Hide file tree
Showing 17 changed files with 1,223 additions and 650 deletions.
4 changes: 2 additions & 2 deletions CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ add_definitions("-DSOURCE_LENGTH=${SOURCE_LENGTH}")
# Version information
#--------------------------------------------------
set(ONNX2TRT_MAJOR 8)
set(ONNX2TRT_MINOR 5)
set(ONNX2TRT_PATCH 1)
set(ONNX2TRT_MINOR 6)
set(ONNX2TRT_PATCH 0)
set(ONNX2TRT_VERSION "${ONNX2TRT_MAJOR}.${ONNX2TRT_MINOR}.${ONNX2TRT_PATCH}" CACHE STRING "ONNX2TRT version")

#--------------------------------------------------
Expand Down
39 changes: 16 additions & 23 deletions ConditionalHelpers.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ using LayerName = std::string;
using InputIndex = int32_t;

// A SubgraphPortsMap maps either the inputs or outputs ports of each node in an ONNX graph.
using SubgraphPortsMap = std::unordered_map<NodeName, std::set<InputIndex>>;
using SubgraphPortsMap = std::unordered_map<nvinfer1::ITensor*, std::set<InputIndex>>;

// An InputsMap tracks which IIfConditionalInputLayer we've added to a layer's inputs,
// so that we can reuse them if needed.
Expand All @@ -27,7 +27,7 @@ using InputsMap = std::unordered_map<LayerName, nvinfer1::IIfConditionalInputLay
SubgraphPortsMap::const_iterator findLayer(const SubgraphPortsMap& inputs, const std::string layerName)
{
return std::find_if(inputs.begin(), inputs.end(), [&](const auto& item) {
const auto& key = item.first;
std::string const key = item.first->getName();
return layerName.compare(0, key.size(), key) == 0;
});
}
Expand Down Expand Up @@ -57,7 +57,7 @@ Status addConditionalInputLayer(IImporterContext* ctx, nvinfer1::IIfConditional*
inputLayer = conditional->addInput(*input);
inputsMap[name] = inputLayer;
const std::string inputLayerName(name);
ctx->registerLayer(inputLayer, inputLayerName + "_InputLayer");
ctx->registerLayer(inputLayer, inputLayerName + "_InputLayer", nullptr);
// Note: Since multiple conditionals may use the same external tensor, check unique names for output tensors of
// IfConditionalInputLayers to avoid tensor name duplication.
ctx->registerTensor(
Expand Down Expand Up @@ -132,9 +132,8 @@ Status addIfInputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditi
{
// Find all of the tensors entering the subgraph.
// The node-names are from the ONNX context.
using NodeName = std::string;
using InputIndex = int32_t;
std::unordered_map<NodeName, std::set<InputIndex>> subgraphInputsMap;
std::unordered_map<nvinfer1::ITensor*, std::set<InputIndex>> subgraphInputsMap;
getSubgraphInputs(newLayers, subgraphInputsMap);

// Add a ConditionalInputLayer in front of each input that is external to the subgraph.
Expand Down Expand Up @@ -166,9 +165,8 @@ Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* condit
}
};

using NodeName = std::string;
std::unordered_map<NodeName, std::set<int32_t>> thenOutputs;
std::unordered_map<NodeName, std::set<int32_t>> elseOutputs;
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);
Expand All @@ -182,14 +180,9 @@ Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* condit
= [](IImporterContext* ctx, std::vector<nvinfer1::ITensor*>& sgOutputs, SubgraphPortsMap& subgraphOutputs,
::ONNX_NAMESPACE::GraphProto const& subgraph, std::vector<nvinfer1::ILayer*> subgraphLayers,
StringMap<TensorOrWeights> const& subgraphTensors) {
for (const auto& layer : subgraphLayers)
for (auto const& pair : subgraphOutputs)
{
const auto layerName = layer->getName();
auto iter = findLayer(subgraphOutputs, layerName);
if (iter != subgraphOutputs.end())
{
sgOutputs.push_back(layer->getOutput(0));
}
sgOutputs.push_back(pair.first);
}

if (sgOutputs.empty())
Expand Down Expand Up @@ -221,15 +214,15 @@ Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* condit
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");
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<std::string, std::set<int32_t>>& externalOutputs, bool extractOutputs,
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalOutputs, bool extractOutputs,
const std::vector<std::string>* reportedOutputs = nullptr)
{
using NodeName = std::string;
Expand Down Expand Up @@ -271,7 +264,7 @@ Status getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
};

// Retrieve the list of tensors either exiting or entering the subgraph.
std::unordered_map<TensorName, std::vector<Port>> externalPortsMap;
std::unordered_map<nvinfer1::ITensor*, std::vector<Port>> externalPortsMap;
auto filterTensors = [&](TensorsSet const& tensors, auto getNodeAccessor) {
for (nvinfer1::ILayer const* l : newLayers)
{
Expand Down Expand Up @@ -307,7 +300,7 @@ Status getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
}
if (!reportedOutputs || prefixFound)
{
externalPortsMap[tensorName].push_back(std::make_pair(nodeName, i));
externalPortsMap[tensor].push_back(std::make_pair(nodeName, i));
}
}
i++;
Expand All @@ -330,23 +323,23 @@ Status getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
{
for (const Port& inPort : input.second)
{
auto const nodeName = inPort.first;
auto* tensor = input.first;
auto const portIndex = inPort.second;
externalOutputs[nodeName].insert(portIndex);
externalOutputs[tensor].insert(portIndex);
}
}
return Status::success();
}

Status getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
std::unordered_map<std::string, std::set<int32_t>>& externalOutputs,
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalOutputs,
const std::vector<std::string>& reportedOutputs)
{
return getSubgraphTensors(newLayers, externalOutputs, true, &reportedOutputs);
}

Status getSubgraphInputs(const std::vector<nvinfer1::ILayer*>& newLayers,
std::unordered_map<std::string, std::set<int32_t>>& externalInputs)
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& externalInputs)
{
return getSubgraphTensors(newLayers, externalInputs, false);
}
Expand Down
15 changes: 7 additions & 8 deletions ConditionalHelpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,16 @@ namespace onnx2trt
{

// Given a subgraph, find all of its external inputs (tensors entering the subgraph).
// The result is returned in `subgraphInputs`, which is a map indexed by layer-name and with values indicating a set
// of external input indices.
Status getSubgraphInputs(
const std::vector<nvinfer1::ILayer*>& newLayers,
std::unordered_map<std::string, std::set<int32_t>>& subgraphInputs);
// The result is returned in `subgraphInputs`, which is a map indexed by ITensor (a tensor entering the subgraph) and
// with values indicating a set of external input indices.
Status getSubgraphInputs(std::vector<nvinfer1::ILayer*> const& newLayers,
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& subgraphInputs);

// Given a subgraph, find all of its external outputs (tensors exiting the subgraph).
// The result is returned in `subgraphInputs`, which is a map indexed by layer-name and with values indicating a set
// of external outputs indices.
// The result is returned in `subgraphInputs`, which is a map indexed by ITensor (a tensor exiting the subgraph) and
// with values indicating a set of external outputs indices.
Status getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
std::unordered_map<std::string, std::set<int32_t>>& subgraphOutputs,
std::unordered_map<nvinfer1::ITensor*, std::set<int32_t>>& subgraphOutputs,
const std::vector<std::string>& reportedOutputs);

// Take a snapshot of the network before and after parsing the subgraph and return a list
Expand Down
154 changes: 153 additions & 1 deletion ImporterContext.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,26 @@
*/

#include "ImporterContext.hpp"
#include "NvInferVersion.h"
#include <sstream>

#if !defined(_WIN32)
#include <dlfcn.h>
#if defined(__linux__)
#include <link.h>
#endif
#else // defined(_WIN32)
#include <windows.h>
#endif // !defined(_WIN32)

#define RT_ASSERT(cond) \
do \
{ \
if (!(cond)) \
{ \
throw std::runtime_error("Assertion " #cond " failed!"); \
} \
} while (0)

namespace onnx2trt
{
Expand Down Expand Up @@ -89,7 +109,7 @@ void ImporterContext::registerTensor(TensorOrWeights tensor, std::string const&
p.first->second = std::move(tensor);
}

void ImporterContext::registerLayer(nvinfer1::ILayer* layer, std::string const& basename)
void ImporterContext::registerLayer(nvinfer1::ILayer* layer, std::string const& basename, ::ONNX_NAMESPACE::NodeProto const* node)
{
// No layer will be added for Constant nodes in ONNX.
if (layer)
Expand All @@ -111,6 +131,138 @@ void ImporterContext::registerLayer(nvinfer1::ILayer* layer, std::string const&
mConstantLayers.insert({uniqueName, static_cast<nvinfer1::IConstantLayer*>(layer)});
}
}
if (node != nullptr)
{
processMetadata(*node, layer);
}
}

void ImporterContext::registerLayer(nvinfer1::ILayer* layer, ::ONNX_NAMESPACE::NodeProto const& node)
{
std::string const& basename = getNodeName(node);
registerLayer(layer, basename, &node);
}

namespace
{

//! Translates a "logical" library name into an OS-dependent DSO or DLL name
std::string getOSLibraryName(char const* logicalName)
{
std::stringstream libName;
#if defined(_WIN32)
libName << logicalName << ".dll";
#else
libName << "lib" << logicalName << ".so." << NV_TENSORRT_SONAME_MAJOR;
#endif
return libName.str();
}

//! Platform-agnostic wrapper around dynamic libraries.
class DynamicLibrary
{
public:
explicit DynamicLibrary(std::string const& name)
: mLibName{name}
{
#if defined(_WIN32)
mHandle = LoadLibraryA(name.c_str());
#else // defined(_WIN32)
int32_t flags{RTLD_LAZY};
mHandle = dlopen(name.c_str(), flags);
#endif // defined(_WIN32)

if (mHandle == nullptr)
{
std::string errorStr{};
#if !defined(_WIN32)
errorStr = std::string{" due to "} + std::string{dlerror()};
#endif
throw std::runtime_error("Unable to open library: " + name + errorStr);
}
}

DynamicLibrary(DynamicLibrary const&) = delete;
DynamicLibrary(DynamicLibrary const&&) = delete;

~DynamicLibrary()
{
try
{
#if defined(_WIN32)
RT_ASSERT(static_cast<bool>(FreeLibrary(static_cast<HMODULE>(mHandle))));
#else
RT_ASSERT(dlclose(mHandle) == 0);
#endif
}
catch (...)
{
std::cerr << "Unable to close library: " << mLibName << std::endl;
}
}

std::string getFullPath() const
{
RT_ASSERT(mHandle != nullptr);
#if defined(__linux__)
link_map* linkMap = nullptr;
auto const err = dlinfo(mHandle, RTLD_DI_LINKMAP, &linkMap);
RT_ASSERT(err == 0 && linkMap != nullptr && linkMap->l_name != nullptr);
return std::string{linkMap->l_name};
#elif defined(_WIN32)
constexpr int32_t kMAX_PATH_LEN{4096};
std::string path(kMAX_PATH_LEN, '\0'); // since C++11, std::string storage is guaranteed to be contiguous
auto const pathLen = GetModuleFileNameA(static_cast<HMODULE>(mHandle), &path[0], kMAX_PATH_LEN);
RT_ASSERT(GetLastError() == ERROR_SUCCESS);
path.resize(pathLen);
path.shrink_to_fit();
return path;
#else
RT_ASSERT(!"Unsupported operation: getFullPath()");
#endif
}

private:
std::string mLibName{}; //!< Name of the DynamicLibrary
void* mHandle{}; //!< Handle to the DynamicLibrary
};

//! Translates an OS-dependent DSO/DLL name into a path on the filesystem
std::string getOSLibraryPath(std::string const& osLibName)
{
DynamicLibrary lib{osLibName};
return lib.getFullPath();
}

} // namespace

void ImporterContext::addUsedVCPluginLibrary(
::ONNX_NAMESPACE::NodeProto const& node, char const* pluginName, char const* pluginLib)
{
auto* ctx = this; // For logging
auto osPluginLibName = getOSLibraryName(pluginLib);
LOG_VERBOSE("Node " << getNodeName(node) << " requires plugin " << pluginName << " which is provided by "
<< osPluginLibName);
mLogicalVCPluginLibraries.insert(osPluginLibName);
}

std::vector<std::string> ImporterContext::getUsedVCPluginLibraries()
{
auto* ctx = this; // For logging
#if defined(_WIN32) || defined(__linux__)
std::vector<std::string> ret;
ret.reserve(mLogicalVCPluginLibraries.size());
for (auto const& l : mLogicalVCPluginLibraries)
{
auto osLibPath = getOSLibraryPath(l);
LOG_VERBOSE("Library " << l << " located on filesystem as " << osLibPath);
ret.emplace_back(std::move(osLibPath));
}
return ret;
#else
LOG_WARNING("getUsedVCPluginLibraries not implemented on platform!");
return {};
#endif
}

} // namespace onnx2trt
Loading

0 comments on commit 6872a94

Please sign in to comment.