Add TRT plugin custom ops support - #13
Merged
Merged
Conversation
Register TensorRT plugins as custom ops under the 'trt.plugins' domain, enabling ORT to recognize and route nodes with TRT plugin op types to the TensorRT EP. Implementation: - Add TensorRTCustomOp/TensorRTCustomKernel structs using Ort::CustomOpBase with variadic inputs/outputs (TRT plugin API doesn't expose I/O counts) - Add CreateTensorRTCustomOpDomainList() that initializes default TRT plugins, loads extra plugin libraries, enumerates all registered plugin creators (V1 and V3One), and creates custom ops in the 'trt.plugins' domain - Wire GetNumCustomOpDomains/GetCustomOpDomains on OrtEpFactory - Add unit test adapted from upstream TRTPluginsCustomOpTest Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
|
Use DISCOVERY_MODE PRE_TEST so test discovery happens at ctest runtime instead of build time. This avoids exit code 0xc0000135 (DLL not found) when onnxruntime.dll and CUDA DLLs are not in PATH during the build step. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The trt_ep_tests unit tests need onnxruntime.dll/libonnxruntime.so at runtime but it was not included in the uploaded artifact. Copy it from the NuGet-fetched ORT package during the build-trtep step so it's available alongside the test exe. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The test binary links against libonnxruntime.so.1 (versioned soname) but the NuGet package file is named libonnxruntime.so. Create a symlink so the dynamic linker can find it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…parser errors The output type mapping in CreateNodeComputeInfoFromGraph was using positional indexing (graph_output[i]) assuming TRT network outputs are in the same order as the ONNX graph proto outputs. This causes out-of- bounds access when TRT reorders outputs, leading to a segfault. Fix by building a name-to-index map and looking up output types by name. Also add error checking for trt_parser->parse() to fail gracefully instead of continuing with an invalid network state. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…Node Match upstream ORT behavior by checking the 'source' attribute of EPContext nodes in GraphHasCtxNode. If the source doesn't match TensorrtExecutionProvider, the node is skipped so the EP doesn't claim it. This prevents segfaults when processing EPContext nodes from other EPs (e.g., OpenVINO) that contain invalid engine data for TRT. Also adds a secondary source check in GetEpContextFromGraph as defense-in-depth. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ach call The static TensorrtLogger in GetTensorrtLogger() captured the OrtLogger reference at first construction and never updated it. When multiple sessions are created sequentially (as in unit tests), the first session's OrtLogger is destroyed, leaving a dangling reference. The next session's TRT operations would trigger logging through the stale reference, causing a segfault. Fix: Change TensorrtLogger to store OrtLogger as a pointer, add update_logger() method, and call it on every GetTensorrtLogger() invocation. Also add null check before logging. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Register TensorRT plugins as custom ops under the
trt.pluginsdomain, enabling ORT to recognize and route nodes with TRT plugin op types to the TensorRT EP.Reference: upstream ORT implementation
Changes
New files
src/tensorrt_execution_provider_custom_ops.h- DefinesTensorRTCustomOp(usingOrt::CustomOpBase) andTensorRTCustomKernelstructs with variadic inputs/outputs, since TRT plugin API doesn't expose I/O counts.src/tensorrt_execution_provider_custom_ops.cc- ImplementsCreateTensorRTCustomOpDomainList():LoadLibraryA/dlopen)initLibNvInferPluginstrt.pluginsdomainModified files
src/tensorrt_provider_factory.h- AddedGetNumCustomOpDomainsImpl/GetCustomOpDomainsImplstatic methods andcustom_op_domain_list_member.src/tensorrt_provider_factory.cc- Wired the two function pointers in the constructor and implemented the static methods.Tests
tests/tensorrt_basic_test.cc- AddedTRTPluginsCustomOpTestadapted from the upstreamTensorrtExecutionProviderTest.TRTPluginsCustomOpTest. Validates that a model withtrt.pluginsdomain ops can be loaded and session creation succeeds.tests/testdata/trt_plugin_custom_op_test.py- Python script to generate the test model.tests/testdata/trt_plugin_custom_op_test.onnx- Generated test model withDisentangledAttention_TRTnode.