[executorch][cuda] Run 1-D convolutions on the implicit-GEMM template

backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py

@@ -230,15 +230,39 @@ 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),
+    # 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),
+    ("1d_k5", "512-1024 ch T8", (1, 512, 8), 1024, 5, 1, 2, True),
+    ("1d_k5", "1024 ch T8", (1, 1024, 8), 1024, 5, 1, 2, True),
+    ("1d_k5", "1024-2048 ch T4", (1, 1024, 4), 2048, 5, 1, 2, True),
+    ("1d_k5", "2048 ch T4", (1, 2048, 4), 2048, 5, 1, 2, True),
+    ("1d_k5", "4096-1024 ch T4", (1, 4096, 4), 1024, 5, 1, 2, True),
+    ("1d_k5", "2048-512 ch T8", (1, 2048, 8), 512, 5, 1, 2, True),
+    ("1d_k5", "512 ch T64", (1, 512, 64), 512, 5, 1, 2, True),
+    ("1d_k5", "1024 ch T32", (1, 1024, 32), 1024, 5, 1, 2, True),
+    ("1d_k5", "2048 ch T16", (1, 2048, 16), 2048, 5, 1, 2, True),
+    ("1d_k1", "2-512 ch T16", (1, 2, 16), 512, 1, 1, 0, True),
+    ("1d_k1", "512-1024 ch T8", (1, 512, 8), 1024, 1, 1, 0, True),
+    ("1d_k1", "1024-2048 ch T4", (1, 1024, 4), 2048, 1, 1, 0, True),
+    ("1d_k1", "4096-1024 ch T4", (1, 4096, 4), 1024, 1, 1, 0, True),
+    ("1d_k1", "2048-512 ch T8", (1, 2048, 8), 512, 1, 1, 0, True),
+    ("1d_k1", "512-2 ch T16", (1, 512, 16), 2, 1, 1, 0, True),
+    ("1d_k1", "1024-2048 ch T16", (1, 1024, 16), 2048, 1, 1, 0, True),
+    ("1d_k3_s2", "512 ch T16", (1, 512, 16), 512, 3, 2, 1, True),
+    ("1d_k3_s2", "1024 ch T8", (1, 1024, 8), 1024, 3, 2, 1, True),
+    ("1d_k3_s2", "512 ch T64", (1, 512, 64), 512, 3, 2, 1, True),
+    ("1d_k3_s2", "1024 ch T32", (1, 1024, 32), 1024, 3, 2, 1, True),
 ]
 
 
-def shape_cases(precision, layout):
+def shape_cases(precision, layout, groups=None):
     return [
         Case(
             **{**asdict(_case(f"{group} {label}", *rest, precision)), "layout": layout}
         )
         for group, label, *rest in SHAPES
+        if groups is None or group in groups
     ]
 
 
@@ -420,10 +444,17 @@ def run():
     return run, schedule
 
 
+def _conv1d(x, w, b, stride, padding=(0,)):
+    """The backend's ``conv1d_to_conv2d``: a 2-D convolution over a unit last dim."""
+    out = F.conv2d(x.unsqueeze(-1), w.unsqueeze(-1), b, (stride[0], 1), (padding[0], 0))
+    return out.squeeze(-1)
+
+
 def aten_conv(x, w, b, case):
-    """What the exported graph holds: a separate pad for asymmetric padding, and
-    a contiguous output where the graph records one."""
-    conv = {1: F.conv1d, 2: F.conv2d, 3: F.conv3d}[len(case.kernel)]
+    """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."""
+    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)
     else:
@@ -554,14 +585,19 @@ def main():
         choices=["nchw", "channels_last", "cl_to_nchw"],
         help="cl_to_nchw: channels-last input, contiguous output (as exported on CUDA)",
     )
+    parser.add_argument(
+        "--groups", help="shapes suite: comma-separated groups, e.g. 3x3,1d_k5"
+    )
     parser.add_argument("--shard", default="0/1", help="run cases i::n of the suite")
     parser.add_argument("--warm", action="store_true", help="do not flush L2")
     parser.add_argument("--threshold", type=float, default=1.25)
     parser.add_argument("--json", help="append one JSON object per case")
     args = parser.parse_args()
 
     cases = {
-        "shapes": lambda: shape_cases(args.precision, args.layout),
+        "shapes": lambda: shape_cases(
+            args.precision, args.layout, args.groups and args.groups.split(",")
+        ),
         "envelope": lambda: envelope_cases(args.precision.split(",")),
         "boundary": lambda: boundary_cases(args.precision, args.layout),
     }[args.suite]()

