CUDA backend: per-program weights file names and antialiased resize

backends/aoti/aoti_backend.py

@@ -28,6 +28,29 @@
 
 class COMPILE_SPEC_KEYS(Enum):
     METHOD_NAME = "method_name"
+    DATA_NAME = "data_name"
+
+
+# The name becomes a file name on every platform, so keep it short, to portable
+# characters, and clear of Windows device names such as NUL. Without re.ASCII,
+# IGNORECASE lets [A-Za-z] match letters like the Kelvin sign. A leading '-'
+# reads as an option on a command line. Some writers add .ptd only when it is
+# missing, so "encoder" and "encoder.ptd" could both write encoder.ptd.
+_DATA_NAME_PATTERN = re.compile(
+    r"(?!(?:CON|PRN|AUX|NUL|COM\d|LPT\d)(?:\.|$))(?!.*\.ptd$)"
+    r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,199}",
+    re.IGNORECASE | re.ASCII,
+)
+
+
+def _checked_data_name(data_name: str) -> str:
+    if not _DATA_NAME_PATTERN.fullmatch(data_name):
+        raise ValueError(
+            f"Invalid data name {data_name!r}: use up to 200 ASCII letters, "
+            "digits, '_', '-' and '.', do not start with '.' or '-', do not end "
+            "with '.ptd', and avoid Windows device names such as NUL"
+        )
+    return data_name
 
 
 @experimental(
@@ -368,7 +391,9 @@ def preprocess(
         # Determine whether to save named data externally based on backend setting
         # External: save to separate .ptd file, otherwise merge with .pte file
         external_tag = (
-            f"aoti_{device_name}_blob" if cls.save_data_externally() else None
+            cls.data_name_from_compile_specs(compile_specs)
+            if cls.save_data_externally()
+            else None
         )
 
         named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag)
@@ -404,6 +429,33 @@ def generate_method_name_compile_spec(
             method_name.encode("utf-8"),
         )
 
+    @classmethod
+    def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec:
+        """
+        Generate a CompileSpec that names the external data file, written as
+        ``<data_name>.ptd``, so several programs can share a folder.
+        """
+        if not cls.save_data_externally():
+            raise ValueError(
+                f"{cls.__name__} keeps its data in the .pte file, so it takes no "
+                "data name"
+            )
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.DATA_NAME.value,
+            _checked_data_name(data_name).encode("utf-8"),
+        )
+
+    @classmethod
+    def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str:
+        """
+        The external data file name, without ``.ptd``, set by
+        generate_data_name_compile_spec, or ``aoti_<device>_blob`` by default.
+        """
+        for spec in compile_specs:
+            if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value:
+                return _checked_data_name(spec.value.decode("utf-8"))
+        return f"aoti_{cls.get_device_name()}_blob"
+
     @classmethod
     def method_name_from_compile_specs(
         cls,

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)
+        # Plain indexing lets Inductor fold the gather into the sum instead of
+        # keeping every tap of every row in memory.
+        gathered = result[(slice(None),) * dim + (index,)]
+        # Taps outside the window already have zero weight, but NaN * 0 is NaN,
+        # so mask them out to keep a NaN or inf local.
+        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."
@@ -704,6 +818,7 @@ def _preprocess_with_weight_collector(
             cls.get_device_name(),
             target_sm=target_sm,
             ptx_compute=ptx_compute,
+            external_tag=cls.data_name_from_compile_specs(compile_specs),
         )
         return result
 
@@ -726,6 +841,17 @@ def preprocess_multimethod(
         compile_specs: Dict[str, List[List[CompileSpec]]],
     ) -> Dict[str, List[PreprocessResult]]:
         """Compile every method against one global ``(device, FQN)`` store."""
+        # The methods share one weights file, so they must agree on its name.
+        data_names = {
+            cls.data_name_from_compile_specs(specs)
+            for method_specs in compile_specs.values()
+            for specs in method_specs
+        }
+        if len(data_names) > 1:
+            raise ValueError(
+                "Every method of a program must use the same data name, got "
+                f"{sorted(data_names)}"
+            )
         collector = CudaWeightCollector()
         results: Dict[str, List[PreprocessResult]] = {}
         for method_name, programs in edge_programs.items():
@@ -852,6 +978,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/cuda_weight_collector.py

@@ -570,6 +570,7 @@ def add_preprocess_result(
         device_name: str,
         target_sm: Optional[int] = None,
         ptx_compute: int = 0,
+        external_tag: Optional[str] = None,
     ) -> None:
         if result.data_store_output is None:
             raise RuntimeError("CUDA AOTI preprocess returned no named data")
@@ -589,7 +590,7 @@ def add_preprocess_result(
             keep_compatibility_blob=not artifact.entries,
         )
 
-        external_tag = f"aoti_{device_name}_blob"
+        external_tag = external_tag or f"aoti_{device_name}_blob"
         serialized_entries = []
         for entry in artifact.entries:
             data = artifact.storages.get(entry.storage_key)
@@ -601,7 +602,7 @@ def add_preprocess_result(
                     ),
                 )
             if data is not None:
-                # Every program writes its weights to the same file name, so the
+                # Programs can write their weights to the same file name, so the
                 # key carries the content: another program's file then fails to
                 # load instead of running this program with its weights.
                 content_key = f"{entry.storage_key}:{data.sha256().hex()}"

backends/cuda/merge_ptes.py

@@ -208,6 +208,11 @@ def _load_artifact(source: CudaPteInput) -> _Artifact:
     return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates)
 
 
+# Compile specs that only shape how a variant was exported, so variants that
+# differ in them still merge.
+_EXPORT_ONLY_SPEC_KEYS = ("cuda_include_ptx", "data_name")
+
+
 def _normalized_program(program: Program) -> Program:
     normalized = copy.deepcopy(program)
     payloads = []
@@ -219,7 +224,7 @@ def _normalized_program(program: Program) -> Program:
                 delegate.compile_specs = [
                     spec
                     for spec in delegate.compile_specs
-                    if spec.key != "cuda_include_ptx"
+                    if spec.key not in _EXPORT_ONLY_SPEC_KEYS
                 ]
             delegate.processed = BackendDelegateDataReference(
                 location=DataLocation.INLINE, index=len(payloads)
@@ -534,7 +539,7 @@ def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStor
                 delegate.compile_specs = [
                     spec
                     for spec in delegate.compile_specs
-                    if spec.key != "cuda_include_ptx"
+                    if spec.key not in _EXPORT_ONLY_SPEC_KEYS
                 ]
 
     merged_store = NamedDataStore()