Decompose antialiased resizes in the CUDA backend

F.interpolate(x, mode="bilinear" or "bicubic", antialias=True) becomes
aten._upsample_bilinear2d_aa or aten._upsample_bicubic2d_aa. Inductor
has no lowering for either, so it calls ATen's kernel, and AOTInductor
has no C shim for them. Lowering to the CUDA backend failed with
"missing fallback kernels". Image preprocessing often does this.

The backend already rewrites ops like this in get_decomposition_table
(conv1d as conv2d, for example), so this adds the bilinear and bicubic
antialiased resizes there, for float32, float16 and bfloat16 inputs.
The backend has no float64, and ATen's CUDA kernel has no integer
version.

An antialiased resize is separable, so the rewrite resizes the width,
then the height. Along each, it computes every output's window and
weights with ATen's formulas: the filter widened by the downscale
factor, cut to its support, and normalized. Then it sums the window's
inputs with those weights. Taps outside the window have zero weight and
are masked out before the sum, so a NaN or inf reaches only the outputs
whose window holds it, and the work grows with the window size, not the
input size. Where a window edge falls exactly on an input, float
rounding can put that input in or out of the window differently from
ATen; its weight there is zero or close to it.

Symbolic shape math keeps dynamic input and output sizes working. The
number of taps per output is a shape, so it is computed with integers:
in float32 the generated kernel can round a ratio such as 21 / 7 just
above 3 and count more taps than the host allocated. For align_corners,
the ratio for an output size of 1 comes from min and max, not a
conditional, because torch.compile cannot handle a conditional size
expression. The result keeps the input's memory format. The .vec
overloads reuse torch's decomposition.

The rewrite computes positions and weights in float32 with its own
rounding, so its float32 results are not bit-identical to ATen's
kernel: for inputs in [0, 1) they stayed within 2e-5 of it on the sizes
tried. Half and bfloat16 inputs are computed in float and rounded once
at the end, so they can differ slightly from ATen, which keeps its
weights in the input type.

The tests are in a new file, test_antialiased_resize.py, because they
run programs through the Python runtime. They lower bilinear and
bicubic resizes whose input and output heights are dynamic, run them
for heights that shrink, grow, go to 1, and have whole ratios, with a
NaN in the input, and compare with eager CUDA. Other cases cover an
uneven scale_factor, float16 and bfloat16 with a channels_last input,
the rewrite's memory format, and torch.compile of an align_corners
resize with dynamic sizes.

Author: shoumikhin

backends/cuda/cuda_backend.py

@@ -9,6 +9,7 @@
 import copy
 import functools
 import logging
+import math
 import os
 import shutil
 import threading
@@ -50,7 +51,9 @@
 from executorch.exir._warnings import experimental
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
+from torch._decomp import decomposition_table
 from torch._inductor.decomposition import conv1d_to_conv2d
+from torch._prims_common import suggest_memory_format
 from torch.nn.attention import SDPBackend
 
 # ---------------------------------------------------------------------------
@@ -548,6 +551,117 @@ def _conv_transpose_as_conv(
     return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups)
 
 
