Skip to content

Commit 7586c16

Browse files
authored
CUDA backend: per-program weights file names and antialiased resize (#24314)
This change fixes several problems in the CUDA backend. Each fix is its own commit. ## Name the weights file of each program The backend saves every program's weights to `aoti_cuda_blob.ptd`. Exporting several programs into one folder overwrites the weights, and the earlier programs then fail to load. Renaming each file after export works, but is easy to miss. A new option names the file: ```python partitioner = CudaPartitioner([ CudaBackend.generate_method_name_compile_spec("forward"), CudaBackend.generate_data_name_compile_spec("encoder"), ]) # writes encoder.ptd ``` Without the option nothing changes. The name is checked: up to 200 ASCII letters, digits, `_`, `-` and `.`, not starting with `.` or `-`, not ending with `.ptd`, and not a Windows device name such as `NUL`. All methods of a program must share it. A backend that keeps its data inside the `.pte`, such as Metal, refuses it. ## Lower antialiased resizes `F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)` failed to lower with "missing fallback kernels". The backend now rewrites it into plain tensor math, with PyTorch's window and weight formulas, for dynamic sizes and float32, float16 and bfloat16. The results are not bit-identical: in float32 they stay within about 2e-5 of PyTorch's kernel for inputs between 0 and 1. ## Tests Ran on an A100 GPU (Linux x86_64): `test_cuda_export.py`, `test_antialiased_resize.py`, `test_missing_weights_blob.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py`. Each test added for a fix fails when that fix is removed.
1 parent fb34f59 commit 7586c16

10 files changed

Lines changed: 511 additions & 17 deletions

File tree

‎backends/aoti/aoti_backend.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,29 @@
2828

2929
class COMPILE_SPEC_KEYS(Enum):
3030
METHOD_NAME = "method_name"
31+
DATA_NAME = "data_name"
32+
33+
34+
# The name becomes a file name on every platform, so keep it short, to portable
35+
# characters, and clear of Windows device names such as NUL. Without re.ASCII,
36+
# IGNORECASE lets [A-Za-z] match letters like the Kelvin sign. A leading '-'
37+
# reads as an option on a command line. Some writers add .ptd only when it is
38+
# missing, so "encoder" and "encoder.ptd" could both write encoder.ptd.
39+
_DATA_NAME_PATTERN = re.compile(
40+
r"(?!(?:CON|PRN|AUX|NUL|COM\d|LPT\d)(?:\.|$))(?!.*\.ptd$)"
41+
r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,199}",
42+
re.IGNORECASE | re.ASCII,
43+
)
44+
45+
46+
def _checked_data_name(data_name: str) -> str:
47+
if not _DATA_NAME_PATTERN.fullmatch(data_name):
48+
raise ValueError(
49+
f"Invalid data name {data_name!r}: use up to 200 ASCII letters, "
50+
"digits, '_', '-' and '.', do not start with '.' or '-', do not end "
51+
"with '.ptd', and avoid Windows device names such as NUL"
52+
)
53+
return data_name
3154

3255

3356
@experimental(
@@ -368,7 +391,9 @@ def preprocess(
368391
# Determine whether to save named data externally based on backend setting
369392
# External: save to separate .ptd file, otherwise merge with .pte file
370393
external_tag = (
371-
f"aoti_{device_name}_blob" if cls.save_data_externally() else None
394+
cls.data_name_from_compile_specs(compile_specs)
395+
if cls.save_data_externally()
396+
else None
372397
)
373398

374399
named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag)
@@ -404,6 +429,33 @@ def generate_method_name_compile_spec(
404429
method_name.encode("utf-8"),
405430
)
406431

432+
@classmethod
433+
def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec:
434+
"""
435+
Generate a CompileSpec that names the external data file, written as
436+
``<data_name>.ptd``, so several programs can share a folder.
437+
"""
438+
if not cls.save_data_externally():
439+
raise ValueError(
440+
f"{cls.__name__} keeps its data in the .pte file, so it takes no "
441+
"data name"
442+
)
443+
return CompileSpec(
444+
COMPILE_SPEC_KEYS.DATA_NAME.value,
445+
_checked_data_name(data_name).encode("utf-8"),
446+
)
447+
448+
@classmethod
449+
def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str:
450+
"""
451+
The external data file name, without ``.ptd``, set by
452+
generate_data_name_compile_spec, or ``aoti_<device>_blob`` by default.
453+
"""
454+
for spec in compile_specs:
455+
if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value:
456+
return _checked_data_name(spec.value.decode("utf-8"))
457+
return f"aoti_{cls.get_device_name()}_blob"
458+
407459
@classmethod
408460
def method_name_from_compile_specs(
409461
cls,

‎backends/cuda/cuda_backend.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import copy
1010
import functools
1111
import logging
12+
import math
1213
import os
1314
import shutil
1415
import threading
@@ -50,7 +51,9 @@
5051
from executorch.exir._warnings import experimental
5152
from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
5253
from executorch.exir.backend.compile_spec_schema import CompileSpec
54+
from torch._decomp import decomposition_table
5355
from torch._inductor.decomposition import conv1d_to_conv2d
56+
from torch._prims_common import suggest_memory_format
5457
from torch.nn.attention import SDPBackend
5558

5659
# ---------------------------------------------------------------------------
@@ -548,6 +551,117 @@ def _conv_transpose_as_conv(
548551
return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups)
549552

550553

554+
def _linear_kernel(x: torch.Tensor) -> torch.Tensor:
555+
return torch.clamp(1 - x.abs(), min=0)
556+
557+
558+
def _cubic_kernel(x: torch.Tensor) -> torch.Tensor:
559+
# Keys' cubic with a = -0.5, the one ATen's antialiased bicubic uses.
560+
a = -0.5
561+
x = x.abs()
562+
near = ((a + 2) * x - (a + 3)) * x * x + 1
563+
far = (((x - 5) * x + 8) * x - 4) * a
564+
return torch.where(x < 1, near, torch.where(x < 2, far, torch.zeros_like(x)))
565+
566+
567+
def _antialias_taps(
568+
in_size: int,
569+
out_size: int,
570+
align_corners: bool,
571+
scale: Optional[float],
572+
support: int,
573+
kernel: typing.Callable[[torch.Tensor], torch.Tensor],
574+
device: torch.device,
575+
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
576+
"""
577+
Taps of an antialiased resize along one dim, computed as ATen does:
578+
output o reads input index[o, k] with weight[o, k], and inside[o, k]
579+
marks the taps in its window. The window is the kernel's support widened
580+
by the downscale factor, and the weights in it are normalized.
581+
"""
582+
if align_corners or scale is None or scale <= 0:
583+
if align_corners:
584+
# ATen uses a ratio of 0 for an output size of 1. min and max give
585+
# that without a conditional, which torch.compile cannot handle for
586+
# dynamic sizes.
587+
num = (in_size - 1) * torch.sym_min(out_size - 1, 1)
588+
den = torch.sym_max(out_size - 1, 1)
589+
else:
590+
num, den = in_size, out_size
591+
ratio = num / den
592+
# The tap count is a shape that both the host and the kernel work out.
593+
# In float32 the kernel can round num / den just past a whole number
594+
# and count more taps than the host allocated, so use integers.
595+
reach = torch.sym_max((support * num + den - 1) // den, support)
596+
else:
597+
ratio = 1.0 / scale
598+
reach = math.ceil(support * max(ratio, 1.0))
599+
widen = torch.sym_max(ratio, 1.0)
600+
dtype = torch.float32
601+
center = ratio * (torch.arange(out_size, dtype=dtype, device=device) + 0.5)
602+
first = torch.clamp(torch.trunc(center - support * widen + 0.5), min=0)
603+
end = torch.clamp(torch.trunc(center + support * widen + 0.5), max=in_size)
604+
taps = 2 * reach + 1
605+
tap = torch.arange(taps, dtype=dtype, device=device)
606+
position = first[:, None] + tap
607+
inside = position < end[:, None]
608+
weights = kernel((tap + (first - center)[:, None] + 0.5) * (1.0 / widen))
609+
weights = torch.where(inside, weights, 0)
610+
total = weights.sum(1, keepdim=True)
611+
weights = torch.where(total != 0, weights / total, weights)
612+
# Clamp after the cast: in float32, in_size - 1 rounds up to in_size for
613+
# very long axes.
614+
index = position.to(torch.int64).clamp(max=in_size - 1)
615+
return index, weights, inside
616+
617+
618+
def _antialiased_resize(
619+
support: int,
620+
kernel: typing.Callable[[torch.Tensor], torch.Tensor],
621+
input: torch.Tensor,
622+
output_size: List[int],
623+
align_corners: bool,
624+
scales_h: Optional[float] = None,
625+
scales_w: Optional[float] = None,
626+
) -> torch.Tensor:
627+
"""
628+
Rewrite an antialiased bilinear or bicubic resize as a weighted sum over
629+
each output's taps, one dim at a time. Inductor has no lowering for these
630+
resizes, so it would call ATen's kernel, which this backend does not ship.
631+
"""
632+
# The backend's tensors have no float64, and ATen's CUDA kernel has no
633+
# integer version.
634+
if input.dtype not in (torch.float32, torch.float16, torch.bfloat16):
635+
return NotImplemented
636+
# Half types are computed in float and rounded once, so they can differ
637+
# slightly from ATen, which keeps its weights in the input type.
638+
result = input.float()
639+
for dim, scale in ((3, scales_w), (2, scales_h)):
640+
index, weights, inside = _antialias_taps(
641+
input.shape[dim],
642+
output_size[dim - 2],
643+
align_corners,
644+
scale,
645+
support,
646+
kernel,
647+
input.device,
648+
)
649+
rows, taps = index.shape
650+
trailing = [1] * (3 - dim)
651+
# Plain indexing lets Inductor fold the gather into the sum instead of
652+
# keeping every tap of every row in memory.
653+
gathered = result[(slice(None),) * dim + (index,)]
654+
# Taps outside the window already have zero weight, but NaN * 0 is NaN,
655+
# so mask them out to keep a NaN or inf local.
656+
weighted = torch.where(
657+
inside.view(rows, taps, *trailing),
658+
gathered * weights.view(rows, taps, *trailing),
659+
0,
660+
)
661+
result = weighted.sum(dim + 1)
662+
return result.to(input.dtype).contiguous(memory_format=suggest_memory_format(input))
663+
664+
551665
@final
552666
@experimental(
553667
"This API and all of cuda backend related functionality are experimental."
@@ -704,6 +818,7 @@ def _preprocess_with_weight_collector(
704818
cls.get_device_name(),
705819
target_sm=target_sm,
706820
ptx_compute=ptx_compute,
821+
external_tag=cls.data_name_from_compile_specs(compile_specs),
707822
)
708823
return result
709824

@@ -726,6 +841,17 @@ def preprocess_multimethod(
726841
compile_specs: Dict[str, List[List[CompileSpec]]],
727842
) -> Dict[str, List[PreprocessResult]]:
728843
"""Compile every method against one global ``(device, FQN)`` store."""
844+
# The methods share one weights file, so they must agree on its name.
845+
data_names = {
846+
cls.data_name_from_compile_specs(specs)
847+
for method_specs in compile_specs.values()
848+
for specs in method_specs
849+
}
850+
if len(data_names) > 1:
851+
raise ValueError(
852+
"Every method of a program must use the same data name, got "
853+
f"{sorted(data_names)}"
854+
)
729855
collector = CudaWeightCollector()
730856
results: Dict[str, List[PreprocessResult]] = {}
731857
for method_name, programs in edge_programs.items():
@@ -852,6 +978,15 @@ def get_decomposition_table(cls) -> Dict[Any, Any]:
852978
# for it and the shim no fallback kernel.
853979
torch.ops.aten._is_all_true.default: torch.ops.aten.all.default,
854980
}
981+
for resize, support, kernel in (
982+
(torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel),
983+
(torch.ops.aten._upsample_bicubic2d_aa, 2, _cubic_kernel),
984+
):
985+
table[resize.default] = functools.partial(
986+
_antialiased_resize, support, kernel
987+
)
988+
# The vec form works out the output size and calls the default form.
989+
table[resize.vec] = decomposition_table[resize.vec]
855990
# Inductor disables its transposed conv2d Triton template only on CUDA
856991
# (https://github.com/pytorch/pytorch/issues/187081).
857992
if torch.version.hip is None:

‎backends/cuda/cuda_weight_collector.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -570,6 +570,7 @@ def add_preprocess_result(
570570
device_name: str,
571571
target_sm: Optional[int] = None,
572572
ptx_compute: int = 0,
573+
external_tag: Optional[str] = None,
573574
) -> None:
574575
if result.data_store_output is None:
575576
raise RuntimeError("CUDA AOTI preprocess returned no named data")
@@ -589,7 +590,7 @@ def add_preprocess_result(
589590
keep_compatibility_blob=not artifact.entries,
590591
)
591592

592-
external_tag = f"aoti_{device_name}_blob"
593+
external_tag = external_tag or f"aoti_{device_name}_blob"
593594
serialized_entries = []
594595
for entry in artifact.entries:
595596
data = artifact.storages.get(entry.storage_key)
@@ -601,7 +602,7 @@ def add_preprocess_result(
601602
),
602603
)
603604
if data is not None:
604-
# Every program writes its weights to the same file name, so the
605+
# Programs can write their weights to the same file name, so the
605606
# key carries the content: another program's file then fails to
606607
# load instead of running this program with its weights.
607608
content_key = f"{entry.storage_key}:{data.sha256().hex()}"

‎backends/cuda/merge_ptes.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,11 @@ def _load_artifact(source: CudaPteInput) -> _Artifact:
208208
return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates)
209209

210210

211+
# Compile specs that only shape how a variant was exported, so variants that
212+
# differ in them still merge.
213+
_EXPORT_ONLY_SPEC_KEYS = ("cuda_include_ptx", "data_name")
214+
215+
211216
def _normalized_program(program: Program) -> Program:
212217
normalized = copy.deepcopy(program)
213218
payloads = []
@@ -219,7 +224,7 @@ def _normalized_program(program: Program) -> Program:
219224
delegate.compile_specs = [
220225
spec
221226
for spec in delegate.compile_specs
222-
if spec.key != "cuda_include_ptx"
227+
if spec.key not in _EXPORT_ONLY_SPEC_KEYS
223228
]
224229
delegate.processed = BackendDelegateDataReference(
225230
location=DataLocation.INLINE, index=len(payloads)
@@ -534,7 +539,7 @@ def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStor
534539
delegate.compile_specs = [
535540
spec
536541
for spec in delegate.compile_specs
537-
if spec.key != "cuda_include_ptx"
542+
if spec.key not in _EXPORT_ONLY_SPEC_KEYS
538543
]
539544

540545
merged_store = NamedDataStore()

0 commit comments

Comments
 (0)