From eaa25a594a14182e924664018cea59ec2d2a6aab Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 28 Jul 2026 18:42:46 -0700 Subject: [PATCH 1/3] Fold constant_pad_nd into convolution via pre/post padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../conversion/custom_ops_converters.py | 33 ++++ .../dynamo/conversion/impl/conv.py | 13 +- .../lowering/passes/_aten_lowering_pass.py | 2 + .../passes/fuse_pad_into_convolution.py | 167 ++++++++++++++++++ 4 files changed, 214 insertions(+), 1 deletion(-) create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/fuse_pad_into_convolution.py diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 7bd70705a0..77fd5731f7 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -20,9 +20,42 @@ tensorrt_fused_nccl_reduce_scatter_op, tensorrt_fused_nccl_scatter_op, ) +from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, +) _LOGGER: logging.Logger = logging.getLogger(__name__) + +@dynamo_tensorrt_converter(tensorrt_conv_asym_pad_op, supports_dynamic_shapes=True) +def conv_asym_pad( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[trt.ITensor, Sequence[trt.ITensor]]: + """Folded pad+conv with independent pre/post padding per spatial dim.""" + del kwargs + source, weight, bias, stride, pre_padding, post_padding, dilation, groups = args + return impl.conv.convNd( + ctx, + target, + SourceIR.ATEN, + name, + is_conv1d=False, + input=source, + weight=weight, + bias=bias, + stride=stride, + padding=[0] * len(pre_padding), + dilation=dilation, + groups=groups, + pre_padding=pre_padding, + post_padding=post_padding, + ) + + if ENABLED_FEATURES.native_trt_collectives: # Use native TensorRT DistCollective API (no TensorRT-LLM dependency) _LOGGER.info("Using native TensorRT DistCollective API for distributed operations") diff --git a/py/torch_tensorrt/dynamo/conversion/impl/conv.py b/py/torch_tensorrt/dynamo/conversion/impl/conv.py index 513346a63b..8d49cf087a 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/conv.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/conv.py @@ -36,6 +36,8 @@ def convNd( output_padding: Union[int, Sequence[int]] = 0, scale: Optional[Union[torch.Tensor, float]] = None, zero_point: Optional[Union[torch.Tensor, float]] = None, + pre_padding: Optional[Sequence[int]] = None, + post_padding: Optional[Sequence[int]] = None, ) -> TRTTensor: if has_dynamic_shape(input.shape): assert input.shape[1] != -1, "Channel dim can't be dynamic for convolution." @@ -159,7 +161,16 @@ def convNd( dilation = (dilation[0], 1) if dilation is not None else dilation # Set relevant attributes of convolution layer - if padding is not None: + if pre_padding is not None or post_padding is not None: + # Asymmetric padding (e.g. causal 3D conv) must use pre/post padding. + # Symmetric padding_nd cannot express (pad_before, 0) on a spatial dim. + if pre_padding is None or post_padding is None: + raise ValueError( + f"Convolution {name}: pre_padding and post_padding must both be set" + ) + conv_layer.pre_padding = tuple(pre_padding) + conv_layer.post_padding = tuple(post_padding) + elif padding is not None: conv_layer.padding_nd = padding if stride is not None: conv_layer.stride_nd = stride diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index 271f7c98b7..d06856a58b 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -14,6 +14,7 @@ from .complex_graph_rewrite import complex_graph_detection from .constant_folding import constant_fold from .force_causal_efficient_attention import force_causal_efficient_attention +from .fuse_pad_into_convolution import fuse_pad_into_convolution from .fuse_prims_broadcast import fuse_prims_broadcast from .pass_manager import DynamoPassManager from .remove_assert_nodes import remove_assert_nodes @@ -38,6 +39,7 @@ repair_input_as_output, fuse_prims_broadcast, replace_max_pool_with_indices, + fuse_pad_into_convolution, remove_assert_nodes, remove_num_users_is_0_nodes, complex_graph_detection, diff --git a/py/torch_tensorrt/dynamo/lowering/passes/fuse_pad_into_convolution.py b/py/torch_tensorrt/dynamo/lowering/passes/fuse_pad_into_convolution.py new file mode 100644 index 0000000000..c4e9dc776c --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/fuse_pad_into_convolution.py @@ -0,0 +1,167 @@ +# mypy: disallow-untyped-decorators=False + +"""Fold ``constant_pad_nd -> convolution`` into a single asymmetric-padded conv. + +Torch-TensorRT's convolution converter only sets symmetric ``padding_nd``. When a +model applies padding itself (common for causal 3D convolutions that need +``(2 * pad_t, 0)`` in time), that pad stays a materialized copy. TensorRT's ONNX +parser folds the same pattern into ``pre_padding``/``post_padding``; this pass +recovers the equivalent for the Dynamo path. +""" + +from __future__ import annotations + +import logging +from typing import List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + clean_up_graph_after_modifications, +) + +logger = logging.getLogger(__name__) + +# Torch-TRT routes conv1d through unsqueeze/squeeze; keep fusion on 2D/3D only. +_SUPPORTED_SPATIAL_RANKS = (2, 3) + + +@torch.library.custom_op("tensorrt::conv_asym_pad", mutates_args=()) +def _conv_asym_pad_impl( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + stride: List[int], + pre_padding: List[int], + post_padding: List[int], + dilation: List[int], + groups: int, +) -> torch.Tensor: + """Convolution with independent leading/trailing padding per spatial dim.""" + flat: List[int] = [] + for pre, post in zip(reversed(pre_padding), reversed(post_padding)): + flat.extend((pre, post)) + padded = F.pad(x, flat) + conv = {2: F.conv2d, 3: F.conv3d}[len(pre_padding)] + return conv(padded, weight, bias, stride, 0, dilation, groups) + + +@_conv_asym_pad_impl.register_fake +def _( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + stride: List[int], + pre_padding: List[int], + post_padding: List[int], + dilation: List[int], + groups: int, +) -> torch.Tensor: + del bias, groups + spatial = [] + for index, (pre, post) in enumerate(zip(pre_padding, post_padding)): + extent = x.shape[2 + index] + pre + post + kernel = (weight.shape[2 + index] - 1) * dilation[index] + 1 + spatial.append((extent - kernel) // stride[index] + 1) + return x.new_empty((x.shape[0], weight.shape[0], *spatial)) + + +# Public alias used as the FX node target and converter key. +tensorrt_conv_asym_pad_op = torch.ops.tensorrt.conv_asym_pad.default + + +def _split_pad( + pad: List[int], num_spatial: int +) -> Optional[Tuple[List[int], List[int]]]: + """Convert ``constant_pad_nd`` amounts into per-spatial-dim pre/post lists. + + ``constant_pad_nd`` pads trailing dimensions first, so ``pad`` is ordered + ``[last_pre, last_post, second_last_pre, ...]``. Returns ``None`` when the + pad reaches past the spatial dims into channels or batch. + """ + if len(pad) % 2 or len(pad) > 2 * num_spatial: + return None + pre = [0] * num_spatial + post = [0] * num_spatial + for index in range(len(pad) // 2): + dim = num_spatial - 1 - index + pre[dim] = pad[2 * index] + post[dim] = pad[2 * index + 1] + return pre, post + + +def fuse_pad_into_convolution( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Rewrite ``constant_pad_nd -> convolution`` into ``tensorrt::conv_asym_pad``.""" + del settings + fused = 0 + + for node in list(gm.graph.nodes): + if ( + node.op != "call_function" + or node.target != torch.ops.aten.convolution.default + ): + continue + + ( + source, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) = node.args + if transposed or any(output_padding): + continue + if len(padding) not in _SUPPORTED_SPATIAL_RANKS: + continue + + pad_node = source + if ( + not isinstance(pad_node, torch.fx.Node) + or pad_node.target != torch.ops.aten.constant_pad_nd.default + or len(pad_node.users) != 1 + ): + continue + + pad = list(pad_node.args[1]) + fill = pad_node.args[2] if len(pad_node.args) > 2 else 0 + # Non-zero fill is not expressible as convolution padding; negatives are crops. + if fill not in (0, 0.0) or any(amount < 0 for amount in pad): + continue + + split = _split_pad(pad, len(padding)) + if split is None: + continue + pre, post = split + pre = [amount + padding[index] for index, amount in enumerate(pre)] + post = [amount + padding[index] for index, amount in enumerate(post)] + + with gm.graph.inserting_before(node): + replacement = gm.graph.call_function( + tensorrt_conv_asym_pad_op, + args=( + pad_node.args[0], + weight, + bias, + list(stride), + pre, + post, + list(dilation), + groups, + ), + ) + replacement.meta.update(node.meta) + node.replace_all_uses_with(replacement) + gm.graph.erase_node(node) + fused += 1 + + if fused: + gm = clean_up_graph_after_modifications(gm) + logger.debug(f"Folded padding into {fused} convolution(s)") + return gm From 373091f6e260efe8b2f8b8224dbbfff7751db7f3 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 28 Jul 2026 18:54:39 -0700 Subject: [PATCH 2/3] Add unit tests for fuse_pad_into_convolution lowering pass 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. --- .../test_fuse_pad_into_convolution.py | 310 ++++++++++++++++++ 1 file changed, 310 insertions(+) create mode 100644 tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py diff --git a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py new file mode 100644 index 0000000000..b495d1aa89 --- /dev/null +++ b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py @@ -0,0 +1,310 @@ +import unittest + +import torch +import torch.nn as nn +import torch.nn.functional as F +from parameterized import parameterized + +from ..conversion.harness import DispatchTestCase + + +def _node_targets(gm: torch.fx.GraphModule) -> list: + return [n.target for n in gm.graph.nodes if n.op == "call_function"] + + +def _build_pad_conv_graph( + *, + pad: list[int], + pad_value: float = 0.0, + stride: list[int], + conv_padding: list[int], + dilation: list[int], + groups: int = 1, + bias: bool = True, + pad_num_users: int = 1, +) -> torch.fx.GraphModule: + """Build ``placeholder -> constant_pad_nd -> convolution -> output``.""" + g = torch.fx.Graph() + x = g.placeholder("x") + weight = g.placeholder("weight") + bias_node = g.placeholder("bias") if bias else None + pad_node = g.call_function( + torch.ops.aten.constant_pad_nd.default, + args=(x, pad, pad_value), + ) + conv = g.call_function( + torch.ops.aten.convolution.default, + args=( + pad_node, + weight, + bias_node, + stride, + conv_padding, + dilation, + False, + [0] * len(stride), + groups, + ), + ) + if pad_num_users > 1: + # A second consumer keeps the pad from being fused. + add = g.call_function(torch.ops.aten.add.Tensor, args=(pad_node, 0.0)) + g.output((conv, add)) + else: + g.output(conv) + return torch.fx.GraphModule({}, g) + + +class TestFusePadIntoConvolutionPass(unittest.TestCase): + """Unit tests for the fuse_pad_into_convolution lowering pass.""" + + def _settings(self): + from torch_tensorrt.dynamo._settings import CompilationSettings + + return CompilationSettings() + + def _run_pass(self, gm: torch.fx.GraphModule) -> torch.fx.GraphModule: + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + fuse_pad_into_convolution, + ) + + return fuse_pad_into_convolution(gm, self._settings()) + + def test_fuses_causal_3d_pad(self) -> None: + """constant_pad_nd + convolution → tensorrt::conv_asym_pad.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[1, 1, 1, 1, 2, 0], + stride=[1, 1, 1], + conv_padding=[0, 0, 0], + dilation=[1, 1, 1], + ) + gm = self._run_pass(gm) + targets = _node_targets(gm) + self.assertNotIn(torch.ops.aten.constant_pad_nd.default, targets) + self.assertNotIn(torch.ops.aten.convolution.default, targets) + self.assertIn(tensorrt_conv_asym_pad_op, targets) + + def test_fused_args_carry_pre_post_padding(self) -> None: + """Fused node stores independent pre/post padding per spatial dim.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[1, 1, 1, 1, 2, 0], + stride=[1, 1, 1], + conv_padding=[0, 0, 0], + dilation=[1, 1, 1], + ) + gm = self._run_pass(gm) + fused = next(n for n in gm.graph.nodes if n.target == tensorrt_conv_asym_pad_op) + # args: (x, weight, bias, stride, pre, post, dilation, groups) + self.assertEqual(fused.args[3], [1, 1, 1]) + self.assertEqual(fused.args[4], [2, 1, 1]) # pre: (t, h, w) + self.assertEqual(fused.args[5], [0, 1, 1]) # post: (t, h, w) + self.assertEqual(fused.args[6], [1, 1, 1]) + self.assertEqual(fused.args[7], 1) + + def test_combines_explicit_pad_with_conv_padding(self) -> None: + """Symmetric conv padding is added into both pre and post.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[0, 0, 0, 0, 1, 0], + stride=[1, 1, 1], + conv_padding=[1, 1, 1], + dilation=[1, 1, 1], + ) + gm = self._run_pass(gm) + fused = next(n for n in gm.graph.nodes if n.target == tensorrt_conv_asym_pad_op) + self.assertEqual(fused.args[4], [2, 1, 1]) + self.assertEqual(fused.args[5], [1, 1, 1]) + + def test_fuses_2d_asymmetric_pad(self) -> None: + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[0, 1, 0, 1], + stride=[1, 1], + conv_padding=[0, 0], + dilation=[1, 1], + ) + gm = self._run_pass(gm) + fused = next(n for n in gm.graph.nodes if n.target == tensorrt_conv_asym_pad_op) + self.assertEqual(fused.args[4], [0, 0]) + self.assertEqual(fused.args[5], [1, 1]) + + def test_skips_multi_user_pad(self) -> None: + """Pad with more than one consumer must not be folded.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[1, 1, 1, 1, 2, 0], + stride=[1, 1, 1], + conv_padding=[0, 0, 0], + dilation=[1, 1, 1], + pad_num_users=2, + ) + gm = self._run_pass(gm) + targets = _node_targets(gm) + self.assertIn(torch.ops.aten.constant_pad_nd.default, targets) + self.assertIn(torch.ops.aten.convolution.default, targets) + self.assertNotIn(tensorrt_conv_asym_pad_op, targets) + + def test_skips_nonzero_pad_value(self) -> None: + """Non-zero fill cannot be expressed as convolution padding.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[1, 1, 1, 1, 2, 0], + pad_value=1.0, + stride=[1, 1, 1], + conv_padding=[0, 0, 0], + dilation=[1, 1, 1], + ) + gm = self._run_pass(gm) + targets = _node_targets(gm) + self.assertIn(torch.ops.aten.constant_pad_nd.default, targets) + self.assertNotIn(tensorrt_conv_asym_pad_op, targets) + + def test_idempotent(self) -> None: + """Running the pass twice is a no-op after the first fusion.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + gm = _build_pad_conv_graph( + pad=[1, 1, 1, 1, 2, 0], + stride=[1, 1, 1], + conv_padding=[0, 0, 0], + dilation=[1, 1, 1], + ) + gm = self._run_pass(gm) + first = _node_targets(gm) + gm = self._run_pass(gm) + self.assertEqual(first, _node_targets(gm)) + self.assertEqual(first.count(tensorrt_conv_asym_pad_op), 1) + + def test_eager_custom_op_matches_pad_then_conv(self) -> None: + """Eager tensorrt::conv_asym_pad matches F.pad + F.conv3d.""" + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + 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) + torch.testing.assert_close(got, expected) + + +class TestFusePadIntoConvolutionConverter(DispatchTestCase): + """End-to-end TRT numerical tests exercising the pad→conv fusion.""" + + @parameterized.expand( + [ + ("causal_t_2_0", (1, 1, 1, 1, 2, 0), 1, 1, 1, True), + ("causal_t_1_0", (1, 1, 1, 1, 1, 0), 1, 1, 1, True), + ("zero_pad", (0, 0, 0, 0, 0, 0), 1, 1, 1, True), + ("no_bias", (1, 1, 1, 1, 2, 0), 1, 1, 1, False), + ("stride_2", (1, 1, 1, 1, 2, 0), 2, 1, 1, True), + ("dilation_2", (2, 2, 2, 2, 4, 0), 1, 2, 1, True), + ("groups_2", (1, 1, 1, 1, 2, 0), 1, 1, 2, True), + ("hw_only", (1, 1, 1, 1), 1, 1, 1, True), + ] + ) + def test_padded_conv3d(self, _, pad, stride, dilation, groups, bias): + class PaddedConv3d(nn.Module): + def __init__(self) -> None: + super().__init__() + self.pad = pad + self.conv = nn.Conv3d( + 8, + 8, + 3, + stride=stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(F.pad(x, self.pad)) + + # Temporal extent must leave room for causal pad + dilation/stride. + t = 8 if dilation > 1 or stride > 1 else 5 + inputs = [torch.randn(1, 8, t, 16, 16)] + self.run_test( + PaddedConv3d(), + inputs, + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("asym_01_01", (0, 1, 0, 1), 1), + ("asym_stride_2", (0, 1, 0, 1), 2), + ] + ) + def test_padded_conv2d(self, _, pad, stride): + class PaddedConv2d(nn.Module): + def __init__(self) -> None: + super().__init__() + self.pad = pad + self.conv = nn.Conv2d(8, 4, 3, stride=stride, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(F.pad(x, self.pad)) + + inputs = [torch.randn(1, 8, 16, 16)] + self.run_test( + PaddedConv2d(), + inputs, + use_dynamo_tracer=True, + enable_passes=True, + ) + + def test_graph_contains_fused_op_after_lowering(self): + """post_lowering should rewrite the pad+conv pair before conversion.""" + from torch_tensorrt.dynamo._settings import CompilationSettings + from torch_tensorrt.dynamo.lowering import get_decompositions, post_lowering + from torch_tensorrt.dynamo.lowering.passes.fuse_pad_into_convolution import ( + tensorrt_conv_asym_pad_op, + ) + + class PaddedConv3d(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv3d(4, 4, 3, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(F.pad(x, (1, 1, 1, 1, 2, 0))) + + model = PaddedConv3d().eval() + x = torch.randn(1, 4, 5, 8, 8) + exported = torch.export.export(model, (x,)) + exported = exported.run_decompositions(get_decompositions(False)) + gm = post_lowering(exported.module(), CompilationSettings()) + targets = _node_targets(gm) + self.assertIn(tensorrt_conv_asym_pad_op, targets) + self.assertNotIn(torch.ops.aten.constant_pad_nd.default, targets) + + +if __name__ == "__main__": + unittest.main() From dfa851cac489e8288913062cd86b4f8be5ab67f0 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 29 Jul 2026 15:22:21 -0700 Subject: [PATCH 3/3] Apply black formatting to satisfy pre-commit lint --- tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py index b495d1aa89..f023aad984 100644 --- a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py +++ b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py @@ -208,7 +208,9 @@ def test_eager_custom_op_matches_pad_then_conv(self) -> None: 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)