diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 8ea87c56a79..80c705bcc0d 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -169,6 +169,11 @@ fbcode_target( name = "triton_kernels", srcs = [ "triton/kernels/__init__.py", + "triton/kernels/conv/__init__.py", + "triton/kernels/conv/conv.py", + "triton/kernels/conv/core/__init__.py", + "triton/kernels/conv/core/common.py", + "triton/kernels/conv/routing.py", "triton/kernels/fused_moe.py", "triton/kernels/int4_quantized_gemm.py", "triton/kernels/int5_quantized_gemm.py", @@ -188,6 +193,7 @@ fbcode_target( ], deps = [ ":autotune", + ":triton_routing", "//caffe2:torch", ], ) diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index b07feeafd87..3829ded98ed 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -52,6 +52,51 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_conv_binding", + srcs = [ + "test_conv_binding.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:triton_kernels", + "//executorch/exir/dialects:lib", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100", + ), + ) + + python_unittest_remote_gpu( + name = "test_conv_routing", + srcs = [ + "test_conv_routing.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", + "//executorch/backends/cuda:triton_kernels", + "//executorch/backends/cuda:triton_replacement_pass", + "//executorch/exir:lib", + "//executorch/exir/backend:backend_details", + "//executorch/exir/backend:compile_spec_schema", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100", + ), + ) + runtime.python_library( name = "autotune_test_utils", srcs = [ diff --git a/backends/cuda/tests/test_conv_binding.py b/backends/cuda/tests/test_conv_binding.py new file mode 100644 index 00000000000..58e5a80a4cb --- /dev/null +++ b/backends/cuda/tests/test_conv_binding.py @@ -0,0 +1,398 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Binding of every conv overload to triton::convolution's arguments, and the +output layout helper, checked against ATen's own fake-tensor convolution. + +No GPU needed: the tensors are fake CUDA tensors. +""" + +import itertools +import unittest + +import torch +from executorch.backends.cuda.triton.kernels.conv.core.common import ( + bind_conv, + ConvBindingError, + ConvCall, + ConvParams, + out_shape_and_stride, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + +aten = torch.ops.aten + + +def _channels_last(t): + """``t`` with channels-last strides (for 3-D tensors: channel stride 1).""" + if t.dim() == 3: + return t.permute(0, 2, 1).contiguous().permute(0, 2, 1) + return t.contiguous( + memory_format=torch.channels_last if t.dim() == 4 else torch.channels_last_3d + ) + + +class ConvBindingTest(unittest.TestCase): + def setUp(self): + self.mode = FakeTensorMode() + + def empty(self, *shape, dtype=torch.float32): + with self.mode: + return torch.empty(*shape, dtype=dtype, device="cuda") + + def assertBinds(self, op, args, kwargs, params): + """``op`` binds to ``params`` and its output matches ATen's.""" + input, weight, bias, bound = bind_conv(op, args, kwargs) + self.assertIs(input, args[0]) + self.assertEqual(bound, params) + with self.mode: + expected = op(*args, **kwargs) + call = ConvCall(input, weight, bias, bound) + self.assertEqual( + out_shape_and_stride(call), + (tuple(expected.shape), tuple(expected.stride())), + ) + edge_op = getattr( + getattr(exir_ops.edge.aten, op._schema.name.removeprefix("aten::")), + op._overloadname, + ) + self.assertEqual(bind_conv(edge_op, args, kwargs)[3], params) + + def test_omitted_defaults(self): + for rank, op in enumerate((aten.conv1d, aten.conv2d, aten.conv3d), 1): + x = self.empty(2, 4, *([9] * rank)) + w = self.empty(6, 4, *([3] * rank)) + ones, zeros = (1,) * rank, (0,) * rank + with self.subTest(rank=rank): + self.assertBinds( + op.default, + (x, w), + {}, + ConvParams(ones, zeros, zeros, ones, False, zeros, 1), + ) + + def test_keywords_and_single_values(self): + x, w, b = self.empty(1, 4, 11, 13), self.empty(8, 2, 3, 5), self.empty(8) + self.assertBinds( + aten.conv2d.default, + (x, w), + {"bias": b, "stride": [2], "padding": [1, 2], "groups": 2}, + ConvParams((2, 2), (1, 2), (1, 2), (1, 1), False, (0, 0), 2), + ) + self.assertBinds( + aten.conv2d.default, + (x, self.empty(8, 4, 3, 5), None, [1, 2]), + {"dilation": [2]}, + ConvParams((1, 2), (0, 0), (0, 0), (2, 2), False, (0, 0), 1), + ) + + def test_padding_modes(self): + x = self.empty(1, 3, 10, 12) + cases = ( + ((4, 4), [1, 1], "same", (1, 1), (2, 2)), + ((3, 4), [2, 1], "same", (2, 1), (2, 2)), + ((2, 5), [3, 1], "same", (1, 2), (2, 2)), + ((3, 3), [1, 1], "valid", (0, 0), (0, 0)), + ) + for kernel, dilation, mode, before, after in cases: + w = self.empty(5, 3, *kernel) + with self.subTest(kernel=kernel, dilation=dilation, mode=mode): + self.assertBinds( + aten.conv2d.padding, + (x, w, None, [1, 1], mode, dilation), + {}, + ConvParams( + (1, 1), before, after, tuple(dilation), False, (0, 0), 1 + ), + ) + self.assertBinds( + aten.conv1d.padding, + (self.empty(2, 3, 7), self.empty(3, 3, 2)), + {"padding": "same"}, + ConvParams((1,), (0,), (1,), (1,), False, (0,), 1), + ) + self.assertBinds( + aten.conv3d.padding, + (self.empty(1, 2, 5, 6, 7), self.empty(4, 1, 2, 3, 4)), + {"padding": "same", "groups": 2}, + ConvParams((1,) * 3, (0, 1, 1), (1, 1, 2), (1,) * 3, False, (0,) * 3, 2), + ) + + def test_same_padding_needs_unit_stride(self): + x, w = self.empty(1, 3, 10, 12), self.empty(5, 3, 3, 3) + with self.assertRaisesRegex(ConvBindingError, "stride 1"): + bind_conv(aten.conv2d.padding, (x, w, None, [2, 2], "same"), {}) + + def test_transposed_argument_order(self): + # conv_transpose takes output_padding before groups and dilation, and an + # [in, out / groups, k] weight. + self.assertBinds( + aten.conv_transpose1d.default, + (self.empty(2, 4, 9), self.empty(4, 3, 4), None, [2], [1], [1], 2, [1]), + {}, + ConvParams((2,), (1,), (1,), (1,), True, (1,), 2), + ) + self.assertBinds( + aten.conv_transpose2d.input, + (self.empty(1, 6, 5, 7), self.empty(6, 2, 3, 3), self.empty(6)), + {"stride": [2, 3], "output_padding": [1, 2], "groups": 3, "dilation": [1]}, + ConvParams((2, 3), (0, 0), (0, 0), (1, 1), True, (1, 2), 3), + ) + self.assertBinds( + aten.conv_transpose3d.input, + (self.empty(1, 2, 3, 4, 5), self.empty(2, 3, 2, 2, 2), None, [2], [1]), + {"dilation": [1, 2, 1]}, + ConvParams((2,) * 3, (1,) * 3, (1,) * 3, (1, 2, 1), True, (0,) * 3, 1), + ) + + def test_convolution(self): + # A forward convolution ignores output_padding. + x, w = self.empty(2, 4, 8, 8), self.empty(4, 2, 3, 3) + for transposed in (False, True): + with self.subTest(transposed=transposed): + self.assertBinds( + aten.convolution.default, + (x, w, None, [2, 1], [1, 0], [1, 1], transposed, [1, 0], 2), + {}, + ConvParams( + (2, 1), + (1, 0), + (1, 0), + (1, 1), + transposed, + (1, 0) if transposed else (0, 0), + 2, + ), + ) + + def test_rejects(self): + x, w = self.empty(4, 8, 8), self.empty(4, 4, 3, 3) + with self.assertRaisesRegex(ConvBindingError, "unbatched"): + bind_conv(aten.conv2d.default, (x, w), {}) + with self.assertRaisesRegex(ConvBindingError, "not a conv"): + bind_conv(aten.add.Tensor, (x, w), {}) + with self.assertRaisesRegex(ConvBindingError, "missing"): + bind_conv(aten.convolution.default, (x, w, None), {}) + with self.assertRaisesRegex(ConvBindingError, "entries"): + bind_conv( + aten.conv2d.default, (self.empty(1, 4, 8, 8), w, None, [1, 1, 1]), {} + ) + + +class CanonicalArgumentsTest(unittest.TestCase): + """A transposed convolution needs 0 <= output_padding < max(stride, + dilation) on every axis, a forward one output_padding 0; both binding and + direct calls of the canonical op check it.""" + + def setUp(self): + self.mode = FakeTensorMode() + + def empty(self, *shape): + with self.mode: + return torch.empty(*shape, device="cuda") + + def cases(self, rank): + """(stride, dilation, output_padding, valid) on the last axis.""" + return ( + (2, 1, 1, True), + (2, 1, 2, False), + (2, 1, -1, False), + # At or past the stride but below the dilation. + (1, 3, 2, True), + (1, 3, 3, False), + (3, 2, 2, True), + ) + + def test_transposed_binding(self): + ops = ( + aten.conv_transpose1d.default, + aten.conv_transpose2d.input, + aten.conv_transpose3d.input, + ) + for rank, op in enumerate(ops, 1): + x, w = self.empty(1, 4, *([6] * rank)), self.empty(4, 2, *([3] * rank)) + for stride, dilation, output_padding, valid in self.cases(rank): + ones = [1] * (rank - 1) + args = ( + x, + w, + None, + ones + [stride], + [0] * rank, + [0] * (rank - 1) + [output_padding], + 1, + ones + [dilation], + ) + convolution_args = (*args[:4], [0] * rank, args[7], True, args[5], 1) + with self.subTest( + rank=rank, stride=stride, output_padding=output_padding + ): + for target, call_args in ( + (op, args), + (aten.convolution.default, convolution_args), + ): + if valid: + params = bind_conv(target, call_args, {})[3] + self.assertEqual(params.output_padding[-1], output_padding) + else: + with self.assertRaisesRegex( + ConvBindingError, "output_padding" + ): + bind_conv(target, call_args, {}) + + def test_canonical_op(self): + import executorch.backends.cuda.triton.kernels # noqa: F401 + + for rank in (1, 2, 3): + x, w = self.empty(1, 4, *([6] * rank)), self.empty(4, 2, *([3] * rank)) + for stride, dilation, output_padding, valid in self.cases(rank): + ones = [1] * (rank - 1) + args = ( + x, + w, + None, + ones + [stride], + [0] * rank, + [0] * rank, + ones + [dilation], + True, + [0] * (rank - 1) + [output_padding], + 1, + ) + with self.subTest( + rank=rank, stride=stride, output_padding=output_padding + ): + with self.mode: + if valid: + out = torch.ops.triton.convolution(*args) + self.assertEqual(out.dim(), rank + 2) + else: + with self.assertRaisesRegex(ValueError, "output_padding"): + torch.ops.triton.convolution(*args) + x, w = self.empty(1, 4, 6, 6), self.empty(4, 4, 3, 3) + forward = (x, w, None, [1, 1], [0, 0], [0, 0], [1, 1], False, [0, 1], 1) + with self.mode, self.assertRaisesRegex(ValueError, "output_padding"): + torch.ops.triton.convolution(*forward) + + +@unittest.skipIf( + torch.version.hip is not None, "compares with the CUDA layouts of aten.convolution" +) +class OutShapeAndStrideTest(unittest.TestCase): + """``out_shape_and_stride`` against FakeTensor ``aten.convolution`` on CUDA.""" + + LAYOUTS = ("contiguous", "input_cl", "weight_cl", "both_cl") + + def check(self, x_shape, w_shape, params, layout, offset=False): + mode = FakeTensorMode() + with mode: + x = torch.empty(x_shape, device="cuda") + w = torch.empty(w_shape, device="cuda") + if offset: + x = torch.empty(x.shape[0], x.shape[1] + 1, *x.shape[2:], device="cuda") + x = x[:, 1:] + if layout in ("input_cl", "both_cl"): + x = _channels_last(x) + if layout in ("weight_cl", "both_cl"): + w = _channels_last(w) + expected = aten.convolution.default( + x, + w, + None, + list(params.stride), + list(params.padding_before), + list(params.dilation), + params.transposed, + list(params.output_padding), + params.groups, + ) + got = out_shape_and_stride(ConvCall(x, w, None, params)) + self.assertEqual(got, (tuple(expected.shape), tuple(expected.stride()))) + + def test_matches_aten(self): + shapes = { + 1: [(2, 4, 9), (1, 1, 7), (3, 4, 1)], + 2: [(2, 4, 9, 7), (1, 4, 1, 1), (2, 1, 5, 6), (1, 4, 1, 6), (2, 4, 6, 1)], + 3: [(2, 4, 5, 6, 7), (1, 4, 1, 1, 1), (1, 1, 3, 1, 4)], + } + for rank, transposed, layout, offset in itertools.product( + (1, 2, 3), (False, True), self.LAYOUTS, (False, True) + ): + for x_shape in shapes[rank]: + channels = x_shape[1] + groups = 2 if channels % 2 == 0 else 1 + kernel = (3,) + (1, 2)[: rank - 1] + w_shape = ( + (channels, 3, *kernel) + if transposed + else (6, channels // groups, *kernel) + ) + params = ConvParams( + (2,) * rank, + (1,) * rank, + (1,) * rank, + (1,) * rank, + transposed, + (1,) * rank if transposed else (0,) * rank, + groups, + ) + with self.subTest( + x=x_shape, transposed=transposed, layout=layout, offset=offset + ): + self.check(x_shape, w_shape, params, layout, offset) + + def test_depthwise_and_dilated(self): + for layout in self.LAYOUTS: + with self.subTest(layout=layout): + self.check( + (2, 8, 9, 9), + (8, 1, 3, 3), + ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 8), + layout, + ) + self.check( + (2, 8, 9, 9), + (16, 1, 3, 3), + ConvParams((1, 1), (2, 2), (2, 2), (2, 2), False, (0, 0), 8), + layout, + ) + + def test_large_output_padding_is_contiguous(self): + # output_padding >= stride (allowed below the dilation) leaves cuDNN. + for layout in self.LAYOUTS: + with self.subTest(layout=layout): + self.check( + (1, 4, 5, 5), + (4, 2, 3, 3), + ConvParams((1, 1), (0, 0), (0, 0), (2, 2), True, (1, 1), 1), + layout, + ) + + def test_asymmetric_forward_matches_explicit_pad(self): + mode = FakeTensorMode() + with mode: + x = torch.empty(1, 3, 10, 11, device="cuda") + w = torch.empty(5, 3, 4, 3, device="cuda") + padded = aten.constant_pad_nd.default(x, [0, 2, 3, 1]) + expected = aten.convolution.default( + padded, w, None, [2, 1], [0, 0], [1, 1], False, [0, 0], 1 + ) + params = ConvParams((2, 1), (3, 0), (1, 2), (1, 1), False, (0, 0), 1) + got = out_shape_and_stride(ConvCall(x, w, None, params)) + self.assertEqual(got, (tuple(expected.shape), tuple(expected.stride()))) + + def test_asymmetric_transposed_crops_the_full_output(self): + # Full transposed length (in - 1) * s + d * (k - 1) + 1, plus + # output_padding, minus before and after. + mode = FakeTensorMode() + with mode: + x = torch.empty(2, 4, 6, 5, device="cuda") + w = torch.empty(4, 3, 3, 4, device="cuda") + params = ConvParams((2, 3), (1, 0), (2, 3), (1, 2), True, (1, 2), 1) + shape, _ = out_shape_and_stride(ConvCall(x, w, None, params)) + full = [(6 - 1) * 2 + 1 * (3 - 1) + 1, (5 - 1) * 3 + 2 * (4 - 1) + 1] + self.assertEqual(shape, (2, 3, full[0] + 1 - 1 - 2, full[1] + 2 - 0 - 3)) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py new file mode 100644 index 00000000000..b594c77e979 --- /dev/null +++ b/backends/cuda/tests/test_conv_routing.py @@ -0,0 +1,407 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Which convolutions ReplaceEdgeOpWithTritonOpPass routes to +triton::convolution, and its failure contract. + +The graphs are the ones the CUDA backend receives from both lowering APIs, +exported on fake CUDA tensors, so no GPU is needed. +""" + +import contextlib +import dataclasses +import importlib.util +import unittest +from unittest import mock + +import torch +import torch.nn as nn + +_PASS_LOGGER = "executorch.backends.cuda.triton.replacement_pass" +_APIS = ("to_edge_transform_and_lower", "to_edge_to_backend") + + +def _backend_graphs(model, *inputs, dynamic_shapes=None, device="cuda"): + """The graph the CUDA backend gets from each lowering API. + + The model is exported on fake CUDA tensors, or on CPU tensors whose fake + values are then moved to CUDA, as the backend does before the pass runs. + """ + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import to_edge, to_edge_transform_and_lower + from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult + from torch._subclasses.fake_tensor import FakeTensorMode + + fake_mode = FakeTensorMode() + with fake_mode if device == "cuda" else contextlib.nullcontext(), torch.device( + device + ): + model = model() + args = tuple(make() for make in inputs) + program = torch.export.export(model, args, dynamic_shapes=dynamic_shapes) + captured = [] + + def capture(cls, edge_program, compile_specs, collector): + captured.append(edge_program.graph_module) + return PreprocessResult(processed_bytes=b"") + + # Lowering looks the backend up by class name; patch every class of that + # name, since the module may have been imported under another name too. + backends = {CudaBackend} + backends.update( + cls for cls in BackendDetails.__subclasses__() if cls.__name__ == "CudaBackend" + ) + with contextlib.ExitStack() as stack: + for backend in backends: + stack.enter_context( + mock.patch.object( + backend, "_preprocess_with_weight_collector", classmethod(capture) + ) + ) + to_edge_transform_and_lower(program, partitioner=[CudaPartitioner([])]) + to_edge(program).to_backend(CudaPartitioner([])) + for graph_module in captured if device == "cpu" else (): + for node in graph_module.graph.nodes: + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + with fake_mode: + node.meta["val"] = torch.empty_strided( + val.shape, val.stride(), dtype=val.dtype, device="cuda" + ) + return dict(zip(_APIS, captured)) + + +def _targets(graph_module): + return [n.target for n in graph_module.graph.nodes if n.op == "call_function"] + + +def _nodes(graph_module): + return [ + (n.op, n.target, n.name, n.args, n.kwargs, sorted(u.name for u in n.users)) + for n in graph_module.graph.nodes + ] + + +def _depthwise(channels=8, kernel=3, padding=1, dtype=torch.float32, **kwargs): + class Depthwise(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d( + channels, + kwargs.pop("out_channels", channels), + kernel, + padding=padding, + groups=channels, + dtype=dtype, + **kwargs, + ) + + def forward(self, x): + return self.conv(x) + + return Depthwise + + +def _runs_on_rocm(test): + """Marks a test of the routing on ROCm: the other tests check the routing + on CUDA, and skip on ROCm, where the router takes no convolution.""" + test.runs_on_rocm = True + return test + + +class _PassTest(unittest.TestCase): + def setUp(self): + if importlib.util.find_spec("triton") is None: + self.skipTest("Triton required") + if torch.version.hip is not None and not getattr( + getattr(self, self._testMethodName), "runs_on_rocm", False + ): + self.skipTest("the router takes no convolution on ROCm") + + def run_pass(self, graph_module): + from executorch.backends.cuda.triton.replacement_pass import ( + ReplaceEdgeOpWithTritonOpPass, + ) + + with self.assertLogs(_PASS_LOGGER, "INFO") as logs: + ReplaceEdgeOpWithTritonOpPass()(graph_module) + graph_module.graph.lint() + return _targets(graph_module), "\n".join(logs.output) + + def assertRouted( + self, model, *inputs, template, canonical_args=None, device="cuda" + ): + for api, graph_module in _backend_graphs(model, *inputs, device=device).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(f"optional, {template} template", logs) + self.assertIn("Replaced 1 nodes", logs) + node = next( + n + for n in graph_module.graph.nodes + if n.target == torch.ops.triton.convolution.default + ) + if canonical_args is not None and api == _APIS[0]: + self.assertEqual(node.args[3:], canonical_args) + + def assertKept(self, model, *inputs, reason, **export_kwargs): + for api, graph_module in _backend_graphs( + model, *inputs, **export_kwargs + ).items(): + with self.subTest(api=api): + targets, logs = self.run_pass(graph_module) + self.assertNotIn(torch.ops.triton.convolution.default, targets) + self.assertIn("Skipping conv replacement", logs) + self.assertIn(reason, logs) + self.assertIn("Replaced 0 nodes", logs) + + +class ConvRoutingTest(_PassTest): + """The routing gates, with a template that supports every call.""" + + def setUp(self): + super().setUp() + from executorch.backends.cuda.triton.kernels.conv import routing + + patch = mock.patch.object(routing, "TEMPLATES", (_permissive_template(),)) + patch.start() + self.addCleanup(patch.stop) + + def test_canonical_arguments(self): + self.assertRouted( + _depthwise(), + lambda: torch.randn(2, 8, 12, 10), + template="permissive", + canonical_args=([1, 1], [1, 1], [1, 1], [1, 1], False, [0, 0], 8), + ) + # An even kernel pads one more element after than before. + self.assertRouted( + _depthwise(kernel=4, padding="same"), + lambda: torch.randn(1, 8, 9, 9), + template="permissive", + canonical_args=([1, 1], [1, 1], [2, 2], [1, 1], False, [0, 0], 8), + ) + + def test_channels_last_cpu_export_is_routed(self): + # torch.export traces with cuDNN off. On CPU it still records the + # channels-last output ATen's CUDA convolution returns. + self.assertRouted( + _depthwise(kernel=7, padding=3, stride=2), + lambda: torch.randn(1, 8, 15, 9).to(memory_format=torch.channels_last), + template="permissive", + device="cpu", + ) + + def channels_last_weight(): + return _depthwise(kernel=5, padding=2)().to( + memory_format=torch.channels_last + ) + + self.assertRouted( + channels_last_weight, + lambda: torch.randn(1, 8, 6, 7), + template="permissive", + device="cpu", + ) + + def test_channels_last_cuda_export_is_kept(self): + # On CUDA, with cuDNN off, the export records a contiguous output. + self.assertKept( + _depthwise(), + lambda: torch.randn(1, 8, 9, 9).to(memory_format=torch.channels_last), + reason="differs from the node's", + ) + + @_runs_on_rocm + def test_rocm_is_kept(self): + for graph_module in _backend_graphs( + _depthwise(), lambda: torch.randn(1, 8, 9, 9) + ).values(): + with mock.patch.object(torch.version, "hip", "6.0"): + targets, logs = self.run_pass(graph_module) + self.assertNotIn(torch.ops.triton.convolution.default, targets) + self.assertIn("Skipping conv replacement", logs) + self.assertIn("ROCm", logs) + + def test_symbolic_shape_is_kept(self): + batch = torch.export.Dim("batch", min=2, max=8) + self.assertKept( + _depthwise(), + lambda: torch.randn(2, 8, 9, 9), + reason="symbolic shapes", + dynamic_shapes=({0: batch},), + ) + + def test_output_layout_mismatch_is_kept(self): + for graph_module in _backend_graphs( + _depthwise(), lambda: torch.randn(1, 8, 9, 9) + ).values(): + conv = next(n for n in graph_module.graph.nodes if "conv" in str(n.target)) + val = conv.meta["val"] + conv.meta["val"] = val.as_strided(val.shape, (648, 1, 72, 8)) + targets, logs = self.run_pass(graph_module) + self.assertNotIn(torch.ops.triton.convolution.default, targets) + self.assertIn("differs from the node's", logs) + + def test_worth_routing_gates_only_optional_calls(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + unprofitable = dataclasses.replace( + _permissive_template(), worth_routing=lambda call: "not worth it" + ) + with mock.patch.object(routing, "TEMPLATES", (unprofitable,)): + self.assertKept( + _depthwise(), lambda: torch.randn(1, 8, 9, 9), reason="not worth it" + ) + # Only a template can compile a dilated convolution. + graph_module = _backend_graphs( + _depthwise(dilation=2, padding=2), lambda: torch.randn(1, 8, 9, 9) + )[_APIS[0]] + targets, logs = self.run_pass(graph_module) + self.assertIn(torch.ops.triton.convolution.default, targets) + self.assertIn("required, permissive template", logs) + + def test_without_a_template_the_node_is_kept(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + with mock.patch.object(routing, "TEMPLATES", ()): + self.assertKept( + _depthwise(), lambda: torch.randn(1, 8, 9, 9), reason="no template" + ) + + +def _permissive_template(): + """A test template that claims every call, so the routing of calls the + regular lowering cannot compile (REQUIRED) can be exercised.""" + from executorch.backends.cuda.triton.kernels.conv.conv import ConvTemplate + + return ConvTemplate( + "permissive", + lambda call: None, + ("any",), + lambda name, call: None, + lambda *args, **kwargs: None, + lambda *args, **kwargs: None, + lambda call: None, + ) + + +class ConvFailureContractTest(unittest.TestCase): + """Faults injected at each step: OPTIONAL calls keep their node, REQUIRED + calls raise; either way the graph is left as it was.""" + + def setUp(self): + if importlib.util.find_spec("triton") is None: + self.skipTest("Triton required") + if torch.version.hip is not None: + self.skipTest("the router takes no convolution on ROCm") + + def graph(self, dilation=1): + """A conv whose output has two users.""" + conv = _depthwise(dilation=dilation, padding=dilation) + + class TwoUsers(nn.Module): + def __init__(self): + super().__init__() + self.conv = conv() + + def forward(self, x): + y = self.conv(x) + return y + 1, y * 2 + + graph_module = _backend_graphs(TwoUsers, lambda: torch.randn(1, 8, 9, 9))[ + _APIS[0] + ] + (conv_node,) = ( + n + for n in graph_module.graph.nodes + if n.op == "call_function" and "conv" in str(n.target) + ) + self.assertEqual(len(conv_node.users), 2) + return graph_module + + def run_pass(self, graph_module): + from executorch.backends.cuda.triton.replacement_pass import ( + ReplaceEdgeOpWithTritonOpPass, + ) + + ReplaceEdgeOpWithTritonOpPass()(graph_module) + + def faults(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + def fail(*args, **kwargs): + raise RuntimeError("injected fault") + + # Faults in the rewrite itself are tested in test_triton_replacement.py. + return { + "binding": mock.patch.object(routing, "bind_conv", fail), + "preflight": mock.patch.object(routing, "out_shape_and_stride", fail), + "template selection": mock.patch.object( + routing, "TEMPLATES", (_failing_template(),) + ), + } + + def assertUnchanged(self, graph_module, before): + graph_module.graph.lint() + self.assertEqual(_nodes(graph_module), before) + self.assertNotIn(torch.ops.triton.convolution.default, _targets(graph_module)) + + def test_optional_failures_keep_the_node(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + for step, fault in self.faults().items(): + graph_module = self.graph() + before = _nodes(graph_module) + with self.subTest(step=step): + with mock.patch.object( + routing, "TEMPLATES", (_permissive_template(),) + ), fault: + self.run_pass(graph_module) + self.assertUnchanged(graph_module, before) + + def test_required_is_replaced(self): + from executorch.backends.cuda.triton.kernels.conv import routing + + graph_module = self.graph(dilation=2) + with self.assertLogs(_PASS_LOGGER, "INFO") as logs: + with mock.patch.object(routing, "TEMPLATES", (_permissive_template(),)): + self.run_pass(graph_module) + self.assertIn("required, permissive template", "\n".join(logs.output)) + self.assertIn(torch.ops.triton.convolution.default, _targets(graph_module)) + + def test_required_failures_raise(self): + from executorch.backends.cuda.triton.kernels.conv import routing + from executorch.backends.cuda.triton.routing import ReplacementError + + faults = self.faults() + # Nothing is classified before binding; a call that cannot be bound + # keeps the regular lowering. + del faults["binding"] + for step, fault in faults.items(): + graph_module = self.graph(dilation=2) + before = _nodes(graph_module) + with self.subTest(step=step): + with mock.patch.object( + routing, "TEMPLATES", (_permissive_template(),) + ), fault: + with self.assertRaisesRegex( + ReplacementError, + r"convolution\(input=float32\[1, 8, 9, 9\].*dilation=\[2, 2\]", + ): + self.run_pass(graph_module) + self.assertUnchanged(graph_module, before) + + +def _failing_template(): + from executorch.backends.cuda.triton.kernels.conv.conv import ConvTemplate + + def fail(*args, **kwargs): + raise RuntimeError("injected fault") + + return ConvTemplate("failing", fail, ("any",), fail, fail, fail, fail) diff --git a/backends/cuda/triton/kernels/__init__.py b/backends/cuda/triton/kernels/__init__.py index 625aa726953..942e8be0f49 100644 --- a/backends/cuda/triton/kernels/__init__.py +++ b/backends/cuda/triton/kernels/__init__.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from executorch.backends.cuda.triton.kernels.conv import convolution from executorch.backends.cuda.triton.kernels.fused_moe import ( fused_moe, fused_moe_batched, @@ -24,6 +25,7 @@ from executorch.backends.cuda.triton.kernels.topk import topk __all__ = [ + "convolution", "dequant_w4_to_bf16", "fused_moe", "fused_moe_batched", diff --git a/backends/cuda/triton/kernels/conv/__init__.py b/backends/cuda/triton/kernels/conv/__init__.py new file mode 100644 index 00000000000..f40db9f455d --- /dev/null +++ b/backends/cuda/triton/kernels/conv/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.cuda.triton.kernels.conv.conv import convolution + +__all__ = ["convolution"] diff --git a/backends/cuda/triton/kernels/conv/conv.py b/backends/cuda/triton/kernels/conv/conv.py new file mode 100644 index 00000000000..7031e79735b --- /dev/null +++ b/backends/cuda/triton/kernels/conv/conv.py @@ -0,0 +1,173 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""``triton::convolution``: the canonical convolution op of the CUDA backend. + +The op body allocates the output once, with the size and stride of +``out_shape_and_stride``, and hands it to the first template whose +``supports()`` accepts the call; the template's schedules only write into it. +""" + +from __future__ import annotations + +import contextlib +import contextvars +from dataclasses import dataclass +from typing import Callable, Iterator, Optional + +import torch +import triton +from executorch.backends.cuda.triton.kernels.conv.core.common import ( + check_params, + ConvCall, + ConvParams, + out_shape_and_stride, +) +from torch.library import triton_op + + +@dataclass(frozen=True) +class ConvTemplate: + name: str + supports: Callable[[ConvCall], Optional[str]] + schedules: tuple[str, ...] + schedule_supports: Callable[[str, ConvCall], Optional[str]] + launch: Callable[..., None] + run_schedule: Callable[..., None] + # Why a call ``supports`` accepts is still left to the regular lowering + # (outside the benchmarked range where the template wins), or None. + worth_routing: Callable[[ConvCall], Optional[str]] + + +TEMPLATES: tuple[ConvTemplate, ...] = () + +_FORCED: contextvars.ContextVar[Optional[tuple[str, Optional[triton.Config]]]] = ( + contextvars.ContextVar("forced_conv_schedule", default=None) +) + + +@contextlib.contextmanager +def force_conv_candidate( + schedule: str, config: Optional[triton.Config] = None +) -> Iterator[None]: + """For tests: ``triton::convolution`` runs ``schedule`` without timing, with + ``config`` instead of its autotuned config if given, and raises if the + schedule cannot run the call.""" + token = _FORCED.set((schedule, config)) + try: + yield + finally: + _FORCED.reset(token) + + +def select_template(call: ConvCall) -> ConvTemplate: + reasons = [] + for template in TEMPLATES: + reason = template.supports(call) + if reason is None: + return template + reasons.append(reason) + raise ValueError(f"no template supports {call.signature()}: {'; '.join(reasons)}") + + +def _canonical_call( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + stride: list[int], + padding_before: list[int], + padding_after: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> ConvCall: + params = ConvParams( + tuple(stride), + tuple(padding_before), + tuple(padding_after), + tuple(dilation), + transposed, + tuple(output_padding), + groups, + ) + check_params(params) + if not input.dim() == weight.dim() == params.rank + 2: + raise ValueError( + f"{params.rank}-D convolution of a {input.dim()}-D input with a " + f"{weight.dim()}-D weight" + ) + return ConvCall(input, weight, bias, params) + + +@triton_op("triton::convolution", mutates_args={}) +def convolution( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + stride: list[int], + padding_before: list[int], + padding_after: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + call = _canonical_call( + input, + weight, + bias, + stride, + padding_before, + padding_after, + dilation, + transposed, + output_padding, + groups, + ) + params = call.params + template = select_template(call) + size, stride_out = out_shape_and_stride(call) + out = torch.empty_strided(size, stride_out, dtype=input.dtype, device=input.device) + forced = _FORCED.get() + if forced is None: + template.launch(input, weight, bias, out, params) + return out + schedule, config = forced + reason = template.schedule_supports(schedule, call) + if reason is not None: + raise ValueError(f"forced schedule cannot run {call.signature()}: {reason}") + template.run_schedule(schedule, input, weight, bias, out, params, config) + return out + + +@convolution.register_fake +def _convolution_fake( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + stride: list[int], + padding_before: list[int], + padding_after: list[int], + dilation: list[int], + transposed: bool, + output_padding: list[int], + groups: int, +) -> torch.Tensor: + call = _canonical_call( + input, + weight, + bias, + stride, + padding_before, + padding_after, + dilation, + transposed, + output_padding, + groups, + ) + size, stride_out = out_shape_and_stride(call) + return input.new_empty_strided(size, stride_out) diff --git a/backends/cuda/triton/kernels/conv/core/__init__.py b/backends/cuda/triton/kernels/conv/core/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/cuda/triton/kernels/conv/core/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/backends/cuda/triton/kernels/conv/core/common.py b/backends/cuda/triton/kernels/conv/core/common.py new file mode 100644 index 00000000000..1ce053cc21f --- /dev/null +++ b/backends/cuda/triton/kernels/conv/core/common.py @@ -0,0 +1,332 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Shared pieces of ``triton::convolution``: the canonical arguments, binding +every ATen conv overload to them, and the output layout. + +Canonical arguments. Every spatial argument has one entry per spatial dim. + +- Forward (``transposed=False``): the input is padded with ``padding_before[i]`` + zeros in front of and ``padding_after[i]`` zeros behind spatial dim ``i``, + then convolved with no padding: + ``out = (in + before + after - dilation * (k - 1) - 1) // stride + 1``. + The weight is ``[C_out, C_in / groups, *k]`` and ``output_padding`` is 0. +- Transposed (``transposed=True``): the full transposed convolution of length + ``(in - 1) * stride + dilation * (k - 1) + 1`` gets ``output_padding[i]`` + positions appended, then ``before`` positions are cropped from the front and + ``after`` from the back: + ``out = (in - 1) * stride - before - after + dilation * (k - 1) + output_padding + 1``. + The weight is ``[C_in, C_out / groups, *k]``. ATen's transposed overloads + always bind ``before == after``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Optional, Sequence + +import torch + +SUPPORTED_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + +# (schema name, overload) -> spatial rank, or None when any rank is accepted. +CONV_OVERLOADS = { + ("aten::conv1d", ""): 1, + ("aten::conv2d", ""): 2, + ("aten::conv3d", ""): 3, + ("aten::conv1d", "padding"): 1, + ("aten::conv2d", "padding"): 2, + ("aten::conv3d", "padding"): 3, + ("aten::conv_transpose1d", ""): 1, + ("aten::conv_transpose2d", "input"): 2, + ("aten::conv_transpose3d", "input"): 3, + ("aten::convolution", ""): None, +} + + +class ConvBindingError(ValueError): + """The arguments of a conv call cannot be bound to the canonical ABI.""" + + +@dataclass(frozen=True) +class ConvParams: + stride: tuple[int, ...] + padding_before: tuple[int, ...] + padding_after: tuple[int, ...] + dilation: tuple[int, ...] + transposed: bool + output_padding: tuple[int, ...] + groups: int + + @property + def rank(self) -> int: + return len(self.stride) + + +@dataclass(frozen=True) +class ConvCall: + input: torch.Tensor + weight: torch.Tensor + bias: Optional[torch.Tensor] + params: ConvParams + + def signature(self) -> str: + def describe(t: Optional[torch.Tensor]) -> str: + if t is None: + return "None" + dtype = str(t.dtype).removeprefix("torch.") + return f"{dtype}{list(t.shape)}@{list(t.stride())}" + + p = self.params + return ( + f"convolution(input={describe(self.input)}, " + f"weight={describe(self.weight)}, bias={describe(self.bias)}, " + f"stride={list(p.stride)}, padding_before={list(p.padding_before)}, " + f"padding_after={list(p.padding_after)}, dilation={list(p.dilation)}, " + f"transposed={p.transposed}, output_padding={list(p.output_padding)}, " + f"groups={p.groups})" + ) + + +def check_params(p: ConvParams) -> None: + """Raises ``ConvBindingError`` for canonical arguments outside the ABI.""" + fields = (p.padding_before, p.padding_after, p.dilation, p.output_padding) + if any(len(f) != p.rank for f in fields) or p.rank not in (1, 2, 3): + raise ConvBindingError(f"spatial arguments of different lengths: {p}") + if min(p.stride) <= 0 or min(p.dilation) <= 0 or p.groups <= 0: + raise ConvBindingError(f"non-positive stride, dilation or groups: {p}") + if min(p.padding_before) < 0 or min(p.padding_after) < 0: + raise ConvBindingError(f"negative padding: {p}") + for o, s, d in zip(p.output_padding, p.stride, p.dilation): + limit = max(s, d) if p.transposed else 1 + if not 0 <= o < limit: + raise ConvBindingError( + f"output_padding {list(p.output_padding)} must be 0 for a forward " + "convolution and in [0, max(stride, dilation)) for a transposed one" + ) + + +def _expand(value: Any, rank: int, name: str) -> tuple[int, ...]: + values = [value] if isinstance(value, int) else list(value) + if len(values) == 1: + values = values * rank + if len(values) != rank: + raise ConvBindingError(f"{name} has {len(values)} entries for {rank} dims") + if not all(type(v) is int for v in values): + raise ConvBindingError(f"{name} is not static: {values}") + return tuple(values) + + +def _bind_schema(schema: Any, args: Sequence[Any], kwargs: dict[str, Any]) -> dict: + names = [a.name for a in schema.arguments] + if len(args) > len(names) or set(kwargs) - set(names[len(args) :]): + raise ConvBindingError(f"unexpected arguments for {schema.name}") + bound = {} + for i, arg in enumerate(schema.arguments): + if i < len(args): + bound[arg.name] = args[i] + elif arg.name in kwargs: + bound[arg.name] = kwargs[arg.name] + elif arg.has_default_value(): + bound[arg.name] = arg.default_value + else: + raise ConvBindingError(f"{schema.name} is missing {arg.name}") + return bound + + +def _padding( + padding: Any, + stride: tuple[int, ...], + dilation: tuple[int, ...], + kernel: Sequence[int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """``padding_before`` and ``padding_after`` for an int list, "valid" or + "same" (ATen puts the odd element of "same" after).""" + rank = len(stride) + if padding == "valid": + return (0,) * rank, (0,) * rank + if padding == "same": + if any(s != 1 for s in stride): + raise ConvBindingError("padding='same' needs stride 1") + total = [d * (k - 1) for d, k in zip(dilation, kernel)] + if not all(type(t) is int for t in total): + raise ConvBindingError("padding='same' with a symbolic kernel size") + before = tuple(t // 2 for t in total) + return before, tuple(t - b for t, b in zip(total, before)) + if isinstance(padding, str): + raise ConvBindingError(f"unknown padding mode {padding!r}") + padding = _expand(padding, rank, "padding") + return padding, padding + + +def bind_conv( + target: Any, + args: Sequence[Any], + kwargs: dict[str, Any], + value: Callable[[Any], Any] = lambda a: a, +) -> tuple[Any, Any, Any, ConvParams]: + """Binds a call of one of ``CONV_OVERLOADS`` to the canonical arguments. + + Returns ``(input, weight, bias, params)``; the first three are the call's + own arguments (e.g. FX nodes) and ``value`` maps them to tensors. + """ + schema = target._schema + key = (schema.name, schema.overload_name) + if key not in CONV_OVERLOADS: + raise ConvBindingError(f"{schema.name}.{schema.overload_name} is not a conv") + bound = _bind_schema(schema, args, kwargs) + input, weight = value(bound["input"]), value(bound["weight"]) + rank = weight.dim() - 2 + if CONV_OVERLOADS[key] not in (None, rank) or rank not in (1, 2, 3): + raise ConvBindingError(f"{schema.name} with a {weight.dim()}-D weight") + if input.dim() != rank + 2: + raise ConvBindingError(f"unbatched {input.dim()}-D input") + stride = _expand(bound["stride"], rank, "stride") + dilation = _expand(bound["dilation"], rank, "dilation") + transposed = schema.name.startswith("aten::conv_transpose") or bool( + bound.get("transposed", False) + ) + # A forward convolution ignores output_padding. + output_padding = ( + _expand(bound["output_padding"], rank, "output_padding") + if transposed + else (0,) * rank + ) + before, after = _padding(bound["padding"], stride, dilation, weight.shape[2:]) + groups = bound["groups"] + if type(groups) is not int: + raise ConvBindingError(f"groups is not static: {groups}") + params = ConvParams( + stride, before, after, dilation, transposed, output_padding, groups + ) + check_params(params) + return bound["input"], bound["weight"], bound["bias"], params + + +def _contiguous_strides(shape: Sequence[int]) -> tuple[int, ...]: + strides, running = [], 1 + for size in reversed(shape): + strides.append(running) + running *= max(size, 1) + return tuple(reversed(strides)) + + +def _channels_last_order(ndim: int) -> list[int]: + return [1, *range(ndim - 1, 1, -1), 0] + + +def channels_last_like(shape: Sequence[int], strides: Sequence[int]) -> bool: + """``Tensor.suggest_memory_format()`` is channels_last(_3d) for these strides.""" + if len(shape) not in (4, 5) or strides[1] == 0: + return False + smallest = 0 + for d in _channels_last_order(len(shape)): + if shape[d] == 0 or strides[d] < smallest: + return False + # ATen falls back to contiguous for ambiguous strides such as N111. + if d == 0 and smallest == strides[1]: + return False + smallest = strides[d] * shape[d] + return True + + +def _channels_last_strides(shape: Sequence[int]) -> tuple[int, ...]: + contiguous = _contiguous_strides(shape) + # Tensor.to(memory_format=...) keeps strides that already suggest it. + if channels_last_like(shape, contiguous): + return contiguous + strides, running = [0] * len(shape), 1 + for d in _channels_last_order(len(shape)): + strides[d] = running + running *= shape[d] + return tuple(strides) + + +def _channels_last_output(call: ConvCall) -> bool: + """Whether ATen's CUDA convolution returns a channels-last output. + + Mirrors the backend choice of a CUDA build with cuDNN: cuDNN picks + channels-last when the input or the weight suggests it; a transposed + convolution with ``output_padding >= stride`` and an empty input take + backends that always return contiguous outputs. A 1-D convolution runs as a + 2-D one, with the input made contiguous and both unsqueezed at dim 2. + """ + x, w, p = call.input, call.weight, call.params + if x.shape[0] == 0 or x.shape[1] == 0: + return False + if p.transposed and any(o >= s for o, s in zip(p.output_padding, p.stride)): + return False + x_shape, x_stride = tuple(x.shape), tuple(x.stride()) + w_shape, w_stride = tuple(w.shape), tuple(w.stride()) + if w.dim() == 3: + if not _is_contiguous(x_shape, x_stride): + x_stride = _contiguous_strides(x_shape) + x_shape, x_stride = _unsqueeze2(x_shape, x_stride) + w_shape, w_stride = _unsqueeze2(w_shape, w_stride) + return channels_last_like(x_shape, x_stride) or channels_last_like( + w_shape, w_stride + ) + + +def _is_contiguous(shape: Sequence[int], strides: Sequence[int]) -> bool: + expected = 1 + for size, stride in zip(reversed(shape), reversed(strides)): + if size != 1: + if stride != expected: + return False + expected *= size + return True + + +def _unsqueeze2( + shape: tuple[int, ...], strides: tuple[int, ...] +) -> tuple[tuple[int, ...], tuple[int, ...]]: + return (*shape[:2], 1, shape[2]), (*strides[:2], strides[2] * shape[2], strides[2]) + + +def out_shape_and_stride(call: ConvCall) -> tuple[tuple[int, ...], tuple[int, ...]]: + """The size and stride of the output, as ATen's CUDA convolution lays it out.""" + x, w, p = call.input, call.weight, call.params + kernel = w.shape[2:] + if p.transposed: + channels = w.shape[1] * p.groups + spatial = [ + (i - 1) * s - b - a + d * (k - 1) + o + 1 + for i, k, s, b, a, d, o in zip( + x.shape[2:], + kernel, + p.stride, + p.padding_before, + p.padding_after, + p.dilation, + p.output_padding, + ) + ] + else: + channels = w.shape[0] + spatial = [ + (i + b + a - d * (k - 1) - 1) // s + 1 + for i, k, s, b, a, d in zip( + x.shape[2:], + kernel, + p.stride, + p.padding_before, + p.padding_after, + p.dilation, + ) + ] + shape = (x.shape[0], channels, *spatial) + if not _channels_last_output(call): + return shape, _contiguous_strides(shape) + if len(shape) == 3: + strides = _channels_last_strides((shape[0], shape[1], 1, shape[2])) + return shape, (strides[0], strides[1], strides[3]) + return shape, _channels_last_strides(shape) + + +def max_offset(t: torch.Tensor) -> int: + """The largest element offset ``t`` addresses, relative to its data pointer.""" + return sum((n - 1) * s for n, s in zip(t.shape, t.stride()) if n > 0) diff --git a/backends/cuda/triton/kernels/conv/routing.py b/backends/cuda/triton/kernels/conv/routing.py new file mode 100644 index 00000000000..f70e39c3845 --- /dev/null +++ b/backends/cuda/triton/kernels/conv/routing.py @@ -0,0 +1,158 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Which conv nodes become ``triton::convolution``. + +``plan(node)`` is the replacement pass's preflight for convolutions +(see ``triton/routing.py``): it binds the node's arguments, checks the output +layout contract and the templates' ``supports()``, and classifies the node +without touching the graph. ``OPTIONAL`` means a template runs it, its +``worth_routing`` (the measured range where it beats the regular lowering) +accepts it, and the regular lowering could run it too; ``REQUIRED`` means a +template runs it and the regular lowering cannot. + +A failure while planning follows the same contract: it raises +``ReplacementError`` when the regular lowering cannot compile the node, and +otherwise leaves the node unchanged. +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +import torch +from executorch.backends.cuda.triton.kernels.conv.conv import TEMPLATES +from executorch.backends.cuda.triton.kernels.conv.core.common import ( + bind_conv, + ConvCall, + ConvParams, + out_shape_and_stride, +) +from executorch.backends.cuda.triton.routing import ( + ReplacementError, + ReplacementPlan, + Route, + unsupported, +) +from torch._subclasses.fake_tensor import FakeTensor +from torch.fx import Node + +logger: logging.Logger = logging.getLogger(__name__) + + +def _unsupported(reason: str) -> ReplacementPlan: + return unsupported("conv", reason) + + +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. + """ + 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", + ) + if p.transposed and not decomposed: + return f"{p.rank}-D transposed convolution has no Inductor template" + depthwise = p.groups > 1 and w.shape[1] == 1 and w.shape[0] == p.groups + if p.rank == 1 and not decomposed and not depthwise: + return "1-D convolution has no Inductor template" + return None + + +def _is_static(t: Optional[torch.Tensor]) -> bool: + return t is None or ( + isinstance(t, FakeTensor) + and all(type(s) is int for s in (*t.shape, *t.stride(), t.storage_offset())) + ) + + +def _check(node: Node, call: ConvCall, gap: Optional[str]) -> ReplacementPlan: + """An UNSUPPORTED plan, or an OPTIONAL one labelled with the template.""" + expected = node.meta.get("val") + if not all(_is_static(t) for t in (call.input, call.weight, call.bias, expected)): + return _unsupported("symbolic shapes or missing fake tensors") + if call.input.device.type != "cuda": + return _unsupported(f"input on {call.input.device}") + size, stride = out_shape_and_stride(call) + if (size, stride) != (tuple(expected.shape), tuple(expected.stride())): + return _unsupported( + f"output {list(size)}@{list(stride)} differs from the node's " + f"{list(expected.shape)}@{list(expected.stride())}" + ) + reasons = [] + for template in TEMPLATES: + reason = template.supports(call) + # A call the regular lowering cannot compile is routed whenever a + # template runs it; otherwise only where the template is worth it. + if reason is None and gap is None: + reason = template.worth_routing(call) + if reason is None: + return ReplacementPlan(Route.OPTIONAL, "", "conv", template.name) + reasons.append(reason) + return _unsupported(f"no template takes it: {'; '.join(reasons)}") + + +def plan(node: Node) -> ReplacementPlan: + """Classifies a conv node; never changes the graph.""" + if torch.version.hip is not None: + return _unsupported("ROCm") + + def value(arg: Any) -> Any: + return arg.meta.get("val") if isinstance(arg, Node) else arg + + try: + input, weight, bias, params = bind_conv( + node.target, node.args, node.kwargs, value + ) + call = ConvCall(value(input), value(weight), value(bias), params) + except Exception as e: # noqa: BLE001 - nothing is classified yet + return _unsupported(f"cannot bind arguments: {e!r}") + gap = regular_lowering_gap(node.target, call) + try: + result = _check(node, call, gap) + except Exception as e: + if gap is not None: + raise ReplacementError( + f"{call.signature()} ({gap}): preflight failed: {e!r}" + ) from e + logger.warning("Conv preflight failed for %s: %r", node.name, e) + return _unsupported(f"preflight failed: {e!r}") + if result.route is Route.UNSUPPORTED: + if gap is not None: + result = _unsupported(f"{result.reason}; regular lowering: {gap}") + return result + route = Route.OPTIONAL if gap is None else Route.REQUIRED + return ReplacementPlan( + route, + gap or "", + "conv", + f"{result.label} template", + torch.ops.triton.convolution.default, + (input, weight, bias, *_canonical_args(params)), + call.signature(), + ) + + +def _canonical_args(p: ConvParams) -> tuple[Any, ...]: + return ( + list(p.stride), + list(p.padding_before), + list(p.padding_after), + list(p.dilation), + p.transposed, + list(p.output_padding), + p.groups, + ) diff --git a/backends/cuda/triton/replacement_pass.py b/backends/cuda/triton/replacement_pass.py index 81625fa6a97..38d719e984e 100644 --- a/backends/cuda/triton/replacement_pass.py +++ b/backends/cuda/triton/replacement_pass.py @@ -15,6 +15,7 @@ import torch from executorch.backends.cuda.triton import routing +from executorch.backends.cuda.triton.kernels.conv import routing as conv_routing from executorch.backends.cuda.triton.kernels.sdpa import ( _prepare_mask_params, _validate_qkv_shapes, @@ -28,15 +29,31 @@ logger = logging.getLogger(__name__) triton = torch.ops.triton +CONV_TARGETS = ( + exir_ops.edge.aten.conv1d.default, + exir_ops.edge.aten.conv2d.default, + exir_ops.edge.aten.conv3d.default, + exir_ops.edge.aten.conv1d.padding, + exir_ops.edge.aten.conv2d.padding, + exir_ops.edge.aten.conv3d.padding, + exir_ops.edge.aten.conv_transpose1d.default, + exir_ops.edge.aten.conv_transpose2d.input, + exir_ops.edge.aten.conv_transpose3d.input, + exir_ops.edge.aten.convolution.default, +) + # Global mapping from edge dialect operators to Triton kernel functions EDGE_TO_TRITON_KERNELS = { exir_ops.edge.aten.scaled_dot_product_attention.default: triton.sdpa, exir_ops.edge.aten.topk.default: triton.topk, + **{target: triton.convolution for target in CONV_TARGETS}, } # Ops whose calls are classified before they are replaced (see routing.py): # edge op -> planner(node) returning a ``routing.ReplacementPlan``. -PLANNERS: dict[object, Callable[[Node], routing.ReplacementPlan]] = {} +PLANNERS: dict[object, Callable[[Node], routing.ReplacementPlan]] = { + target: conv_routing.plan for target in CONV_TARGETS +} _SPLITK_LKV_THRESHOLD = 256