resolve comments

Author: bowang007

examples/dynamo/auto_generate_plugin.py

@@ -1,10 +1,10 @@
 """
 .. _auto_generate_converters:
 
-Automatically Generate a Converter for a Custom Kernel
+Automatically Generate a Plugin for a Custom Kernel
 ===================================================================
 
-We are going to demonstrate how to automatically generate a converter for a custom kernel using Torch-TensorRT using
+We are going to demonstrate how to automatically generate a plugin for a custom kernel using Torch-TensorRT using
 the new Python based plugin system in TensorRT 10.7.
 
 Torch-TensorRT supports falling back to PyTorch implementations of operations in the case that Torch-TensorRT
@@ -102,16 +102,21 @@ def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Ten
 torch_tensorrt.dynamo.conversion.plugins.generate_plugin("torchtrt_ex::elementwise_mul")
 
 
-# %%
-# Generating the Converter
-# -------------------------------------------------------------------
-# Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation.
-# As long as the namespace and names match, the following function will automatically generate the converter for the operation.
+# # %%
+# # Generating the Converter
+# # -------------------------------------------------------------------
+# # Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation.
+# # As long as the namespace and names match, the following function will automatically generate the converter for the operation.
 torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
     "torchtrt_ex::elementwise_mul", supports_dynamic_shapes=True
 )
 
 
+# # %%
+# # Above two commands can be replaced with the following single one line:
+# torch_tensorrt.dynamo.conversion.plugins.custom_op("torchtrt_ex::elementwise_mul", supports_dynamic_shapes=True)
+
+
 # %%
 # Using our converter with a model
 # -------------------------------------------------------------------

py/torch_tensorrt/dynamo/conversion/plugins/__init__.py

@@ -1,3 +1,4 @@
+from torch_tensorrt.dynamo.conversion.plugins._custom_op import custom_op
 from torch_tensorrt.dynamo.conversion.plugins._generate_plugin import generate_plugin
 from torch_tensorrt.dynamo.conversion.plugins._generate_plugin_converter import (
     generate_plugin_converter,

py/torch_tensorrt/dynamo/conversion/plugins/_custom_op.py

@@ -0,0 +1,21 @@
+from typing import Callable, Optional
+
+from torch.fx.node import Node
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.conversion._ConverterRegistry import ConverterPriority
+from torch_tensorrt.dynamo.conversion.plugins._generate_plugin import generate_plugin
+from torch_tensorrt.dynamo.conversion.plugins._generate_plugin_converter import (
+    generate_plugin_converter,
+)
+
+
+def custom_op(
+    op_name: str,
+    capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
+    priority: ConverterPriority = ConverterPriority.STANDARD,
+    supports_dynamic_shapes: bool = False,
+):
+    generate_plugin(op_name)
+    generate_plugin_converter(
+        op_name, capability_validator, priority, supports_dynamic_shapes
+    )

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py

@@ -1,3 +1,4 @@
+import logging
 from types import FunctionType
 from typing import Tuple
 
@@ -8,6 +9,8 @@
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.fx.experimental.symbolic_shapes import DimDynamic, ShapeEnv
 
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
 
 def mksym(shape_env, value, source, dynamic_dim):
     return shape_env.create_symintnode(
@@ -42,6 +45,16 @@ def generate_signature(torch_op):
         for arg in schema.arguments:
             arg_list.append(arg.name)
 
+            # TODO: Torch types need to be converted to python primitive types here
+            # Some other types are not handled:
+            # - torch._C.ListType.ofT(<type>)
+            # - torch._C.TupleType.get()
+            # - torch._C.DictType.get(<key_type>, <value_type>)
+            # - torch._C.OptionalType.ofT(<type>)
+            # - torch._C.DeviceObjType.get()
+            # - torch._C.FunctionType.get()
+            # - torch._C.ClassType
+
             if arg.type.isSubtypeOf(torch._C.TensorType.get()):
                 tensor_args.append(arg)
                 register_func_annotation[arg.name] = trtp.TensorDesc
@@ -52,6 +65,12 @@ def generate_signature(torch_op):
             elif arg.type.isSubtypeOf(torch._C.IntType.get()):
                 register_func_annotation[arg.name] = int
                 impl_func_annotation[arg.name] = int
+            elif arg.type.isSubtypeOf(torch._C.Booltype.get()):
+                register_func_annotation[arg.name] = bool
+                impl_func_annotation[arg.name] = bool
+            elif arg.type.isSubtypeOf(torch._C.Stringtype.get()):
+                register_func_annotation[arg.name] = str
+                impl_func_annotation[arg.name] = str
             else:
                 raise ValueError("arg type is not handled")
 
@@ -94,12 +113,6 @@ def generate_signature(torch_op):
         register_func_annotation,
         impl_func_annotation,
     ) = generate_signature(torch_op)
-    print(args_input)
-    print(kwargs_input)
-    print(plugin_signature)
-    print(plugin_impl_signature)
-    print(register_func_annotation)
-    print(impl_func_annotation)
 
     def _generic_plugin_desc(*args, **kwargs) -> Tuple[trtp.TensorDesc]:
         shape_env = ShapeEnv()
@@ -141,6 +154,8 @@ def _generic_plugin_desc(*args, **kwargs) -> Tuple[trtp.TensorDesc]:
     return _generic_plugin_desc({args_input}, {kwargs_input})
     """
 
+    _LOGGER.warning(f"Plugin registration function: \n{codegen_plugin}")
+
     plugin_code = compile(codegen_plugin, "<string>", "exec")
 
     globals()["_generic_plugin_desc"] = _generic_plugin_desc
@@ -167,6 +182,8 @@ def _generic_plugin_impl(outputs, stream, *args, **kwargs):
     _generic_plugin_impl(outputs, stream, {args_input}, {kwargs_input})
     """
 
+    _LOGGER.warning(f"Plugin implementation function: \n{plugin_impl_func}")
+
     plugin_impl_code = compile(plugin_impl_func, "<string>", "exec")
 
     globals()["_generic_plugin_impl"] = _generic_plugin_impl
@@ -175,15 +192,6 @@ def _generic_plugin_impl(outputs, stream, *args, **kwargs):
 
     plugin_impl.__annotations__ = impl_func_annotation
 
-    import inspect
-
-    sig = inspect.signature(plugin_impl)
-
-    # input arg annotations are optional, but we will validate if provided
-    for name, param in sig.parameters.items():
-        print(name)
-        print(param.annotation)
-
     trtp.impl(plugin_name)(plugin_impl)
 
     return plugin

pyproject.toml

@@ -12,6 +12,7 @@ requires = [
     "torch>=2.7.0.dev,<2.8.0",
     "pybind11==2.6.2",
     "numpy",
+    "sympy",
 ]
 build-backend = "setuptools.build_meta"