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..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." @@ -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: 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_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 + ) 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