backends/cuda/tests/test_conv_implicit_gemm.py

@@ -28,7 +28,7 @@
 
 def _reference(x, w, b, stride, before, after):
     pad = [p for pair in zip(reversed(before), reversed(after)) for p in pair]
-    conv = F.conv2d if x.dim() == 4 else F.conv3d
+    conv = {3: F.conv1d, 4: F.conv2d, 5: F.conv3d}[x.dim()]
     bias = None if b is None else b.double()
     return conv(F.pad(x.double(), pad), w.double(), bias, stride)
 
@@ -43,6 +43,8 @@ def _dense_strides(shape, order):
 
 
 def _channels_last(t):
+    if t.dim() == 3:
+        return t.transpose(1, 2).contiguous().transpose(1, 2)
     return t.to(
         memory_format=torch.channels_last if t.dim() == 4 else torch.channels_last_3d
     )
@@ -185,6 +187,12 @@ def check_geometries(self, dtype, layouts=("contiguous", "channels_last")):
             ((1, 20, 9, 7), 3, (2, 4), (1, 3), (2, 0), (0, 5)),
             ((1, 17, 4, 9, 8), 18, (3, 3, 3), (1, 2, 1), (2, 1, 1), (0, 1, 1)),
             ((1, 2, 1, 1), 1, (1, 1), (1, 1), (0, 0), (0, 0)),
+            # 1-D: k5, k3 s2 and k1 convolutions over short sequences, and an
+            # asymmetric pad.
+            ((1, 64, 16), 96, (5,), (1,), (2,), (2,)),
+            ((2, 32, 8), 32, (3,), (2,), (1,), (1,)),
+            ((1, 48, 4), 64, (1,), (1,), (0,), (0,)),
+            ((1, 5, 13), 7, (4,), (3,), (2,), (1,)),
         )
         cases = [
             (geometry, layout, layout == "weight_cl" or i % 2 == 0)
@@ -236,6 +244,7 @@ def test_output_layouts(self):
                 (0, 0, 0),
                 [(0, 1, 2, 3, 4), (0, 2, 3, 4, 1), (1, 0, 2, 3, 4)],
             ),
+            ((2, 16, 12), 24, (5,), (1,), (2,), [(0, 1, 2), (0, 2, 1), (2, 1, 0)]),
         ):
             x, w, b = self.tensors(shape, c_out, kernel)
             out_shape = (
@@ -430,30 +439,55 @@ def forward(self, x):
         return self.proj(self.down(F.relu(self.padded(x)))), self.unpadded(x)
 
 
+class _Conv1dBlock(nn.Module):
+    """1-D convolutions: k5 p2, a k1 residual and a k3 s2 p1 downsample; and a
+    "same"-padded even kernel."""
+
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv1d(16, 32, 5, padding=2)
+        self.res = nn.Conv1d(16, 32, 1)
+        self.same = nn.Conv1d(32, 32, 4, padding="same")
+        self.down = nn.Conv1d(32, 32, 3, stride=2, padding=1)
+
+    def forward(self, x):
+        return self.down(self.same(F.mish(self.conv(x)) + self.res(x)))
+
+
 class ImplicitGemmLoweringTest(unittest.TestCase):
     def setUp(self):
         if not torch.cuda.is_available() or torch.version.hip is not None:
             self.skipTest("CUDA required")
 
-    def lower_and_run(self, model, inputs, routed, tolerance):
+    def lower_and_run(
+        self, model, inputs, routed, tolerance, route="optional", to_backend=False
+    ):
+        """Lowers ``model`` with to_edge_transform_and_lower, or with to_edge
+        and to_backend, checks that ``routed`` convolutions take the template
+        on ``route``, and runs the program against eager (outside fbcode)."""
         from executorch.backends.cuda.cuda_backend import CudaBackend
         from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
-        from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+        from executorch.exir import (
+            EdgeCompileConfig,
+            to_edge,
+            to_edge_transform_and_lower,
+        )
 
         with torch.no_grad():
             program = torch.export.export(model, inputs, strict=True)
+        partitioner = CudaPartitioner(
+            [CudaBackend.generate_method_name_compile_spec("forward")]
+        )
+        config = EdgeCompileConfig(_check_ir_validity=False)
         with self.assertLogs(_PASS_LOGGER, "INFO") as logs:
-            edge = to_edge_transform_and_lower(
-                program,
-                partitioner=[
-                    CudaPartitioner(
-                        [CudaBackend.generate_method_name_compile_spec("forward")]
-                    )
-                ],
-                compile_config=EdgeCompileConfig(_check_ir_validity=False),
-            )
+            if to_backend:
+                edge = to_edge(program, compile_config=config).to_backend(partitioner)
+            else:
+                edge = to_edge_transform_and_lower(
+                    program, partitioner=[partitioner], compile_config=config
+                )
         output = "\n".join(logs.output)
-        self.assertEqual(output.count("(optional, implicit_gemm template)"), routed)
+        self.assertEqual(output.count(f"({route}, implicit_gemm template)"), routed)
         self.assertIn(f"Replaced {routed} nodes with Triton kernels", output)
 
         from executorch.export.utils import is_fbcode
@@ -476,7 +510,9 @@ def lower_and_run(self, model, inputs, routed, tolerance):
             enabled=True, allow_tf32=False
         ):
             expected = model(*inputs)
