Skip to content

feat: support in-place plugins#4279

Open
bowang007 wants to merge 2 commits into
mainfrom
inplace_plugin
Open

feat: support in-place plugins#4279
bowang007 wants to merge 2 commits into
mainfrom
inplace_plugin

Conversation

@bowang007

Copy link
Copy Markdown
Collaborator

Description

End-to-end support for in-place QDP plugins in torch-tensorrt

Fixes #4240

Type of change

Please delete options that are not relevant and/or add your own.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

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

@meta-cla meta-cla Bot added the cla signed label May 19, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: runtime component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels May 19, 2026
@github-actions github-actions Bot requested a review from zewenli98 May 19, 2026 12:26
@bowang007 bowang007 requested a review from narendasan May 19, 2026 12:28
@bowang007 bowang007 requested a review from cehongwang May 19, 2026 14:56

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/unfunctionalize_qdp_inplace.py	2026-05-20 14:19:18.477550+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/unfunctionalize_qdp_inplace.py	2026-05-20 14:19:37.985186+00:00
@@ -131,13 +131,11 @@
                gm.graph.erase_node(user)
        else:
            return_getitems: List[Any] = []
            with gm.graph.inserting_after(new_call):
                for i in range(n_outputs):
-                    g = gm.graph.call_function(
-                        operator.getitem, args=(new_call, i)
-                    )
+                    g = gm.graph.call_function(operator.getitem, args=(new_call, i))
                    if isinstance(hop_val, tuple) and i < len(hop_val):
                        g.meta["val"] = hop_val[i]
                    return_getitems.append(g)
            for user in getitem_users:
                idx = user.args[1]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you explain this lowering pass? Auto functionalization is a bit more broader reaching that just for QDP nodes. Maybe it might be useful elsewhere?

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/unfunctionalize_qdp_inplace.py	2026-07-14 09:13:16.073785+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/unfunctionalize_qdp_inplace.py	2026-07-14 09:13:43.216625+00:00
@@ -131,13 +131,11 @@
                gm.graph.erase_node(user)
        else:
            return_getitems: List[Any] = []
            with gm.graph.inserting_after(new_call):
                for i in range(n_outputs):
-                    g = gm.graph.call_function(
-                        operator.getitem, args=(new_call, i)
-                    )
+                    g = gm.graph.call_function(operator.getitem, args=(new_call, i))
                    if isinstance(hop_val, tuple) and i < len(hop_val):
                        g.meta["val"] = hop_val[i]
                    return_getitems.append(g)
            for user in getitem_users:
                idx = user.args[1]

…r converter

Convert in-place QDP plugins by registering a converter on the
auto_functionalized(_v2) higher-order wrapper (gated by a capability_validator
to the ops that have a plugin) instead of reversing functionalization in a
lowering pass. The graph stays functional and scoping is per-node.

- Add _auto_functionalized_converter: wrapper converter that reconstructs the
  inner op's args, delegates to the inner op's converter, and returns the
  [*returns, *bases] tuple with a faithful base->return alias map; plus a
  scoped no-op converter for the aliased write-back copy_.
- Add qdp_inplace_metadata pass: metadata-only (no graph reversal) so
  pre-conversion symbolic-shape extraction handles the direct-return case.
- Remove unfunctionalize_qdp_inplace: its global copy_ deletion corrupted
  mixed graphs (a non-plugin in-place op's write-back was dropped).

All in-place tests pass (5/5); full plugin+kernels suite green (57 passed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: runtime component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

✨[Feature] Support in-place custom plugins (and multiple outputs)

2 participants