[executorch][cuda] Test 3-D convolutions on the implicit-GEMM template

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
 
 

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)

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"
@@ -726,14 +727,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
 
@@ -926,8 +930,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",
@@ -1210,6 +1239,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
@@ -1242,3 +1281,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))