diff --git a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py index 7ea984abe43..330df91c82d 100644 --- a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py +++ b/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,6 +585,9 @@ 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) @@ -561,7 +595,9 @@ def main(): 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]() diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py index 3a42bca57b8..ce629106018 100644 --- a/backends/cuda/tests/test_conv_implicit_gemm.py +++ b/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) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index 36b7f31961b..fa822a8fea0 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/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) diff --git a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py index 659e1c0e6fb..4273941a31e 100644 --- a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py +++ b/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): diff --git a/backends/cuda/triton/kernels/conv/routing.py b/backends/cuda/triton/kernels/conv/routing.py index 45227e3981b..65d7d727241 100644 --- a/backends/cuda/triton/kernels/conv/routing.py +++ b/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"