-        for got, want in zip(outputs, expected):
+        if isinstance(expected, torch.Tensor):
+            expected = (expected,)
+        for got, want in zip(outputs, expected, strict=True):
             torch.testing.assert_close(
                 got.float(), want.float().cpu(), atol=tolerance, rtol=tolerance
             )
@@ -517,3 +553,17 @@ def test_three_channel_block_lowers_and_runs(self):
                 )
                 # The 1x1 convolutions keep Inductor's lowering.
                 self.assertEqual(output.count("Skipping conv replacement"), 2)
+
+    def test_conv1d_block_lowers_and_runs(self):
+        torch.manual_seed(0)
+        model = _Conv1dBlock().to("cuda").eval()
+        inputs = (torch.randn(1, 16, 16, device="cuda"),)
+        # to_edge keeps aten.convolution, which Inductor cannot lower in 1-D:
+        # the template is required. to_edge_transform_and_lower keeps
+        # aten.conv1d, which the backend lowers as a 2-D convolution, except
+        # for the "same" overload.
+        self.lower_and_run(model, inputs, 4, 1e-4, route="required", to_backend=True)
+        output = self.lower_and_run(
+            model, inputs, 1, self.fp32_tolerance(), route="required"
+        )
+        self.assertEqual(output.count("Skipping conv replacement"), 3)

backends/cuda/tests/test_conv_routing.py

@@ -926,6 +926,33 @@ def admitted(input):
     )
 
 
+# 1-D convolutions over short sequences: (row, model, input shape, dtype,
+# why Inductor keeps them).
+_CONV1D_ROWS = (
+    (
+        "k5",
+        _dense(nn.Conv1d, 512, 512, 5, torch.float32, padding=2),
+        (1, 512, 16),
+        torch.float32,
+        "512 input channels",
+    ),
+    (
+        "k1",
+        _dense(nn.Conv1d, 512, 1024, 1, torch.float32),
+        (1, 512, 8),
+        torch.float32,
+        "1x1 convolution",
+    ),
+    (
+        "k3 s2",
+        _dense(nn.Conv1d, 512, 512, 3, torch.float32, stride=2, padding=1),
+        (1, 512, 16),
+        torch.float32,
+        "512 input channels",
+    ),
+)
+
+
 def _steps_past(families, shape, c_out, kernel, stride, padding):
     """The calls one step past each limit a benchmarked corner sits on."""
     from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm
@@ -1182,3 +1209,36 @@ def test_only_benchmarked_layouts_are_routed(self):
                             self.assertIsNone(reason)
                         else:
                             self.assertIn("not a benchmarked layout", 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
+        # leaves aten.conv1d, which the backend lowers as a 2-D convolution;
+        # no row is benchmarked faster, so Inductor keeps it.
+        for row, model, shape, dtype, reason in _CONV1D_ROWS:
+            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):
+                    targets, logs = self.run_pass(graph_module)
+                    if api == "to_edge_to_backend":
+                        self.assertEqual(
+                            targets.count(torch.ops.triton.convolution.default), 1
+                        )
+                        self.assertIn("required, implicit_gemm template", logs)
+                    else:
+                        self.assertNotIn(torch.ops.triton.convolution.default, targets)
+                        self.assertIn(reason, logs)
+
+    def test_same_padded_conv1d_is_required(self):
+        # The backend decomposes only conv1d.default into a 2-D convolution.
+        graphs = _backend_graphs(
+            _dense(nn.Conv1d, 16, 32, 4, torch.float32, padding="same"),
+            lambda: torch.randn(1, 16, 16),
+        )
+        for api, graph_module in graphs.items():
+            with self.subTest(api=api):
+                targets, logs = self.run_pass(graph_module)
+                self.assertEqual(targets.count(torch.ops.triton.convolution.default), 1)
+                self.assertIn("required, implicit_gemm template", logs)

backends/cuda/triton/kernels/conv/core/implicit_gemm.py

@@ -4,13 +4,13 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Implicit-GEMM template: dense (groups 1) forward convolution.
+"""Implicit-GEMM template: dense (groups 1) 1-D, 2-D and 3-D forward convolution.
 
 The convolution is a GEMM of the output pixels (``M = N * OD * OH * OW``) by the
 output channels, reducing over ``K = C_in * KD * KH * KW``; input tiles are
 gathered from the input as the GEMM walks K. The kernel is written once for
-(D, H, W): a 2-D convolution has a unit depth, which, like every shape, stride
-and geometry argument, is constexpr and compiles away; so does every bounds
+(D, H, W): lower ranks have unit leading dims, which, like every shape, stride
+and geometry argument, are constexpr and compile away; so does every bounds
 check the geometry cannot fail (all of them when kernel == stride and
 there is no padding). Tiles are autotuned. K is walked in
 one of two orders:
@@ -80,8 +80,6 @@ def _output_size(call: ConvCall) -> list[int]:
 
 def _geometry_reason(call: ConvCall) -> Optional[str]:
     x, w, p = call.input, call.weight, call.params
-    if p.rank not in (2, 3):
-        return f"implicit GEMM: {p.rank}-D convolution, needs 2-D or 3-D"
     if p.transposed:
         return "implicit GEMM: transposed convolution"
     if any(d != 1 for d in p.dilation):

backends/cuda/triton/kernels/conv/routing.py

@@ -55,17 +55,18 @@ def regular_lowering_gap(target: Any, call: ConvCall) -> Optional[str]:
     """Why the backend's regular lowering cannot compile this conv, or None.
 
     Inductor's Triton templates need dilation 1 and a forward convolution of
-    rank 2 or 3 (rank 1 only for depthwise). The backend decomposes ``conv1d``,
-    ``conv_transpose1d`` and ``conv_transpose2d`` into 2-D forward convolutions
-    before Inductor sees them.
+    rank 2 or 3 (rank 1 only for depthwise). The backend decomposes
+    ``conv1d.default``, ``conv_transpose1d.default`` and
+    ``conv_transpose2d.input`` into 2-D forward convolutions before Inductor
+    sees them; ``conv1d.padding`` stays a 1-D convolution.
     """
     p, w = call.params, call.weight
     if any(d != 1 for d in p.dilation):
         return f"dilation {list(p.dilation)} has no Inductor template"
-    decomposed = target._schema.name in (
-        "aten::conv1d",
-        "aten::conv_transpose1d",
-        "aten::conv_transpose2d",
+    decomposed = (target._schema.name, target._schema.overload_name) in (
+        ("aten::conv1d", ""),
+        ("aten::conv_transpose1d", ""),
+        ("aten::conv_transpose2d", "input"),
     )
     if p.transposed and not decomposed:
         return f"{p.rank}-D transposed convolution has no Inductor template"