From b1ff931b386c813cc32683e4530059cac311703e Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 13 Nov 2026 23:58:45 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- .../benchmarks/conv/benchmark_depthwise.py | 56 +++- .../conv/benchmark_implicit_gemm.py | 32 +- backends/cuda/tests/test_conv_binding.py | 98 ++++++ backends/cuda/tests/test_conv_depthwise.py | 67 +++- .../cuda/tests/test_conv_implicit_gemm.py | 84 ++++- backends/cuda/tests/test_conv_routing.py | 309 +++++++++++++++--- backends/cuda/triton/kernels/conv/conv.py | 15 +- .../cuda/triton/kernels/conv/core/common.py | 74 ++++- .../cuda/triton/kernels/conv/core/direct.py | 30 ++ .../triton/kernels/conv/core/implicit_gemm.py | 35 +- backends/cuda/triton/kernels/conv/routing.py | 26 +- 11 files changed, 710 insertions(+), 116 deletions(-) diff --git a/backends/cuda/benchmarks/conv/benchmark_depthwise.py b/backends/cuda/benchmarks/conv/benchmark_depthwise.py index a3e6a90e87d..83cd8cc4677 100644 --- a/backends/cuda/benchmarks/conv/benchmark_depthwise.py +++ b/backends/cuda/benchmarks/conv/benchmark_depthwise.py @@ -30,10 +30,14 @@ Suites: ``representative`` (model shapes, one dtype and layout per run), ``envelope`` (each axis of what the template supports varied around base shapes, every dtype and layout), ``extent`` (the largest shapes and the -remaining weight layouts) and ``grouped`` (1 to 8 input channels per group, -channel multipliers 1 to 4). The router's ``direct.worth_routing`` takes the -range these suites cover. For grouped shapes a case is significantly faster -when it is against both the template and the tap sum. +remaining weight layouts), ``layouts`` (those shapes with channels-last +inputs or weights and the contiguous output a graph exported on CUDA records, +which triton::convolution produces through ``output_strides`` and the other +paths through a ``contiguous()`` that Inductor may fuse) and ``grouped`` (1 to +8 input channels per group, channel multipliers 1 to 4). The router's +``direct.worth_routing`` takes the range these suites cover. For grouped shapes +a case is significantly faster when it is against both the template and the +tap sum. Usage: python benchmark_depthwise.py @@ -181,6 +185,21 @@ def envelope_cases(): return cases +def layout_cases(): + """The envelope and extent shapes whose input or weight is channels-last, + with the contiguous output a graph exported on CUDA records for them.""" + follow = { + "channels_last": "cl_to_nchw", + "weight_cl": "weight_cl_to_nchw", + "both_cl": "both_cl_to_nchw", + } + return [ + Case(**{**asdict(case), "layout": follow[case.layout]}) + for case in envelope_cases() + extent_cases() + if case.layout in follow + ] + + def extent_cases(): """The largest shapes and the remaining weight layouts the router takes: 4096 channels, 2048x2048 images, batch 256, ~50M-element inputs, and @@ -351,7 +370,11 @@ def _pad(x, case): return F.pad(x, [case.before[1], case.after[1], case.before[0], case.after[0]]) -def triton_runner(x, w, b, case): +def _contiguous_output(case): + return case.layout.endswith("_to_nchw") + + +def triton_runner(x, w, b, case, output_strides=None): args = ( list(case.stride), list(case.before), @@ -360,6 +383,7 @@ def triton_runner(x, w, b, case): False, [0, 0], case.groups, + output_strides, ) stream = torch.cuda.Stream() with torch.cuda.stream(stream), autotune_launch_params() as stats: @@ -375,10 +399,13 @@ def run(): def aten_conv(x, w, b, case): - """What the exported graph holds: a separate pad for asymmetric padding.""" + """What the exported graph holds: a separate pad for asymmetric padding, and + a contiguous output where the graph records one.""" if case.before == case.after: - return F.conv2d(x, w, b, case.stride, case.before, 1, case.groups) - return F.conv2d(_pad(x, case), w, b, case.stride, 0, 1, case.groups) + out = F.conv2d(x, w, b, case.stride, case.before, 1, case.groups) + else: + out = F.conv2d(_pad(x, case), w, b, case.stride, 0, 1, case.groups) + return out.contiguous() if _contiguous_output(case) else out def tap_sum(x, w, b, case): @@ -415,11 +442,15 @@ def run_case(case, flush, threshold): ) w = w / (case.group_channels * case.kernel[0] * case.kernel[1]) ** 0.5 b = torch.randn(out_channels, device="cuda", dtype=dtype) - if case.layout in ("channels_last", "both_cl"): + if case.layout in ("channels_last", "both_cl", "cl_to_nchw", "both_cl_to_nchw"): x = x.to(memory_format=torch.channels_last) - if case.layout in ("weight_cl", "both_cl"): + if case.layout in ("weight_cl", "both_cl", "weight_cl_to_nchw", "both_cl_to_nchw"): w = w.to(memory_format=torch.channels_last) - ours, schedule = triton_runner(x, w, b, case) + output_strides = None + if _contiguous_output(case): + shape = aten_conv(x[:1].double(), w.double(), None, case).shape + output_strides = list(torch.empty((case.shape[0], *shape[1:])).stride()) + ours, schedule = triton_runner(x, w, b, case, output_strides) runners = [ ("ours", ours), ("inductor", inductor_runner(x, w, b, case)), @@ -462,7 +493,7 @@ def main(): parser.add_argument( "--suite", default="representative", - choices=["representative", "envelope", "extent", "grouped"], + choices=["representative", "envelope", "extent", "layouts", "grouped"], ) parser.add_argument( "--dtype", default="float32", choices=["float32", "float16", "bfloat16"] @@ -480,6 +511,7 @@ def main(): "representative": lambda: representative_cases(args.dtype, args.layout), "envelope": envelope_cases, "extent": extent_cases, + "layouts": layout_cases, "grouped": grouped_cases, }[args.suite]() index, count = map(int, args.shard.split("/")) diff --git a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py index ea2abe1aedf..7ea984abe43 100644 --- a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py +++ b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py @@ -21,6 +21,10 @@ Precisions: ``fp32`` (TF32 off everywhere), ``tf32`` (TF32 allowed; needs a PyTorch whose Inductor conv templates can use TF32), ``fp16`` and ``bf16``. +Layouts: NCHW, channels-last input with channels-last output, and +channels-last input with the contiguous output a graph exported on CUDA +records (``cl_to_nchw``: triton::convolution gets ``output_strides``, the +other paths a ``contiguous()`` that Inductor may fuse). Suites: ``shapes`` (dense and kernel == stride convolutions at representative shapes), ``envelope`` (each axis of what the @@ -398,7 +402,11 @@ def _args(case): def triton_runner(x, w, b, case): - args = _args(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) stream = torch.cuda.Stream() with torch.cuda.stream(stream), autotune_launch_params() as stats: convolution(x, w, b, *args) @@ -413,12 +421,17 @@ def run(): def aten_conv(x, w, b, case): - """What the exported graph holds: a separate pad for asymmetric padding.""" + """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)] if case.before == case.after: - return conv(x, w, b, case.stride, case.before) - pad = [p for pair in zip(reversed(case.before), reversed(case.after)) for p in pair] - return conv(F.pad(x, pad), w, b, case.stride) + out = conv(x, w, b, case.stride, case.before) + 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) + return out.contiguous() if case.layout == "cl_to_nchw" else out def inductor_runner(x, w, b, case): @@ -470,7 +483,7 @@ def _tensors(case): else None ) memory_format = torch.channels_last if x.dim() == 4 else torch.channels_last_3d - if case.layout == "channels_last": + if case.layout in ("channels_last", "cl_to_nchw"): x = x.to(memory_format=memory_format) if case.layout == "weight_cl": w = w.to(memory_format=memory_format) @@ -535,7 +548,12 @@ def main(): default="fp32", help="fp32, tf32, fp16 or bf16; the envelope suite takes a comma-separated list", ) - parser.add_argument("--layout", default="nchw", choices=["nchw", "channels_last"]) + parser.add_argument( + "--layout", + default="nchw", + choices=["nchw", "channels_last", "cl_to_nchw"], + help="cl_to_nchw: channels-last input, contiguous output (as exported on CUDA)", + ) 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) diff --git a/backends/cuda/tests/test_conv_binding.py b/backends/cuda/tests/test_conv_binding.py index e73e62cfccc..ad927d4e1eb 100644 --- a/backends/cuda/tests/test_conv_binding.py +++ b/backends/cuda/tests/test_conv_binding.py @@ -19,6 +19,8 @@ ConvBindingError, ConvCall, ConvParams, + dense_strides, + memory_formats, out_shape_and_stride, ) from executorch.exir.dialects._ops import ops as exir_ops @@ -393,3 +395,99 @@ def test_asymmetric_transposed_crops_the_full_output(self): 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)) + + +class OutputStridesTest(unittest.TestCase): + """Explicit ``output_strides``: the output follows the strides the graph + recorded, which must be dense.""" + + def test_dense_strides(self): + self.assertTrue(dense_strides((2, 3, 4, 5), (60, 20, 5, 1))) + self.assertTrue(dense_strides((2, 3, 4, 5), (60, 1, 15, 3))) + self.assertTrue(dense_strides((2, 3, 4, 5), (5, 40, 10, 1))) + # Size-1 dims may have any stride. + self.assertTrue(dense_strides((1, 3, 1, 5), (7, 5, 999, 1))) + self.assertFalse(dense_strides((2, 3, 4, 5), (61, 20, 5, 1))) + self.assertFalse(dense_strides((2, 3, 4, 5), (60, 20, 0, 1))) + self.assertFalse(dense_strides((2, 3, 4, 5), (60, 20, 5, 2))) + self.assertFalse(dense_strides((1, 3, 1, 5), (7, 5, -1, 1))) + self.assertFalse(dense_strides((2, 3), (3,))) + + def test_memory_formats(self): + both = {"contiguous", "channels_last"} + for shape, strides, formats in ( + ((2, 3, 4, 5), (60, 20, 5, 1), {"contiguous"}), + ((2, 3, 4, 5), (60, 1, 15, 3), {"channels_last"}), + ((2, 3, 4, 5, 6), (360, 1, 90, 18, 3), {"channels_last"}), + ((2, 3, 4, 5), (5, 40, 10, 1), set()), + ((2, 3, 4, 5), (60, 20, 5, 2), set()), + ((8, 1, 3, 3), (9, 9, 3, 1), both), + ((2, 3, 1, 1), (3, 1, 1, 1), both), + ((1, 3, 4, 5), (999, 20, 5, 1), {"contiguous"}), + ((2, 3, 4), (12, 4, 1), {"contiguous"}), + ): + with self.subTest(shape=shape, strides=strides): + self.assertEqual(memory_formats(shape, strides), formats) + + def setUp(self): + import executorch.backends.cuda.triton.kernels # noqa: F401 + + self.mode = FakeTensorMode() + + def call(self, x_shape, w_shape, output_strides, channels_last=False): + with self.mode: + x = torch.empty(x_shape, device="cuda") + w = torch.empty(w_shape, device="cuda") + if channels_last: + x = _channels_last(x) + params = ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 1) + return ConvCall(x, w, None, params, output_strides) + + def test_output_follows_the_strides(self): + for x_shape, strides in ( + ((2, 3, 4, 5), (120, 20, 5, 1)), + ((2, 3, 4, 5), (120, 1, 30, 6)), + ((2, 3, 4, 5), (20, 40, 5, 1)), + ((2, 3, 4, 5), (1, 2, 60, 12)), + ((1, 3, 1, 5), (4321, 5, 77, 1)), + ((1, 3, 1, 5), (0, 1, 0, 6)), + ): + for channels_last in (False, True): + call = self.call(x_shape, (6, 3, 3, 3), strides, channels_last) + shape = (x_shape[0], 6, *x_shape[2:]) + with self.subTest(x=x_shape, strides=strides, cl=channels_last): + self.assertEqual(out_shape_and_stride(call), (shape, strides)) + with self.mode: + out = torch.ops.triton.convolution( + call.input, + call.weight, + None, + [1, 1], + [1, 1], + [1, 1], + [1, 1], + False, + [0, 0], + 1, + list(strides), + ) + self.assertEqual((tuple(out.shape), out.stride()), (shape, strides)) + + def test_non_dense_strides_are_rejected(self): + call = self.call((2, 3, 4, 5), (6, 3, 3, 3), (121, 20, 5, 1)) + with self.assertRaisesRegex(ConvBindingError, "not a dense layout"): + out_shape_and_stride(call) + with self.mode, self.assertRaisesRegex(ConvBindingError, "not a dense layout"): + torch.ops.triton.convolution( + call.input, + call.weight, + None, + [1, 1], + [1, 1], + [1, 1], + [1, 1], + False, + [0, 0], + 1, + [121, 20, 5, 1], + ) diff --git a/backends/cuda/tests/test_conv_depthwise.py b/backends/cuda/tests/test_conv_depthwise.py index de9507d6ab8..4c03b7e35ae 100644 --- a/backends/cuda/tests/test_conv_depthwise.py +++ b/backends/cuda/tests/test_conv_depthwise.py @@ -9,6 +9,7 @@ against an FP64 reference, and models lowered through the CUDA backend.""" import glob +import itertools import os import tempfile import unittest @@ -17,11 +18,21 @@ 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" _TOLERANCE = {torch.float32: 2e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2} +def _dense_strides(shape, order): + """Dense strides of ``shape`` with the dims in ``order``, outermost first.""" + strides, running = [0] * len(shape), 1 + for d in reversed(order): + strides[d] = running + running *= shape[d] + return strides + + def _reference(x, w, b, stride, before, after): padded = F.pad(x.double(), [before[1], after[1], before[0], after[0]]) bias = None if b is None else b.double() @@ -34,7 +45,7 @@ def setUp(self): self.skipTest("CUDA required") torch.manual_seed(0) - def convolution(self, x, w, b, stride, before, after): + def convolution(self, x, w, b, stride, before, after, output_strides=None): from executorch.backends.cuda.triton.kernels.conv.conv import convolution return convolution( @@ -48,11 +59,15 @@ def convolution(self, x, w, b, stride, before, after): False, [0, 0], x.shape[1] // w.shape[1], + output_strides, ) - def check(self, x, w, b, stride, before, after, every_config=False): + def check( + self, x, w, b, stride, before, after, every_config=False, output_strides=None + ): """Runs every schedule that supports the call, with its first config or - every config, and checks the result and the layout.""" + every config, and checks the result and the layout (the fake + implementation's, and ``output_strides`` if given).""" from executorch.backends.cuda.triton.kernels.conv.conv import ( force_conv_candidate, ) @@ -72,10 +87,23 @@ def check(self, x, w, b, stride, before, after, every_config=False): (0, 0), x.shape[1] // w.shape[1], ) - call = ConvCall(x, w, b, params) + call = ConvCall( + x, w, b, params, None if output_strides is None else tuple(output_strides) + ) self.assertIsNone(direct.supports(call)) expected = _reference(x, w, b, stride, before, after).to(x.dtype) layout = out_shape_and_stride(call) + with FakeTensorMode() as mode: + fake = self.convolution( + *(None if t is None else mode.from_tensor(t) for t in (x, w, b)), + stride, + before, + after, + output_strides, + ) + self.assertEqual((tuple(fake.shape), fake.stride()), layout) + if output_strides is not None: + self.assertEqual(layout[1], tuple(output_strides)) ran = [] for schedule, all_configs in ( ("dw_flat", direct.FLAT_CONFIGS), @@ -86,11 +114,13 @@ def check(self, x, w, b, stride, before, after, every_config=False): ): if direct.schedule_supports(schedule, call) is not None: with force_conv_candidate(schedule), self.assertRaises(ValueError): - self.convolution(x, w, b, stride, before, after) + self.convolution(x, w, b, stride, before, after, output_strides) continue for config in all_configs if every_config else all_configs[:1]: with force_conv_candidate(schedule, config): - out = self.convolution(x, w, b, stride, before, after) + out = self.convolution( + x, w, b, stride, before, after, output_strides + ) self.assertEqual((tuple(out.shape), out.stride()), layout) tolerance = _TOLERANCE[x.dtype] torch.testing.assert_close( @@ -293,6 +323,31 @@ def test_unsupported_groupings(self): ) with self.subTest(groups=groups): self.assertIn(reason, direct.supports(call)) + def test_output_layouts(self): + # Every schedule writes the output strides the graph recorded, whatever + # the input layout; dims of size 1 keep their recorded strides. + for shape, kernel, stride, padding in ( + ((2, 24, 11, 13), (3, 3), (1, 1), (1, 1)), + ((1, 16, 9, 1), (5, 1), (2, 1), (2, 0)), + ): + x, w, b = self.tensors(*shape, *kernel) + out_shape = ( + shape[0], + shape[1], + (shape[2] + 2 * padding[0] - kernel[0]) // stride[0] + 1, + (shape[3] + 2 * padding[1] - kernel[1]) // stride[1] + 1, + ) + orders = [(0, 1, 2, 3), (0, 2, 3, 1), (1, 0, 2, 3), (2, 3, 1, 0)] + inputs = (x, x.to(memory_format=torch.channels_last)) + for x, order in itertools.product(inputs, orders): + strides = [ + 12345 if n == 1 else s + for n, s in zip(out_shape, _dense_strides(out_shape, order)) + ] + with self.subTest(shape=shape, input=x.stride(), strides=strides): + self.check( + x, w, b, stride, padding, padding, output_strides=strides + ) def test_sliced_inputs_with_storage_offset(self): base = torch.randn(2, 20, 15, 18, device="cuda") diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py index a727595715a..3a42bca57b8 100644 --- a/backends/cuda/tests/test_conv_implicit_gemm.py +++ b/backends/cuda/tests/test_conv_implicit_gemm.py @@ -9,6 +9,7 @@ through the CUDA backend.""" import glob +import itertools import os import tempfile import unittest @@ -18,6 +19,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" _TOLERANCE = {torch.float32: 2e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2} @@ -31,6 +33,15 @@ def _reference(x, w, b, stride, before, after): return conv(F.pad(x.double(), pad), w.double(), bias, stride) +def _dense_strides(shape, order): + """Dense strides of ``shape`` with the dims in ``order``, outermost first.""" + strides, running = [0] * len(shape), 1 + for d in reversed(order): + strides[d] = running + running *= shape[d] + return strides + + def _channels_last(t): return t.to( memory_format=torch.channels_last if t.dim() == 4 else torch.channels_last_3d @@ -44,7 +55,7 @@ def setUp(self): torch.manual_seed(0) self.tensor_cores = torch.cuda.get_device_capability() >= (8, 0) - def convolution(self, x, w, b, stride, before, after): + def convolution(self, x, w, b, stride, before, after, output_strides=None): from executorch.backends.cuda.triton.kernels.conv.conv import convolution rank = x.dim() - 2 @@ -59,12 +70,16 @@ def convolution(self, x, w, b, stride, before, after): False, [0] * rank, 1, + output_strides, ) - def check(self, x, w, b, stride, before, after, every_config=False): + def check( + self, x, w, b, stride, before, after, every_config=False, output_strides=None + ): """Runs every schedule and K order on the call, with the first config of - each or every config, and checks the result and the layout; the - tensor-core schedule runs FP32 operands in TF32.""" + each or every config, and checks the result and the layout (the fake + implementation's, and ``output_strides`` if given); the tensor-core + schedule runs FP32 operands in TF32.""" from executorch.backends.cuda.triton.kernels.conv.conv import ( force_conv_candidate, ) @@ -85,11 +100,24 @@ def check(self, x, w, b, stride, before, after, every_config=False): (0,) * rank, 1, ) - call = ConvCall(x, w, b, params) + call = ConvCall( + x, w, b, params, None if output_strides is None else tuple(output_strides) + ) tf32 = mock.patch.object(implicit_gemm, "conv_allows_tf32", return_value=True) self.assertIsNone(implicit_gemm.supports(call)) expected = _reference(x, w, b, stride, before, after) layout = out_shape_and_stride(call) + with FakeTensorMode() as mode: + fake = self.convolution( + *(None if t is None else mode.from_tensor(t) for t in (x, w, b)), + stride, + before, + after, + output_strides, + ) + self.assertEqual((tuple(fake.shape), fake.stride()), layout) + if output_strides is not None: + self.assertEqual(layout[1], tuple(output_strides)) schedules = ["igemm_ffma"] + (["igemm_tc"] if self.tensor_cores else []) for (schedule, _), configs in implicit_gemm.CONFIGS.items(): if schedule not in schedules: @@ -101,7 +129,9 @@ def check(self, x, w, b, stride, before, after, every_config=False): tolerance = _TF32_TOLERANCE for config in configs if every_config else configs[:1]: with tf32, force_conv_candidate(schedule, config): - out = self.convolution(x, w, b, stride, before, after) + out = self.convolution( + x, w, b, stride, before, after, output_strides + ) self.assertEqual((tuple(out.shape), out.stride()), layout) torch.testing.assert_close( out.double(), @@ -185,6 +215,48 @@ def test_geometries_fp16(self): def test_geometries_bf16(self): self.check_geometries(torch.bfloat16) + def test_output_layouts(self): + # Every schedule writes the output strides the graph recorded, whatever + # the input layout; dims of size 1 keep their recorded strides. + for shape, c_out, kernel, stride, padding, orders in ( + ( + (2, 3, 30, 26), + 40, + (7, 7), + (2, 2), + (3, 3), + [(0, 1, 2, 3), (0, 2, 3, 1), (1, 0, 2, 3), (2, 3, 1, 0)], + ), + ((1, 3, 56, 14), 48, (14, 14), (14, 14), (0, 0), [(0, 1, 2, 3)]), + ( + (2, 3, 2, 28, 32), + 32, + (2, 16, 16), + (2, 16, 16), + (0, 0, 0), + [(0, 1, 2, 3, 4), (0, 2, 3, 4, 1), (1, 0, 2, 3, 4)], + ), + ): + x, w, b = self.tensors(shape, c_out, kernel) + out_shape = ( + shape[0], + c_out, + *( + (e + 2 * p - k) // s + 1 + for e, k, s, p in zip(shape[2:], kernel, stride, padding) + ), + ) + inputs = (x, _channels_last(x)) + for x, order in itertools.product(inputs, orders): + strides = [ + 12345 if n == 1 else s + for n, s in zip(out_shape, _dense_strides(out_shape, order)) + ] + with self.subTest(shape=shape, input=x.stride(), strides=strides): + self.check( + x, w, b, stride, padding, padding, output_strides=strides + ) + def test_sliced_inputs_with_storage_offset(self): base = torch.randn(2, 20, 15, 18, device="cuda") _, w, b = self.tensors((1, 16, 1, 1), 24, (3, 3)) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index 345b8383bd2..97bfa42e683 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/backends/cuda/tests/test_conv_routing.py @@ -129,8 +129,11 @@ def assertRouted( *inputs, template, canonical_args=None, + output_strides=None, device="cuda", ): + """Routed, with the output strides the conv node recorded (and, if + given, ``output_strides``).""" for api, graph_module in _backend_graphs(model, *inputs, device=device).items(): with self.subTest(api=api): targets, logs = self.run_pass(graph_module) @@ -142,8 +145,11 @@ def assertRouted( for n in graph_module.graph.nodes if n.target == torch.ops.triton.convolution.default ) + self.assertEqual(node.args[-1], list(node.meta["val"].stride())) + if output_strides is not None: + self.assertEqual(node.args[-1], output_strides) if canonical_args is not None and api == _APIS[0]: - self.assertEqual(node.args[3:], canonical_args) + self.assertEqual(node.args[3:-1], canonical_args) def assertKept(self, model, *inputs, reason, **export_kwargs): for api, graph_module in _backend_graphs( @@ -205,12 +211,13 @@ def channels_last_weight(): device="cpu", ) - def test_channels_last_cuda_export_is_kept(self): + def test_channels_last_cuda_export_follows_the_graph(self): # On CUDA, with cuDNN off, the export records a contiguous output. - self.assertKept( + self.assertRouted( _depthwise(), lambda: torch.randn(1, 8, 9, 9).to(memory_format=torch.channels_last), - reason="differs from the node's", + template="permissive", + output_strides=[648, 81, 9, 1], ) def test_rocm_is_kept(self): @@ -232,16 +239,42 @@ def test_symbolic_shape_is_kept(self): 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_recorded_output_layout(self): + # (size, strides) recorded for the [2, 8, 9, 9] output, and why the + # node is kept; every dense layout is followed. + for size, strides, reason in ( + ((2, 8, 9, 9), (648, 81, 9, 1), None), + ((2, 8, 9, 9), (648, 1, 72, 8), None), + ((2, 8, 9, 9), (81, 162, 9, 1), None), + ((2, 8, 9, 9), (700, 81, 9, 1), "are not a dense layout"), + ((2, 8, 9, 9), (648, 81, 9, 2), "are not a dense layout"), + ((2, 8, 9, 8), (576, 72, 8, 1), "differs from the node's"), + ): + for graph_module in _backend_graphs( + _depthwise(), lambda: torch.randn(2, 8, 9, 9) + ).values(): + conv = next( + n + for n in graph_module.graph.nodes + if n.op == "call_function" and "conv" in str(n.target) + ) + val = conv.meta["val"] + with val.fake_mode: + conv.meta["val"] = torch.empty_strided( + size, strides, dtype=val.dtype, device=val.device + ) + with self.subTest(size=size, strides=strides): + targets, logs = self.run_pass(graph_module) + if reason is None: + (node,) = [ + n + for n in graph_module.graph.nodes + if n.target == torch.ops.triton.convolution.default + ] + self.assertEqual(node.args[-1], list(strides)) + else: + self.assertNotIn(torch.ops.triton.convolution.default, targets) + self.assertIn(reason, logs) def test_worth_routing_gates_only_optional_calls(self): from executorch.backends.cuda.triton.kernels.conv import routing @@ -447,6 +480,97 @@ def test_grouped_is_routed(self): lambda channels=channels: torch.randn(2, channels, 12, 10), template="direct", ) + def test_channels_last_cuda_export_is_routed(self): + # The export on CUDA records a contiguous output for a channels-last + # input; the converted conv writes exactly that layout. + for dtype in (torch.float32, torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + self.assertRouted( + _depthwise(channels=64, dtype=dtype), + lambda dtype=dtype: torch.randn(2, 64, 28, 28, dtype=dtype).to( + memory_format=torch.channels_last + ), + template="direct", + output_strides=[50176, 784, 28, 1], + ) + + def test_reordering_size_edge(self): + # Channels-last input, contiguous output: routed from the smallest + # measured input that wins on; the same-order layouts are not limited. + from executorch.backends.cuda.triton.kernels.conv.core import direct + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + ) + + params = ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 8) + for batch, output, routed in ( + (2, "contiguous", True), + (1, "contiguous", False), + (1, "channels_last", True), + ): + with FakeTensorMode(): + x = torch.empty(batch, 8, 28, 28, device="cuda").contiguous( + memory_format=torch.channels_last + ) + w = torch.empty(8, 1, 3, 3, device="cuda") + out = torch.empty(batch, 8, 28, 28, device="cuda") + if output == "channels_last": + out = out.contiguous(memory_format=torch.channels_last) + self.assertEqual(2 * 8 * 28 * 28, direct.MIN_ROUTED_REORDERING_ELEMENTS) + reason = direct.worth_routing( + ConvCall(x, w, None, params, tuple(out.stride())) + ) + with self.subTest(batch=batch, output=output): + if routed: + self.assertIsNone(reason) + else: + self.assertIn("contiguous output is not benchmarked", reason) + self.assertKept( + _depthwise(), + lambda: torch.randn(1, 8, 9, 9).to(memory_format=torch.channels_last), + reason="contiguous output is not benchmarked faster", + ) + + def test_only_benchmarked_layouts_are_routed(self): + from executorch.backends.cuda.triton.kernels.conv.core import direct + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + ) + + params = ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 64) + shape = (2, 64, 28, 28) + contiguous, channels_last = (50176, 784, 28, 1), (50176, 1, 1792, 64) + channel_outermost = (784, 1568, 28, 1) + with FakeTensorMode(): + x = torch.empty(shape, device="cuda") + w = torch.empty(64, 1, 3, 3, device="cuda") + x_cl = x.contiguous(memory_format=torch.channels_last) + w_cl = w.contiguous(memory_format=torch.channels_last) + sliced = x.as_strided(shape, x.stride(), 1) + padded = x.as_strided(shape, (64 * 785, 785, 28, 1)) + # (input, weight, output strides, routed); a depthwise weight matches + # both memory formats. + for x_, w_, out, routed in ( + (x, w, contiguous, True), + (x, w, channels_last, True), + (x, w_cl, channels_last, True), + (x_cl, w, channels_last, True), + (x_cl, w_cl, contiguous, True), + (x_cl, w, contiguous, True), + (x, w, channel_outermost, False), + (x_cl, w, channel_outermost, False), + (sliced, w, contiguous, False), + (padded, w, contiguous, False), + ): + call = ConvCall(x_, w_, None, params, out) + with self.subTest(input=x_.stride(), weight=w_.stride(), output=out): + reason = direct.worth_routing(call) + if routed: + self.assertIsNone(reason) + else: + self.assertIn("not a benchmarked layout", reason) def test_other_convolutions_are_kept(self): x = lambda: torch.randn(1, 8, 9, 9) # noqa: E731 @@ -705,6 +829,88 @@ def forward(self, x): ) +def _layout_variants(routed_layouts): # noqa: C901 + """(layout, whether ``routed_layouts`` admits it): functions that give a + dense NCHW call with a contiguous weight and the default output another + input, weight or output layout. Where size-1 dims make the output strides + match another format too, whether that one is admitted decides.""" + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + memory_formats, + out_shape_and_stride, + ) + + def output_in(*formats): + def routed(call): + return bool(memory_formats(*out_shape_and_stride(call)) & set(formats)) + + return routed + + def channels_last(t): + fmt = torch.channels_last if t.dim() == 4 else torch.channels_last_3d + return t.contiguous(memory_format=fmt) + + def output(call, order): + shape, _ = out_shape_and_stride(call) + strides, running = [0] * len(shape), 1 + for d in reversed(order(len(shape))): + strides[d] = running + running *= shape[d] + return dataclasses.replace(call, output_strides=tuple(strides)) + + def nchw(call): + return call + + def channels_last_input(call): + return dataclasses.replace(call, input=channels_last(call.input)) + + def channels_last_input_contiguous_output(call): + return output(channels_last_input(call), lambda n: range(n)) + + def contiguous_input_channels_last_output(call): + return output(call, lambda n: (0, *range(2, n), 1)) + + def channels_last_input_channel_outermost_output(call): + return output(channels_last_input(call), lambda n: (1, 0, *range(2, n))) + + def sliced_input(call): + x = call.input + return dataclasses.replace(call, input=x.as_strided(x.shape, x.stride(), 1)) + + def padded_input(call): + x = call.input + n, c, *spatial = x.shape + plane = math.prod(spatial) + 1 + strides = (c * plane, plane, *x.stride()[2:]) + return dataclasses.replace(call, input=x.as_strided(x.shape, strides)) + + def channels_last_weight(call): + return dataclasses.replace(call, weight=channels_last(call.weight)) + + def sliced_weight(call): + w = call.weight + return dataclasses.replace(call, weight=w.as_strided(w.shape, w.stride(), 1)) + + contiguous, cl = "contiguous", "channels_last" + + def admitted(input): + return [o for i, w, o in routed_layouts if (i, w) == (input, contiguous)] + + return ( + (nchw, True), + (channels_last_input, (cl, contiguous, cl) in routed_layouts), + ( + channels_last_input_contiguous_output, + (cl, contiguous, contiguous) in routed_layouts, + ), + (contiguous_input_channels_last_output, output_in(*admitted(contiguous))), + (channels_last_input_channel_outermost_output, output_in(*admitted(cl))), + (sliced_input, False), + (padded_input, False), + (channels_last_weight, False), + (sliced_weight, False), + ) + + 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 @@ -790,6 +996,23 @@ def test_benchmarked_rows(self): reason=reason, ) + def test_channels_last_cuda_export_is_routed(self): + # The export on CUDA records a contiguous output for a channels-last + # input; the converted conv writes exactly that layout. + for (row, model, shape), strides in ( + (_ROUTED_ROWS[6], [4915200, 76800, 320, 1]), + (_ROUTED_ROWS[7], [48672, 1521, 39, 1]), + ): + with self.subTest(row=row): + self.assertRouted( + model(torch.float32), + lambda shape=shape: torch.randn(shape).to( + memory_format=torch.channels_last + ), + template="implicit_gemm", + output_strides=strides, + ) + def test_canonical_arguments(self): # padding="valid" binds to zero padding on both sides. self.assertRouted( @@ -817,19 +1040,18 @@ def test_precision_without_tensor_cores_is_kept(self): template="implicit_gemm", ) - def worth_routing(self, shape, out_channels, kernel, stride, before, after, **kw): + def worth_routing(self, *call, dtype, tf32, layout=None): from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + with self.tf32(tf32): + return implicit_gemm.worth_routing(self.call(*call, dtype, layout)) + + def call(self, shape, out_channels, kernel, stride, before, after, dtype, layout): from executorch.backends.cuda.triton.kernels.conv.core.common import ( ConvCall, ConvParams, ) - with FakeTensorMode(): - x = torch.empty(shape, dtype=kw["dtype"], device="cuda") - w = torch.empty( - out_channels, shape[1], *kernel, dtype=kw["dtype"], device="cuda" - ) - x, w = kw.get("layout", lambda x, w: (x, w))(x, w) rank = len(kernel) params = ConvParams( tuple(stride), @@ -840,8 +1062,11 @@ def worth_routing(self, shape, out_channels, kernel, stride, before, after, **kw (0,) * rank, 1, ) - with self.tf32(kw["tf32"]): - return implicit_gemm.worth_routing(ConvCall(x, w, None, params)) + with FakeTensorMode(): + x = torch.empty(shape, dtype=dtype, device="cuda") + w = torch.empty(out_channels, shape[1], *kernel, dtype=dtype, device="cuda") + call = ConvCall(x, w, None, params) + return call if layout is None else layout(call) def precisions(self): """(dtype, TF32 allowed) of the FMA and tensor-core cases.""" @@ -911,41 +1136,15 @@ def test_one_step_past_the_benchmarked_calls_is_kept(self): ) def test_only_benchmarked_layouts_are_routed(self): - """The benchmark measures dense NCHW and channels-last inputs with - contiguous weights; sliced or padded inputs and channels-last or sliced - weights keep the regular lowering.""" + """The (input, weight, output) memory formats the benchmark measures are + routed; sliced or padded inputs, channels-last or sliced weights and + other output layouts keep the regular lowering.""" from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm - def channels_last(x, w): - fmt = torch.channels_last if x.dim() == 4 else torch.channels_last_3d - return x.contiguous(memory_format=fmt), w - - def sliced_input(x, w): - return x.as_strided(x.shape, x.stride(), 1), w - - def padded_input(x, w): - n, c, *spatial = x.shape - plane = math.prod(spatial) + 1 - strides = (c * plane, plane, *x.stride()[2:]) - return x.as_strided(x.shape, strides), w - - def channels_last_weight(x, w): - fmt = torch.channels_last if w.dim() == 4 else torch.channels_last_3d - return x, w.contiguous(memory_format=fmt) - - def sliced_weight(x, w): - return x, w.as_strided(w.shape, w.stride(), 1) - for dtype, tf32 in self.precisions(): for call in implicit_gemm.boundary_calls(): shape, c_out, kernel, stride, padding = call - for layout, routed in ( - (channels_last, True), - (sliced_input, False), - (padded_input, False), - (channels_last_weight, False), - (sliced_weight, False), - ): + for layout, routed in _layout_variants(implicit_gemm.ROUTED_LAYOUTS): with self.subTest( call=call, layout=layout.__name__, dtype=dtype, tf32=tf32 ): @@ -960,6 +1159,10 @@ def sliced_weight(x, w): tf32=tf32, layout=layout, ) + if callable(routed): + routed = routed( + self.call(*call[:4], padding, padding, dtype, layout) + ) if routed: self.assertIsNone(reason) else: diff --git a/backends/cuda/triton/kernels/conv/conv.py b/backends/cuda/triton/kernels/conv/conv.py index 66709311cd0..6ad299cd714 100644 --- a/backends/cuda/triton/kernels/conv/conv.py +++ b/backends/cuda/triton/kernels/conv/conv.py @@ -9,6 +9,8 @@ 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. +``output_strides`` fixes the output's strides (a dense layout); the router +passes the strides the source node recorded, so the output follows the graph. """ from __future__ import annotations @@ -104,6 +106,7 @@ def _canonical_call( transposed: bool, output_padding: list[int], groups: int, + output_strides: Optional[list[int]], ) -> ConvCall: params = ConvParams( tuple(stride), @@ -120,7 +123,13 @@ def _canonical_call( f"{params.rank}-D convolution of a {input.dim()}-D input with a " f"{weight.dim()}-D weight" ) - return ConvCall(input, weight, bias, params) + return ConvCall( + input, + weight, + bias, + params, + None if output_strides is None else tuple(output_strides), + ) @triton_op("triton::convolution", mutates_args={}) @@ -135,6 +144,7 @@ def convolution( transposed: bool, output_padding: list[int], groups: int, + output_strides: Optional[list[int]] = None, ) -> torch.Tensor: call = _canonical_call( input, @@ -147,6 +157,7 @@ def convolution( transposed, output_padding, groups, + output_strides, ) params = call.params template = select_template(call) @@ -176,6 +187,7 @@ def _convolution_fake( transposed: bool, output_padding: list[int], groups: int, + output_strides: Optional[list[int]] = None, ) -> torch.Tensor: call = _canonical_call( input, @@ -188,6 +200,7 @@ def _convolution_fake( transposed, output_padding, groups, + output_strides, ) 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/common.py b/backends/cuda/triton/kernels/conv/core/common.py index 909aa82d687..8567071ad1c 100644 --- a/backends/cuda/triton/kernels/conv/core/common.py +++ b/backends/cuda/triton/kernels/conv/core/common.py @@ -25,6 +25,7 @@ from __future__ import annotations +import itertools from dataclasses import dataclass from typing import Any, Callable, NamedTuple, Optional, Sequence @@ -91,6 +92,9 @@ class ConvCall: weight: torch.Tensor bias: Optional[torch.Tensor] params: ConvParams + # The output strides the caller requires (a dense layout), or None for the + # layout ATen's CUDA convolution returns. + output_strides: Optional[tuple[int, ...]] = None def signature(self) -> str: def describe(t: Optional[torch.Tensor]) -> str: @@ -106,7 +110,13 @@ def describe(t: Optional[torch.Tensor]) -> str: 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})" + f"groups={p.groups}" + + ( + "" + if self.output_strides is None + else f", output_strides={list(self.output_strides)}" + ) + + ")" ) @@ -306,8 +316,61 @@ def _unsqueeze2( return (*shape[:2], 1, shape[2]), (*strides[:2], strides[2] * shape[2], strides[2]) +def dense_strides(shape: Sequence[int], strides: Sequence[int]) -> bool: + """Whether ``strides`` lay ``shape`` out densely: each element once, with no + gaps, in some order of the dimensions. Dimensions of size 1 may have any + non-negative stride.""" + if len(strides) != len(shape) or any(s < 0 for s in strides): + return False + expected = 1 + for size, stride in sorted( + ((n, s) for n, s in zip(shape, strides) if n != 1), key=lambda d: d[1] + ): + if stride != expected: + return False + expected *= size + return True + + +def memory_formats(shape: Sequence[int], strides: Sequence[int]) -> frozenset[str]: + """The dense layouts ("contiguous", "channels_last") ``strides`` match; both + for strides that are ambiguous because of dimensions of size 1.""" + formats = set() + if _is_contiguous(shape, strides): + formats.add("contiguous") + if len(shape) in (4, 5): + expected, running = [0] * len(shape), 1 + for d in _channels_last_order(len(shape)): + expected[d] = running + running *= shape[d] + if all(n == 1 or s == e for n, s, e in zip(shape, strides, expected)): + formats.add("channels_last") + return frozenset(formats) + + +def unbenchmarked_layout( + call: ConvCall, routed: frozenset[tuple[str, str, str]] +) -> Optional[str]: + """Why the (input, weight, output) memory formats of ``call`` are not one of + ``routed``, or None. A sliced input or weight matches no format.""" + x, w = call.input, call.weight + shape, out = out_shape_and_stride(call) + formats = [ + memory_formats(t.shape, t.stride()) if t.storage_offset() == 0 else () + for t in (x, w) + ] + if routed.intersection(itertools.product(*formats, memory_formats(shape, out))): + return None + return ( + f"input strides {list(x.stride())} (offset {x.storage_offset()}), weight " + f"strides {list(w.stride())} (offset {w.storage_offset()}) and output " + f"strides {list(out)} are not a benchmarked layout" + ) + + 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.""" + """The size and stride of the output: ``call.output_strides`` if given (which + must be dense), otherwise the strides ATen's CUDA convolution returns.""" x, w, p = call.input, call.weight, call.params kernel = w.shape[2:] if p.transposed: @@ -338,6 +401,13 @@ def out_shape_and_stride(call: ConvCall) -> tuple[tuple[int, ...], tuple[int, .. ) ] shape = (x.shape[0], channels, *spatial) + if call.output_strides is not None: + if not dense_strides(shape, call.output_strides): + raise ConvBindingError( + f"output strides {list(call.output_strides)} are not a dense " + f"layout of {list(shape)}" + ) + return shape, tuple(call.output_strides) if not _channels_last_output(call): return shape, _contiguous_strides(shape) if len(shape) == 3: diff --git a/backends/cuda/triton/kernels/conv/core/direct.py b/backends/cuda/triton/kernels/conv/core/direct.py index d74023937d7..b9985e793d9 100644 --- a/backends/cuda/triton/kernels/conv/core/direct.py +++ b/backends/cuda/triton/kernels/conv/core/direct.py @@ -54,6 +54,7 @@ max_offset, Schedule, SUPPORTED_DTYPES, + unbenchmarked_layout, ) from executorch.backends.cuda.triton.routing import device_supports from torch.library import wrap_triton @@ -151,6 +152,24 @@ def supports(call: ConvCall) -> Optional[str]: MAX_ROUTED_GROUPED_STRIDE = 2 MIN_TEMPLATE_GROUP_CHANNELS = 4 MIN_TEMPLATE_HALF_REDUCTION = 4 * 7 * 7 +# The (input, weight, output) memory formats measured. A depthwise weight has +# one input channel, so it matches both formats. +SAME_ORDER_LAYOUTS = frozenset( + { + ("contiguous", "contiguous", "contiguous"), + ("contiguous", "channels_last", "channels_last"), + ("channels_last", "contiguous", "channels_last"), + ("channels_last", "channels_last", "channels_last"), + } +) +ROUTED_LAYOUTS = SAME_ORDER_LAYOUTS | { + ("channels_last", "contiguous", "contiguous"), + ("channels_last", "channels_last", "contiguous"), +} +# A channels-last input written to a contiguous output (what an export on CUDA +# records) loses to Inductor on the smallest, launch-bound inputs; it is +# measured faster from this many input elements on. +MIN_ROUTED_REORDERING_ELEMENTS = 2 * 8 * 28 * 28 def worth_routing(call: ConvCall) -> Optional[str]: @@ -173,6 +192,17 @@ def worth_routing(call: ConvCall) -> Optional[str]: return f"direct stencil: spatial size {list(x.shape[2:])} is not benchmarked" if x.numel() > MAX_ROUTED_INPUT_ELEMENTS: return f"direct stencil: a {x.numel()}-element input is not benchmarked" + reason = unbenchmarked_layout(call, ROUTED_LAYOUTS) + if reason is not None: + return f"direct stencil: {reason}" + if ( + x.numel() < MIN_ROUTED_REORDERING_ELEMENTS + and unbenchmarked_layout(call, SAME_ORDER_LAYOUTS) is not None + ): + return ( + f"direct stencil: a {x.numel()}-element channels-last input with a " + "contiguous output is not benchmarked faster" + ) if w.shape[1] == 1 and w.shape[0] == p.groups: return None group_channels, multiplier = w.shape[1], w.shape[0] // p.groups diff --git a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py index a2c9b7a2a20..659e1c0e6fb 100644 --- a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py +++ b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py @@ -54,6 +54,7 @@ max_offset, Schedule, SUPPORTED_DTYPES, + unbenchmarked_layout, ) from executorch.backends.cuda.triton.routing import device_supports from torch.library import wrap_triton @@ -127,15 +128,22 @@ class RoutedFamily: # measured call significantly faster than Inductor's Triton conv template: # 3-channel inputs, which the flat K order walks without the empty # lanes of the template's channel blocks, in the geometry families below, with -# a dense NCHW or channels-last input and a contiguous weight (the layouts the -# benchmark measures). Its ``boundary`` suite measures ``boundary_calls()``: -# each geometry at every corner of these limits. +# the (input, weight, output) memory formats in ROUTED_LAYOUTS. Its +# ``boundary`` suite measures ``boundary_calls()``: each geometry at every +# corner of these limits. ROUTED_CHANNELS = 3 MIN_ROUTED_OUT_CHANNELS = 32 MAX_ROUTED_OUT_CHANNELS = 1280 MIN_ROUTED_OUTPUT_PIXELS = 256 MAX_ROUTED_INPUT_PIXELS = 1024 * 1024 MAX_ROUTED_OUTPUT_ELEMENTS = 1 << 24 +ROUTED_LAYOUTS = frozenset( + { + ("contiguous", "contiguous", "contiguous"), + ("channels_last", "contiguous", "channels_last"), + ("channels_last", "contiguous", "contiguous"), + } +) ROUTED_FAMILIES = ( RoutedFamily( "window", @@ -164,18 +172,6 @@ class RoutedFamily: ) -def _benchmarked_layout(x: torch.Tensor, w: torch.Tensor) -> bool: - """A dense, unsliced NCHW or channels-last input and a contiguous weight.""" - channels_last = torch.channels_last if x.dim() == 4 else torch.channels_last_3d - dense_input = x.is_contiguous() or x.is_contiguous(memory_format=channels_last) - return ( - dense_input - and x.storage_offset() == 0 - and w.is_contiguous() - and w.storage_offset() == 0 - ) - - def _family(call: ConvCall) -> Optional[RoutedFamily]: p = call.params if p.padding_before != p.padding_after: @@ -197,12 +193,9 @@ def worth_routing(call: ConvCall) -> Optional[str]: ) if x.shape[1] != ROUTED_CHANNELS: return f"implicit GEMM: {x.shape[1]} input channels are not benchmarked faster" - if not _benchmarked_layout(x, w): - return ( - f"implicit GEMM: input strides {list(x.stride())} (offset " - f"{x.storage_offset()}) or weight strides {list(w.stride())} are not " - "a benchmarked layout" - ) + reason = unbenchmarked_layout(call, ROUTED_LAYOUTS) + if reason is not None: + return f"implicit GEMM: {reason}" family = _family(call) if family is None: return ( diff --git a/backends/cuda/triton/kernels/conv/routing.py b/backends/cuda/triton/kernels/conv/routing.py index f70e39c3845..45227e3981b 100644 --- a/backends/cuda/triton/kernels/conv/routing.py +++ b/backends/cuda/triton/kernels/conv/routing.py @@ -7,9 +7,10 @@ """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 +(see ``triton/routing.py``): it binds the node's arguments, checks the +templates' ``supports()``, and classifies the node without touching the graph. +The replacement produces exactly the output strides the node recorded, which +must be a dense layout. ``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. @@ -21,6 +22,7 @@ from __future__ import annotations +import dataclasses import logging from typing import Any, Optional @@ -30,6 +32,7 @@ bind_conv, ConvCall, ConvParams, + dense_strides, out_shape_and_stride, ) from executorch.backends.cuda.triton.routing import ( @@ -86,12 +89,17 @@ def _check(node: Node, call: ConvCall, gap: Optional[str]) -> ReplacementPlan: 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())): + if not dense_strides(expected.shape, expected.stride()): return _unsupported( - f"output {list(size)}@{list(stride)} differs from the node's " - f"{list(expected.shape)}@{list(expected.stride())}" + f"the node's output strides {list(expected.stride())} are not a " + "dense layout" ) + size, _ = out_shape_and_stride(call) + if size != tuple(expected.shape): + return _unsupported( + f"output size {list(size)} differs from the node's {list(expected.shape)}" + ) + call = dataclasses.replace(call, output_strides=tuple(expected.stride())) reasons = [] for template in TEMPLATES: reason = template.supports(call) @@ -135,13 +143,15 @@ def value(arg: Any) -> Any: result = _unsupported(f"{result.reason}; regular lowering: {gap}") return result route = Route.OPTIONAL if gap is None else Route.REQUIRED + # The output follows the strides the node recorded. + call = dataclasses.replace(call, output_strides=tuple(node.meta["val"].stride())) return ReplacementPlan( route, gap or "", "conv", f"{result.label} template", torch.ops.triton.convolution.default, - (input, weight, bias, *_canonical_args(params)), + (input, weight, bias, *_canonical_args(params), list(call.output_strides)), call.signature(), )