From 0f000d5a2958e3f9035d2606c5c9ba799063260e Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 13 Nov 2026 23:58:55 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- .../conv/benchmark_implicit_gemm.py | 146 ++++++++++++++++-- .../cuda/tests/test_conv_implicit_gemm.py | 38 +++++ backends/cuda/tests/test_conv_routing.py | 74 ++++++++- 3 files changed, 239 insertions(+), 19 deletions(-) diff --git a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py index 330df91c82d..9bd576ad0e5 100644 --- a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py +++ b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py @@ -67,6 +67,17 @@ class Case: bias: bool precision: str = "fp32" layout: str = "nchw" + # The padding is a separate pad op before a padding-0 convolution, as the + # model writes it; both paths run that pad. + explicit_pad: bool = False + + +@dataclass(frozen=True) +class Pad: + """An explicit pad before a padding-0 convolution, per spatial dim.""" + + before: tuple[int, ...] + after: tuple[int, ...] def _case(label, shape, out_channels, kernel, stride, padding, bias, precision): @@ -75,20 +86,28 @@ def _case(label, shape, out_channels, kernel, stride, padding, bias, precision): def expand(v): return tuple(v) if isinstance(v, tuple) else (v,) * rank - padding = expand(padding) + if isinstance(padding, Pad): + before, after, explicit = padding.before, padding.after, True + else: + before = after = expand(padding) + explicit = False return Case( label, shape, out_channels, expand(kernel), expand(stride), - padding, - padding, + before, + after, bias, precision, + explicit_pad=explicit, ) +_CAUSAL = Pad((2, 1, 1), (0, 1, 1)) +_CACHED = Pad((0, 1, 1), (0, 1, 1)) + # (group, label, input shape, C_out, kernel, stride, padding, bias) SHAPES = [ ( @@ -230,6 +249,98 @@ def expand(v): ("32ch_s2", "32 ch 5x5 125x125 x2", (2, 32, 125, 125), 32, 5, 2, 0, True), ("32ch_s2", "32 ch 3x3 61x61 x2", (2, 32, 61, 61), 32, 3, 2, 0, True), ("32ch_s2", "32 ch 3x3 30x30 x2", (2, 32, 30, 30), 32, 3, 2, 0, True), + # Causal 3x3x3 convolutions pad time in front and space by 1 in a + # separate pad; chunked inputs carry two cached time steps in front. + ( + "3x3x3", + "12-160 ch 1x112x224", + (1, 12, 1, 112, 224), + 160, + 3, + 1, + _CAUSAL, + True, + ), + ( + "3x3x3", + "160 ch 1x112x224", + (1, 160, 1, 112, 224), + 160, + 3, + 1, + _CAUSAL, + True, + ), + ("3x3x3", "320 ch 1x56x112", (1, 320, 1, 56, 112), 320, 3, 1, _CAUSAL, True), + ("3x3x3", "640 ch 1x28x56", (1, 640, 1, 28, 56), 640, 3, 1, _CAUSAL, True), + ("3x3x3", "640 ch 1x14x28", (1, 640, 1, 14, 28), 640, 3, 1, _CAUSAL, True), + ( + "3x3x3", + "160 ch 18x64x64 x2", + (2, 160, 18, 64, 64), + 160, + 3, + 1, + _CACHED, + True, + ), + ( + "3x3x3", + "320 ch 18x32x32 x2", + (2, 320, 18, 32, 32), + 320, + 3, + 1, + _CACHED, + True, + ), + ("3x3x3", "640 ch 6x8x8 x2", (2, 640, 6, 8, 8), 640, 3, 1, _CACHED, True), + ( + "3x3x3", + "160 ch 10x128x160 x3", + (3, 160, 10, 128, 160), + 160, + 3, + 1, + _CACHED, + True, + ), + ("1x1x1", "160-320 ch 1x56x112", (1, 160, 1, 56, 112), 320, 1, 1, 0, True), + ("1x1x1", "320-640 ch 1x28x56", (1, 320, 1, 28, 56), 640, 1, 1, 0, True), + ("1x1x1", "96 ch 1x14x28", (1, 96, 1, 14, 28), 96, 1, 1, 0, True), + ("1x1x1", "160-320 ch 16x32x32 x2", (2, 160, 16, 32, 32), 320, 1, 1, 0, True), + ("1x1x1", "320-640 ch 8x16x16 x2", (2, 320, 8, 16, 16), 640, 1, 1, 0, True), + ("1x1x1", "96 ch 4x8x8 x2", (2, 96, 4, 8, 8), 96, 1, 1, 0, True), + ( + "3x1x1_s2", + "320 ch 17x16x16 x2", + (2, 320, 17, 16, 16), + 320, + (3, 1, 1), + (2, 1, 1), + 0, + True, + ), + ( + "3x1x1_s2", + "640 ch 9x8x8 x2", + (2, 640, 9, 8, 8), + 640, + (3, 1, 1), + (2, 1, 1), + 0, + True, + ), + ( + "3x1x1_s2", + "320 ch 9x16x20 x3", + (3, 320, 9, 16, 20), + 320, + (3, 1, 1), + (2, 1, 1), + 0, + True, + ), # 1-D convolutions over sequences of 4 to 64 steps at batch 1. ("1d_k5", "2-512 ch T16", (1, 2, 16), 512, 5, 1, 2, True), ("1d_k5", "512 ch T16", (1, 512, 16), 512, 5, 1, 2, True), @@ -425,21 +536,29 @@ def _args(case): ) +def _pad(x, case): + pad = [p for pair in zip(reversed(case.before), reversed(case.after)) for p in pair] + return F.pad(x, pad) + + def triton_runner(x, w, b, case): output_strides = None if case.layout == "cl_to_nchw": shape = aten_conv(x.to("meta"), w.to("meta"), None, case).shape output_strides = list(torch.empty(shape, device="meta").stride()) - args = (*_args(case), output_strides) + args = list(_args(case)) + [output_strides] + padded = case.explicit_pad or case.before != case.after + if padded: + args[1] = args[2] = [0] * len(case.kernel) stream = torch.cuda.Stream() with torch.cuda.stream(stream), autotune_launch_params() as stats: - convolution(x, w, b, *args) + convolution(_pad(x, case) if padded else x, w, b, *args) stream.synchronize() schedule = stats.picks[0].choice if stats.picks else "igemm_ffma" def run(): with force_conv_candidate(schedule): - return convolution(x, w, b, *args) + return convolution(_pad(x, case) if padded else x, w, b, *args) return run, schedule @@ -451,17 +570,14 @@ def _conv1d(x, w, b, stride, padding=(0,)): def aten_conv(x, w, b, case): - """What the CUDA backend gives Inductor: a separate pad for asymmetric - padding, a 1-D convolution as a 2-D one, and a contiguous output where the - graph records one.""" + """What the CUDA backend gives Inductor: a separate pad for explicit or + asymmetric padding, a 1-D convolution as a 2-D one, and a contiguous output + where the graph records one.""" conv = {1: _conv1d, 2: F.conv2d, 3: F.conv3d}[len(case.kernel)] - if case.before == case.after: - out = conv(x, w, b, case.stride, case.before) + if case.explicit_pad or case.before != case.after: + out = conv(_pad(x, case), w, b, case.stride) else: - pad = [ - p for pair in zip(reversed(case.before), reversed(case.after)) for p in pair - ] - out = conv(F.pad(x, pad), w, b, case.stride) + out = conv(x, w, b, case.stride, case.before) return out.contiguous() if case.layout == "cl_to_nchw" else out diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py index ce629106018..5fc5392b8ad 100644 --- a/backends/cuda/tests/test_conv_implicit_gemm.py +++ b/backends/cuda/tests/test_conv_implicit_gemm.py @@ -8,6 +8,7 @@ the implicit-GEMM template against an FP64 reference, and models lowered through the CUDA backend.""" +import dataclasses import glob import itertools import os @@ -193,6 +194,11 @@ def check_geometries(self, dtype, layouts=("contiguous", "channels_last")): ((2, 32, 8), 32, (3,), (2,), (1,), (1,)), ((1, 48, 4), 64, (1,), (1,), (0,), (0,)), ((1, 5, 13), 7, (4,), (3,), (2,), (1,)), + # 3-D: a causal 3x3x3 on an explicitly padded input, a + # temporal (3,1,1) s(2,1,1) and a 1x1x1. + ((1, 24, 3, 10, 12), 16, (3, 3, 3), (1, 1, 1), (0,) * 3, (0,) * 3), + ((2, 32, 9, 4, 5), 32, (3, 1, 1), (2, 1, 1), (0,) * 3, (0,) * 3), + ((1, 16, 2, 5, 6), 24, (1, 1, 1), (1, 1, 1), (0,) * 3, (0,) * 3), ) cases = [ (geometry, layout, layout == "weight_cl" or i % 2 == 0) @@ -454,6 +460,22 @@ def forward(self, x): return self.down(self.same(F.mish(self.conv(x)) + self.res(x))) +class _ExplicitPadConv3dBlock(nn.Module): + """3-D convolutions as models write them: a causal + pad and a 3x3x3, a 1x1x1 shortcut, and a (3,1,1) s(2,1,1) time + convolution.""" + + def __init__(self, dtype): + super().__init__() + self.causal = nn.Conv3d(12, 16, 3, dtype=dtype) + self.shortcut = nn.Conv3d(12, 16, 1, dtype=dtype) + self.time = nn.Conv3d(16, 16, (3, 1, 1), stride=(2, 1, 1), dtype=dtype) + + def forward(self, x): + y = self.causal(F.pad(x, (1, 1, 1, 1, 2, 0))) + self.shortcut(x) + return self.time(y) + + class ImplicitGemmLoweringTest(unittest.TestCase): def setUp(self): if not torch.cuda.is_available() or torch.version.hip is not None: @@ -567,3 +589,19 @@ def test_conv1d_block_lowers_and_runs(self): model, inputs, 1, self.fp32_tolerance(), route="required" ) self.assertEqual(output.count("Skipping conv replacement"), 3) + + def test_conv3d_block_lowers_and_runs(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + torch.manual_seed(0) + model = _ExplicitPadConv3dBlock(torch.bfloat16).to("cuda").eval() + inputs = (torch.randn(1, 12, 5, 16, 16, device="cuda", dtype=torch.bfloat16),) + # Outside the benchmarked range: Inductor lowers all three. + self.lower_and_run(model, inputs, 0, 3e-2) + # Every one is correct through the template; the pad stays separate. + take_all = tuple( + dataclasses.replace(t, worth_routing=lambda call: None) + for t in routing.TEMPLATES + ) + with mock.patch.object(routing, "TEMPLATES", take_all): + self.lower_and_run(model, inputs, 3, 3e-2) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index f39ad6bce1f..6750a7e9a61 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/backends/cuda/tests/test_conv_routing.py @@ -20,6 +20,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from torch._subclasses.fake_tensor import FakeTensorMode _PASS_LOGGER = "executorch.backends.cuda.triton.replacement_pass" @@ -711,14 +712,17 @@ def test_benchmarked_range(self): ) -def _dense(conv, in_channels, out_channels, kernel, dtype, view=None, **kwargs): +def _dense( + conv, in_channels, out_channels, kernel, dtype, view=None, pad=None, **kwargs +): class Dense(nn.Module): def __init__(self): super().__init__() self.conv = conv(in_channels, out_channels, kernel, dtype=dtype, **kwargs) def forward(self, x): - return self.conv(x if view is None else x.view(view)) + x = x if view is None else x.view(view) + return self.conv(x if pad is None else F.pad(x, pad)) return Dense @@ -911,8 +915,33 @@ def admitted(input): ) -# 1-D convolutions over short sequences: (row, model, input shape, dtype, -# why Inductor keeps them). +# 3-D convolutions: a causal 3x3x3 after an explicit pad, a 1x1x1 and a +# (3,1,1) stride (2,1,1): (row, model, input shape, dtype, why Inductor +# keeps them). +_CONV3D_ROWS = ( + ( + "causal 3x3x3", + _dense(nn.Conv3d, 160, 160, 3, torch.bfloat16, pad=(1, 1, 1, 1, 2, 0)), + (1, 160, 1, 112, 224), + torch.bfloat16, + "160 input channels", + ), + ( + "1x1x1", + _dense(nn.Conv3d, 160, 320, 1, torch.bfloat16), + (2, 160, 16, 32, 32), + torch.bfloat16, + "1x1 convolution", + ), + ( + "3x1x1 s2", + _dense(nn.Conv3d, 320, 320, (3, 1, 1), torch.bfloat16, stride=(2, 1, 1)), + (2, 320, 17, 16, 16), + torch.bfloat16, + "320 input channels", + ), +) +# 1-D convolutions over short sequences, likewise. _CONV1D_ROWS = ( ( "k5", @@ -1195,6 +1224,16 @@ def test_only_benchmarked_layouts_are_routed(self): else: self.assertIn("not a benchmarked layout", reason) + def test_conv3d_rows(self): + # Inductor keeps every row: none is benchmarked faster. + for row, model, shape, dtype, reason in _CONV3D_ROWS: + with self.subTest(row=row): + self.assertKept( + model, + lambda shape=shape, dtype=dtype: torch.randn(shape, dtype=dtype), + reason=reason, + ) + def test_conv1d_rows(self): # to_edge leaves a 1-D convolution as aten.convolution, which Inductor # cannot lower, so the template takes it. to_edge_transform_and_lower @@ -1227,3 +1266,30 @@ def test_same_padded_conv1d_is_required(self): targets, logs = self.run_pass(graph_module) self.assertEqual(targets.count(torch.ops.triton.convolution.default), 1) self.assertIn("required, implicit_gemm template", logs) + + def test_explicit_pad_stays_its_own_node(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + take_all = tuple( + dataclasses.replace(t, worth_routing=lambda call: None) + for t in routing.TEMPLATES + ) + for row, model, shape, dtype, _ in _CONV3D_ROWS[:1]: + graphs = _backend_graphs( + model, lambda shape=shape, dtype=dtype: torch.randn(shape, dtype=dtype) + ) + for api, graph_module in graphs.items(): + with self.subTest(row=row, api=api), mock.patch.object( + routing, "TEMPLATES", take_all + ): + self.run_pass(graph_module) + (conv,) = [ + n + for n in graph_module.graph.nodes + if n.target == torch.ops.triton.convolution.default + ] + self.assertRegex( + str(conv.args[0].target), r"aten\.(constant_pad_nd|pad)\." + ) + rank = len(shape) - 2 + self.assertEqual(conv.args[4:6], ([0] * rank, [0] * rank))