Skip to content

Add TRT plugin custom ops support#13

Open
chilo-ms wants to merge 5 commits into
mainfrom
chi/add_custom_ops
Open

Add TRT plugin custom ops support#13
chilo-ms wants to merge 5 commits into
mainfrom
chi/add_custom_ops

Conversation

@chilo-ms

Copy link
Copy Markdown
Collaborator

Summary

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.

Reference: upstream ORT implementation

Changes

New files

  • src/tensorrt_execution_provider_custom_ops.h - Defines TensorRTCustomOp (using Ort::CustomOpBase) and TensorRTCustomKernel structs with variadic inputs/outputs, since TRT plugin API doesn't expose I/O counts.
  • src/tensorrt_execution_provider_custom_ops.cc - Implements CreateTensorRTCustomOpDomainList():
    • Loads extra TRT plugin libraries (LoadLibraryA/dlopen)
    • Initializes default TRT plugins via initLibNvInferPlugins
    • Enumerates all registered plugin creators (V1 and V3One types)
    • Creates custom ops under the trt.plugins domain
    • Uses static storage with mutex for thread-safe, process-lifetime persistence

Modified files

  • src/tensorrt_provider_factory.h - Added GetNumCustomOpDomainsImpl/GetCustomOpDomainsImpl static methods and custom_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 - Added TRTPluginsCustomOpTest adapted from the upstream TensorrtExecutionProviderTest.TRTPluginsCustomOpTest. Validates that a model with trt.plugins domain 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 with DisentangledAttention_TRT node.

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>
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

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>
- Run 'Run unit tests' (trt_ep_tests) first
- Rename 'Run tests' to 'Run onnxruntime_provider_test' for clarity

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The ORT build artifact may have onnxruntime.dll in a subdirectory.
Use Get-ChildItem -Recurse to find it rather than assuming a fixed path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants