chore: import upstream snapshot with attribution

python/packaging/bindings_wheel/tensorrt/plugin/__init__.py

@@ -0,0 +1,46 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorrt as trt
+
+logger = trt.Logger()
+logger.log(trt.Logger.WARNING, "Functionality provided through tensorrt.plugin module is experimental.")
+
+# export.public_api() will expose things here. To make sure that happens, we just need to
+# import all the submodules so that the decorator is actually executed (__discover_modules() below).
+__all__ = []
+
+def __discover_modules():
+    import importlib
+    import pkgutil
+
+    mods = [importlib.import_module(__package__)]
+    while mods:
+        mod = mods.pop(0)
+
+        yield mod
+
+        if hasattr(mod, "__path__"):
+            mods.extend(
+                [
+                    importlib.import_module(f"{mod.__name__}.{submod.name}")
+                    for submod in pkgutil.iter_modules(mod.__path__)
+                ]
+            )
+
+
+_ = list(__discover_modules())

python/packaging/bindings_wheel/tensorrt/plugin/_autotune.py

@@ -0,0 +1,270 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import builtins
+import tensorrt as trt
+from typing import List, Iterable
+import copy
+
+from ._utils import _str_to_data_type
+from ._export import public_api
+
+
+# "onesided" means either type or format combinations. After combinations for each are separately generated, we will combine them later.
+# e.g. io_variants = ["FP32|FP16", "FP32|FP16", "FP32*FP16"] for a plugin with 3 I/Os. i.e. I/O indices 0 and 1 are dependently either FP32/FP16 and index 2 is independently FP32/FP16.
+# There will be 2 * 2 = 4 combinations here: ["FP32", "FP32", "FP32"], ["FP16", "FP16", "FP32"], ["FP32", "FP32", "FP16"], ["FP16", "FP16", "FP16"]
+def _gen_onesided_combinations(io_variants):
+
+    # Algorithm:
+    # (1) Ignore independent variants and count the (max) number of dependent variants `mx_poly`
+    # (2) Compile initial list of #`mx_poly` combinations using the first option (option 0) for any independent variants
+    # (3) For each independent variant IO index, add combinations with that index replaced by option 1, 2, ...
+
+    combinations = []
+    mx_poly = 0  # This is the number of dependent variants
+
+    for io_variant in io_variants:
+        io_variant_list = io_variant.split("|")
+
+        if len(io_variant_list) > 1:
+            if "*" in io_variant:
+                raise ValueError(
+                    f"Type/Format '{io_variant}' contains both '|' and '*'"
+                )
+            if mx_poly > 1:
+                if mx_poly != len(io_variant_list):
+                    raise ValueError(
+                        f"Type/Format combinations {io_variants} contain illegal dependent lengths"
+                    )
+
+        mx_poly = builtins.max(mx_poly, len(io_variant_list))
+
+    for _ in range(mx_poly):
+        combinations.append([None] * len(io_variants))
+
+    for j, io_variant in enumerate(io_variants):
+        io_variant_list = io_variant.split("|")
+
+        if len(io_variant_list) == 1:
+            if "*" in io_variant:
+                io_variant_list = io_variant.split("*")
+            for i in range(len(combinations)):
+                combinations[i][j] = io_variant_list[0]
+        else:
+            for k in range(len(io_variant_list)):
+                combinations[k][j] = io_variant_list[k]
+
+    for j, io_variant in enumerate(io_variants):
+        new_combs = []
+        if "*" in io_variant:
+            io_variant_list = io_variant.split("*")
+            for k in range(1, len(io_variant_list)):
+                for c in combinations:
+                    new_c = copy.deepcopy(c)
+                    new_c[j] = io_variant_list[k]
+                    new_combs.append(new_c)
+            combinations.extend(new_combs)
+
+    return combinations
+
+
+class _TypeFormatCombination:
+    def __init__(self, num=0):
+        self.types = [None] * num
+        self.layouts = [None] * num
+        self.tactics = []
+
+    def set_types(self, types):
+        self.types = types
+
+    def set_layouts(self, layouts=None):
+        if isinstance(layouts, List):
+            self.layouts = layouts
+        else:
+            self.layouts = [layouts] * len(self.types)
+
+    def __hash__(self):
+        return hash((tuple(self.types), tuple(self.layouts)))
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, _TypeFormatCombination)
+            and self.types == other.types
+            and self.layouts == other.layouts
+        )
+
+    def __str__(self) -> str:
+        return "{" + str(self.types) + ", " + str(self.layouts) + "}"
+
+
+@public_api()
+class AutoTuneCombination:
+    def __init__(
+        self, io_types: str = None, layouts: str = None, tactics: Iterable[int] = None
+    ):
+        """
+        Construct a set of supported type/format combinations of a plugin's I/O.
+
+        Any custom *tactic* s per each such type/format combination can also be advertised. A tactic is simply another way to
+        calculate the output of a plugin for the same type/format combination of the I/O (e.g. if there are multiple kernels available).
+
+        Args:
+            io_types (str, optional): A string representation of a type combination.
+
+                Valid format is "type0,type1,...,type#io" where 'type' is of the form "TYPE0[sep]TYPE1[sep]...".
+
+                TYPE is a valid string representation of a `trt.DataType`. These include "FP32" for trt.float32, "FP16" for trt.float16. The string representation of other data types is the same as their name in the trt.DataType enum.
+
+
+                [sep] is a valid separator, which is either '|' or '*'. Only one of these separators can appear in a given `io_types`.
+
+                (1). '|' indicates a dependent combination: the dependence of the type of one I/O to another I/O. e.g. "FP32|FP16,FP32|FP16" indicates the IO can only be both FP32 or both FP16.
+
+                (2). '*' indicates an independent combination. e.g. "FP32*FP16,FP32|FP16,FP32|FP16" indicates that the first input is independently either FP32 or FP16 regardless of the rest of the IO.
+
+            layouts (str, optional): A string representation of a format combination.
+
+                Valid format is "format0,format1,...,format#io" where 'format' is of the form "FORMAT0[sep]FORMAT1[sep]...".
+
+                FORMAT is a valid string representation of a `trt.TensorFormat`. These are string versions for the enum values of `trt.TensorFormat`. e.g. "LINEAR" for `trt.TensorFormat.LINEAR`.
+
+                [sep] is a valid separator, which is either '|' or '*'. The rules are the same as for `io_types`.
+
+            tactics (Iterable[int], optional): Custom tactics for this type/format combination. Each custom tactic must be a positive integer. Defaults to default tactic (0).
+
+        .. code-block:: python
+            :linenos:
+            :caption: For a plugin with 3 I/Os, I/O indices 0 and 1 are dependently either FP32/FP16 and index 2 is independently FP32/FP16.
+
+            @trtp.autotune("my::plugin")
+            def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+                # The following would result in the following type combinations:
+                # [FP32, FP32, FP32], [FP16, FP16, FP32], [FP32, FP32, FP16], [FP16, FP16, FP16]
+                return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16, FP32|FP16", "LINEAR", [1, 2])]
+
+        .. code-block:: python
+            :linenos:
+            :caption: For a plugin with 2 I/Os, the input/output supports either LINEAR or HWC format for FP32 and LINEAR format for FP16.
+
+            @trtp.autotune("my::plugin")
+            def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+                # Even though (FP16, HWC) is not a valid combination (see next example), TRT should intelligently reject those
+                # and pass the following combinations to the impl function:
+                # [{FP32, FP32}, {LINEAR, LINEAR}], [{FP32, FP32}, {HWC, LINEAR}], [{FP16, FP32}, {LINEAR, LINEAR}]
+                return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR*HWC, LINEAR", [1, 2])]
+
+        .. code-block:: python
+            :linenos:
+            :caption: For a plugin with 2 I/Os, the input/output supports either LINEAR or HWC format for FP32 and LINEAR format for FP16 (second method).
+
+            @trtp.autotune("my::plugin")
+            def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+                # We can use two AutoTuneCombination objects to avoid communicating illegal combinations
+                return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR, LINEAR", [1, 2]), trtp.AutoTuneCombination("FP32, FP32", "HWC, LINEAR", [1, 2])]
+        """
+
+        if io_types is not None:
+            self.io_types = [s.strip() for s in io_types.split(",")]
+            if layouts is None:
+                layouts = "LINEAR"
+            self.layouts = [s.strip() for s in layouts.split(",")]
+
+            if len(self.layouts) > 1:
+                assert len(self.io_types) == len(self.layouts)
+
+            if len(self.io_types) > len(self.layouts):
+                assert len(self.layouts) == 1
+                self.layouts = [self.layouts[0]] * len(self.io_types)
+        else:
+            self.io_types = []
+            self.layouts = []
+
+        self.combinations = []
+        self._tactics = tactics
+
+    def pos(self, pos: Iterable[int], io_types: str, layouts: str = "LINEAR") -> None:
+        """
+        Specify I/O types and formats for a specified set of I/O indices.
+
+        Args:
+            pos (Iterable[int]): I/O indices. Input indices are [0, 1, ..., #inputs - 1] and output indices are [#inputs, #inputs + 1, ..., #inputs + #outputs - 1].
+            io_types (str): Data types for these I/O indices.
+            layouts (str, optional): Tensor format(s) for these I/O indices. Defaults to "LINEAR".
+        Raises:
+            ValueError: If types or layouts for any of these I/O indices is already specified.
+
+        .. code-block:: python
+            :linenos:
+            :caption: For a plugin with 3 I/Os, I/O indices 0 and 1 are dependently either FP32/FP16 and index 2 is independently FP32/FP16.
+
+            @trtp.autotune("my::plugin")
+            def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+                c = trtp.AutoTuneCombination()
+                c.pos([0, 1], "FP32|FP16", "LINEAR")
+                c.pos(2, "FP32*FP16") # Omitting format is the same as declaring it to be LINEAR.
+                c.tactics([1, 2])
+                return [c]
+        """
+        if max(pos) >= len(self.io_types):
+            self.io_types.extend([None] * (max(pos) + 1 - len(self.io_types)))
+            self.layouts.extend([None] * (max(pos) + 1 - len(self.layouts)))
+            assert len(self.io_types) == len(self.layouts)
+
+        for p in pos:
+            if self.io_types[p] is not None:
+                raise ValueError(f"Type(s) for position {p} already specified")
+            if self.layouts[p] is not None:
+                raise ValueError(f"Layout(s) for position {p} already specified")
+            self.io_types[p] = io_types
+            self.layouts[p] = layouts
+
+    def tactics(self, tactics: Iterable[int]) -> None:
+        """
+        Specify custom tactics for this type/format combination
+
+        Args:
+            tactics (Iterable[int]): Custom tactics. These must be positive integers.
+        """
+        self._tactics = tactics
+
+    def _generate_combinations(self):
+
+        self.combinations = []
+
+        type_combinations = _gen_onesided_combinations(self.io_types)
+        layout_combinations = _gen_onesided_combinations(self.layouts)
+
+        for t in type_combinations:
+            for l in layout_combinations:
+                c = _TypeFormatCombination(len(self.io_types))
+                c.types = [_str_to_data_type(tt) for tt in t]
+                c.layouts = [getattr(trt.TensorFormat, ff) for ff in l]
+                c.tactics = self._tactics
+                self.combinations.append(c)
+
+    def _get_combinations(self):
+        self._generate_combinations()
+        return self.combinations
+
+    def _check(self, pos, type, layout):
+        for i in range(len(self.combinations)):
+            if (
+                self.combinations[i].types[pos] == _str_to_data_type(type)
+                and self.combinations[i].layouts[pos] == layout.name
+            ):
+                return True
+        return False

python/packaging/bindings_wheel/tensorrt/plugin/_export.py

@@ -0,0 +1,39 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorrt as trt
+from types import ModuleType
+import importlib
+
+def public_api(module: ModuleType = None, symbol: str = None):
+    def export_impl(obj):
+        nonlocal module, symbol
+
+        module = module or importlib.import_module(__package__)
+        symbol = symbol or obj.__name__
+
+        if not hasattr(module, "__all__"):
+            module.__all__ = []
+
+        module.__all__.append(symbol)
+        setattr(module, symbol, obj)
+
+        return obj
+
+    return export_impl
+
+IS_AOT_ENABLED = hasattr(trt, "QuickPluginCreationRequest")

python/packaging/bindings_wheel/tensorrt/plugin/_lib.py

@@ -0,0 +1,695 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorrt as trt
+import types
+import typing
+from typing import Callable, Tuple, List
+import numpy as np
+from ._plugin_class import _TemplateJITPlugin
+from ._export import IS_AOT_ENABLED
+if IS_AOT_ENABLED:
+    from ._plugin_class import _TemplateAOTPlugin
+from ._validate import (
+    _parse_register_inputs,
+    _parse_register_return,
+    _validate_autotune,
+    _validate_impl,
+    _validate_aot_impl,
+    _validate_name_and_namespace,
+)
+from ._utils import (
+    _built_in_to_plugin_field_type,
+    _join_with,
+    _numpy_to_plugin_field_type,
+    _is_numpy_array,
+    _infer_numpy_type,
+)
+
+from ._export import public_api
+
+# Namespace to which plugins are dynamically bound
+# A namespace can be thought of as a library of plugins from the same author/common objective
+class _PluginNamespace(types.ModuleType):
+    def __init__(self, namespace):
+        super().__init__("tensorrt.plugin.op." + namespace)
+        self._namespace = namespace
+
+    def define(self, name, plugin_def):
+        assert not hasattr(self, name)
+        setattr(self, name, plugin_def)
+
+    def __getattr__(self, name):
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object '{self._namespace}' has no attribute '{name}'"
+        )
+
+    def __repr__(self):
+        return f'_PluginNamespace(namespace="{self._namespace}")'
+
+
+# `tensorrt.plugin.op` module to which plugin namespaces are dynamically bound
+class _Op(types.ModuleType):
+    def __init__(self):
+        super().__init__("tensorrt.plugin.op")
+
+    def define_or_get(self, namespace):
+        if hasattr(self, namespace):
+            return getattr(self, namespace)
+
+        ns = _PluginNamespace(namespace)
+        setattr(self, namespace, ns)
+
+        return ns
+
+    def __getattr__(self, name):
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{name}'"
+        )
+
+
+op = _Op()
+public_api(symbol="op")(op)
+
+QDP_CREATORS = {}
+QDP_REGISTRY = {}
+
+# Contains metadata about a registered plugin and `__call__()`` that allows for a plugin instance to be created
+class PluginDef:
+    def __init__(self):
+        self.plugin_id = None  # includes namespace (format is ns::name)
+        self.register_func = None
+        self.impl_func = None
+        self.aot_impl_func = None
+        self.autotune_func = None
+        self.autotune_attr_names = None
+        self.input_tensor_names = None
+        self.input_attrs = None  # map name -> type
+        self.impl_attr_names = None
+        self.aot_impl_attr_names = None
+        self.num_outputs = None
+        self.input_arg_schema = None
+        self.expects_tactic = None
+
+    def __call__(
+        self, *args, **kwargs
+    ) -> Tuple[List[trt.ITensor], List[trt.ITensor], trt.IPluginV3]:
+        namespace, name = self.plugin_id.split("::")
+
+        input_tensors = []
+        schema_chunks = []
+
+        for t in args:
+            if not isinstance(t, trt.ITensor):
+                raise ValueError(
+                    f"Expected trt.ITensor but got input of type {type(t)}"
+                )
+
+            schema_chunks.append("ITensor")
+            input_tensors.append(t)
+
+        attrs = {}
+        for key, value in kwargs.items():
+            if key not in self.input_attrs:
+                raise ValueError(
+                    f"Unexpected attribute {key} provided. Expected one of {self.input_attrs.keys()}."
+                )
+            attrs[key] = value
+            attr_annotation = self.input_attrs[key]
+            if isinstance(value, np.ndarray):
+                if typing.get_origin(attr_annotation) == np.ndarray:
+                    np_dtype = typing.get_args(typing.get_args(attr_annotation)[1])[0]
+                    if np.dtype(np_dtype) != np.dtype(value.dtype):
+                        raise ValueError(
+                            f"Unexpected dtype '{np.dtype(value.dtype)}' for attribute '{key}'. Expected '{np_dtype}'."
+                        )
+            else:
+                if attr_annotation is not type(value):
+                    raise ValueError(
+                        f"Unexpected type '{type(value)}' for attribute '{key}'. Expected '{attr_annotation}'."
+                    )
+
+            schema_chunks.append(key)
+
+        expected_schema = (
+            f"({_join_with(['ITensor'] * len(self.input_tensor_names))}"
+            + _join_with(self.input_attrs.keys(), True)
+            + ")"
+        )
+        schema = f"({', '.join(schema_chunks)})"
+
+        if schema != expected_schema:
+            raise ValueError(
+                f"Unexpected schema {schema} received. Expected {expected_schema}."
+            )
+
+        if self.plugin_id in QDP_CREATORS:
+            plg_creator = trt.get_plugin_registry().get_creator(name, "1", namespace)
+        else:
+            attrs_types = {}
+            for key, value in kwargs.items():
+                if isinstance(value, np.ndarray):
+                    attrs_types[key] = (False, value.dtype)  # (builtin?, type)
+                else:
+                    attrs_types[key] = (True, type(value))  # (builtin?, type)
+
+            plg_creator = _register_plugin_creator(name, namespace, attrs_types)
+
+        fields = []
+        for key, value in attrs.items():
+            if isinstance(value, np.ndarray):
+                np_type = np.dtype(value.dtype)
+                if np_type == np.float16:
+                    fields.append(
+                        trt.PluginField(
+                            key, value.tobytes(), trt.PluginFieldType.UNKNOWN
+                        )
+                    )
+                else:
+                    fields.append(
+                        trt.PluginField(
+                            key, value, _numpy_to_plugin_field_type[np_type]
+                        )
+                    )
+            elif isinstance(value, str):
+                fields.append(
+                    trt.PluginField(key, value.encode(), trt.PluginFieldType.CHAR)
+                )
+            elif isinstance(value, bytes):
+                fields.append(trt.PluginField(key, value, trt.PluginFieldType.UNKNOWN))
+            else:
+                fields.append(
+                    trt.PluginField(
+                        key,
+                        np.array([value]),
+                        _built_in_to_plugin_field_type[type(value)],
+                    )
+                )
+
+        def create_plugin_instance(quick_plugin_creation_request: "trt.QuickPluginCreationRequest" = None):
+            if quick_plugin_creation_request is None:
+                plg = plg_creator.create_plugin(
+                    name,
+                    namespace,
+                    trt.PluginFieldCollection(fields),
+                    trt.TensorRTPhase.BUILD
+                )
+            else:
+                plg = plg_creator.create_plugin(
+                    name,
+                    namespace,
+                    trt.PluginFieldCollection(fields),
+                    trt.TensorRTPhase.BUILD,
+                    quick_plugin_creation_request
+                )
+
+            return input_tensors, [], plg
+
+        return create_plugin_instance
+
+class _TemplatePluginCreator(trt.IPluginCreatorV3Quick):
+    def __init__(self, name, namespace, attrs):
+        trt.IPluginCreatorV3Quick.__init__(self)
+        self.name = name
+        self.plugin_namespace = namespace
+        self.plugin_version = "1"
+        field_names = []
+        for name, (builtin, type_) in attrs.items():
+            if builtin:
+                if type_ is str:
+                    field_names.append(
+                        trt.PluginField(name, b"", trt.PluginFieldType.CHAR)
+                    )
+                elif type_ is bytes:
+                    field_names.append(
+                        trt.PluginField(name, b"", trt.PluginFieldType.UNKNOWN)
+                    )
+                else:
+                    field_names.append(
+                        trt.PluginField(
+                            name, np.array([]), _built_in_to_plugin_field_type[type_]
+                        )
+                    )
+            else:
+                field_names.append(
+                    trt.PluginField(
+                        name, np.array([]), _numpy_to_plugin_field_type[np.dtype(type_)]
+                    )
+                )
+
+        self.field_names = trt.PluginFieldCollection(field_names)
+
+    def create_plugin(self, name, namespace, fc, phase, qpcr: "trt.QuickPluginCreationRequest" = None):
+        desc = QDP_REGISTRY[f"{namespace}::{name}"]
+        name = name
+        namespace = namespace
+
+        attrs = {}
+        for f in fc:
+            if f.name not in desc.input_attrs:
+                raise AssertionError(
+                    f"Unexpected attribute {f.name} provided to create_plugin. Expected one of {desc.input_attrs.keys()}."
+                )
+
+            attr_type_annot = desc.input_attrs[f.name]
+            if _is_numpy_array(attr_type_annot):
+                np_type = _infer_numpy_type(attr_type_annot)
+                if np_type == np.float16:
+                    attrs[f.name] = np.frombuffer(f.data.tobytes(), dtype=np.float16)
+                else:
+                    attrs[f.name] = f.data.astype(np_type)
+            else:
+                if issubclass(attr_type_annot, str):
+                    attrs[f.name] = f.data.tobytes().decode("utf-8")
+                else:
+                    attrs[f.name] = attr_type_annot(f.data)
+
+        jit_or_aot = None # True if JIT is to be created, False if AOT. Not None will be asserted before plugin creation.
+
+        if qpcr is None:
+            plg = _TemplateJITPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.impl_attr_names,
+                desc.impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func,
+                desc.expects_tactic,
+            )
+
+            return plg
+
+        # If there is a strict preference, that takes precedence
+        if qpcr == trt.QuickPluginCreationRequest.STRICT_AOT:
+            if desc.aot_impl_func is None:
+                raise ValueError(f"AOT implementation requested, but not defined for '{desc.plugin_id}'. Was @trt.plugin.aot_impl defined?")
+            jit_or_aot = False
+        elif qpcr == trt.QuickPluginCreationRequest.STRICT_JIT:
+            if desc.impl_func is None:
+                raise ValueError(f"JIT implementation requested, but not defined for '{desc.plugin_id}'. Was @trt.plugin.impl defined?")
+            jit_or_aot = True
+        else:
+            aot_defined = desc.aot_impl_func is not None
+            jit_defined = desc.impl_func is not None
+
+            # A preferemce must be indicated if both AOT and JIT implementations are defined
+            if aot_defined and jit_defined:
+                if qpcr == trt.QuickPluginCreationRequest.PREFER_AOT:
+                    jit_or_aot = False
+                elif qpcr == trt.QuickPluginCreationRequest.PREFER_JIT:
+                    jit_or_aot = True
+                else:
+                    raise ValueError(f"Plugin '{desc.plugin_id}' has both AOT and JIT implementations. NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS or NetworkDefinitionCreationFlag.PREFER_JIT_PYTHON_PLUGINS should be specified.")
+            else:
+                # If only one implementation is defined, use that.
+                # Any preference specified is ignored. If the preference is strong, a strict flag should have been specified.
+                if aot_defined:
+                    jit_or_aot = False
+                elif jit_defined:
+                    jit_or_aot = True
+                else:
+                    raise ValueError(f"Plugin '{desc.plugin_id}' does not have either a AOT or JIT implementation.")
+
+        assert jit_or_aot is not None
+
+        if jit_or_aot:
+            plg = _TemplateJITPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.impl_attr_names,
+                desc.impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func,
+                desc.expects_tactic,
+            )
+
+        else:
+            plg = _TemplateAOTPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.aot_impl_attr_names,
+                desc.aot_impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func
+            )
+
+        # the caller can determine if the created plugin is an AOT or JIT plugin by inspecting the interface info
+        return plg
+
+def _register_plugin_creator(name: str, namespace: str, attrs_types):
+    plg_registry = trt.get_plugin_registry()
+    plg_creator = _TemplatePluginCreator(name, namespace, attrs_types)
+    plg_registry.register_creator(plg_creator, namespace)
+    plg_creator = plg_registry.get_creator(name, "1", namespace)
+    QDP_CREATORS[f"{namespace}::{name}"] = plg_creator
+    return plg_creator
+
+
+# Decorator for `tensorrt.plugin.register`
+# By default, the plugin will be immediately registered in the TRT plugin registry
+# During plugin development/when building engine, lazy registration may be used to delay plugin registration until the plugin is explicitly instantiated using `trt.plugin.op.ns.plugin_name(...)`
+@public_api()
+def register(plugin_id: str, lazy_register: bool = False) -> Callable:
+    """
+    Wraps a function to register and describe a TensorRT plugin's IO characteristics. In addition, a complete plugin at least needs an `trt.plugin.impl` function to be registered.
+
+    This API is only intended to be used as a decorator. The decorated function must have type hints for all inputs as well as return value.
+
+    .. code-block:: text
+
+        (inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, ...) -> Union[TensorDesc, Tuple[TensorDesc]]
+
+    * Input tensors are declared first, each described by a tensor descriptor TensorDesc.
+    * Plugin attributes are declared next. "SupportedAttrType" must be one of:
+       * Supported built-in types: int, float, str, bool, bytes (Note: Lists/tuples of these types are not supported)
+       * 1-D Numpy arrays of the following types: int8, int16, int32, int64, float16, float32, float64, bool. These must be annotated with 'numpy.typing.NDArray[dtype]', where 'dtype' is the expected numpy dtype.
+    * If the plugin has only one output, the return annotation could be TensorDesc. Tuple[TensorDesc] could be used for any number of outputs.
+
+    By default, the plugin will be immediately registered in the TRT plugin registry. Use the lazy_register argument to change this.
+
+    Args:
+        plugin_id: An ID for the plugin in the form "{namespace}::{name}",
+            e.g. "my_project::add_plugin". The namespace is used to avoid collisions
+            so using your product/project name is recommended.
+
+        lazy_register: During plugin development/when building engine, lazy registration may be used to delay plugin registration until the plugin is explicitly instantiated using `trt.plugin.op.ns.plugin_name(...)`
+
+    .. code-block:: python
+        :linenos:
+        :caption: Registration of an elementwise plugin (output has same characteristics as the input)
+
+        import tensorrt.plugin as trtp
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+    """
+
+    def decorator(register_func: Callable):
+
+        plugin_ns, plugin_name = plugin_id.split("::")
+        _validate_name_and_namespace(plugin_ns, plugin_name)
+
+        op_namespace = op.define_or_get(plugin_ns)
+
+        if hasattr(op_namespace, plugin_name):
+            raise ValueError(
+                f"'{op.__class__.__name__}' already has a defintion for '{plugin_name}'"
+            )
+
+        (
+            tensor_names,
+            input_attrs,
+            input_arg_schema,
+            attrs_types,
+        ) = _parse_register_inputs(register_func, lazy_register)
+
+        plugin_def = PluginDef()
+        plugin_def.plugin_id = plugin_id
+        plugin_def.register_func = register_func
+        plugin_def.input_tensor_names = tensor_names
+        plugin_def.input_attrs = input_attrs
+        plugin_def.input_arg_schema = input_arg_schema
+
+        num_outputs = _parse_register_return(register_func)
+
+        plugin_def.num_outputs = num_outputs
+        QDP_REGISTRY[plugin_id] = plugin_def
+
+        if not lazy_register:
+            _register_plugin_creator(plugin_name, plugin_ns, attrs_types)
+
+        op_namespace.define(plugin_name, plugin_def)
+
+        return register_func
+
+    return decorator
+
+
+# Decorator for `tensorrt.plugin.impl`
+@public_api()
+def impl(plugin_id: str) -> Callable:
+    """
+    Wraps a function to define an implementation for a plugin already registered through `trt.plugin.register`.
+
+    This API is only intended to be used as a decorator. The decorated function is not required to have type hints for input arguments or return value;
+    however, any type hints specified will be validated against the `trt.plugin.register` signature for consistency.
+
+    The schema for the function is as follows:
+
+    .. code-block:: text
+
+        (inp0: Tensor, inp1: Tensor, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[Tensor], stream: int, tactic: Optional[int]) -> None
+
+    * Input tensors are passed first, each described by a `Tensor`.
+    * Plugin attributes are declared next.
+       * Not all attributes included in `trt.plugin.register` must be specified here -- they could be a subset.
+       * Included attributes will be serialized to the TRT engine. Therefore, only attributes the plugin actually needs to perform inference (within the body of `trt.plugin.impl`) should be included.
+    * `tactic` is an optional argument. If the plugin is using custom tactics, it must be specified to receive the tactic value to use for the current execution of the plugin.
+
+    Args:
+        plugin_id: The ID for the plugin in the form "{namespace}::{name}", which must match that used during `trt.plugin.register`
+
+    .. code-block:: python
+        :linenos:
+        :caption: Implementation of an elementwise plugin with an OpenAI Triton kernel
+
+        import tensorrt.plugin as trtp
+        import triton
+        import triton.language as tl
+
+        @triton.jit
+        def add_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < n_elements
+            x = tl.load(x_ptr + offsets, mask=mask)
+            tl.store(y_ptr + offsets, x + 1, mask=mask)
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+        @trtp.impl("my::add_plugin")
+        def add_plugin_impl(inp0: trtp.Tensor, block_size: int, outputs: Tuple[trtp.Tensor], stream: int) -> None:
+
+            n = inp0.numel()
+            inp0_t = torch.as_tensor(inp0, device="cuda")
+            out_t = torch.as_tensor(outputs[0], device="cuda")
+
+            add_kernel[(triton.cdiv(n, block_size),)](inp0_t, out_t, n, BLOCK_SIZE = block_size)
+    """
+
+    def decorator(impl_func: Callable):
+        if plugin_id not in QDP_REGISTRY:
+            raise ValueError(
+                f"Plugin {plugin_id} is not registered. Did you register it with tensorrt.plugin.register API?"
+            )
+
+        plugin_def = QDP_REGISTRY[plugin_id]
+        impl_attr_names, found_tactic = _validate_impl(impl_func, plugin_def)
+
+        plugin_def.impl_func = impl_func
+        plugin_def.impl_attr_names = impl_attr_names
+        plugin_def.expects_tactic = found_tactic
+        return impl_func
+
+    return decorator
+
+# Decorator for `tensorrt.plugin.aot_impl`
+@public_api()
+def aot_impl(plugin_id: str) -> Callable:
+    """
+    Wraps a function to define an Ahead-of-Time (AOT) implementation for a plugin already registered through `trt.plugin.register`.
+
+    This API is only intended to be used as a decorator. The decorated function is not required to have type hints for input arguments or return value;
+    however, any type hints specified will be validated against the `trt.plugin.register` signature for consistency.
+
+    The schema for the function is as follows:
+    .. code-block:: text
+
+        (inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[TensorDesc], tactic: Optional[int]) -> Tuple[str, str, KernelLaunchParams, SymExprs]
+
+    * Input tensors are passed first, each described by a `TensorDesc`.
+    * Plugin attributes are declared next.
+       * Not all attributes included in `trt.plugin.register` must be specified here -- they could be a subset.
+       * NOTE: Plugin attributes are not serialized into the engine when using an AOT implementation.
+    * `tactic` is an optional argument. If the plugin is using custom tactics, it must be specified to receive the tactic value to use for the current execution of the plugin.
+
+    Args:
+        plugin_id: The ID for the plugin in the form "{namespace}::{name}", which must match that used during `trt.plugin.register`
+
+    :returns:
+        - kernel_name: The name of the kernel.
+        - compiled_kernel: Compiled form of the kernel. Presently, only PTX is supported.
+        - launch_params: The launch parameters for the kernel
+        - extra_args: Symbolic expressions for scalar inputs to the kernel, located after the tensor inputs and before the tensor outputs
+
+    .. code-block:: python
+        :linenos:
+        :caption: Implementation of an elementwise plugin with an OpenAI Triton kernel
+
+        import tensorrt.plugin as trtp
+        import triton
+        import triton.language as tl
+
+        @triton.jit
+        def add_kernel(x_ptr, n_elements, y_ptr, BLOCK_SIZE: tl.constexpr):
+            pid = tl.program_id(0)
+            offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < n_elements
+            x = tl.load(x_ptr + offsets, mask=mask)
+            tl.store(y_ptr + offsets, x + 1, mask=mask)
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+        @trtp.aot_impl("my::elemwise_add_plugin")
+        def add_plugin_aot_impl(
+            inp0: trtp.TensorDesc, block_size: int, single_tactic: bool, outputs: Tuple[trtp.TensorDesc], tactic: int
+        ) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
+
+            type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"
+
+            src = triton.compiler.ASTSource(
+                fn=add_kernel,
+                signature={
+                    "x_ptr": f"*{type_str}",
+                    "n_elements": "i32",
+                    "y_ptr": f"*{type_str}",
+                },
+                constexprs={
+                    "BLOCK_SIZE": block_size,
+                },
+            )
+
+            compiled_kernel = triton.compile(src)
+
+            N = inp0.shape_expr.numel()
+            launch_params = trtp.KernelLaunchParams()
+
+            # grid dims
+            launch_params.grid_x = trtp.cdiv(N, block_size)
+            # block dims
+            launch_params.block_x = compiled_kernel.metadata.num_warps * 32
+            # shared memory
+            launch_params.shared_mem = compiled_kernel.metadata.shared
+
+            extra_args = trtp.SymIntExprs(1)
+            extra_args[0] = trtp.SymInt32(N)
+
+            return compiled_kernel.metadata.name, compiled_kernel.asm["ptx"], launch_params, extra_args
+    """
+    def decorator(aot_impl_func: Callable):
+        if plugin_id not in QDP_REGISTRY:
+            raise ValueError(
+                f"Plugin {plugin_id} is not registered. Did you register it with tensorrt.plugin.register API?"
+            )
+
+        plugin_def = QDP_REGISTRY[plugin_id]
+        aot_impl_attr_names = _validate_aot_impl(aot_impl_func, plugin_def)
+
+        plugin_def.aot_impl_func = aot_impl_func
+        plugin_def.aot_impl_attr_names = aot_impl_attr_names
+        return aot_impl_func
+
+    return decorator
+
+
+# Decorator for `tensorrt.plugin.autotune`
+@public_api()
+def autotune(plugin_id: str) -> Callable:
+    """
+    Wraps a function to define autotune logic for a plugin already registered through `trt.plugin.register`.
+
+    Autotuning is the process by which TensorRT executes the plugin over IO type/format combinations, and any custom tactics advertised as being supported by the plugin.
+    The (type, format, tactic) combination with the lowest latency is used to execute the plugin once the engine is built.
+
+    .. note:: An autotune function is optional. If not specified, TensorRT will assume the plugin only supports input types specified at network creation, output types specifeid through `trt.plugin.register`, and linear formats for all I/O.
+
+    This API is only intended to be used as a decorator. The decorated function is not required to have type hints for input arguments or return value; however, any type hints specified will be validated against the `trt.plugin.register` signature for consistency.
+
+    The schema for the function is as follows:
+
+    .. code-block:: text
+
+        (inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[TensorDesc]) -> List[AutoTuneCombination]
+
+    * Input tensors are passed first, each described by a :class:`TensorDesc`.
+    * Plugin attributes are declared next. Not all attributes included in `trt.plugin.register` must be specified here -- they could be a subset.
+    * The function should return a list of :class:`AutoTuneCombination`\s.
+
+    Args:
+        plugin_id: The ID for the plugin in the form "{namespace}::{name}", which must match that used during `trt.plugin.register`
+
+    .. code-block:: python
+        :linenos:
+        :caption: An elementwise add plugin which supports both FP32 and FP16 linear I/O and wants to be tuned over 2 custom tactics.
+
+        import tensorrt.plugin as trtp
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+        @trtp.autotune("my::add_plugin")
+        def add_plugin_autotune(inp0: trtp.TensorDesc, block_size: int, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+
+            return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR", [1, 2])]
+
+    .. code-block:: python
+        :linenos:
+        :caption: Same as above example but using index-by-index construction of an `AutoTuneCombination`
+
+        import tensorrt.plugin as trtp
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+        @trtp.autotune("my::add_plugin")
+        def add_plugin_autotune(inp0: trtp.TensorDesc, block_size: int, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
+            c = trtp.AutoTuneCombination()
+            c.pos(0, "FP32|FP16", "LINEAR")
+            c.pos(1, "FP32|FP16") # index 1 is the output. Omitting format is the same as declaring it to be LINEAR.
+            c.tactics([1, 2])
+            return [c]
+    """
+
+    def decorator(autotune_func: Callable):
+        if plugin_id not in QDP_REGISTRY:
+            raise ValueError(
+                f"Plugin {plugin_id} is not registered. Did you register it with tensorrt.plugin.register API?"
+            )
+
+        plugin_def = QDP_REGISTRY[plugin_id]
+        autotune_attr_names = _validate_autotune(autotune_func, plugin_def)
+
+        plugin_def.autotune_func = autotune_func
+        plugin_def.autotune_attr_names = autotune_attr_names
+
+        return autotune_func
+
+    return decorator

python/packaging/bindings_wheel/tensorrt/plugin/_plugin_class.py

@@ -0,0 +1,445 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import tensorrt as trt
+from typing import Tuple, Union
+
+import numpy as np
+from ._utils import _numpy_to_plugin_field_type, _built_in_to_plugin_field_type
+from ._tensor import TensorDesc, Tensor, Shape, ShapeExpr, ShapeExprs, SymIntExpr, SymExprs, SymInt32
+from ._export import IS_AOT_ENABLED
+
+if IS_AOT_ENABLED:
+    from ._tensor import KernelLaunchParams
+from ._autotune import _TypeFormatCombination
+
+from ._export import public_api
+
+
+class _TemplatePluginBase(
+    trt.IPluginV3,
+    trt.IPluginV3QuickCore,
+    trt.IPluginV3QuickBuild,
+):
+    def __init__(self, name, namespace, num_outputs):
+        trt.IPluginV3.__init__(self)
+        trt.IPluginV3QuickCore.__init__(self)
+        trt.IPluginV3QuickBuild.__init__(self)
+
+        self.plugin_version = "1"
+        self.input_types = []
+        self.aliased_map = {}  # output index -> input index
+
+        self.plugin_namespace = namespace
+        self.plugin_name = name
+        self.num_outputs = num_outputs
+
+        self.autotune_combs = []
+        self.supported_combs = {}
+        self.curr_comb = None
+
+    def get_num_outputs(self):
+        return self.num_outputs
+
+    def get_output_data_types(self, input_types, ranks):
+        self.input_types = input_types
+
+        input_descs = [None] * len(input_types)
+        input_desc_map = {}
+        for i in range(len(input_types)):
+            input_descs[i] = TensorDesc()
+            input_descs[i].dtype = input_types[i]
+            input_descs[i].shape_expr = ShapeExprs(ranks[i], _is_dummy=True)
+            input_descs[i]._immutable = True
+            input_desc_map[id(input_descs[i])] = i
+
+        output_descs = self.register_function(*input_descs, **self.attrs)
+        if not isinstance(output_descs, Tuple):
+            output_descs = tuple([output_descs])
+
+        self.output_types = []
+
+        for i in range(len(output_descs)):
+            self.output_types.append(output_descs[i].dtype)
+
+            if output_descs[i].get_aliased() is not None:
+                self.aliased_map[i] = input_desc_map[id(output_descs[i].get_aliased())]
+            else:
+                self.aliased_map[i] = -1
+
+        return self.output_types
+
+    def get_fields_to_serialize(self):
+        fields = []
+        for key, value in self.attrs.items():
+            if key in self.impl_attr_names:
+                if isinstance(value, np.ndarray):
+                    if np.dtype(value.dtype) == np.float16:
+                        fields.append(trt.PluginField(key, value.tobytes(), trt.PluginFieldType.UNKNOWN))
+                    else:
+                        fields.append(
+                            trt.PluginField(
+                                key,
+                                value,
+                                _numpy_to_plugin_field_type[np.dtype(value.dtype)],
+                            )
+                        )
+                elif isinstance(value, str):
+                    fields.append(trt.PluginField(key, value.encode(), trt.PluginFieldType.CHAR))
+                elif isinstance(value, bytes):
+                    fields.append(trt.PluginField(key, value, trt.PluginFieldType.UNKNOWN))
+                else:
+                    fields.append(
+                        trt.PluginField(
+                            key,
+                            np.array([value]),
+                            _built_in_to_plugin_field_type[type(value)],
+                        )
+                    )
+
+        return trt.PluginFieldCollection(fields)
+
+    def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
+        assert len(shape_inputs) == 0  # Shape inputs are not yet supported for QDPs
+        SymIntExpr._exprBuilder = exprBuilder
+        self.input_descs = []
+        for i in range(len(inputs)):
+            desc = TensorDesc()
+            inp = inputs[i]
+
+            desc.dtype = self.input_types[i]
+            desc.shape_expr = ShapeExprs(len(inp))
+            for j in range(len(inp)):
+                desc.shape_expr[j] = ShapeExpr(inp[j])
+            desc._immutable = True
+
+            self.input_descs.append(desc)
+
+        self.output_descs = self.register_function(*self.input_descs, **self.attrs)
+        if not isinstance(self.output_descs, Tuple):
+            self.output_descs = tuple([self.output_descs])
+
+        for idx, desc in enumerate(self.output_descs):
+            if desc.is_size_tensor:
+                desc._set_index(idx)
+
+        output_exprs = []
+        for i in range(len(self.output_descs)):
+            exprs = trt.DimsExprs(len(self.output_descs[i].shape_expr))
+            for j in range(len(exprs)):
+                exprs[j] = self.output_descs[i].shape_expr[j]._expr
+
+            output_exprs.append(exprs)
+
+        SymIntExpr._exprBuilder = None
+        return output_exprs
+
+    def configure_plugin(self, inputs, outputs):
+        self.curr_comb = _TypeFormatCombination()
+        self.curr_comb.types = [inp.desc.type for inp in inputs] + [out.desc.type for out in outputs]
+        self.curr_comb.layouts = [inp.desc.format for inp in inputs] + [out.desc.format for out in outputs]
+
+    def get_supported_format_combinations(self, in_out, num_inputs):
+        if self.autotune_function is not None:
+            if len(self.autotune_attr_names) > 0:
+                val = [self.attrs[k] for k in self.autotune_attr_names]
+            else:
+                val = ()
+
+            for i, desc in enumerate(in_out):
+                if i < num_inputs:
+                    self.input_descs[i]._immutable = False
+                    self.input_descs[i].shape = Shape(desc)
+                    self.input_descs[i].format = desc.desc.format
+                    self.input_descs[i].scale = desc.desc.scale
+                    self.input_descs[i]._immutable = True
+                else:
+                    self.output_descs[i - num_inputs]._immutable = False
+                    self.output_descs[i - num_inputs].shape = Shape(desc)
+                    self.output_descs[i - num_inputs].format = desc.desc.format
+                    self.output_descs[i - num_inputs].scale = desc.desc.scale
+                    self.output_descs[i - num_inputs]._immutable = True
+
+            self.autotune_combs = self.autotune_function(*self.input_descs, *val, self.output_descs)
+
+        if len(self.autotune_combs) == 0:
+            default_comb = [None] * len(in_out)
+            comb = _TypeFormatCombination(len(in_out))
+            for j in range(len(in_out)):
+                default_comb[j] = trt.PluginTensorDesc()
+                default_comb[j].type = (
+                    self.input_types[j] if j < num_inputs else self.output_descs[j - num_inputs].dtype
+                )
+                default_comb[j].format = trt.TensorFormat.LINEAR
+                comb.types[j] = default_comb[j].type
+                comb.layouts[j] = default_comb[j].format
+
+            self.supported_combs[comb] = set()
+
+            return default_comb
+
+        all_combs = []
+        for comb in self.autotune_combs:
+            all_combs.extend(comb._get_combinations())
+
+        ret_supported_combs = []
+        self.supported_combs = {}
+
+        for i, comb in enumerate(all_combs):
+            value = self.supported_combs.get(comb)
+            if value is not None:
+                value.update(set(comb.tactics) if comb.tactics is not None else set())
+            else:
+                self.supported_combs[comb] = set(comb.tactics) if comb.tactics is not None else set()
+                for j in range(len(in_out)):
+                    curr_comb = trt.PluginTensorDesc()
+                    curr_comb.type = comb.types[j]
+                    curr_comb.format = comb.layouts[j]
+                    ret_supported_combs.append(curr_comb)
+
+        return ret_supported_combs
+
+    def get_aliased_input(self, output_index: int):
+        return self.aliased_map[output_index]
+
+    def get_valid_tactics(self):
+        tactics = self.supported_combs.get(self.curr_comb)
+        assert tactics is not None
+        return list(tactics)
+
+    def set_tactic(self, tactic):
+        self._tactic = tactic
+
+
+class _TemplateJITPlugin(_TemplatePluginBase, trt.IPluginV3QuickRuntime):
+    def __init__(self, name, namespace, num_outputs):
+        super().__init__(name, namespace, num_outputs)
+        trt.IPluginV3QuickRuntime.__init__(self)
+
+        self.expects_tactic = False
+
+    def init(
+        self,
+        register_function,
+        attrs,
+        impl_attr_names,
+        impl_function,
+        autotune_attr_names,
+        autotune_function,
+        expects_tactic,
+    ):
+        self.register_function = register_function
+        self.impl_function = impl_function
+        self.attrs = attrs
+        self.impl_attr_names = impl_attr_names
+        self.autotune_attr_names = autotune_attr_names
+        self.autotune_function = autotune_function
+        self.expects_tactic = expects_tactic
+
+    def get_capability_interface(self, type):
+        return self
+
+    def enqueue(
+        self,
+        input_desc,
+        output_desc,
+        inputs,
+        outputs,
+        in_strides,
+        out_strides,
+        stream,
+    ):
+        input_tensors = [None] * (len(inputs))
+        aliased_input_idxs = list(self.aliased_map.values())
+
+        for i in range(len(inputs)):
+            input_tensors[i] = Tensor()
+            input_tensors[i].dtype = input_desc[i].type
+            input_tensors[i].shape = Shape(input_desc[i])
+            input_tensors[i].format = input_desc[i].format
+            input_tensors[i].scale = input_desc[i].scale
+            input_tensors[i].data_ptr = inputs[i]
+            input_tensors[i]._stream = stream
+            input_tensors[i]._read_only = i not in aliased_input_idxs
+            input_tensors[i].strides = in_strides[i]
+
+        output_tensors = [None] * (len(outputs))
+        for i in range(len(outputs)):
+            output_tensors[i] = Tensor()
+            output_tensors[i].dtype = output_desc[i].type
+            output_tensors[i].shape = Shape(output_desc[i])
+            output_tensors[i].format = output_desc[i].format
+            output_tensors[i].scale = output_desc[i].scale
+            output_tensors[i].data_ptr = outputs[i]
+            output_tensors[i]._stream = stream
+            output_tensors[i]._read_only = False
+            output_tensors[i].strides = out_strides[i]
+
+        for i, j in self.aliased_map.items():
+            output_tensors[i]._aliased_to = input_tensors[j]
+            input_tensors[j]._aliased_to = output_tensors[i]
+
+        for t in input_tensors:
+            t._immutable = True
+
+        for t in output_tensors:
+            t._immutable = True
+
+        if len(self.impl_attr_names) > 0:
+            val = [self.attrs[k] for k in self.impl_attr_names]
+        else:
+            val = ()
+
+        if self.expects_tactic:
+            self.impl_function(*input_tensors, *val, output_tensors, stream, self._tactic)
+        else:
+            self.impl_function(*input_tensors, *val, output_tensors, stream=stream)
+
+    def clone(self):
+        cloned_plugin = _TemplateJITPlugin(self.plugin_name, self.plugin_namespace, self.num_outputs)
+        cloned_plugin.__dict__.update(self.__dict__)
+        return cloned_plugin
+
+
+if IS_AOT_ENABLED:
+
+    class _TemplateAOTPlugin(
+        _TemplatePluginBase,
+        trt.IPluginV3QuickAOTBuild,
+    ):
+        def __init__(self, name, namespace, num_outputs):
+            _TemplatePluginBase.__init__(self, name, namespace, num_outputs)
+            trt.IPluginV3QuickAOTBuild.__init__(self)
+            self.kernel_map = {}
+
+        def set_tactic(self, tactic):
+            self._tactic = tactic
+
+        def init(
+            self,
+            register_function,
+            attrs,
+            aot_impl_attr_names,
+            aot_impl_function,
+            autotune_attr_names,
+            autotune_function,
+        ):
+            self.register_function = register_function
+            self.aot_impl_function = aot_impl_function
+            self.attrs = attrs
+            self.aot_impl_attr_names = aot_impl_attr_names
+            self.autotune_attr_names = autotune_attr_names
+            self.autotune_function = autotune_function
+
+        def get_capability_interface(self, type):
+            return self
+
+        def get_kernel(self, inputDesc, outputDesc):
+            io_types = []
+            io_formats = []
+
+            for i, desc in enumerate(inputDesc):
+                io_types.append(desc.type)
+                io_formats.append(desc.format)
+
+            for i, desc in enumerate(outputDesc):
+                io_types.append(desc.type)
+                io_formats.append(desc.format)
+
+            key = (tuple(io_types), tuple(io_formats), self._tactic)
+
+            assert key in self.kernel_map, "key {} not in kernel_map".format(key)
+
+            kernel_name, ptx = self.kernel_map[key]
+
+            return kernel_name, ptx.encode() if isinstance(ptx, str) else ptx
+
+        def get_launch_params(self, inDimsExprs, in_out, num_inputs, launchParams, symExprSetter, exprBuilder):
+
+            SymIntExpr._exprBuilder = exprBuilder
+
+            if len(self.attrs) > 0:
+                _, val = zip(*self.attrs.items())
+            else:
+                val = ()
+
+            io_types = []
+            io_formats = []
+
+            for i, desc in enumerate(in_out):
+                if i < num_inputs:
+                    self.input_descs[i]._immutable = False
+                    self.input_descs[i].shape = Shape(desc)
+                    self.input_descs[i].dtype = desc.desc.type
+                    self.input_descs[i].format = desc.desc.format
+                    self.input_descs[i].scale = desc.desc.scale
+                    io_types.append(desc.desc.type)
+                    io_formats.append(desc.desc.format)
+                    self.input_descs[i]._immutable = True
+                else:
+                    self.output_descs[i - num_inputs]._immutable = False
+                    self.output_descs[i - num_inputs].shape = Shape(desc)
+                    self.output_descs[i - num_inputs].dtype = desc.desc.type
+                    self.output_descs[i - num_inputs].format = desc.desc.format
+                    self.output_descs[i - num_inputs].scale = desc.desc.scale
+                    io_types.append(desc.desc.type)
+                    io_formats.append(desc.desc.format)
+                    self.output_descs[i - num_inputs]._immutable = True
+
+            kernel_name, ptx, launch_params, extra_args = self.aot_impl_function(
+                *self.input_descs, *val, self.output_descs, self._tactic
+            )
+
+            if not isinstance(kernel_name, str) and not isinstance(kernel_name, bytes):
+                raise TypeError(f"Kernel name must be a 'str' or 'bytes'.  Got: {type(kernel_name)}.")
+
+            if not isinstance(ptx, str) and not isinstance(ptx, bytes):
+                raise TypeError(f"PTX/CUBIN must be a 'str' or 'bytes'.  Got: {type(ptx)}.")
+
+            if not isinstance(launch_params, KernelLaunchParams):
+                raise TypeError(
+                    f"Launch params must be a 'tensorrt.plugin.KernelLaunchParams'.  Got: {type(launch_params)}."
+                )
+
+            if not isinstance(extra_args, SymExprs):
+                raise TypeError(f"Extra args must be a 'tensorrt.plugin.SymIntExprs'.  Got: {type(extra_args)}.")
+
+            launchParams.grid_x = launch_params.grid_x()
+            launchParams.grid_y = launch_params.grid_y()
+            launchParams.grid_z = launch_params.grid_z()
+            launchParams.block_x = launch_params.block_x()
+            launchParams.block_y = launch_params.block_y()
+            launchParams.block_z = launch_params.block_z()
+            launchParams.shared_mem = launch_params.shared_mem()
+
+            self.kernel_map[(tuple(io_types), tuple(io_formats), self._tactic)] = (kernel_name, ptx)
+
+            symExprSetter.nbSymExprs = len(extra_args)
+
+            for i, arg in enumerate(extra_args):
+                if not isinstance(arg, SymInt32):
+                    raise TypeError(f"Extra args must be a 'tensorrt.plugin.SymInt32'.  Got: {type(arg)}.")
+                symExprSetter[i] = arg()
+
+            SymIntExpr._exprBuilder = None
+
+        def get_timing_cache_id(self):
+            return ""
+
+        def clone(self):
+            cloned_plugin = _TemplateAOTPlugin(self.plugin_name, self.plugin_namespace, self.num_outputs)
+            cloned_plugin.__dict__.update(self.__dict__)
+            return cloned_plugin

python/packaging/bindings_wheel/tensorrt/plugin/_top_level.py

@@ -0,0 +1,132 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Union, Tuple
+import tensorrt as trt
+from ._tensor import ShapeExpr, TensorDesc, ShapeExprs, SizeTensorDesc
+from ._export import public_api
+
+# Miscellaneous top-level functions accessible through `tensorrt.plugin`
+
+# Performs `trt.DimensionOperation.CEIL_DIV`
+@public_api()
+def cdiv(first: Union[int, ShapeExpr], second: Union[int, ShapeExpr]) -> ShapeExpr:
+    """
+    Computes symbolic ceiling division of `first` by `second`
+
+    Args:
+        first (Union[int, ShapeExpr]): Dividend
+        second (Union[int, ShapeExpr]): Divisor
+
+    Raises:
+        ValueError: If both arguments are `int`\s or if `second` evaluates to 0
+
+    Returns:
+        ShapeExpr: Symbolic expression for the ceiling division of `first` by `second`
+    """
+    if isinstance(first, int):
+        if isinstance(second, int):
+            raise ValueError("Both arguments cannot be 'int's")
+        first = ShapeExpr(first)
+
+    return first._op(trt.DimensionOperation.CEIL_DIV, second)
+
+
+# Performs `trt.DimensionOperation.MAX`
+@public_api()
+def max(first: Union[int, ShapeExpr], second: Union[int, ShapeExpr]) -> ShapeExpr:
+    """
+    Computes the maximum of `first` and `second`
+
+    Args:
+        first (Union[int, ShapeExpr]): First operand
+        second (Union[int, ShapeExpr]): Second operand
+
+    Raises:
+        ValueError: If both arguments are `int`\s
+
+    Returns:
+        ShapeExpr: Symbolic expression for the maximum of `first` and `second`
+    """
+    if isinstance(first, int):
+        if isinstance(second, int):
+            raise ValueError("Both arguments cannot be 'int's")
+        first = ShapeExpr(first)
+
+    return first._op(trt.DimensionOperation.MAX, second)
+
+
+# Performs `trt.DimensionOperation.MIN`
+@public_api()
+def min(first: Union[int, ShapeExpr], second: Union[int, ShapeExpr]) -> ShapeExpr:
+    """
+    Computes the minimum of `first` and `second`
+
+    Args:
+        first (Union[int, ShapeExpr]): First operand
+        second (Union[int, ShapeExpr]): Second operand
+
+    Raises:
+        ValueError: If both arguments are `int`\s
+
+    Returns:
+        ShapeExpr: Symbolic expression for the minimum of `first` and `second`
+    """
+    if isinstance(first, int):
+        if isinstance(second, int):
+            raise ValueError("Both arguments cannot be 'int's")
+        first = ShapeExpr(first)
+
+    return first._op(trt.DimensionOperation.MIN, second)
+
+
+# Declare a size tensor descriptor with the specified autotune shape expression `opt` and `upper-bound` shape expression
+@public_api()
+def size_tensor(opt: ShapeExpr, upper_bound: ShapeExpr) -> SizeTensorDesc:
+    """
+    Constructs a size tensor with the specified autotune shape expression `opt` and `upper_bound`
+
+    Args:
+        opt (ShapeExpr): Symbolic expression for the extent of this size tensor to use in the autotune process of the engine build
+        upper_bound (ShapeExpr): Symbolic expression for the upper-bound of this size tensor
+
+    Returns:
+        SizeTensorDesc: A tensor descriptor for a size tensor with the specified autotune extent and upper-bound
+    """
+    return SizeTensorDesc(opt, upper_bound)
+
+# Create a TensorDesc using shape expressions and a dtype
+@public_api()
+def from_shape_expr(shape_expr: Union[Tuple[Union[ShapeExpr, int]], ShapeExprs], dtype: trt.DataType) -> TensorDesc:
+    """
+    Constructs a tensor descriptor with the specified shape expression and data type
+
+    Args:
+        shape_expr (Union[Tuple[Union[ShapeExpr, int]], ShapeExprs]): Expressions or constants denoting the shape of the tensor
+        dtype (trt.DataType): Data type of the tensor
+
+    Returns:
+        TensorDesc: Tensor descriptor with the specified shape expression and data type
+    """
+    if isinstance(shape_expr, tuple):
+        shape_expr_ = ShapeExprs.from_tuple(shape_expr)
+    else:
+        shape_expr_ = shape_expr
+    
+    return TensorDesc(shape_expr_, dtype)
+
+    

python/packaging/bindings_wheel/tensorrt/plugin/_utils.py

@@ -0,0 +1,77 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorrt as trt
+import numpy as np
+import typing
+
+_numpy_to_plugin_field_type = {
+    np.dtype('int32'): trt.PluginFieldType.INT32,
+    np.dtype('int16'): trt.PluginFieldType.INT16,
+    np.dtype('int8'): trt.PluginFieldType.INT8,
+    np.dtype('bool'): trt.PluginFieldType.INT8,
+    np.dtype('int64'): trt.PluginFieldType.INT64,
+    np.dtype('float32'): trt.PluginFieldType.FLOAT32,
+    np.dtype('float64'): trt.PluginFieldType.FLOAT64,
+    np.dtype('float16'): trt.PluginFieldType.FLOAT16
+}
+
+_built_in_to_plugin_field_type = {
+    int: trt.PluginFieldType.INT64,
+    float: trt.PluginFieldType.FLOAT64,
+    bool: trt.PluginFieldType.INT8,
+    # str is handled separately, so not needed here
+}
+
+def _str_to_data_type(dtype: str) -> trt.DataType:
+    if dtype == "FP32":
+        return trt.DataType.FLOAT
+    if dtype == "FP16":
+        return trt.DataType.HALF
+    try:
+        return getattr(trt.DataType, dtype)
+    except KeyError:
+        raise ValueError(f"Unknown data type string '{dtype}'") from None
+
+
+def _join_with(lst, middle = False, delim = ", "):
+    if len(lst) == 0:
+        return ""
+    
+    ret = ""
+    if middle:
+        ret += ", "
+
+    ret += delim.join(lst)
+    
+    return ret
+
+def _is_npt_ndarray(annotation):
+    return (typing.get_origin(annotation) == np.ndarray) or (hasattr(annotation, "__origin__") and annotation.__origin__ == np.ndarray)
+
+def _is_numpy_array(annotation):
+    return (annotation == np.ndarray) or _is_npt_ndarray(annotation)
+
+def _infer_numpy_type(annotation):
+    assert _is_npt_ndarray(annotation)
+    annot_args = typing.get_args(annotation) or annotation.__args__
+    if len(annot_args) >= 2:
+        np_type = typing.get_args(annot_args[1]) or annot_args[1].__args__
+        if len(np_type) >= 1:
+            return np_type[0]
+
+    raise AttributeError("Improper annotation for numpy array. Annotate numpy array attributes using 'numpy.typing.NDArray[dtype]', where 'dtype' is the expected numpy dtype of the array.")

python/packaging/bindings_wheel/tensorrt/plugin/_validate.py

@@ -0,0 +1,475 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import inspect
+import numpy as np
+import typing
+import types
+
+from ._utils import _is_numpy_array, _join_with, _infer_numpy_type, _is_npt_ndarray
+from ._tensor import TensorDesc, Tensor, SymExprs
+from ._export import IS_AOT_ENABLED
+if IS_AOT_ENABLED:
+    from ._tensor import KernelLaunchParams
+from ._autotune import AutoTuneCombination
+
+SERIALIZABLE_BUILTIN_TYPES = (int, float, bytes, bool, str)
+SERIALIZABLE_NP_DTYPES = (
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+    bool,
+    np.bool_,
+)
+
+# Reserve some namespaces for future use/avoid confusion
+RESERVED_NAMESPACES = {
+    "",
+    "trt",
+    "tensorrt",
+    "std",
+}
+
+DISALLOWED_ATTR_NAMES = {
+    "outputs",
+    "stream",
+    "tactic",
+}
+
+def _validate_name_and_namespace(ns: str, name: str):
+    if "." in ns:
+        raise ValueError(
+            f"Provided namespace {ns} cannot have any '.' in trt.plugin.register(\"{ns}::{name}\", ...)"
+        )
+
+    if "." in name:
+        raise ValueError(
+            f"Provided name {name} cannot have any '.' in trt.plugin.register(\"{ns}::{name}\", ...)"
+        )
+
+    if ns in RESERVED_NAMESPACES:
+        raise ValueError(
+            f"Provided namespace {ns} is a reserved namespace"
+        )
+
+
+# Parse `tensorrt.plugin.register` schema
+def _parse_register_inputs(register_func, lazy_register):
+    tensor_names = []
+    input_attrs = (
+        dict()
+    )  # order is important here but for Python >= 3.7, dict respects key order
+
+    schema_chunks = []
+
+    # TensorDescs and attribute args cannot be interspersed, so remember when we saw the first attribute arg
+    saw_first_attr = False
+
+    # Map of (attr_name: str) -> (is_builtin_type?: bool, type annotation: str)
+    attrs_types = {}
+
+    sig = inspect.signature(register_func)
+
+    for idx, (name, param) in enumerate(sig.parameters.items()):
+
+        if param.kind not in (
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        ):
+            raise ValueError(
+                f"Argument {name} is not a positional-or-keyword or keyword-only arg"
+            )
+
+        # Type annotations are manadatory for `tensorrt.plugin.register` args
+        if param.annotation == inspect.Parameter.empty:
+            raise ValueError(
+                f"Argument {name} does not have a type annotation. Please mark as TensorDesc or one of the serializable attribute types."
+            )
+
+        # Presently, we do not support default values for attributes
+        if param.default is not inspect.Parameter.empty:
+            raise ValueError(
+                f"Argument {name} has a default value. Default values are not supported yet."
+            )
+
+
+        if issubclass(param.annotation, TensorDesc):
+            if saw_first_attr:
+                raise ValueError(
+                    f"TensorDescs args and attribute args cannot be interspersed. Received function with signature {sig}."
+                )
+
+            tensor_names.append(name)
+            schema_chunks.append(f"TensorDesc {name}")
+        # At this point, we don't validate attribute types since we only care about the types of serializable attributes
+        # However, we memorize name and type so that we may validate that the autotune function maintains consistency
+        else:
+            if idx == 0:
+                raise ValueError(
+                    f"TensorDescs args should come first, followed by attributes. Received function with signature {sig}."
+                )
+
+            if name in DISALLOWED_ATTR_NAMES:
+                raise ValueError(
+                    f"'{name}' is not allowed as a plugin attribute name."
+                )
+
+            if param.annotation not in SERIALIZABLE_BUILTIN_TYPES:
+                if _is_numpy_array(param.annotation):
+                    if not lazy_register:
+                        if param.annotation == np.ndarray:
+                            raise ValueError(
+                                "If using non-lazy registration, annotate numpy array attributes using 'numpy.typing.NDArray[dtype]', where 'dtype' is the expected numpy dtype of the array."
+                            )
+
+                        if _is_npt_ndarray(param.annotation):
+                            np_dtype = _infer_numpy_type(param.annotation)
+                            if np_dtype not in SERIALIZABLE_NP_DTYPES:
+                                raise ValueError(
+                                    f"Attribute '{name}' is not a supported numpy array type. Supported numpy arrays type are {SERIALIZABLE_NP_DTYPES}."
+                                )
+                            attrs_types[name] = (False, np_dtype)
+
+                else:
+                    raise ValueError(
+                        f"Attribute '{name}' of type {param.annotation} is not a supported serializable type. Supported types are {SERIALIZABLE_BUILTIN_TYPES} or numpy arrays of type {SERIALIZABLE_NP_DTYPES}."
+                    )
+            else:
+                attrs_types[name] = (True, param.annotation)
+
+            saw_first_attr = True
+
+            schema_chunks.append(f"{param.annotation} {name}")
+            input_attrs[name] = param.annotation
+
+    return (
+        tensor_names,
+        input_attrs,
+        f"({_join_with(schema_chunks)})",
+        attrs_types,
+    )
+
+
+def _parse_register_return(register_func):
+    sig = inspect.signature(register_func)
+
+    ret_annotation = sig.return_annotation
+
+    if ret_annotation == inspect.Parameter.empty:
+        raise ValueError(
+            f"No return annotation found for register function. Received signature {sig}."
+        )
+
+    if typing.get_origin(ret_annotation) is not tuple:
+        if not inspect.isclass(ret_annotation) or not issubclass(
+            ret_annotation, TensorDesc
+        ):
+            raise ValueError(
+                f"Return argument is of type {ret_annotation}. Return types can only be TensorDesc or Tuple[TensorDesc]."
+            )
+
+        num_outputs = 1
+    else:
+        args = typing.get_args(ret_annotation)
+
+        for arg in args:
+            if not issubclass(arg, TensorDesc):
+                raise ValueError(
+                    f"Return argument is of type {ret_annotation}. Return types can only be TensorDesc or Tuple[TensorDesc]."
+                )
+
+        num_outputs = len(args)
+
+    return num_outputs
+
+
+def _validate_impl(impl_func, plugin_def):
+    impl_attr_names = []
+    found_tactic = False
+
+    sig = inspect.signature(impl_func)
+    registered_attr_names = plugin_def.input_attrs.keys()
+
+    # input arg annotations are optional, but we will validate if provided
+    for name, param in sig.parameters.items():
+        # tactic arg is optional in impl function. If specified, remember so that we can pass it during enqueue.
+        if name == "tactic":
+            found_tactic = True
+        if param.annotation != inspect.Parameter.empty:
+            if name == "outputs":
+                if typing.get_origin(param.annotation) is not tuple:
+                    raise ValueError(
+                        f"'outputs' should be of type Tuple[Tensor]. Received {param.annotation}."
+                    )
+                args = typing.get_args(param.annotation)
+                for arg in args:
+                    if not issubclass(arg, Tensor):
+                        raise ValueError(
+                            f"Argument for receiving output Tensor, '{name}' contains a {param.annotation}. '{name}' should be a Tuple[Tensor]."
+                        )
+            elif name == "stream":
+                if not issubclass(param.annotation, int):
+                    raise ValueError("'stream' input argument should be an int")
+            elif name == "tactic":
+                if not issubclass(param.annotation, int):
+                    raise ValueError("'tactic' input argument should be an int")
+            elif issubclass(param.annotation, Tensor):
+                if name not in plugin_def.input_tensor_names:
+                    raise ValueError(
+                        f"Unexpected tensor '{name}' specified in autotune function. Expected one of {plugin_def.input_tensor_names}."
+                    )
+            else:
+                if name not in plugin_def.input_attrs:
+                    raise ValueError(
+                        f"Unexpected attribute '{name}' specified in impl function. Expected one of {list(registered_attr_names)}."
+                    )
+
+                if param.annotation != plugin_def.input_attrs[name]:
+                    raise ValueError(
+                        f"Attribute '{name}' has a type annotation different from the one specified at registration. Expected '{plugin_def.input_attrs[name]}'."
+                    )
+
+                impl_attr_names.append(name)
+        else:
+            if name in plugin_def.input_attrs:
+                impl_attr_names.append(name)
+
+    # Expected attribute schema should be constructed in the order they appeared in the register function
+    expected_attr_schema_chunks = [
+        n for n in registered_attr_names if n in impl_attr_names
+    ]
+
+    expected_schema = (
+        "("
+        + _join_with(plugin_def.input_tensor_names)
+        + _join_with(expected_attr_schema_chunks, True)
+        + ", outputs, stream"
+    )
+    if found_tactic:
+        expected_schema += ", tactic)"
+    else:
+        expected_schema += ")"
+
+    if f"({', '.join(sig.parameters.keys())})" != expected_schema:
+        raise ValueError(
+            f"Signature of the impl function '{sig}' does not match the expected input arg schema: {expected_schema}"
+        )
+
+    # Return annotation is optional, but we will validate if one is specified
+    if sig.return_annotation != inspect.Parameter.empty and sig.return_annotation is not None:
+        raise ValueError("Return annotation should be None.")
+
+    return impl_attr_names, found_tactic
+
+def _validate_aot_impl(aot_impl_func, plugin_def):
+    aot_impl_attr_names = []
+
+    sig = inspect.signature(aot_impl_func)
+    registered_attr_names = plugin_def.input_attrs.keys()
+
+    # input arg annotations are optional, but we will validate if provided
+    for name, param in sig.parameters.items():
+        if param.annotation != inspect.Parameter.empty:
+            if name == "outputs":
+                if typing.get_origin(param.annotation) is not tuple:
+                    raise ValueError(
+                        f"'outputs' should be of type Tuple[TensorDesc]. Received {param.annotation}."
+                    )
+                args = typing.get_args(param.annotation)
+                for arg in args:
+                    if not issubclass(arg, TensorDesc):
+                        raise ValueError(
+                            f"Argument for receiving output TensorDesc, '{name}' contains a {param.annotation}. '{name}' should be a Tuple[TensorDesc]."
+                        )
+            elif name == "tactic":
+                if not issubclass(param.annotation, int):
+                    raise ValueError("'tactic' input argument should be an int")
+            elif issubclass(param.annotation, TensorDesc):
+                if name not in plugin_def.input_tensor_names:
+                    raise ValueError(
+                        f"Unexpected tensor '{name}' specified in autotune function. Expected one of {plugin_def.input_tensor_names}."
+                    )
+            else:
+                if name not in plugin_def.input_attrs:
+                    raise ValueError(
+                        f"Unexpected attribute '{name}' specified in aot_impl function. Expected one of {list(registered_attr_names)}."
+                    )
+
+                if param.annotation != plugin_def.input_attrs[name]:
+                    raise ValueError(
+                        f"Attribute '{name}' has a type annotation different from the one specified at registration. Expected '{plugin_def.input_attrs[name]}'."
+                    )
+
+                aot_impl_attr_names.append(name)
+        else:
+            if name in plugin_def.input_attrs:
+                aot_impl_attr_names.append(name)
+
+    # Expected attribute schema should be constructed in the order they appeared in the register function
+    expected_attr_schema_chunks = [
+        n for n in registered_attr_names if n in aot_impl_attr_names
+    ]
+
+    expected_schema = (
+        "("
+        + _join_with(plugin_def.input_tensor_names)
+        + _join_with(expected_attr_schema_chunks, True)
+        + ", outputs, tactic)"
+    )
+
+    if f"({', '.join(sig.parameters.keys())})" != expected_schema:
+        raise ValueError(
+            f"Signature of the aot_impl function '{sig}' does not match the expected input arg schema: {expected_schema}"
+        )
+
+    ret_annotation = sig.return_annotation
+
+    if ret_annotation == inspect.Parameter.empty:
+        raise ValueError(
+            f"No return annotation found for aot_impl function. Received signature {sig}."
+        )
+
+    expected_return_schema = "tuple[str | bytes, str | bytes, tensorrt.plugin.KernelLaunchParams, tensorrt.plugin.SymIntExprs]"
+
+    # Return annotation is optional, but we will validate if one is specified
+    if ret_annotation != inspect.Parameter.empty:
+        if typing.get_origin(ret_annotation) is not tuple:
+            raise ValueError(
+                f"Return annotation is {ret_annotation}. Expected {expected_return_schema}."
+            )
+        else:
+            args = typing.get_args(ret_annotation)
+
+            if len(args) != 4:
+                raise ValueError(
+                    f"Return annotation is {ret_annotation}. Expected {expected_return_schema}."
+                )
+
+            def validate_union_str_or_bytes(index):
+                def validate_str_or_bytes(arg_):
+                    if (arg_ is not str) and (arg_ is not bytes):
+                        raise ValueError(
+                            f"Return annotation for argument at {index} is '{arg_}'. Expected 'str' or 'bytes'."
+                        )
+
+                orig = typing.get_origin(args[index])
+                # orig is `typing.Union` when annotation uses typing module (e.g, Union[str, bytes])
+                # orig is `types.UnionType` when annotation is of the new (3.10+) native syntax (e.g, str | bytes)
+                if orig is typing.Union or orig is types.UnionType:
+                    for a in typing.get_args(args[index]):
+                        validate_str_or_bytes(a)
+                else:
+                # when annoted with `str` or `bytes`
+                    validate_str_or_bytes(args[index])
+
+            # kernel name should be str or bytes encoding
+            validate_union_str_or_bytes(0)
+            # kernel PTX should be str or bytes encoding
+            validate_union_str_or_bytes(1)
+
+            if not issubclass(args[2], KernelLaunchParams):
+                raise ValueError(f"Argument at index 2 of return annotation is '{args[2]}'. Expected 'tensorrt.plugin.KernelLaunchParams'.")
+
+            if not issubclass(args[3], SymExprs):
+                raise ValueError(f"Argument at index 3 of return annotation is '{args[3]}'. Expected a descendent of tensorrt.plugin.SymExprs.")
+
+    return aot_impl_attr_names
+
+
+def _validate_autotune(autotune_func, plugin_def):
+
+    sig = inspect.signature(autotune_func)
+    registered_attr_names = plugin_def.input_attrs.keys()
+
+    autotune_attr_names = []
+
+    # input arg annotations are optional, but we will validate if provided
+    for name, param in sig.parameters.items():
+        if param.annotation != inspect.Parameter.empty:
+            if name == "outputs":
+                if typing.get_origin(param.annotation) is not tuple:
+                    raise ValueError(
+                        f"'outputs' should be of type Tuple[TensorDesc]. Received {param.annotation}."
+                    )
+                args = typing.get_args(param.annotation)
+                for arg in args:
+                    if not issubclass(arg, TensorDesc):
+                        raise ValueError(
+                            f"Argument for receiving output TensorDescs, '{name}' contains a {param.annotation}. '{name}' should be a Tuple[TensorDesc]."
+                        )
+            elif issubclass(param.annotation, TensorDesc):
+                if name not in plugin_def.input_tensor_names:
+                    raise ValueError(
+                        f"Unexpected tensor '{name}' specified in autotune function. Expected one of {plugin_def.input_tensor_names}."
+                    )
+            else:
+                if name not in plugin_def.input_attrs:
+                    raise ValueError(
+                        f"Unexpected attribute '{name}' specified in autotune function. Expected one of {list(registered_attr_names)}."
+                    )
+                if param.annotation != plugin_def.input_attrs[name]:
+                    raise ValueError(
+                        f"Attribute '{name}' has a type annotation different from the one specified at registration. Expected '{plugin_def.input_attrs[name]}'."
+                    )
+
+                autotune_attr_names.append(name)
+        else:
+            if name in plugin_def.input_attrs:
+                autotune_attr_names.append(name)
+
+    # Expected attribute schema should be constructed in the order they appeared in the register function
+    expected_attr_schema_chunks = [
+        n for n in registered_attr_names if n in autotune_attr_names
+    ]
+
+    expected_schema = (
+        "("
+        + _join_with(plugin_def.input_tensor_names)
+        + _join_with(expected_attr_schema_chunks, True)
+        + ", outputs)"
+    )
+
+    if f"({', '.join(sig.parameters.keys())})" != expected_schema:
+        raise ValueError(
+            f"Specified autotune function signature {sig} is not consistent with the expected input arg schema {expected_schema}."
+        )
+
+    ret_annotation = sig.return_annotation
+
+    # Return annotation is optional, but we will validate if one is specified
+    if ret_annotation != inspect.Parameter.empty:
+        if typing.get_origin(ret_annotation) is not list:
+            if not inspect.isclass(ret_annotation) or not issubclass(
+                ret_annotation, AutoTuneCombination
+            ):
+                raise ValueError(
+                    f"Return argument is of type {ret_annotation}. Return types can only be AutoTuneCombination or List[AutoTuneCombination]."
+                )
+        else:
+            args = typing.get_args(ret_annotation)
+
+            for arg in args:
+                if not issubclass(arg, AutoTuneCombination):
+                    raise ValueError(
+                        f"Return argument is of type {ret_annotation}. Return types can only be AutoTuneCombination or List[AutoTuneCombination]."
+                    )
+
+    return autotune_attr_names