feat: Automatically generate QDP plugins #3370

bowang007 · 2025-01-29T23:35:53Z

Description

This PR introduces a new feature which enables generating automatic plugin generation using TensorRT QDP feature.

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

github-actions

outdated

examples/dynamo/auto_generate_plugin.py

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

narendasan · 2025-01-30T22:52:06Z

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

+
+    # Use the helper function to get the required signatures
+    args_input, kwargs_input, plugin_signature, plugin_impl_signature, register_func_annotation, impl_func_annotation = generate_signature(torch_op)
+    print(args_input)


Make this debug info

@narendasan when I use:

_LOGGER.debug(f"Plugin registration function: \n{codegen_plugin}")

It won't print anything. How to resolve this?

narendasan

Figure out if QDP will hold the reference to the plugin for us
Add test cases for 1. (Tensor, Tensor) -> (Tensor) 2. (Tensor, int, float) -> (Tensor) 3. (Tensor, Tensor) -> (Tensor, Tensor)

narendasan · 2025-02-07T02:07:30Z

Rebase as well

examples/dynamo/auto_generate_plugin.py

narendasan · 2025-02-22T00:14:58Z

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

+    capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
+    priority: ConverterPriority = ConverterPriority.STANDARD,
+    supports_dynamic_shapes: bool = False,
+):


Add a docstring as this is a user API

narendasan · 2025-02-22T00:15:07Z

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

+    )
+
+
+def generate_plugin(plugin_name: str):


Add a docstring

narendasan

LGTM after docstrings are added

peri044

Added minor comments. LGTM

peri044 · 2025-02-24T21:25:03Z

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

+
+        for tensor_arg in tensor_args:
+
+            sample = {f"{i}": 5 for i in range(tensor_arg.ndim)}


what is 5 here ?

If it's a default value or something, consider storing it in a global variable to make it more clear ?

peri044 · 2025-02-24T21:36:07Z

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

+        outputs: Tuple[trtp.Tensor], stream: int, *args: Any, **kwargs: Any
+    ) -> None:
+        tensor_args = [elem for elem in args if isinstance(elem, trtp.Tensor)]
+        print(args)


Is this necessary ? If so, can you make this message more effective for users ?

looks like I forgot to delete the debugging lines. Thanks!

zewenli98

Added some minor comments

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

zewenli98 · 2025-02-24T23:03:50Z

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

+        tensor_args = [elem for elem in args if isinstance(elem, trtp.Tensor)]
+        print(args)
+        non_tensor_args = [elem for elem in args if not isinstance(elem, trtp.Tensor)]


Is it necessary to go for loops twice?

I can write it as:

tensor_args, non_tensor_args = [], [] for elem in args: (tensor_args if isinstance(elem, trtp.Tensor) else non_tensor_args).append(elem)

Since args won't be long I think the first one will be easier to understand?

I think it's clearer and faster in this way:

tensor_args, non_tensor_args = [], [] for elem in args: if isinstance(elem, trtp.Tensor): tensor_args.append(elem) else: non_tensor_args.append(elem)

zewenli98 · 2025-02-24T23:29:32Z

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin_converter.py

