[executorch][cuda] Converted convolutions follow the output layout the graph recorded

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
@@ -348,7 +367,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),
@@ -357,6 +380,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:
@@ -372,10 +396,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):
@@ -412,11 +439,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)),
@@ -460,7 +491,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"]
@@ -478,6 +509,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("/"))

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)

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
@@ -396,3 +398,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],
+            )