Fuse pad into convolution - #4442
Open
micwill755 wants to merge 3 commits into
Open
Conversation
Torch-TensorRT only set symmetric padding_nd on convolutions, so models that apply asymmetric padding themselves (e.g. causal 3D conv with (2*pad_t, 0) in time) left the pad as a materialized copy. Add a post-lowering pass that rewrites constant_pad_nd -> convolution into tensorrt::conv_asym_pad, and convert that op with TRT pre_padding / post_padding — matching what the ONNX parser already does.
Cover graph rewrite (causal 3D/2D pads, conv-padding combine, multi-user and non-zero fill skips, idempotence, eager custom-op parity) and end-to-end TRT numerical checks for padded conv2d/conv3d across stride, dilation, groups, and bias configurations.
github-actions
Bot
added
component: tests
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:49:49.445228+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:50:08.060518+00:00
@@ -206,11 +206,13 @@
x = torch.randn(1, 4, 5, 8, 8)
weight = torch.randn(4, 4, 3, 3, 3)
bias = torch.randn(4)
pre, post = [2, 1, 1], [0, 1, 1]
expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
- got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+ got = tensorrt_conv_asym_pad_op(
+ x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+ )
torch.testing.assert_close(got, expected)
class TestFusePadIntoConvolutionConverter(DispatchTestCase):
"""End-to-end TRT numerical tests exercising the pad→conv fusion."""
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:49:49.054179+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:50:15.101134+00:00
@@ -206,11 +206,13 @@
x = torch.randn(1, 4, 5, 8, 8)
weight = torch.randn(4, 4, 3, 3, 3)
bias = torch.randn(4)
pre, post = [2, 1, 1], [0, 1, 1]
expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
- got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+ got = tensorrt_conv_asym_pad_op(
+ x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+ )
torch.testing.assert_close(got, expected)
class TestFusePadIntoConvolutionConverter(DispatchTestCase):
"""End-to-end TRT numerical tests exercising the pad→conv fusion."""
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:50:10.421070+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 02:50:40.458077+00:00
@@ -206,11 +206,13 @@
x = torch.randn(1, 4, 5, 8, 8)
weight = torch.randn(4, 4, 3, 3, 3)
bias = torch.randn(4)
pre, post = [2, 1, 1], [0, 1, 1]
expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
- got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+ got = tensorrt_conv_asym_pad_op(
+ x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+ )
torch.testing.assert_close(got, expected)
class TestFusePadIntoConvolutionConverter(DispatchTestCase):
"""End-to-end TRT numerical tests exercising the pad→conv fusion."""
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 03:56:45.218617+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py 2026-07-29 03:57:10.275509+00:00
@@ -206,11 +206,13 @@
x = torch.randn(1, 4, 5, 8, 8)
weight = torch.randn(4, 4, 3, 3, 3)
bias = torch.randn(4)
pre, post = [2, 1, 1], [0, 1, 1]
expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
- got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+ got = tensorrt_conv_asym_pad_op(
+ x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+ )
torch.testing.assert_close(got, expected)
class TestFusePadIntoConvolutionConverter(DispatchTestCase):
"""End-to-end TRT numerical tests exercising the pad→conv fusion."""
Collaborator
|
@micwill755 run precommit for the linting |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fold asymmetric constant_pad_nd → convolution patterns into TensorRT convolution pre_padding / post_padding via a new Dynamo post-lowering pass (fuse_pad_into_convolution) and tensorrt::conv_asym_pad converter.
Motivation: Wan-style causal 3D convolutions (used by Cosmos3 / Cosmos3-Edge VAE) zero their own padding and apply it themselves, often asymmetrically in time (e.g. (2 * pad_t, 0)). TensorRT’s ONNX parser already folds Pad → Conv into the convolution’s explicit pre/post padding. Torch-TRT previously only set symmetric padding_nd, so those pads stayed as materialized copy kernels (extra work on the VAE encode path).
Change:
Post-lowering pass rewrites eligible constant_pad_nd → convolution into tensorrt::conv_asym_pad
Converter emits IConvolutionLayer with explicit pre_padding / post_padding (via extended convNd)
Unit tests in tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py
Validation on Cosmos3-Edge policy (Torch-TRT → cosmos3_policy_inference):
Geometry: 256×256, 5 frames, und_len 121, batch 1, 4 denoise steps
Pad fusion removed ~45 extra VAE pad kernels and improved VAE encode 7.46 → 5.50 ms (~1.96 ms)
Full policy steady-state median improved ~23.85 → 21.81 ms (~2.0 ms) after enabling this pass (with MoT builder settings already applied)
No new runtime dependencies beyond existing Torch-TensorRT / TensorRT.
Fixes # (n/a — performance / parity with ONNX Pad→Conv folding)
Type of change
New feature (non-breaking change which adds functionality)
Performance optimization for models with asymmetric Pad→Conv (e.g. Cosmos3 VAE)
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