Skip to content

Commit a605ff5

Browse files
authored
build: Extend TRT Plugin Handling to Support Windows (#7924)
1 parent 48e615d commit a605ff5

File tree

1 file changed

+28
-3
lines changed

1 file changed

+28
-3
lines changed

qa/common/gen_qa_trt_plugin_models.py

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/usr/bin/env python3
22

3-
# Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
44
#
55
# Redistribution and use in source and binary forms, with or without
66
# modification, are permitted provided that the following conditions
@@ -27,6 +27,7 @@
2727
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2828

2929
import argparse
30+
import ctypes
3031
import os
3132

3233
import numpy as np
@@ -38,13 +39,13 @@
3839
TRT_LOGGER = trt.Logger()
3940

4041
trt.init_libnvinfer_plugins(TRT_LOGGER, "")
41-
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
4242

4343

4444
def get_trt_plugin(plugin_name):
4545
plugin = None
4646
field_collection = None
47-
for plugin_creator in PLUGIN_CREATORS:
47+
plugin_creators = trt.get_plugin_registry().plugin_creator_list
48+
for plugin_creator in plugin_creators:
4849
if (plugin_creator.name == "CustomHardmax") and (
4950
plugin_name == "CustomHardmax"
5051
):
@@ -272,13 +273,37 @@ def create_plugin_models(models_dir):
272273
)
273274

274275

276+
def windows_load_plugin_lib(win_plugin_dll):
277+
if os.path.isfile(win_plugin_dll):
278+
try:
279+
ctypes.CDLL(win_plugin_dll, winmode=0)
280+
except TypeError:
281+
# winmode only introduced in python 3.8
282+
ctypes.CDLL(win_plugin_dll)
283+
return
284+
285+
raise IOError('Failed to load library: "{}".'.format(win_plugin_dll))
286+
287+
275288
if __name__ == "__main__":
276289
parser = argparse.ArgumentParser()
277290
parser.add_argument(
278291
"--models_dir", type=str, required=True, help="Top-level model directory"
279292
)
293+
parser.add_argument(
294+
"--win_plugin_dll",
295+
type=str,
296+
required=False,
297+
default="",
298+
help="Path to Windows plugin .dll",
299+
)
280300
FLAGS, unparsed = parser.parse_known_args()
281301

282302
import test_util as tu
283303

304+
# Linux can leverage LD_PRELOAD. We must load the Windows plugin manually
305+
# in order for it to be discovered in the registry.
306+
if os.name == "nt":
307+
windows_load_plugin_lib(FLAGS.win_plugin_dll)
308+
284309
create_plugin_models(FLAGS.models_dir)

0 commit comments

Comments
 (0)