+def _linear_kernel(x: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(1 - x.abs(), min=0)
+
+
+def _cubic_kernel(x: torch.Tensor) -> torch.Tensor:
+    # Keys' cubic with a = -0.5, the one ATen's antialiased bicubic uses.
+    a = -0.5
+    x = x.abs()
+    near = ((a + 2) * x - (a + 3)) * x * x + 1
+    far = (((x - 5) * x + 8) * x - 4) * a
+    return torch.where(x < 1, near, torch.where(x < 2, far, torch.zeros_like(x)))
+
+
+def _antialias_taps(
+    in_size: int,
+    out_size: int,
+    align_corners: bool,
+    scale: Optional[float],
+    support: int,
+    kernel: typing.Callable[[torch.Tensor], torch.Tensor],
+    device: torch.device,
+) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Taps of an antialiased resize along one dim, computed as ATen does:
+    output o reads input index[o, k] with weight[o, k], and inside[o, k]
+    marks the taps in its window. The window is the kernel's support widened
+    by the downscale factor, and the weights in it are normalized.
+    """
+    if align_corners or scale is None or scale <= 0:
+        if align_corners:
+            # ATen uses a ratio of 0 for an output size of 1. min and max give
+            # that without a conditional, which torch.compile cannot handle for
+            # dynamic sizes.
+            num = (in_size - 1) * torch.sym_min(out_size - 1, 1)
+            den = torch.sym_max(out_size - 1, 1)
+        else:
+            num, den = in_size, out_size
+        ratio = num / den
+        # The tap count is a shape that both the host and the kernel work out.
+        # In float32 the kernel can round num / den just past a whole number
+        # and count more taps than the host allocated, so use integers.
+        reach = torch.sym_max((support * num + den - 1) // den, support)
+    else:
+        ratio = 1.0 / scale
+        reach = math.ceil(support * max(ratio, 1.0))
+    widen = torch.sym_max(ratio, 1.0)
+    dtype = torch.float32
+    center = ratio * (torch.arange(out_size, dtype=dtype, device=device) + 0.5)
+    first = torch.clamp(torch.trunc(center - support * widen + 0.5), min=0)
+    end = torch.clamp(torch.trunc(center + support * widen + 0.5), max=in_size)
+    taps = 2 * reach + 1
+    tap = torch.arange(taps, dtype=dtype, device=device)
+    position = first[:, None] + tap
+    inside = position < end[:, None]
+    weights = kernel((tap + (first - center)[:, None] + 0.5) * (1.0 / widen))
+    weights = torch.where(inside, weights, 0)
+    total = weights.sum(1, keepdim=True)
+    weights = torch.where(total != 0, weights / total, weights)
+    # Clamp after the cast: in float32, in_size - 1 rounds up to in_size for
+    # very long axes.
+    index = position.to(torch.int64).clamp(max=in_size - 1)
+    return index, weights, inside
+
+
+def _antialiased_resize(
+    support: int,
+    kernel: typing.Callable[[torch.Tensor], torch.Tensor],
+    input: torch.Tensor,
+    output_size: List[int],
+    align_corners: bool,
+    scales_h: Optional[float] = None,
+    scales_w: Optional[float] = None,
+) -> torch.Tensor:
+    """
+    Rewrite an antialiased bilinear or bicubic resize as a weighted sum over
+    each output's taps, one dim at a time. Inductor has no lowering for these
+    resizes, so it would call ATen's kernel, which this backend does not ship.
+    """
+    # The backend's tensors have no float64, and ATen's CUDA kernel has no
+    # integer version.
+    if input.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+        return NotImplemented
+    # Half types are computed in float and rounded once, so they can differ
+    # slightly from ATen, which keeps its weights in the input type.
+    result = input.float()
+    for dim, scale in ((3, scales_w), (2, scales_h)):
+        index, weights, inside = _antialias_taps(
+            input.shape[dim],
+            output_size[dim - 2],
+            align_corners,
+            scale,
+            support,
+            kernel,
+            input.device,
+        )
+        rows, taps = index.shape
+        trailing = [1] * (3 - dim)
+        gathered = result.index_select(dim, index.flatten()).unflatten(
+            dim, (rows, taps)
+        )
+        # Taps outside the window already have zero weight, but NaN * 0 is NaN,
+        # so mask them out to keep a NaN or inf local, as in ATen.
+        weighted = torch.where(
+            inside.view(rows, taps, *trailing),
+            gathered * weights.view(rows, taps, *trailing),
+            0,
+        )
+        result = weighted.sum(dim + 1)
+    return result.to(input.dtype).contiguous(memory_format=suggest_memory_format(input))
+
+
 @final
 @experimental(
     "This API and all of cuda backend related functionality are experimental."
@@ -862,6 +976,15 @@ def get_decomposition_table(cls) -> Dict[Any, Any]:
             # for it and the shim no fallback kernel.
             torch.ops.aten._is_all_true.default: torch.ops.aten.all.default,
         }
+        for resize, support, kernel in (
+            (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel),
+            (torch.ops.aten._upsample_bicubic2d_aa, 2, _cubic_kernel),
+        ):
+            table[resize.default] = functools.partial(
+                _antialiased_resize, support, kernel
+            )
+            # The vec form works out the output size and calls the default form.
+            table[resize.vec] = decomposition_table[resize.vec]
         # Inductor disables its transposed conv2d Triton template only on CUDA
         # (https://github.com/pytorch/pytorch/issues/187081).
         if torch.version.hip is None:

backends/cuda/tests/test_antialiased_resize.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.
+
+"""Antialiased bilinear and bicubic resizes on the CUDA backend.
+
+Inductor has no lowering for these resizes, so the backend rewrites them into
+plain tensor math. These tests lower real programs, run them through the Python
+runtime, and compare with eager PyTorch.
+"""
+
+import functools
+import os
+import tempfile
+import unittest
+from typing import Optional, Sequence, Tuple
+
+import torch
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+from torch.export import Dim, export
+
+
+class Resize(torch.nn.Module):
+    """Resizes x by scale_factor, or else to like's height and a width of 56."""
+
+    def __init__(
+        self,
+        mode: str,
+        align_corners: bool = False,
+        scale_factor: Optional[Tuple[float, float]] = None,
+    ) -> None:
+        super().__init__()
+        self.mode = mode
+        self.align_corners = align_corners
+        self.scale_factor = scale_factor
+
+    def forward(self, x: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
+        size = None if self.scale_factor else (like.shape[2], 56)
+        return torch.nn.functional.interpolate(
+            x,
+            size=size,
+            scale_factor=self.scale_factor,
+            mode=self.mode,
+            align_corners=self.align_corners,
+            antialias=True,
+        )
+
+
+# Dynamic heights let one program shrink, grow, and resize to 1.
+DYNAMIC_HEIGHTS = (
+    {2: Dim("in_height", min=4, max=64)},
+    {2: Dim("out_height", min=1, max=32)},
+)
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+@unittest.skipIf(
+    torch.version.hip is not None,
+    "loads through the Python runtime, which the ROCm build does not include",
+)
+class TestAntialiasedResize(unittest.TestCase):
+    def _run(
+        self,
+        model: Resize,
+        cases: Sequence[Tuple[torch.Tensor, torch.Tensor]],
+        dynamic_shapes=None,
+        **tolerance: float,
+    ) -> None:
+        """Lower model for cases[0], then run every case and compare with eager."""
+        exported = export(model, cases[0], dynamic_shapes=dynamic_shapes)
+        spec = CudaBackend.generate_method_name_compile_spec("forward")
+        lowered = to_edge_transform_and_lower(
+            exported,
+            partitioner=[CudaPartitioner([spec])],
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+        from executorch.runtime import Runtime
+
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.pte")
+            with open(path, "wb") as f:
+                lowered.to_executorch().write_to_file(f)
+            method = Runtime.get().load_program(path).load_method("forward")
+            for x, like in cases:
+                (output,) = method.execute([x, like])
+                # The rewrite stands in for ATen's CUDA kernel.
+                expected = model(x.cuda().float(), like.cuda()).to(x.dtype).cpu()
+                torch.testing.assert_close(
+                    output.cpu(), expected, equal_nan=True, **tolerance
+                )
+
+    def _check_dynamic_heights(self, mode: str, align_corners: bool) -> None:
+        cases = []
+        # 21 to 7 and 22 to 8 give whole ratios, with and without align_corners.
+        heights = ((48, 20), (64, 20), (10, 20), (10, 1), (21, 7), (22, 8))
+        for in_height, out_height in heights:
+            x = torch.randn(1, 3, in_height, 32)
+            # As in ATen, a NaN reaches only the outputs whose window has it.
+            x[0, 0, -1, -1] = float("nan")
+            cases.append((x, torch.empty(1, 1, out_height)))
+        self._run(
+            Resize(mode, align_corners),
+            cases,
+            DYNAMIC_HEIGHTS,
+            atol=1e-5,
+            rtol=1e-4,
+        )
+
+    def test_bilinear_antialias_resize(self):
+        self._check_dynamic_heights("bilinear", align_corners=False)
+
+    def test_bicubic_antialias_resize(self):
+        self._check_dynamic_heights("bicubic", align_corners=True)
+
+    def test_antialias_resize_with_uneven_scale_factor(self):
+        # 37 * 1.75 and 51 * 0.45 are not whole, so the scale the resize uses
+        # differs from output size / input size.
+        x = torch.randn(1, 3, 37, 51)
+        self._run(
+            Resize("bicubic", scale_factor=(1.75, 0.45)),
+            [(x, torch.empty(0))],
+            atol=1e-5,
+            rtol=1e-4,
+        )
+
+    def test_half_precision_antialias_resize(self):
+        for dtype in (torch.float16, torch.bfloat16):
+            with self.subTest(dtype=dtype):
+                x = torch.randn(1, 3, 48, 32).to(
+                    dtype=dtype, memory_format=torch.channels_last
+                )
+                self._run(
+                    Resize("bilinear"),
+                    [(x, torch.empty(1, 1, 20))],
+                    atol=1e-2,
+                    rtol=1e-2,
+                )
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestAntialiasedResizeDecomposition(unittest.TestCase):
+    def test_keeps_memory_format(self):
+        # The program's output layout is fixed at export, so only a direct call
+        # shows the rewrite's own layout.
+        resize = CudaBackend.get_decomposition_table()[
+            torch.ops.aten._upsample_bilinear2d_aa.default
+        ]
+        x = torch.randn(1, 3, 48, 32, device="cuda")
+        output = resize(x.to(memory_format=torch.channels_last), [20, 56], False)
+        self.assertTrue(output.is_contiguous(memory_format=torch.channels_last))
+
+    def test_align_corners_compiles_with_dynamic_sizes(self):
+        from torch._inductor.compile_fx import compile_fx
+
+        model = Resize("bicubic", align_corners=True)
+        compiled = torch.compile(
+            model,
+            backend=functools.partial(
+                compile_fx, decompositions=CudaBackend.get_decomposition_table()
+            ),
+            dynamic=True,
+            fullgraph=True,
+        )
+        for in_height, out_height in ((48, 20), (10, 1)):
+            x = torch.randn(1, 3, in_height, 32, device="cuda")
+            like = torch.empty(1, 1, out_height, device="cuda")
+            torch.testing.assert_close(
+                compiled(x, like), model(x, like), atol=1e-5, rtol=1e-4
+            )