@@ -58,13 +57,24 @@ def custom_kernel_converter(
        # Assuming TensorRT preserves kwargs order like PyTorch does
        non_tensor_inputs = plugin.input_attrs

+        kwargs = {}


Do you intend to override kwargs here? If yes, it seems kwargs is not necessary in the arguments (line 46).

yes, this kwargs cannot be erased since it's required in TorchTRT converter here.

bowang007 requested a review from narendasan January 29, 2025 23:35

facebook-github-bot added the cla signed label Jan 29, 2025

github-actions bot added component: conversion Issues re: Conversion stage component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Jan 29, 2025

github-actions bot requested a review from peri044 January 29, 2025 23:36

github-actions bot requested changes Jan 29, 2025

View reviewed changes

This comment was marked as outdated.

Sign in to view

narendasan reviewed Jan 30, 2025

View reviewed changes

examples/dynamo/auto_generate_plugin.py Show resolved Hide resolved

narendasan reviewed Jan 30, 2025

View reviewed changes

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py Outdated Show resolved Hide resolved

narendasan reviewed Jan 30, 2025

View reviewed changes

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py Outdated Show resolved Hide resolved

narendasan reviewed Jan 30, 2025

View reviewed changes

This comment was marked as outdated.

Sign in to view

github-actions bot added the component: build system Issues re: Build system label Feb 7, 2025

This comment was marked as outdated.

Sign in to view

narendasan reviewed Feb 7, 2025

View reviewed changes

github-actions bot added the component: tests Issues re: Tests label Feb 20, 2025

This comment was marked as outdated.

Sign in to view

narendasan reviewed Feb 21, 2025

View reviewed changes

examples/dynamo/auto_generate_plugin.py Outdated Show resolved Hide resolved

bowang007 force-pushed the automatic_plugin branch from 5381375 to a35a9ec Compare February 21, 2025 22:18

This comment was marked as outdated.

Sign in to view

bowang007 force-pushed the automatic_plugin branch from a35a9ec to 56dacec Compare February 21, 2025 22:20

bowang007 requested a review from narendasan February 21, 2025 22:21

narendasan reviewed Feb 22, 2025

View reviewed changes

narendasan approved these changes Feb 22, 2025

View reviewed changes

bowang007 force-pushed the automatic_plugin branch from 56dacec to 6683559 Compare February 22, 2025 01:53

This comment was marked as outdated.

Sign in to view

bowang007 force-pushed the automatic_plugin branch from 6683559 to d5a787d Compare February 22, 2025 01:58

peri044 reviewed Feb 24, 2025

View reviewed changes

zewenli98 reviewed Feb 24, 2025

View reviewed changes

feat: Automatically generate QDP plugins

b60c4c6

bowang007 force-pushed the automatic_plugin branch from d5a787d to b60c4c6 Compare February 25, 2025 22:27

bowang007 merged commit 31bdf77 into main Feb 25, 2025
50 of 68 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Automatically generate QDP plugins #3370

feat: Automatically generate QDP plugins #3370

bowang007 commented Jan 29, 2025

github-actions bot left a comment •

edited by bowang007

Loading

This comment was marked as outdated.

This comment was marked as outdated.

narendasan Jan 30, 2025

bowang007 Feb 7, 2025

This comment was marked as outdated.

This comment was marked as outdated.

narendasan left a comment

narendasan commented Feb 7, 2025

This comment was marked as outdated.

This comment was marked as outdated.

narendasan Feb 22, 2025

narendasan Feb 22, 2025

narendasan left a comment

This comment was marked as outdated.

peri044 left a comment

peri044 Feb 24, 2025

peri044 Feb 24, 2025

peri044 Feb 24, 2025

bowang007 Feb 25, 2025

zewenli98 left a comment

zewenli98 Feb 24, 2025

bowang007 Feb 25, 2025

zewenli98 Feb 27, 2025

zewenli98 Feb 24, 2025

bowang007 Feb 25, 2025 •

edited

Loading


		for tensor_arg in tensor_args:

		sample = {f"{i}": 5 for i in range(tensor_arg.ndim)}

feat: Automatically generate QDP plugins #3370

feat: Automatically generate QDP plugins #3370

Conversation

bowang007 commented Jan 29, 2025

Description

Checklist:

github-actions bot left a comment • edited by bowang007 Loading

Choose a reason for hiding this comment

This comment was marked as outdated.

This comment was marked as outdated.

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment was marked as outdated.

This comment was marked as outdated.

narendasan left a comment

Choose a reason for hiding this comment

narendasan commented Feb 7, 2025

This comment was marked as outdated.

This comment was marked as outdated.

Choose a reason for hiding this comment

Choose a reason for hiding this comment

narendasan left a comment

Choose a reason for hiding this comment

This comment was marked as outdated.

peri044 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zewenli98 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bowang007 Feb 25, 2025 • edited Loading

Choose a reason for hiding this comment

github-actions bot left a comment •

edited by bowang007

Loading

bowang007 Feb 25, 2025 •

edited

Loading