diff --git a/docsrc/tutorials/model_zoo.rst b/docsrc/tutorials/model_zoo.rst index 11a50f968a..eaad2da65a 100644 --- a/docsrc/tutorials/model_zoo.rst +++ b/docsrc/tutorials/model_zoo.rst @@ -10,6 +10,7 @@ and benchmark against eager PyTorch. :caption: Vision Example: Compiling ResNet with dynamic shapes <_rendered_examples/dynamo/torch_compile_resnet_example> + Example: Global Performance Tuner <_rendered_examples/dynamo/global_perf_tuner_attention_example> Example: Compiling BERT with torch.compile <_rendered_examples/dynamo/torch_compile_transformers_example> Example: Engine Caching (BERT) <_rendered_examples/dynamo/engine_caching_bert_example> diff --git a/docsrc/user_guide/performance_tuning.rst b/docsrc/user_guide/performance_tuning.rst index 745ad8752f..29a7d6f624 100644 --- a/docsrc/user_guide/performance_tuning.rst +++ b/docsrc/user_guide/performance_tuning.rst @@ -339,3 +339,76 @@ Benchmarking Checklist - For latency workloads: enable CUDA graphs * - ☐ - For large models: try weight streaming or INT8 quantization + +---- + +Global Performance Tuner +------------------------ + +TensorRT's `Global Performance Tuner `_ +searches internal builder knobs (a *build route*) for faster engines. Torch-TensorRT +exposes the same capability in-process for Dynamo TRT subgraphs. + +**Requirements:** TensorRT with Global Performance Tuner enabled (enterprise Linux +builds >= 11.1; currently unavailable on TensorRT-RTX / Windows). Probe with +``torch_tensorrt.dynamo.is_global_perf_tuner_available()``. + +**Discover knobs** (``trtexec --helpBuildRoute`` equivalent):: + + from torch_tensorrt.dynamo import get_all_build_routes + + knobs = get_all_build_routes() + print(knobs["tuner_version"], len(knobs["tuner_options"])) + +**Apply a known route** (``trtexec --setBuildRoute`` equivalent):: + + trt_model = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=inputs, + build_route="-slice_fusion=off -kgen:codegen:cuda_tile=3", + ) + +**Sweep routes** (``trtexec --tuneBuildRoutes`` equivalent):: + + trt_model = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=inputs, + tune_build_routes="-match_ragged_mha=[on|off] -copy_ppg=[on|off]", + tuning_search="fast", # or "full" / "mixed" + accuracy_threshold=0.01, + accuracy_algorithm="cos", + tuning_cache_file="/tmp/torch_trt_tune.jsonl", + ) + +See also the runnable attention walkthrough: +:ref:`global_perf_tuner_attention_example` +(``examples/dynamo/global_perf_tuner_attention_example.py``). + +Notes: + +- Tuning runs **per TRT partition**. Prefer ``require_full_compilation=True`` when you + want a single-engine sweep similar to ``trtexec``. +- ``tuning_cache_file`` is treated as a **base path**. Each TRT partition writes + ``..jsonl`` so multi-subgraph models do not overwrite one + another. Use the same base path with ``tuning_continue=True`` to resume; Torch-TensorRT + resolves the partition file from the subgraph fingerprint. +- ``tuning_search="fast"`` is linear in the number of variable knobs; ``full`` is a + Cartesian product; ``mixed`` runs ``fast`` search and then ``full`` search only on knobs that improved performance. + Prefer ``fast`` / ``mixed`` over large ``full`` expressions. +- Some knob values are model-dependent and may fail to build (for example certain + ``-kgen:codegen:cuda_tile`` settings). Those trials are recorded with ``crash=true`` + (``error_message`` may include recent TensorRT ERROR log lines) and are skipped when + selecting the winner. +- ``tuning_dry_run=True`` enumerates routes without building TRT engines (incompatible with ``mixed``). +- Accuracy metrics match trtexec: ``l0`` / ``l1`` / ``l2`` / ``lInf`` / ``cos`` (lower is better). + References are eager Torch outputs on the compile example inputs. + ``accuracy_atol`` / ``accuracy_rtol`` apply only to ``l0``. +- Resume an interrupted sweep with ``tuning_continue=True`` and the same + ``tuning_cache_file`` base path (Torch-TensorRT JSONL header; not interchangeable with + ``trtexec`` cache files). +- Gains are model-, GPU-, and TensorRT-version-dependent; re-tune after hardware or + TensorRT upgrades. Non-default routes have no cross-release performance guarantee. + Engines are not bit-deterministic across builds. Sweeps multiply compile time, + especially with multi-partition graphs and ``full`` search. Use ``tuning_timeout_s`` to limit the sweep duration. \ No newline at end of file diff --git a/examples/dynamo/global_perf_tuner_attention_example.py b/examples/dynamo/global_perf_tuner_attention_example.py new file mode 100644 index 0000000000..659088a416 --- /dev/null +++ b/examples/dynamo/global_perf_tuner_attention_example.py @@ -0,0 +1,175 @@ +""" +.. _global_perf_tuner_attention_example: + +Global Performance Tuner (Attention) +==================================== + +This example shows how to use TensorRT's Global Performance Tuner (GPT) from +Torch-TensorRT: discover knobs, sweep a small route space on a simple +multi-head attention module, and apply the winning ``build_route``. + +**Requirements:** TensorRT with GPT enabled. This feature is available since +TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows. The +script exits early if GPT is not available. +""" + +# %% +# Imports and model +# ^^^^^^^^^^^^^^^^^ + +import json +import os +import tempfile + +import torch +import torch.nn as nn +import torch_tensorrt +from torch_tensorrt.dynamo import get_all_build_routes, is_global_perf_tuner_available + +torch.manual_seed(0) + +if not torch.cuda.is_available(): + raise SystemExit("CUDA is required for this example.") + +if not is_global_perf_tuner_available(): + raise SystemExit( + "Global Performance Tuner is not available on this TensorRT build (IBuilderConfig.build_route / all_build_routes). " + "This feature is available since TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows." + ) + + +class SimpleAttention(nn.Module): + """Minimal self-attention block suitable for a short GPT demo.""" + + def __init__(self, embed_dim: int = 64, num_heads: int = 4) -> None: + super().__init__() + self.attn = nn.MultiheadAttention( + embed_dim, num_heads, batch_first=True, bias=True + ) + self.norm = nn.LayerNorm(embed_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y, _ = self.attn(x, x, x, need_weights=False) + return self.norm(x + y) + + +batch, seq, embed = 1, 32, 64 +model = SimpleAttention(embed_dim=embed, num_heads=4).eval().cuda().half() +example_inputs = [torch.randn((batch, seq, embed), device="cuda", dtype=torch.float16)] + +# %% +# Discover knobs (``trtexec --helpBuildRoute``) +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +knobs = get_all_build_routes() +print(f"tuner_version={knobs.get('tuner_version')}") +print(f"num_knobs={len(knobs.get('tuner_options', []))}") + +# Tune selected knobs. Some values (for example certain ``cuda_tile`` settings) may +# fail to build for a given graph; those trials are recorded as ``crash=True`` in the +# tuning cache and skipped when picking the winner. +tune_expr = "-match_ragged_mha=[on|off] -slice_fusion=[on|off] -copy_ppg=[on|off] -reshape_ppg=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]" +print(f"tune_build_routes={tune_expr}") + +# %% +# Sweep routes (``trtexec --tuneBuildRoutes``) +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# +# Use ``require_full_compilation=True`` so the sweep targets a single TRT engine +# (closer to whole-network ``trtexec``). ``tuning_search="fast"`` is linear in the +# number of variable knobs and is the recommended starting point. +# +# ``tuning_cache_file`` is a *base* path; Torch-TensorRT writes a per-partition file +# ``..jsonl`` so multi-subgraph models do not overwrite each other. + +cache_base = os.path.join(tempfile.gettempdir(), "torch_trt_attention_tune.jsonl") +# Remove any leftover partition caches from prior runs of this example. +cache_dir = os.path.dirname(cache_base) or "." +cache_stem = os.path.splitext(os.path.basename(cache_base))[0] +for name in os.listdir(cache_dir): + if name.startswith(cache_stem) and name.endswith(".jsonl"): + os.remove(os.path.join(cache_dir, name)) +print(f"tuning_cache_file base={cache_base}") + +optimized = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=example_inputs, + min_block_size=1, + tune_build_routes=tune_expr, + tuning_search="fast", + accuracy_threshold=0.01, + accuracy_algorithm="cos", + tuning_cache_file=cache_base, +) + +# %% +# Check accuracy and inspect the tuning cache +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +with torch.no_grad(): + ref = model(*example_inputs) + out = optimized(*example_inputs) + +cos = torch.nn.functional.cosine_similarity( + ref.flatten().float(), out.flatten().float(), dim=0 +).item() +print(f"cosine_similarity(torch_eager, tuned_trt)={cos:.6f}") + +partition_caches = sorted( + os.path.join(cache_dir, name) + for name in os.listdir(cache_dir) + if name.startswith(cache_stem + ".") and name.endswith(".jsonl") +) +assert partition_caches, f"expected per-partition cache under {cache_base}" +cache_path = partition_caches[0] +print(f"partition cache_path={cache_path}") + +with open(cache_path, "r", encoding="utf-8") as f: + lines = [ln.strip() for ln in f if ln.strip()] + +header = json.loads(lines[0]) +print("cache header keys:", sorted(header.keys())) +print(f"recorded iterations: {len(lines) - 1}") +for line in lines[1:]: + row = json.loads(line) + print( + f" iter={row['iter']} crash={row['crash']} " + f"gpu_time={row.get('gpu_time')} route={row.get('build_route')}" + ) + +# %% +# Re-apply a known winning route +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# +# After a sweep (or from ``trtexec``), pin the best route with ``build_route`` +# for subsequent compiles without re-running the search. + +best_route = "" +best_time = None +for line in lines[1:]: + row = json.loads(line) + t = row.get("gpu_time") + if row.get("crash") or t is None: + continue + if best_time is None or t < best_time: + best_time = t + best_route = row["build_route"] + +print(f"best cached route={best_route!r} gpu_time_ms={best_time}") + +if best_route: + pinned = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=example_inputs, + require_full_compilation=True, + min_block_size=1, + build_route=best_route, + ) + with torch.no_grad(): + pinned_out = pinned(*example_inputs) + print( + "pinned route max abs err vs torch eager:", + (pinned_out.float() - ref.float()).abs().max().item(), + ) diff --git a/py/torch_tensorrt/dynamo/__init__.py b/py/torch_tensorrt/dynamo/__init__.py index 607dca76bf..bc3995ff84 100644 --- a/py/torch_tensorrt/dynamo/__init__.py +++ b/py/torch_tensorrt/dynamo/__init__.py @@ -20,3 +20,4 @@ from ._SourceIR import SourceIR from ._tracer import trace from .debug._Debugger import Debugger + from .tuning import get_all_build_routes, is_global_perf_tuner_available diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..fc8e9478d7 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -5,7 +5,6 @@ import os import platform import warnings - from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union import sympy @@ -484,6 +483,18 @@ def compile( decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS, + build_route: str = _defaults.BUILD_ROUTE, + tune_build_routes: str = _defaults.TUNE_BUILD_ROUTES, + tune_build_route_file: Optional[str] = _defaults.TUNE_BUILD_ROUTE_FILE, + tuning_search: str = _defaults.TUNING_SEARCH, + tuning_timeout_s: int = _defaults.TUNING_TIMEOUT_S, + tuning_cache_file: Optional[str] = _defaults.TUNING_CACHE_FILE, + tuning_continue: bool = _defaults.TUNING_CONTINUE, + tuning_dry_run: bool = _defaults.TUNING_DRY_RUN, + accuracy_threshold: Optional[float] = _defaults.ACCURACY_THRESHOLD, + accuracy_algorithm: str = _defaults.ACCURACY_ALGORITHM, + accuracy_atol: float = _defaults.ACCURACY_ATOL, + accuracy_rtol: float = _defaults.ACCURACY_RTOL, **kwargs: Any, ) -> torch.fx.GraphModule: """Compile an ExportedProgram module for NVIDIA GPUs using TensorRT @@ -578,6 +589,18 @@ def compile( the final output back to FP16. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. + build_route (str): TensorRT Global Performance Tuner build route string (space-separated "-knob=value" tokens). Empty uses the default route. This feature requires TensorRT with Global Performance Tuner enabled; currently unavailable on TensorRT-RTX / Windows. + tune_build_routes (str): Build-route expression for an in-process autotuning sweep (e.g. "-slice_fusion=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"). Empty disables tuning. + tune_build_route_file (Optional[str]): Path to a file with one route token per line (same as "tune_build_routes"). Mutually exclusive with "tune_build_routes". + tuning_search (str): Search algorithm: "fast", "full", or "mixed". Default is "fast". "fast" runs a baseline with all knobs at default, then varies one knob at a time. "full" runs a full grid search. "mixed" runs "fast" first, and then "full" only on knobs that improved performance. + tuning_timeout_s (int): Time budget for the entire tuning process. The current iteration finishes before the loop stops. Use -1 (default) to disable the timeout. Helpful for capping large "full" sweeps. + tuning_cache_file (Optional[str]): JSONL path for tuning results / resume. + tuning_continue (bool): Resume an interrupted sweep from "tuning_cache_file". When it is True, "tuning_cache_file" must be provided; "tune_build_routes", "tune_build_route_file", and "tuning_dry_run" must not be provided. + tuning_dry_run (bool): Enumerate candidate routes without building engines. + accuracy_threshold (Optional[float]): Max allowed accuracy loss between the engine outputs and eager Torch reference outputs; routes with accuracy loss above this threshold are excluded from best-engine selection. "None" skips accuracy checks. + accuracy_algorithm (str): Loss metric: "l0", "l1", "l2", "lInf", or "cos" (lower is better). + accuracy_atol (float): Absolute tolerance for "accuracy_algorithm="l0"". + accuracy_rtol (float): Relative tolerance for "accuracy_algorithm="l0"". **kwargs: Any, Returns: torch.fx.GraphModule: Compiled FX Module, when run it will execute via TensorRT @@ -767,6 +790,18 @@ def compile( "decompose_attention": decompose_attention, "attn_bias_is_causal": attn_bias_is_causal, "fallback_data_dependent_ops": fallback_data_dependent_ops, + "build_route": build_route, + "tune_build_routes": tune_build_routes, + "tune_build_route_file": tune_build_route_file, + "tuning_search": tuning_search, + "tuning_timeout_s": tuning_timeout_s, + "tuning_cache_file": tuning_cache_file, + "tuning_continue": tuning_continue, + "tuning_dry_run": tuning_dry_run, + "accuracy_threshold": accuracy_threshold, + "accuracy_algorithm": accuracy_algorithm, + "accuracy_atol": accuracy_atol, + "accuracy_rtol": accuracy_rtol, } logger.debug(f"CPU memory usage before lowering: {get_cpu_memory_usage()} MB") settings = CompilationSettings(**compilation_options) diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 9c8a1f9f90..e8974a35af 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -68,6 +68,20 @@ ATTN_BIAS_IS_CAUSAL = True FALLBACK_DATA_DEPENDENT_OPS = False +# Global Performance Tuner (TensorRT build-route knobs) +BUILD_ROUTE = "" +TUNE_BUILD_ROUTES = "" +TUNE_BUILD_ROUTE_FILE = None +TUNING_SEARCH = "fast" +TUNING_TIMEOUT_S = -1 +TUNING_CACHE_FILE = None +TUNING_CONTINUE = False +TUNING_DRY_RUN = False +ACCURACY_THRESHOLD = None +ACCURACY_ALGORITHM = "l0" +ACCURACY_ATOL = 1e-5 +ACCURACY_RTOL = 1e-5 + if platform.system() == "Linux": import pwd diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 7ff5f6bdff..ad7849e107 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -7,6 +7,10 @@ from torch_tensorrt._Device import Device from torch_tensorrt._enums import EngineCapability, dtype from torch_tensorrt.dynamo._defaults import ( + ACCURACY_ALGORITHM, + ACCURACY_ATOL, + ACCURACY_RTOL, + ACCURACY_THRESHOLD, ASSUME_DYNAMIC_SHAPE_SUPPORT, ATTN_BIAS_IS_CAUSAL, AUTOCAST_CALIBRATION_DATALOADER, @@ -15,6 +19,7 @@ AUTOCAST_LOW_PRECISION_TYPE, AUTOCAST_MAX_DEPTH_OF_REDUCTION, AUTOCAST_MAX_OUTPUT_THRESHOLD, + BUILD_ROUTE, CACHE_BUILT_ENGINES, CPU_MEMORY_BUDGET, DECOMPOSE_ATTENTION, @@ -49,6 +54,13 @@ TILING_OPTIMIZATION_LEVEL, TIMING_CACHE_PATH, TRUNCATE_DOUBLE, + TUNE_BUILD_ROUTE_FILE, + TUNE_BUILD_ROUTES, + TUNING_CACHE_FILE, + TUNING_CONTINUE, + TUNING_DRY_RUN, + TUNING_SEARCH, + TUNING_TIMEOUT_S, USE_DISTRIBUTED_MODE_TRACE, USE_FAST_PARTITIONER, USE_FP32_ACC, @@ -121,6 +133,18 @@ class CompilationSettings: the final output back to FP16. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. + build_route (str): TensorRT Global Performance Tuner build route string (space-separated "-knob=value" tokens). Empty uses the default route. This feature requires TensorRT with Global Performance Tuner enabled; currently unavailable on TensorRT-RTX / Windows. + tune_build_routes (str): Build-route expression for an in-process autotuning sweep (e.g. "-slice_fusion=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"). Empty disables tuning. + tune_build_route_file (Optional[str]): Path to a file with one route token per line (same as "tune_build_routes"). Mutually exclusive with "tune_build_routes". + tuning_search (str): Search algorithm: "fast", "full", or "mixed". Default is "fast". "fast" runs a baseline with all knobs at default, then varies one knob at a time. "full" runs a full grid search. "mixed" runs "fast" first, and then "full" only on knobs that improved performance. + tuning_timeout_s (int): Time budget for the entire tuning process. The current iteration finishes before the loop stops. Use -1 (default) to disable the timeout. Helpful for capping large "full" sweeps. + tuning_cache_file (Optional[str]): JSONL path for tuning results / resume. + tuning_continue (bool): Resume an interrupted sweep from "tuning_cache_file". When it is True, "tuning_cache_file" must be provided; "tune_build_routes", "tune_build_route_file", and "tuning_dry_run" must not be provided. + tuning_dry_run (bool): Enumerate candidate routes without building engines. + accuracy_threshold (Optional[float]): Max allowed accuracy loss between the engine outputs and eager Torch reference outputs; routes with accuracy loss above this threshold are excluded from best-engine selection. "None" skips accuracy checks. + accuracy_algorithm (str): Loss metric: "l0", "l1", "l2", "lInf", or "cos" (lower is better). + accuracy_atol (float): Absolute tolerance for "accuracy_algorithm="l0"". + accuracy_rtol (float): Relative tolerance for "accuracy_algorithm="l0"". """ workspace_size: int = WORKSPACE_SIZE @@ -180,6 +204,18 @@ class CompilationSettings: decompose_attention: bool = DECOMPOSE_ATTENTION attn_bias_is_causal: bool = ATTN_BIAS_IS_CAUSAL fallback_data_dependent_ops: bool = FALLBACK_DATA_DEPENDENT_OPS + build_route: str = BUILD_ROUTE + tune_build_routes: str = TUNE_BUILD_ROUTES + tune_build_route_file: Optional[str] = TUNE_BUILD_ROUTE_FILE + tuning_search: str = TUNING_SEARCH + tuning_timeout_s: int = TUNING_TIMEOUT_S + tuning_cache_file: Optional[str] = TUNING_CACHE_FILE + tuning_continue: bool = TUNING_CONTINUE + tuning_dry_run: bool = TUNING_DRY_RUN + accuracy_threshold: Optional[float] = ACCURACY_THRESHOLD + accuracy_algorithm: str = ACCURACY_ALGORITHM + accuracy_atol: float = ACCURACY_ATOL + accuracy_rtol: float = ACCURACY_RTOL def __getstate__(self) -> dict[str, Any]: from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( @@ -196,6 +232,18 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: state.pop("use_python_runtime", None) state.setdefault("fallback_data_dependent_ops", FALLBACK_DATA_DEPENDENT_OPS) + state.setdefault("build_route", BUILD_ROUTE) + state.setdefault("tune_build_routes", TUNE_BUILD_ROUTES) + state.setdefault("tune_build_route_file", TUNE_BUILD_ROUTE_FILE) + state.setdefault("tuning_search", TUNING_SEARCH) + state.setdefault("tuning_timeout_s", TUNING_TIMEOUT_S) + state.setdefault("tuning_cache_file", TUNING_CACHE_FILE) + state.setdefault("tuning_continue", TUNING_CONTINUE) + state.setdefault("tuning_dry_run", TUNING_DRY_RUN) + state.setdefault("accuracy_threshold", ACCURACY_THRESHOLD) + state.setdefault("accuracy_algorithm", ACCURACY_ALGORITHM) + state.setdefault("accuracy_atol", ACCURACY_ATOL) + state.setdefault("accuracy_rtol", ACCURACY_RTOL) self.__dict__.update(state) @@ -222,6 +270,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: "autocast_calibration_dataloader", "decompose_attention", "attn_bias_is_causal", + "build_route", } diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py index 448d96a0ce..a7f43d3c33 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py @@ -15,6 +15,7 @@ ) import numpy as np +import tensorrt as trt import torch import torch.fx from torch.fx.experimental.proxy_tensor import unset_fake_temporarily @@ -52,8 +53,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - _LOGGER: logging.Logger = logging.getLogger(__name__) TRT_INTERPRETER_CALL_PRE_OBSERVER: Observer[Callable[[torch.fx.GraphModule], None]] = ( @@ -385,6 +384,27 @@ def _populate_trt_builder_config( self.compilation_settings.l2_limit_for_tiling ) + build_route = getattr(self.compilation_settings, "build_route", "") or "" + if build_route: + if not hasattr(builder_config, "build_route"): + raise RuntimeError( + f"build_route={build_route} was requested, but this TensorRT " + "build does not expose IBuilderConfig.build_route " + "(Global Performance Tuner unavailable). This feature is available " + "since TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows." + ) + all_routes = getattr(builder_config, "all_build_routes", "") or "" + if not all_routes.strip(): + raise RuntimeError( + f"build_route={build_route} was requested, but " + "IBuilderConfig.all_build_routes is empty " + "(Global Performance Tuner disabled on this platform/build). " + "This feature is available since TensorRT 11.1 and is currently " + "not available in TensorRT-RTX or Windows." + ) + _LOGGER.info(f"Using TensorRT build route: {build_route}") + builder_config.build_route = build_route + return builder_config def _create_timing_cache( diff --git a/py/torch_tensorrt/dynamo/conversion/_conversion.py b/py/torch_tensorrt/dynamo/conversion/_conversion.py index f644c3e2ba..2650707a13 100644 --- a/py/torch_tensorrt/dynamo/conversion/_conversion.py +++ b/py/torch_tensorrt/dynamo/conversion/_conversion.py @@ -4,6 +4,7 @@ import logging from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple +import tensorrt as trt import torch from torch_tensorrt._enums import dtype from torch_tensorrt._features import ENABLED_FEATURES @@ -25,8 +26,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - logger = logging.getLogger(__name__) @@ -217,6 +216,48 @@ def interpret_module_to_result( Returns: SerializedInterpreterResult """ + from torch_tensorrt.dynamo.tuning import ( + gpt_settings_requested, + require_global_perf_tuner, + should_run_tuning, + tune_subgraph, + validate_tuning_options, + ) + + if gpt_settings_requested(settings) or should_run_tuning(settings): + require_global_perf_tuner("Requested Global Performance Tuner settings") + validate_tuning_options(settings) + + if should_run_tuning(settings): + return tune_subgraph( + module, + inputs, + settings, + engine_cache=engine_cache, + input_binding_names=input_binding_names, + output_binding_names=output_binding_names, + ) + + return _interpret_module_to_result_impl( + module, + inputs, + settings, + engine_cache, + input_binding_names=input_binding_names, + output_binding_names=output_binding_names, + ) + + +def _interpret_module_to_result_impl( + module: torch.fx.GraphModule, + inputs: Sequence[Input], + settings: CompilationSettings = CompilationSettings(), + engine_cache: Optional[BaseEngineCache] = None, + *, + input_binding_names: Optional[Sequence[str]] = None, + output_binding_names: Optional[Sequence[str]] = None, +) -> SerializedInterpreterResult: + """Interpret an FX module to a TRTInterpreterResult (single build, no GPT sweep).""" symbolic_shape_expressions = extract_symbolic_shape_expressions(module) if symbolic_shape_expressions is None: diff --git a/py/torch_tensorrt/dynamo/tuning/__init__.py b/py/torch_tensorrt/dynamo/tuning/__init__.py new file mode 100644 index 0000000000..84c4fe9fda --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/__init__.py @@ -0,0 +1,61 @@ +"""TensorRT Global Performance Tuner support for Torch-TensorRT Dynamo.""" + +from torch_tensorrt.dynamo.tuning._capability import ( + get_all_build_routes, + get_all_build_routes_raw, + gpt_settings_requested, + is_global_perf_tuner_available, + require_global_perf_tuner, +) +from torch_tensorrt.dynamo.tuning.accuracy import ( + compute_output_losses, + compute_tensor_loss, + loss_cos, + loss_l0, + loss_l1, + loss_l2, + loss_linf, +) +from torch_tensorrt.dynamo.tuning.cache import ( + resolve_partition_tuning_cache_path, + subgraph_partition_key, +) +from torch_tensorrt.dynamo.tuning.routes import ( + BuildRouteExprParser, + BuildRouteKnobDatabase, + expand_build_routes, + expand_routes_fast, + expand_routes_full, + expand_routes_mixed, +) +from torch_tensorrt.dynamo.tuning.sweeper import ( + should_run_tuning, + tune_subgraph, + validate_tuning_options, +) + +__all__ = [ + "BuildRouteExprParser", + "BuildRouteKnobDatabase", + "compute_output_losses", + "compute_tensor_loss", + "expand_build_routes", + "expand_routes_fast", + "expand_routes_full", + "expand_routes_mixed", + "get_all_build_routes", + "get_all_build_routes_raw", + "gpt_settings_requested", + "is_global_perf_tuner_available", + "loss_cos", + "loss_l0", + "loss_l1", + "loss_l2", + "loss_linf", + "require_global_perf_tuner", + "resolve_partition_tuning_cache_path", + "should_run_tuning", + "subgraph_partition_key", + "tune_subgraph", + "validate_tuning_options", +] diff --git a/py/torch_tensorrt/dynamo/tuning/_capability.py b/py/torch_tensorrt/dynamo/tuning/_capability.py new file mode 100644 index 0000000000..e7f7a117fc --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/_capability.py @@ -0,0 +1,102 @@ +"""Global Performance Tuner helpers for Torch-TensorRT Dynamo.""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, Optional + +import tensorrt as trt + +_LOGGER = logging.getLogger(__name__) + + +def is_global_perf_tuner_available() -> bool: + """Return True if TensorRT exposes Global Performance Tuner build-route APIs.""" + if not hasattr(trt.IBuilderConfig, "build_route") or not hasattr( + trt.IBuilderConfig, "all_build_routes" + ): + return False + try: + builder = trt.Builder(trt.Logger(trt.Logger.WARNING)) + config = builder.create_builder_config() + routes = getattr(config, "all_build_routes", "") or "" + return bool(routes.strip()) + except Exception as exc: # pragma: no cover - depends on local TRT/CUDA + _LOGGER.debug(f"Global Performance Tuner probe failed: {exc}") + return False + + +def require_global_perf_tuner(reason: str) -> None: + """Raise if GPT is unavailable when the user requested a GPT feature.""" + if not is_global_perf_tuner_available(): + raise RuntimeError( + f"{reason} requires TensorRT Global Performance Tuner " + "(IBuilderConfig.build_route / all_build_routes). " + "This feature is available since TensorRT 11.1 and is currently not available in TensorRT-RTX or Windows." + ) + + +def get_all_build_routes_raw() -> str: + """Return the raw JSON string from ``IBuilderConfig.all_build_routes``.""" + require_global_perf_tuner("Querying build routes") + builder = trt.Builder(trt.Logger(trt.Logger.WARNING)) + config = builder.create_builder_config() + return config.all_build_routes or "" + + +def get_all_build_routes(knob: Optional[str] = None) -> Dict[str, Any]: + """Parse ``all_build_routes`` JSON, optionally filtering to one knob. + + Args: + knob: Optional knob name (with or without leading "-"), matching "trtexec --helpBuildRoute[=knob]". + + Returns: + Parsed knob database dict with "tuner_version" and "tuner_options". + """ + raw = get_all_build_routes_raw() + try: + root: Dict[str, Any] = json.loads(raw) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Failed to parse all_build_routes JSON from TensorRT: {exc}" + ) from exc + + if knob is None or knob == "": + return root + + wanted = knob[1:] if knob.startswith("-") else knob + options = root.get("tuner_options", []) + filtered = [] + for opt in options: + name = opt.get("option", "") + bare = name[1:] if isinstance(name, str) and name.startswith("-") else name + if bare == wanted: + filtered.append(opt) + + if not filtered: + raise ValueError( + f"No such knob in the Global Performance Tuner database: {knob}. " + "Call get_all_build_routes() without a filter to list knobs." + ) + + out: Dict[str, Any] = {} + if "tuner_version" in root: + out["tuner_version"] = root["tuner_version"] + out["tuner_options"] = filtered + return out + + +def gpt_settings_requested(settings: Any) -> bool: + """True if CompilationSettings requests any GPT feature.""" + if getattr(settings, "build_route", ""): + return True + if getattr(settings, "tune_build_routes", ""): + return True + if getattr(settings, "tune_build_route_file", None): + return True + if getattr(settings, "tuning_continue", False): + return True + if getattr(settings, "tuning_dry_run", False): + return True + return False diff --git a/py/torch_tensorrt/dynamo/tuning/accuracy.py b/py/torch_tensorrt/dynamo/tuning/accuracy.py new file mode 100644 index 0000000000..0191ff51c7 --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/accuracy.py @@ -0,0 +1,136 @@ +"""Accuracy loss metrics matching trtexec Global Performance Tuner validators.""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Sequence, Tuple, Union + +import torch +from torch_tensorrt.dynamo.utils import cosine_similarity + +TensorTree = Union[torch.Tensor, Sequence["TensorTree"], Dict[str, "TensorTree"]] + + +def _as_float_tensor(t: torch.Tensor) -> torch.Tensor: + return t.detach().float().flatten() + + +def loss_l0( + actual: torch.Tensor, + reference: torch.Tensor, + atol: float = 1e-5, + rtol: float = 1e-5, +) -> float: + """Fraction of elements outside ``atol + rtol * abs(ref)`` (PyTorch allclose).""" + a = _as_float_tensor(actual) + b = _as_float_tensor(reference) + if a.numel() == 0: + raise ValueError("Cannot compute L0 accuracy on empty tensors") + if a.shape != b.shape: + raise ValueError(f"Shape mismatch for L0: {tuple(a.shape)} vs {tuple(b.shape)}") + outside = torch.abs(a - b) > (atol + rtol * torch.abs(b)) + return float(outside.float().mean().item()) + + +def loss_l1(actual: torch.Tensor, reference: torch.Tensor) -> float: + a = _as_float_tensor(actual) + b = _as_float_tensor(reference) + if a.numel() == 0: + raise ValueError("Cannot compute L1 accuracy on empty tensors") + return float(torch.mean(torch.abs(a - b)).item()) + + +def loss_l2(actual: torch.Tensor, reference: torch.Tensor) -> float: + a = _as_float_tensor(actual) + b = _as_float_tensor(reference) + if a.numel() == 0: + raise ValueError("Cannot compute L2 accuracy on empty tensors") + return float(torch.mean((a - b) ** 2).item()) + + +def loss_linf(actual: torch.Tensor, reference: torch.Tensor) -> float: + a = _as_float_tensor(actual) + b = _as_float_tensor(reference) + if a.numel() == 0: + raise ValueError("Cannot compute LInf accuracy on empty tensors") + return float(torch.max(torch.abs(a - b)).item()) + + +def loss_cos(actual: torch.Tensor, reference: torch.Tensor) -> float: + """``1 - cosine_similarity`` (lower is better; 0 = perfect match).""" + return 1.0 - float(cosine_similarity(reference, actual)) + + +def compute_tensor_loss( + actual: torch.Tensor, + reference: torch.Tensor, + algorithm: str = "l0", + atol: float = 1e-5, + rtol: float = 1e-5, +) -> float: + algo = algorithm.lower() + if algo == "l0": + return loss_l0(actual, reference, atol=atol, rtol=rtol) + if algo == "l1": + return loss_l1(actual, reference) + if algo == "l2": + return loss_l2(actual, reference) + if algo in {"linf", "linfinity"}: + return loss_linf(actual, reference) + if algo in {"cos", "cosine"}: + return loss_cos(actual, reference) + raise ValueError( + f"Unknown accuracy_algorithm={algorithm}; " + "expected one of 'l0', 'l1', 'l2', 'linf', 'cos'." + ) + + +def _flatten_named_tensors( + tree: TensorTree, prefix: str = "output" +) -> List[Tuple[str, torch.Tensor]]: + """Flatten a nested tree of tensors into a list of (name, tensor) tuples.""" + if isinstance(tree, torch.Tensor): + return [(prefix, tree)] + if isinstance(tree, dict): + out: List[Tuple[str, torch.Tensor]] = [] + for k, v in tree.items(): + out.extend(_flatten_named_tensors(v, f"{prefix}.{k}")) + return out + if isinstance(tree, (list, tuple)): + seq_out: List[Tuple[str, torch.Tensor]] = [] + for i, v in enumerate(tree): + seq_out.extend(_flatten_named_tensors(v, f"{prefix}.{i}")) + return seq_out + raise TypeError(f"Unsupported output type for accuracy: {type(tree)}") + + +def compute_output_losses( + actual: TensorTree, + reference: TensorTree, + algorithm: str = "l0", + atol: float = 1e-5, + rtol: float = 1e-5, +) -> Dict[str, float]: + """Compute per-tensor accuracy loss for nested outputs.""" + actual_flat = _flatten_named_tensors(actual) + reference_flat = _flatten_named_tensors(reference) + if len(actual_flat) != len(reference_flat): + raise ValueError( + f"Output arity mismatch: actual={len(actual_flat)} ref={len(reference_flat)}" + ) + losses: Dict[str, float] = {} + for (aname, at), (rname, rt) in zip(actual_flat, reference_flat): + name = aname if aname == rname else f"{aname}/{rname}" + losses[name] = compute_tensor_loss( + at, rt, algorithm=algorithm, atol=atol, rtol=rtol + ) + return losses + + +def accuracy_failed( + losses: Dict[str, float], + threshold: Optional[float], +) -> bool: + """Return True if any tensor accuracy loss exceeds the threshold.""" + if threshold is None: + return False + return any(v > threshold for v in losses.values()) diff --git a/py/torch_tensorrt/dynamo/tuning/cache.py b/py/torch_tensorrt/dynamo/tuning/cache.py new file mode 100644 index 0000000000..93f149efe1 --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/cache.py @@ -0,0 +1,122 @@ +"""JSONL tuning cache (trtexec-inspired header + per-iteration lines).""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import torch + +_LOGGER = logging.getLogger(__name__) + + +@dataclass +class TuningCacheHeader: + """Header of a tuning cache file.""" + + argv_like: Dict[str, Any] + tuning_expr: str + completed_iterations: int + tuner_version: str = "unknown" + + +def subgraph_partition_key(module: torch.fx.GraphModule) -> str: + digest = hashlib.sha256() + + for node in module.graph.nodes: + payload = { + "op": node.op, + "target": str(node.target), + "name": node.name, + "args": str(node.args), + "kwargs": str(node.kwargs), + } + digest.update(repr(payload).encode("utf-8")) + digest.update(b"\n") + + return digest.hexdigest()[:16] + + +def resolve_partition_tuning_cache_path( + base_path: Optional[str], + module: torch.fx.GraphModule, +) -> Optional[str]: + """Derive a per-partition cache path so multi-subgraph sweeps do not clobber. + + ``/tmp/tune.jsonl`` for partition ``abcd1234ef56`` becomes + ``/tmp/tune.abcd1234ef56.jsonl``. + """ + if not base_path: + return None + root, ext = os.path.splitext(base_path) + if not ext: + ext = ".jsonl" + key = subgraph_partition_key(module) + return f"{root}.{key}{ext}" + + +def write_header(path: str, header: Dict[str, Any]) -> None: + """Write the header of a tuning cache file.""" + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + f.write(json.dumps(header, sort_keys=False) + "\n") + + +def append_iteration( + path: str, + *, + iter_idx: int, + build_route: str, + crashed: bool, + error_message: str = "", + accuracy_loss: Optional[Dict[str, float]] = None, + gpu_time_ms: Optional[float] = None, +) -> None: + """Append an iteration to a tuning cache file.""" + row: Dict[str, Any] = { + "iter": iter_idx, + "build_route": build_route, + "crash": crashed, + "error_message": error_message, + "accuracy_loss": None if crashed or accuracy_loss is None else accuracy_loss, + "gpu_time": None if crashed or gpu_time_ms is None else gpu_time_ms, + } + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(row) + "\n") + + +def read_cache(path: str) -> TuningCacheHeader: + """Read the header of a tuning cache file.""" + if not os.path.isfile(path): + raise FileNotFoundError(f"tuning_cache_file not found: {path}") + with open(path, "r", encoding="utf-8") as f: + lines = [ln.strip() for ln in f.readlines() if ln.strip()] + if not lines: + raise ValueError(f"Empty tuning cache file: {path}") + header = json.loads(lines[0]) + return TuningCacheHeader( + argv_like=header, + tuning_expr=header.get("tuning_expr", ""), + completed_iterations=max(0, len(lines) - 1), + tuner_version=header.get("tuner_version", "unknown"), + ) + + +def read_iteration_gpu_times(path: str, max_iters: int) -> List[Optional[float]]: + """Read the GPU times of a tuning cache file.""" + times: List[Optional[float]] = [] + with open(path, "r", encoding="utf-8") as f: + lines = [ln.strip() for ln in f.readlines() if ln.strip()] + for line in lines[1 : 1 + max_iters]: + row = json.loads(line) + if row.get("crash"): + times.append(None) + else: + times.append(row.get("gpu_time")) + return times diff --git a/py/torch_tensorrt/dynamo/tuning/routes.py b/py/torch_tensorrt/dynamo/tuning/routes.py new file mode 100644 index 0000000000..b786e114f3 --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/routes.py @@ -0,0 +1,366 @@ +"""Build-route expression parsing and expansion (trtexec sampleTuning parity).""" + +from __future__ import annotations + +import itertools +import json +import re +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Tuple + +_BRACKET_VALUES_RE = re.compile(r"\[([^\]]*)\]") + + +@dataclass +class BuildRouteKnobDef: + option: str + allowed_values: str + default_value: str + help: str + values: List[str] = field(default_factory=list) + is_bounded: bool = False + + +@dataclass +class BuildRouteParsedExpr: + knob_name: str + values: List[str] + is_fixed: bool = False + + +class BuildRouteKnobDatabase: + """Knob definitions loaded from ``IBuilderConfig.all_build_routes`` JSON.""" + + def __init__(self) -> None: + self._knobs: Dict[str, BuildRouteKnobDef] = {} + self._knob_order: List[str] = [] + self.tuner_version: str = "unknown" + + def load_from_json(self, json_str: str) -> bool: + self._knobs.clear() + self._knob_order.clear() + self.tuner_version = "unknown" + if not json_str or not json_str.strip(): + return False + try: + root = json.loads(json_str) + except json.JSONDecodeError: + return False + + if isinstance(root.get("tuner_version"), str): + self.tuner_version = root["tuner_version"] + + options = root.get("tuner_options") + if not isinstance(options, list): + return False + + for item in options: + if not isinstance(item, dict): + continue + option = item.get("option", "") + if not option: + continue + allowed = item.get("allowed_values", "") or "" + values = self.parse_allowed_values(allowed) + knob = BuildRouteKnobDef( + option=option, + allowed_values=allowed, + default_value=str(item.get("default_value", "")), + help=str(item.get("help", "")), + values=values, + is_bounded=bool(values), + ) + self._knob_order.append(option) + self._knobs[option] = knob + return bool(self._knobs) + + @staticmethod + def parse_allowed_values(allowed_str: str) -> List[str]: + match = _BRACKET_VALUES_RE.search(allowed_str) + if not match: + return [] + values_str = match.group(1) + if "..." in values_str: # "..." is open-ended range marker + return [] + return [v.strip() for v in values_str.split("|") if v.strip()] + + def has_knob(self, knob_name: str) -> bool: + return knob_name in self._knobs + + def get_knob(self, knob_name: str) -> Optional[BuildRouteKnobDef]: + return self._knobs.get(knob_name) + + def get_default_value(self, knob_name: str) -> str: + knob = self.get_knob(knob_name) + return knob.default_value if knob is not None else "" + + def validate_values(self, knob_name: str, values: Sequence[str]) -> bool: + knob = self.get_knob(knob_name) + if knob is None: + return False + if not knob.is_bounded: + return True + allowed = set(knob.values) + return all(v in allowed for v in values) + + def build_default_path(self) -> str: + parts = [] + for name in self._knob_order: + knob = self._knobs[name] + parts.append(f"{knob.option}={knob.default_value}") + return " ".join(parts) + + +class BuildRouteExprParser: + """Parse ``-knob=[a|b] -fixed=on`` expressions against a knob database.""" + + def __init__(self, db: BuildRouteKnobDatabase) -> None: + self._db = db + self.error: str = "" + + def parse(self, input_str: str) -> Optional[List[BuildRouteParsedExpr]]: + self.error = "" + if not input_str or not input_str.strip(): + self.error = "Empty input" + return None + tokens = self.tokenize(input_str) + if not tokens: + self.error = "No expressions found" + return None + result: List[BuildRouteParsedExpr] = [] + for token in tokens: + expr = self._parse_expr(token) + if expr is None: + return None + result.append(expr) + return result + + @staticmethod + def tokenize(input_str: str) -> List[str]: + tokens: List[str] = [] + current: List[str] = [] + bracket_depth = 0 + for c in input_str: + if c == "[": + bracket_depth += 1 + current.append(c) + elif c == "]": + bracket_depth -= 1 + current.append(c) + elif c == " " and bracket_depth == 0: + if current: + tokens.append("".join(current)) + current = [] + else: + current.append(c) + if current: + tokens.append("".join(current)) + return tokens + + def _parse_expr(self, expr: str) -> Optional[BuildRouteParsedExpr]: + eq_pos = expr.find("=") + if eq_pos < 0: + self.error = f"Invalid expression (no '='): {expr}" + return None + knob_name = expr[:eq_pos].strip() + if not self._db.has_knob(knob_name): + self.error = f"Unknown knob: {knob_name}" + return None + value_str = expr[eq_pos + 1 :].strip() + if value_str.startswith("[") and value_str.endswith("]"): + inner = value_str[1:-1] + values = [v.strip() for v in inner.split("|") if v.strip()] + if not values: + self.error = f"Empty value list for knob: {knob_name}" + return None + if not self._db.validate_values(knob_name, values): + self.error = f"Invalid values for knob {knob_name}: {values}" + return None + return BuildRouteParsedExpr( + knob_name=knob_name, values=values, is_fixed=False + ) + if not self._db.validate_values(knob_name, [value_str]): + # Fixed values may still be legal defaults; allow if knob unbounded + knob = self._db.get_knob(knob_name) + if knob is None or (knob.is_bounded and value_str not in knob.values): + self.error = f"Invalid fixed value for knob {knob_name}: {value_str}" + return None + return BuildRouteParsedExpr( + knob_name=knob_name, values=[value_str], is_fixed=True + ) + + +def _route_from_assignment(names: Sequence[str], values: Sequence[str]) -> str: + return " ".join(f"{n}={v}" for n, v in zip(names, values)) + + +def expand_routes_full(exprs: Sequence[BuildRouteParsedExpr]) -> List[str]: + """Cartesian product over variable knobs (trtexec ``full``).""" + names = [e.knob_name for e in exprs] + value_lists = [e.values for e in exprs] + return [ + _route_from_assignment(names, combo) + for combo in itertools.product(*value_lists) + ] + + +def expand_routes_fast( + exprs: Sequence[BuildRouteParsedExpr], db: BuildRouteKnobDatabase +) -> List[str]: + """Baseline (defaults) + one-knob-at-a-time variants (trtexec ``fast``).""" + names = [e.knob_name for e in exprs] + defaults: List[str] = [] + for e in exprs: + if e.is_fixed: + defaults.append(e.values[0]) + else: + default = db.get_default_value(e.knob_name) + if default in e.values: + defaults.append(default) + else: + defaults.append(e.values[0]) + + routes = [_route_from_assignment(names, defaults)] + seen = set(routes) + + for i, e in enumerate(exprs): + if e.is_fixed: + continue + for val in e.values: + if val == defaults[i]: + continue + combo = list(defaults) + combo[i] = val + route = _route_from_assignment(names, combo) + if route not in seen: + seen.add(route) + routes.append(route) + return routes + + +def expand_routes_mixed( + exprs: Sequence[BuildRouteParsedExpr], + db: BuildRouteKnobDatabase, + positive_knob_indices: Sequence[int], +) -> List[str]: + """Exhaustive sweep over knobs that improved latency in the fast phase.""" + names = [e.knob_name for e in exprs] + defaults: List[str] = [] + for e in exprs: + if e.is_fixed: + defaults.append(e.values[0]) + else: + default = db.get_default_value(e.knob_name) + if default in e.values: + defaults.append(default) + else: + defaults.append(e.values[0]) + + positive = set(positive_knob_indices) + value_lists: List[List[str]] = [] + for i, e in enumerate(exprs): + if e.is_fixed or i not in positive: + value_lists.append([defaults[i]]) + else: + value_lists.append(list(e.values)) + + routes = [ + _route_from_assignment(names, combo) + for combo in itertools.product(*value_lists) + ] + # Prefer stable unique order + seen = set() + unique: List[str] = [] + for r in routes: + if r not in seen: + seen.add(r) + unique.append(r) + return unique + + +def load_tuning_expr_from_file(path: str) -> str: + with open(path, "r", encoding="utf-8") as f: + lines = [line.strip() for line in f.readlines()] + return " ".join(line for line in lines if line) + + +def resolve_tuning_expression( + tune_build_routes: str = "", + tune_build_route_file: Optional[str] = None, +) -> str: + if tune_build_routes and tune_build_route_file: + raise ValueError( + "Cannot specify both tune_build_routes and tune_build_route_file." + ) + if tune_build_route_file: + return load_tuning_expr_from_file(tune_build_route_file) + return tune_build_routes or "" + + +def expand_build_routes( + expression: str, + search: str, + db: BuildRouteKnobDatabase, + *, + dry_run: bool = False, +) -> Tuple[List[BuildRouteParsedExpr], List[str]]: + """Parse and expand a tuning expression. + + Returns: + (parsed_exprs, route_strings). For ``mixed``, only the fast phase is + expanded here; the sweeper runs the second phase after measuring. + """ + search = search.lower() + if search not in {"fast", "full", "mixed"}: + raise ValueError( + f"Unknown tuning_search={search}; expected 'fast', 'full', or 'mixed'." + ) + if dry_run and search == "mixed": + raise ValueError("tuning_dry_run is incompatible with tuning_search='mixed'.") + + parser = BuildRouteExprParser(db) + exprs = parser.parse(expression) + if exprs is None: + raise ValueError(f"Failed to parse tune_build_routes: {parser.error}") + + if search == "full": + return exprs, expand_routes_full(exprs) + # fast and mixed phase-1 use the same expansion + return exprs, expand_routes_fast(exprs, db) + + +def identify_positive_knobs( + exprs: Sequence[BuildRouteParsedExpr], + gpu_times: Sequence[Optional[float]], + db: BuildRouteKnobDatabase, +) -> List[int]: + """Return indices of knobs whose one-off variants beat the baseline.""" + if not gpu_times or gpu_times[0] is None: + return [] + baseline = gpu_times[0] + defaults: List[str] = [] + for e in exprs: + if e.is_fixed: + defaults.append(e.values[0]) + else: + default = db.get_default_value(e.knob_name) + if default in e.values: + defaults.append(default) + else: + defaults.append(e.values[0]) + + positive: List[int] = [] + idx = 1 + for i, e in enumerate(exprs): + if e.is_fixed: + continue + for val in e.values: + if val == defaults[i]: + continue + if idx >= len(gpu_times): + return positive + t = gpu_times[idx] + if t is not None and t < baseline and i not in positive: + positive.append(i) + idx += 1 + return positive diff --git a/py/torch_tensorrt/dynamo/tuning/sweeper.py b/py/torch_tensorrt/dynamo/tuning/sweeper.py new file mode 100644 index 0000000000..c90b14a6b4 --- /dev/null +++ b/py/torch_tensorrt/dynamo/tuning/sweeper.py @@ -0,0 +1,365 @@ +"""In-process Global Performance Tuner sweep for Dynamo TRT subgraphs.""" + +from __future__ import annotations + +import logging +import statistics +import time +from dataclasses import replace +from typing import Any, List, Optional, Sequence, Tuple + +import torch +from torch_tensorrt._Input import Input +from torch_tensorrt.dynamo._engine_cache import BaseEngineCache +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.conversion._conversion import SerializedInterpreterResult +from torch_tensorrt.dynamo.tuning import cache as tuning_cache +from torch_tensorrt.dynamo.tuning._capability import ( + get_all_build_routes_raw, + require_global_perf_tuner, +) +from torch_tensorrt.dynamo.tuning.accuracy import ( + accuracy_failed, + compute_output_losses, +) +from torch_tensorrt.dynamo.tuning.routes import ( + BuildRouteKnobDatabase, + expand_build_routes, + expand_routes_mixed, + identify_positive_knobs, + resolve_tuning_expression, +) + +_LOGGER = logging.getLogger(__name__) + +_WARMUP_ITERS = 3 +_BENCH_ITERS = 10 + + +def _inputs_to_tensors( + inputs: Sequence[Input], device: torch.device +) -> List[torch.Tensor]: + tensors: List[torch.Tensor] = [] + for inp in inputs: + if getattr(inp, "torch_tensor", None) is not None: + t = inp.torch_tensor + elif inp.shape_mode == Input._ShapeMode.STATIC: + t = inp.example_tensor() + else: + t = inp.example_tensor("opt_shape") + tensors.append(t.to(device)) + return tensors + + +def _benchmark_callable( + fn: Any, + args: Sequence[torch.Tensor], + *, + warmup: int = _WARMUP_ITERS, + iters: int = _BENCH_ITERS, +) -> float: + """Return median GPU latency in milliseconds.""" + for _ in range(max(0, warmup)): + fn(*args) + torch.cuda.synchronize() + times: List[float] = [] + for _ in range(max(1, iters)): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn(*args) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + return float(statistics.median(times)) + + +def validate_tuning_options(settings: CompilationSettings) -> None: + """Validate GPT-related CompilationSettings (trtexec-like rules).""" + if settings.tune_build_routes and settings.tune_build_route_file: + raise ValueError( + "Cannot specify both tune_build_routes and tune_build_route_file." + ) + if settings.tuning_continue: + if not settings.tuning_cache_file: + raise ValueError("tuning_continue requires tuning_cache_file.") + if ( + settings.tune_build_routes + or settings.tune_build_route_file + or settings.tuning_dry_run + ): + raise ValueError( + "tuning_continue cannot be combined with tune_build_routes, " + "tune_build_route_file, or tuning_dry_run; recover the sweep " + "from tuning_cache_file." + ) + if settings.tuning_dry_run and settings.tuning_search == "mixed": + raise ValueError("tuning_dry_run is incompatible with tuning_search='mixed'.") + if settings.accuracy_atol != 1e-5 or settings.accuracy_rtol != 1e-5: + if settings.accuracy_algorithm.lower() != "l0": + raise ValueError( + "accuracy_atol/accuracy_rtol are only valid when accuracy_algorithm='l0'." + ) + + +def should_run_tuning(settings: CompilationSettings) -> bool: + if settings.tuning_continue: + return True + if settings.tune_build_routes or settings.tune_build_route_file: + return True + return False + + +def tune_subgraph( + module: torch.fx.GraphModule, + inputs: Sequence[Input], + settings: CompilationSettings, + engine_cache: Optional[BaseEngineCache] = None, + *, + input_binding_names: Optional[Sequence[str]] = None, + output_binding_names: Optional[Sequence[str]] = None, +) -> SerializedInterpreterResult: + """Sweep build routes and return the best SerializedInterpreterResult.""" + from torch_tensorrt.dynamo.conversion._conversion import ( + _interpret_module_to_result_impl, + ) + from torch_tensorrt.dynamo.runtime import TorchTensorRTModule + + require_global_perf_tuner("Global Performance Tuner sweep") + validate_tuning_options(settings) + + # Per-partition cache path so multi-subgraph compiles do not overwrite each other. + user_cache_file = settings.tuning_cache_file + cache_file = tuning_cache.resolve_partition_tuning_cache_path( + user_cache_file, module + ) + if cache_file and cache_file != user_cache_file: + _LOGGER.info( + f"Using per-partition tuning cache {cache_file} (from {user_cache_file})", + ) + + db = BuildRouteKnobDatabase() + raw = get_all_build_routes_raw() + if not db.load_from_json(raw): + raise RuntimeError("Failed to load Global Performance Tuner knob database.") + + start_iter = 0 + if settings.tuning_continue: + assert cache_file is not None + header = tuning_cache.read_cache(cache_file) + expression = header.tuning_expr + start_iter = header.completed_iterations + _LOGGER.info( + f"Resuming tuning from {cache_file} at iteration {start_iter}", + ) + else: + expression = resolve_tuning_expression( + settings.tune_build_routes, settings.tune_build_route_file + ) + if not expression: + raise ValueError("tune_build_routes expression is empty.") + + exprs, routes = expand_build_routes( + expression, + settings.tuning_search, + db, + dry_run=settings.tuning_dry_run, + ) + + if settings.tuning_dry_run: + for i, route in enumerate(routes): + _LOGGER.info(f"[Tuning Dry Run] iter={i} BuildRoute = '{route}'") + raise RuntimeError( + f"tuning_dry_run enumerated {len(routes)} build routes; " + "no engines were built." + ) + + if cache_file and not settings.tuning_continue: + tuning_cache.write_header( + cache_file, + { + "tuner_version": db.tuner_version, + "accuracy_algorithm": settings.accuracy_algorithm, + "accuracy_parameter": { + "atol": settings.accuracy_atol, + "rtol": settings.accuracy_rtol, + "epsilon": settings.accuracy_threshold, + }, + "searching_algorithm": settings.tuning_search, + "tuning_expr": expression, + "default_build_route": db.build_default_path(), + "partition_key": tuning_cache.subgraph_partition_key(module), + "user_tuning_cache_file": user_cache_file, + }, + ) + + device = torch.device( + f"cuda:{settings.device.gpu_id}" + if settings.device.gpu_id is not None + else "cuda" + ) + sample_args = _inputs_to_tensors(inputs, device) + module = module.to(device).eval() + with torch.no_grad(): + ref_outputs = module(*sample_args) + + best_result = None + best_time: Optional[float] = None + best_route = "" + gpu_times: List[Optional[float]] = [] + sweep_start = time.monotonic() + + def _trial( + route: str, iter_idx: int, *, record_cache: bool = True + ) -> Tuple[Optional[Any], Optional[float]]: + nonlocal best_result, best_time, best_route + _LOGGER.info(f"&&&& TASK_BEGIN [iter={iter_idx}] BuildRoute = '{route}'") + trial_settings = replace( + settings, + build_route=route, + tune_build_routes="", + tune_build_route_file=None, + tuning_continue=False, + tuning_dry_run=False, + reuse_cached_engines=False, + cache_built_engines=False, + ) + crashed = False + error_message = "" + accuracy_loss = None + gpu_time: Optional[float] = None + result = None + try: + result = _interpret_module_to_result_impl( + module, + inputs, + trial_settings, + engine_cache=None, + input_binding_names=input_binding_names, + output_binding_names=output_binding_names, + ) + trt_mod = TorchTensorRTModule( + serialized_engine=result.serialized_engine, + input_binding_names=list(result.input_names), + output_binding_names=list(result.output_names), + name=f"tune_iter_{iter_idx}", + settings=trial_settings, + requires_output_allocator=result.requires_output_allocator, + requires_native_multidevice=result.requires_native_multidevice, + symbolic_shape_expressions=result.symbolic_shape_expressions, + aliased_io=result.aliased_io, + ) + trt_mod.eval() + with torch.no_grad(): + actual = trt_mod(*sample_args) + if settings.accuracy_threshold is not None: + accuracy_loss = compute_output_losses( + actual, + ref_outputs, + algorithm=settings.accuracy_algorithm, + atol=settings.accuracy_atol, + rtol=settings.accuracy_rtol, + ) + if accuracy_failed(accuracy_loss, settings.accuracy_threshold): + error_message = f"accuracy threshold exceeded: {accuracy_loss}" + _LOGGER.warning( + "iter=%d route failed accuracy: %s", + iter_idx, + accuracy_loss, + ) + else: + gpu_time = _benchmark_callable(trt_mod, sample_args) + else: + gpu_time = _benchmark_callable(trt_mod, sample_args) + + if ( + gpu_time is not None + and not error_message + and (best_time is None or gpu_time < best_time) + ): + best_time = gpu_time + best_result = result + best_route = route + del trt_mod + torch.cuda.empty_cache() + _LOGGER.info(f"&&&& TASK_END [iter={iter_idx}] BuildRoute = '{route}'") + except Exception as exc: + crashed = True + _LOGGER.warning( + f"&&&& TASK_ABORT [iter={iter_idx}] BuildRoute = '{route}': {str(exc)}" + ) + + if record_cache and cache_file: + tuning_cache.append_iteration( + cache_file, + iter_idx=iter_idx, + build_route=route, + crashed=crashed, + error_message=error_message, + accuracy_loss=accuracy_loss, + gpu_time_ms=gpu_time, + ) + return result, gpu_time + + if settings.tuning_continue and start_iter > 0 and cache_file: + cached_times = tuning_cache.read_iteration_gpu_times(cache_file, start_iter) + gpu_times.extend(cached_times) + best_cached_idx = None + best_cached_time: Optional[float] = None + for idx, t in enumerate(cached_times): + if t is not None and (best_cached_time is None or t < best_cached_time): + best_cached_time = t + best_cached_idx = idx + if best_cached_idx is not None and best_cached_idx < len(routes): + _LOGGER.info( + f"Rebuilding best cached route from iter {best_cached_idx} for resume", + ) + # TODO (@Evan): Consider rebuilding the cached best engine later if it's still the best after the remaining trials. + # Need to think about how to deal with timeout. + _trial(routes[best_cached_idx], best_cached_idx, record_cache=False) + + for i, route in enumerate(routes): + if i < start_iter: + continue + if settings.tuning_timeout_s >= 0: + elapsed = time.monotonic() - sweep_start + if elapsed >= settings.tuning_timeout_s: + _LOGGER.info( + f"Tuning timeout reached after {elapsed:.1f}s; stopping before iter {i}", + ) + break + _, gpu_time = _trial(route, i) + gpu_times.append(gpu_time) + + if settings.tuning_search == "mixed" and start_iter < len(routes): + # Only run phase-2 if we completed (or nearly) phase-1 + if len(gpu_times) >= len(routes): + positive = identify_positive_knobs(exprs, gpu_times[: len(routes)], db) + phase2 = expand_routes_mixed(exprs, db, positive) + # Skip routes already evaluated in phase 1 + phase1_set = set(routes) + phase2_new = [r for r in phase2 if r not in phase1_set] + base_idx = len(routes) + for j, route in enumerate(phase2_new): + iter_idx = base_idx + j + if settings.tuning_timeout_s >= 0: + elapsed = time.monotonic() - sweep_start + if elapsed >= settings.tuning_timeout_s: + _LOGGER.info( + f"Tuning timeout reached during mixed phase-2 at iter {iter_idx}", + ) + break + _trial(route, iter_idx) + + if best_result is None: + raise RuntimeError( + "Global Performance Tuner sweep completed without a valid engine " + "(all routes crashed or failed accuracy checks)." + ) + + _LOGGER.info( + f"Selected best build route '{best_route}' with gpu_time={best_time if best_time is not None else float('nan'):.3f} ms" + ) + # TODO (@Evan): Consider persisting winner on settings for engine-cache hashing / introspection + settings.build_route = best_route + return best_result diff --git a/tests/py/dynamo/runtime/test_global_perf_tuner.py b/tests/py/dynamo/runtime/test_global_perf_tuner.py new file mode 100644 index 0000000000..93bcefa8aa --- /dev/null +++ b/tests/py/dynamo/runtime/test_global_perf_tuner.py @@ -0,0 +1,264 @@ +"""Unit tests for Global Performance Tuner route parsing and accuracy metrics.""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest + +import torch +from torch_tensorrt.dynamo._settings import CompilationSettings, settings_are_compatible +from torch_tensorrt.dynamo.tuning.accuracy import ( + compute_tensor_loss, + loss_cos, + loss_l0, + loss_l1, + loss_l2, + loss_linf, +) +from torch_tensorrt.dynamo.tuning.routes import ( + BuildRouteKnobDatabase, + expand_build_routes, + expand_routes_mixed, + identify_positive_knobs, + resolve_tuning_expression, +) +from torch_tensorrt.dynamo.tuning.sweeper import validate_tuning_options + + +def _sample_knob_db_json() -> str: + return json.dumps( + { + "tuner_version": "test-1.0", + "tuner_options": [ + { + "option": "-slice_fusion", + "allowed_values": "-slice_fusion=[on|off]", + "default_value": "on", + "help": "slice fusion", + }, + { + "option": "-copy_ppg", + "allowed_values": "-copy_ppg=[on|off]", + "default_value": "on", + "help": "copy ppg", + }, + { + "option": "-kgen:codegen:cuda_tile", + "allowed_values": "-kgen:codegen:cuda_tile=[0|1|2|3]", + "default_value": "1", + "help": "cuda tile", + }, + ], + } + ) + + +class TestBuildRouteParsing(unittest.TestCase): + def setUp(self) -> None: + self.db = BuildRouteKnobDatabase() + assert self.db.load_from_json(_sample_knob_db_json()) + + def test_full_expansion_two_binary(self) -> None: + exprs, routes = expand_build_routes( + "-slice_fusion=[on|off] -copy_ppg=[on|off]", "full", self.db + ) + self.assertEqual(len(exprs), 2) + self.assertEqual(len(routes), 4) + self.assertIn("-slice_fusion=on -copy_ppg=on", routes) + self.assertIn("-slice_fusion=off -copy_ppg=off", routes) + + def test_fast_expansion_linear(self) -> None: + exprs, routes = expand_build_routes( + "-slice_fusion=[on|off] -copy_ppg=[on|off]", "fast", self.db + ) + # baseline + one off for each binary knob = 3 + self.assertEqual(len(routes), 3) + self.assertEqual(routes[0], "-slice_fusion=on -copy_ppg=on") + + def test_dry_run_rejects_mixed(self) -> None: + with self.assertRaises(ValueError): + expand_build_routes( + "-slice_fusion=[on|off]", "mixed", self.db, dry_run=True + ) + + def test_unknown_knob(self) -> None: + with self.assertRaises(ValueError): + expand_build_routes("-not_a_real_knob=[on|off]", "fast", self.db) + + def test_fixed_and_variable(self) -> None: + exprs, routes = expand_build_routes( + "-slice_fusion=off -copy_ppg=[on|off]", "full", self.db + ) + self.assertEqual(len(routes), 2) + self.assertTrue(all(r.startswith("-slice_fusion=off") for r in routes)) + + def test_identify_positive_knobs(self) -> None: + exprs, routes = expand_build_routes( + "-slice_fusion=[on|off] -copy_ppg=[on|off]", "fast", self.db + ) + self.assertEqual(len(routes), 3) + # baseline slow, first one-off faster -> positive slice_fusion + gpu_times = [10.0, 5.0, 11.0] + positive = identify_positive_knobs(exprs, gpu_times, self.db) + self.assertEqual(positive, [0]) + mixed = expand_routes_mixed(exprs, self.db, positive) + self.assertTrue(any("slice_fusion=off" in r for r in mixed)) + + def test_identify_positive_knobs_2(self) -> None: + exprs, routes = expand_build_routes( + "-slice_fusion=[on|off] -copy_ppg=[on|off]", "fast", self.db + ) + self.assertEqual(len(routes), 3) + # baseline slow, first one-off faster -> positive slice_fusion + gpu_times = [10.0, 5.0, 3.0] + positive = identify_positive_knobs(exprs, gpu_times, self.db) + self.assertEqual(positive, [0, 1]) + mixed = expand_routes_mixed(exprs, self.db, positive) + self.assertTrue(any("slice_fusion=off" in r for r in mixed)) + self.assertTrue(any("copy_ppg=off" in r for r in mixed)) + + +class TestAccuracyMetrics(unittest.TestCase): + def test_perfect_match_zero_loss(self) -> None: + t = torch.randn(8, 8) + self.assertEqual(loss_l0(t, t), 0.0) + self.assertEqual(loss_l1(t, t), 0.0) + self.assertEqual(loss_l2(t, t), 0.0) + self.assertEqual(loss_linf(t, t), 0.0) + self.assertAlmostEqual(loss_cos(t, t), 0.0, places=5) + + def test_l0_fraction(self) -> None: + ref = torch.zeros(4) + actual = torch.tensor([0.0, 0.0, 1.0, 1.0]) + # atol=rtol=0 => half the elements differ + self.assertAlmostEqual(loss_l0(actual, ref, atol=0.0, rtol=0.0), 0.5) + + def test_algorithm_dispatch(self) -> None: + a = torch.ones(3) + b = torch.zeros(3) + self.assertGreater(compute_tensor_loss(a, b, "l1"), 0.0) + self.assertGreater(compute_tensor_loss(a, b, "lInf"), 0.0) + + +class TestSettingsAndValidation(unittest.TestCase): + def test_build_route_engine_invariant(self) -> None: + a = CompilationSettings(build_route="") + b = CompilationSettings(build_route="-slice_fusion=off") + ok, incompatible = settings_are_compatible(a, b) + self.assertFalse(ok) + self.assertIn("build_route", incompatible) + + def test_validate_mutually_exclusive_exprs(self) -> None: + with self.assertRaises(ValueError): + validate_tuning_options( + CompilationSettings( + tune_build_routes="-a=[on|off]", + tune_build_route_file="/tmp/x.txt", + ) + ) + + def test_validate_continue_requires_cache(self) -> None: + with self.assertRaises(ValueError): + validate_tuning_options(CompilationSettings(tuning_continue=True)) + + def test_validate_dry_run_mixed(self) -> None: + with self.assertRaises(ValueError): + validate_tuning_options( + CompilationSettings(tuning_dry_run=True, tuning_search="mixed") + ) + + def test_partition_cache_path(self) -> None: + from torch_tensorrt.dynamo.tuning.cache import ( + resolve_partition_tuning_cache_path, + subgraph_partition_key, + ) + + g = torch.fx.symbolic_trace(torch.nn.ReLU()) + key = subgraph_partition_key(g) + path = resolve_partition_tuning_cache_path("/tmp/tune.jsonl", g) + self.assertEqual(path, f"/tmp/tune.{key}.jsonl") + self.assertIsNone(resolve_partition_tuning_cache_path(None, g)) + + def test_tune_expr_from_file(self) -> None: + with tempfile.NamedTemporaryFile("w", delete=False) as f: + f.write("-slice_fusion=[on|off]\n") + f.write("-copy_ppg=[on|off]\n") + path = f.name + expr = resolve_tuning_expression(tune_build_route_file=path) + self.assertIn("-slice_fusion=[on|off]", expr) + self.assertIn("-copy_ppg=[on|off]", expr) + + +class TestGPTAvailabilityAndIntegration(unittest.TestCase): + def test_capability_probe(self) -> None: + from torch_tensorrt.dynamo.tuning import is_global_perf_tuner_available + + # Should not raise; result depends on local TensorRT build. + available = is_global_perf_tuner_available() + self.assertIsInstance(available, bool) + + def test_small_tune_sweep(self) -> None: + from torch_tensorrt.dynamo.tuning import ( + get_all_build_routes, + is_global_perf_tuner_available, + ) + + if not torch.cuda.is_available() or not is_global_perf_tuner_available(): + self.skipTest("Global Performance Tuner or CUDA unavailable") + + import torch_tensorrt + + knobs = get_all_build_routes() + options = knobs.get("tuner_options", []) + binary = None + for opt in options: + allowed = opt.get("allowed_values", "") + if "=[on|off]" in allowed: + binary = opt["option"] + break + if binary is None: + self.skipTest("No binary on/off knob found in tuner database") + + class Tiny(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x + 1.0) + + model = Tiny().eval().cuda() + x = torch.randn(1, 8, device="cuda") + expr = f"{binary}=[on|off]" + with tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) as f: + cache_base = f.name + + compiled = torch_tensorrt.compile( + model, + ir="dynamo", + arg_inputs=[x], + min_block_size=1, + tune_build_routes=expr, + tuning_search="full", + accuracy_threshold=0.5, + accuracy_algorithm="cos", + tuning_cache_file=cache_base, + ) + with torch.no_grad(): + out = compiled(x) + ref = model(x) + self.assertTrue(torch.allclose(out, ref, rtol=1e-2, atol=1e-2)) + + cache_dir = os.path.dirname(cache_base) or "." + cache_stem = os.path.splitext(os.path.basename(cache_base))[0] + partition_caches = [ + os.path.join(cache_dir, name) + for name in os.listdir(cache_dir) + if name.startswith(cache_stem + ".") and name.endswith(".jsonl") + ] + self.assertTrue(partition_caches, "expected per-partition tuning cache file") + with open(partition_caches[0], "r", encoding="utf-8") as f: + lines = [ln for ln in f.readlines() if ln.strip()] + self.assertGreaterEqual(len(lines), 3) # header + 2 iters + + +if __name__ == "__main__": + unittest.main()