From 375860a3aaac2adb0c11ae518f751fa02a632611 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 23:20:55 -0700 Subject: [PATCH 1/2] Let each CUDA program name its own weights file The CUDA backend writes every program's weights to a file named aoti_cuda_blob.ptd. Exporting two programs into one folder therefore makes the second overwrite the first one's weights. Since the weights are keyed by their content, the first program then fails to load instead of running with the wrong weights. Renaming each file right after its export works, because the program does not store the file name, but it is an extra step that is easy to miss. Add a data name option. A program lowered with CudaPartitioner([..., CudaBackend.generate_data_name_compile_spec("encoder")]) writes its weights to encoder.ptd. Without the option the file is still aoti_cuda_blob.ptd, so existing exports, scripts and runners are unchanged, and existing files load as before. The name becomes a file name, so it may use only up to 200 ASCII letters, digits, '_', '-' and '.'. It may not start with '.' or '-' (a leading '-' reads as an option, for example in merge_ptes --input-ptd), may not be a Windows device name such as NUL, and may not end with '.ptd'. One of the two data file writers adds '.ptd' only when it is missing, so "encoder" and "encoder.ptd" would write the same file. The check also runs where the name is used, so a hand-written compile spec cannot point outside the output folder. All methods of one program share one weights file, so they must use the same name. A backend that keeps its data in the .pte, such as Metal, refuses the option instead of ignoring it. merge_ptes ignores the name when it compares variants, like cuda_include_ptx, so variants exported with different names still merge. The CUDA backend docs describe the option, its rules, and how to load a program with its data file. The TensorRT page no longer says the name is fixed or that a clash gives wrong answers, and its tip on sharing the CUDA data file with the engines covers a custom name. A new test in test_missing_weights_blob.py exports two programs with their own names, and one with no constants, into one folder, checks the data file names, then loads and runs the first two against eager PyTorch. test_cuda_export.py checks the name rules, including non-ASCII letters, path separators, a leading '-' and a '.ptd' ending, checks that good names and the default are accepted, that a backend which keeps its data in the .pte refuses a name, and the one name per program rule. The existing merge test now merges variants with different names. --- backends/aoti/aoti_backend.py | 54 ++++++++++++++++++- backends/cuda/cuda_backend.py | 12 +++++ backends/cuda/cuda_weight_collector.py | 5 +- backends/cuda/merge_ptes.py | 9 +++- backends/cuda/tests/test_cuda_export.py | 54 +++++++++++++++++++ backends/cuda/tests/test_merge_ptes.py | 15 ++++-- .../cuda/tests/test_missing_weights_blob.py | 53 ++++++++++++++++-- docs/source/backends/cuda/cuda-overview.md | 19 +++++++ .../backends/tensorrt/tensorrt-overview.md | 11 ++-- 9 files changed, 215 insertions(+), 17 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 4a2ed0be168..0d0670b7318 100644 --- a/backends/aoti/aoti_backend.py +++ b/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 + ``.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__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, diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1f55a6099a9..e8e7f0947d9 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -704,6 +704,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 +727,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(): diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 19e9f8c8bb6..f826a34d718 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/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()}" diff --git a/backends/cuda/merge_ptes.py b/backends/cuda/merge_ptes.py index 14456dda54f..184bbfd3b86 100644 --- a/backends/cuda/merge_ptes.py +++ b/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() diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 96f637a829c..637ff4036b1 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -237,6 +237,60 @@ def test_invalid_cuda_include_ptx_compile_spec(self): [CompileSpec(key="cuda_include_ptx", value=b"MAYBE")] ) + def test_invalid_data_name_compile_spec(self): + for name in ( + "", + "..", + ".hidden", + "../outside", + "a/b", + "a\\b", + "a/../../outside", + "C:model", + "model?", + "nul", + "COM1.txt", + "a" * 201, + # Non-ASCII letters that IGNORECASE alone would let through. + "\u212a" * 200, + "\u017fizes", + "\u0130d", + "\u0131d", + "-enc", + # One writer turns "encoder.ptd" into the same file as "encoder". + # The check ignores case, like the rest of the pattern. + "encoder.ptd", + "encoder.PTD", + ): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.generate_data_name_compile_spec(name) + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.data_name_from_compile_specs( + [CompileSpec(key="data_name", value=name.encode())] + ) + # A backend that keeps its data in the .pte would ignore the name. + with patch.object(CudaBackend, "save_data_externally", return_value=False): + with self.assertRaisesRegex(ValueError, "takes no data name"): + CudaBackend.generate_data_name_compile_spec("encoder") + + def test_valid_data_name_compile_spec(self): + self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") + for name in ("encoder", "model_sm80", "v1.2-final", "a" * 200, "x.ptd.bak"): + with self.subTest(name=name): + spec = CudaBackend.generate_data_name_compile_spec(name) + self.assertEqual(CudaBackend.data_name_from_compile_specs([spec]), name) + + def test_methods_of_one_program_must_share_a_data_name(self): + def specs(name): + return [[CudaBackend.generate_data_name_compile_spec(name)]] + + with self.assertRaisesRegex(ValueError, "same data name"): + CudaBackend.preprocess_multimethod( + {"forward": [None], "second": [None]}, + {"forward": specs("first"), "second": specs("second")}, + ) + class TestCudaExport(unittest.TestCase): """Test CUDA export functionality for various operations using to_edge_transform_and_lower.""" diff --git a/backends/cuda/tests/test_merge_ptes.py b/backends/cuda/tests/test_merge_ptes.py index f0bb773070f..8ca8b8ecee5 100644 --- a/backends/cuda/tests/test_merge_ptes.py +++ b/backends/cuda/tests/test_merge_ptes.py @@ -58,6 +58,7 @@ def _write_artifact( fqn: str = "model.weight", weight_key: str = "cuda_fqn_weight:cuda:model.weight", ptx_compute: int | None = None, + data_name: str | None = None, ) -> CudaPteInput: if ptx_compute is None: ptx_compute = target_sm @@ -79,6 +80,8 @@ def _write_artifact( compile_specs.append( CompileSpec("cuda_include_ptx", b"ON" if ptx_compute else b"OFF") ) + if data_name is not None: + compile_specs.append(CompileSpec("data_name", data_name.encode())) delegate = BackendDelegate( id="CudaBackend", processed=BackendDelegateDataReference(DataLocation.INLINE, 0), @@ -148,8 +151,12 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: sm80.mkdir() sm120.mkdir() inputs = [ - self._write_artifact(sm80, 80, b"sm80-so", b"shared-weight"), - self._write_artifact(sm120, 120, b"sm120-so", b"shared-weight"), + self._write_artifact( + sm80, 80, b"sm80-so", b"shared-weight", data_name="model_sm80" + ), + self._write_artifact( + sm120, 120, b"sm120-so", b"shared-weight", data_name="model_sm120" + ), ] merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) @@ -176,14 +183,14 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: hashlib.sha256(b"sm120-so").hexdigest() + "_so_blob", }, ) - self.assertNotIn( - "cuda_include_ptx", + self.assertEqual( { spec.key for spec in merged.program.execution_plan[0] .delegates[0] .compile_specs }, + {"method_name"}, ) def test_uses_ptx_only_from_explicit_fallback(self) -> None: diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index fb2556cb3f8..6f1308b20c3 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -22,6 +22,7 @@ import os import tempfile import unittest +from typing import Optional import torch from executorch.backends.cuda.cuda_backend import CudaBackend @@ -51,14 +52,25 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: "loads through the Python runtime, which the ROCm build does not include", ) class TestMissingWeightsBlob(unittest.TestCase): - def _lower(self, outdir: str, seed: int = 0, name: str = "model.pte") -> str: + def _lower( + self, + outdir: str, + seed: int = 0, + name: str = "model.pte", + data_name: Optional[str] = None, + module: Optional[torch.nn.Module] = None, + ) -> str: torch.manual_seed(seed) - module = ModuleWithConstants().eval().cuda() + if module is None: + module = ModuleWithConstants() + module = module.eval().cuda() inputs = (torch.randn(4, 64, device="cuda"),) exported = export(module, inputs) - spec = CudaBackend.generate_method_name_compile_spec("forward") + specs = [CudaBackend.generate_method_name_compile_spec("forward")] + if data_name is not None: + specs.append(CudaBackend.generate_data_name_compile_spec(data_name)) lowered = to_edge_transform_and_lower( - exported, partitioner=[CudaPartitioner([spec])] + exported, partitioner=[CudaPartitioner(specs)] ) program = lowered.to_executorch() path = os.path.join(outdir, name) @@ -177,6 +189,39 @@ def test_load_fails_with_another_programs_blob(self) -> None: ): program.load_method("forward") + def test_programs_with_their_own_data_names_share_a_folder(self) -> None: + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + for seed, name in ((0, "first"), (1, "second")): + self._lower(outdir, seed=seed, name=f"{name}.pte", data_name=name) + # A program with no constants still writes a data file. + self._lower( + outdir, name="third.pte", data_name="third", module=torch.nn.ReLU() + ) + self.assertEqual( + sorted(f for f in os.listdir(outdir) if f.endswith(".ptd")), + ["first.ptd", "second.ptd", "third.ptd"], + ) + for seed, name in ((0, "first"), (1, "second")): + # The same draws _lower made, so the same weights and input. + torch.manual_seed(seed) + module = ModuleWithConstants().eval().cuda() + value = torch.randn(4, 64, device="cuda") + method = ( + Runtime.get() + .load_program( + os.path.join(outdir, f"{name}.pte"), + data_path=os.path.join(outdir, f"{name}.ptd"), + ) + .load_method("forward") + ) + out = method.execute([value.cpu()]) + torch.testing.assert_close( + (out[0] if isinstance(out, (list, tuple)) else out).cpu(), + module(value).cpu(), + ) + def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob. diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index 518073101d9..a78a4b831fa 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -81,6 +81,25 @@ save_pte_program(exec_program, model_name, "./output_dir") ``` This generates `.pte` and `.ptd` files that can be executed on CUDA devices. +By default, the weights go to a data file named `aoti_cuda_blob.ptd`. To export several programs into one folder, give each program its own +data file name. This partitioner writes `encoder.ptd`: + +```python +partitioner = CudaPartitioner([ + CudaBackend.generate_method_name_compile_spec("forward"), + CudaBackend.generate_data_name_compile_spec("encoder"), +]) +``` + +The name can have up to 200 ASCII letters, digits, `_`, `-` and `.`. It cannot start with `.` or `-`, cannot end with `.ptd` +(the export adds it), and cannot be a Windows device name such as `NUL`. Load each program with its own data file: + +```python +from executorch.runtime import Runtime + +program = Runtime.get().load_program("encoder.pte", data_path="encoder.ptd") +``` + For a complete working example, see the [CUDA export script](https://github.com/pytorch/executorch/blob/main/examples/cuda/scripts/export.py). diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index 411cc262961..9af3a20a2da 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -170,9 +170,11 @@ The Torch-TensorRT docs call such a program *coalesced*. The export writes the weights of the CUDA part to a data file, `aoti_cuda_blob.ptd`, next to the `.pte`. Pass that file when you load the program: if the CUDA part has weights, loading fails -without it. The file has the same name for every model, so a second export into the same directory -overwrites the first one's weights, and the first model can then return wrong answers with no -error. Save each model into its own directory. +without it. By default the file has the same name for every model, so a second export into the +same directory overwrites the first one's weights, and the first model then fails to load. To keep +several models in one directory, give each its own data file name with +`CudaBackend.generate_data_name_compile_spec`, as the +[CUDA backend page](../cuda/cuda-overview.md) shows, or save each model into its own directory. ### Keeping inputs and outputs on the GPU @@ -432,7 +434,8 @@ shows how to set the options that apply at run time, from C++ or from Python. Pass `external_engine_data="engines"` to `torch_tensorrt.save` to write the engines to `engines.ptd` instead of the `.pte`, and pass that file when you load the program. The Python loader takes one data file, so for a program that also uses the CUDA backend, pass - `external_engine_data="aoti_cuda_blob"` to share its file. + the name of the CUDA data file to share it: `external_engine_data="aoti_cuda_blob"` by default, + or the name given to `CudaBackend.generate_data_name_compile_spec`. - [**Optimization profiles.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#optimization-profiles) One engine can be tuned for several ranges of input shapes, for example one for decoding one token at a time and one for a long prompt. A C++ application picks the profile for each call From 3c71632ddecbfb563da3b9a5b74b0903b7bbca9f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 08:21:38 -0700 Subject: [PATCH 2/2] 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 about 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. --- backends/cuda/cuda_backend.py | 123 +++++++++++++ .../cuda/tests/test_antialiased_resize.py | 173 ++++++++++++++++++ 2 files changed, 296 insertions(+) create mode 100644 backends/cuda/tests/test_antialiased_resize.py diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index e8e7f0947d9..891fb81bd06 100644 --- a/backends/cuda/cuda_backend.py +++ b/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." @@ -864,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: diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py new file mode 100644 index 00000000000..963444f80e7 --- /dev/null +++ b/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 + )