from __future__ import annotations import argparse from pathlib import Path import tensorrt as trt def _parse_enum_list(value: str | None, enum_cls): if not value: return [] out = [] for raw in value.split(","): name = raw.strip().upper() if not name: continue if not hasattr(enum_cls, name): valid = ", ".join(x for x in dir(enum_cls) if x.isupper()) raise ValueError(f"Unknown {enum_cls.__name__} value {raw!r}; valid: {valid}") out.append(getattr(enum_cls, name)) return out def _enum_choice(value: str | None, enum_cls): if value is None: return None name = value.strip().upper() if not hasattr(enum_cls, name): valid = ", ".join(x.lower() for x in dir(enum_cls) if x.isupper()) raise ValueError(f"Unknown {enum_cls.__name__} value {value!r}; valid: {valid}") return getattr(enum_cls, name) def main(): parser = argparse.ArgumentParser() parser.add_argument("--onnx_path", required=True) parser.add_argument("--engine_path", required=True) parser.add_argument("--precision", choices=("none", "fp16", "bf16", "bf16_fp16"), default="bf16_fp16") parser.add_argument("--workspace_gb", type=float, default=64.0) parser.add_argument("--builder_optimization_level", type=int, default=5) parser.add_argument("--max_aux_streams", type=int, default=None) parser.add_argument("--strongly_typed", action="store_true") parser.add_argument("--sparse_weights", action="store_true") parser.add_argument("--prefer_precision_constraints", action="store_true") parser.add_argument("--obey_precision_constraints", action="store_true") parser.add_argument("--direct_io", action="store_true") parser.add_argument("--no_tf32", action="store_true") parser.add_argument("--tactic_sources", type=str, default=None, help="Comma-separated TensorRT tactic sources, e.g. CUBLAS_LT,CUDNN,JIT_CONVOLUTIONS") parser.add_argument("--preview_features", type=str, default=None, help="Comma-separated TensorRT preview features to enable.") parser.add_argument("--tiling_optimization_level", type=str, default=None, help="TensorRT tiling optimization level: none, fast, moderate, or full.") parser.add_argument("--hardware_compatibility_level", type=str, default=None, help="TensorRT hardware compatibility level: none, ampere_plus, same_compute_capability.") parser.add_argument("--profiling_verbosity", type=str, default=None, help="TensorRT profiling verbosity: none, layer_names_only, detailed.") parser.add_argument("--timing_cache", type=str, default=None, help="Target-specific TensorRT timing-cache file to read and update.") parser.add_argument("--version_compatible", action="store_true", help="Build a version-compatible plan. Native exact-version plans remain the default.") parser.add_argument("--exclude_lean_runtime", action="store_true", help="Exclude the lean runtime from a version-compatible plan.") args = parser.parse_args() onnx_path = Path(args.onnx_path) engine_path = Path(args.engine_path) engine_path.parent.mkdir(parents=True, exist_ok=True) logger = trt.Logger(trt.Logger.INFO) trt.init_libnvinfer_plugins(logger, "") builder = trt.Builder(logger) network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) if args.strongly_typed: network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) network = builder.create_network(network_flags) parser = trt.OnnxParser(network, logger) if not parser.parse_from_file(str(onnx_path)): for idx in range(parser.num_errors): print(parser.get_error(idx)) raise RuntimeError(f"Failed to parse ONNX: {onnx_path}") config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(args.workspace_gb * (1024 ** 3))) if hasattr(config, "builder_optimization_level"): config.builder_optimization_level = args.builder_optimization_level if args.max_aux_streams is not None and hasattr(config, "max_aux_streams"): config.max_aux_streams = args.max_aux_streams if not args.strongly_typed: if args.precision in {"fp16", "bf16_fp16"}: config.set_flag(trt.BuilderFlag.FP16) if args.precision in {"bf16", "bf16_fp16"}: config.set_flag(trt.BuilderFlag.BF16) if args.sparse_weights: config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) if args.prefer_precision_constraints: config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS) if args.obey_precision_constraints: config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) if args.direct_io: config.set_flag(trt.BuilderFlag.DIRECT_IO) if args.no_tf32 and hasattr(config, "clear_flag"): config.clear_flag(trt.BuilderFlag.TF32) tactic_sources = _parse_enum_list(args.tactic_sources, trt.TacticSource) if tactic_sources: mask = 0 for source in tactic_sources: mask |= 1 << int(source) config.set_tactic_sources(mask) preview_features = _parse_enum_list(args.preview_features, trt.PreviewFeature) for feature in preview_features: config.set_preview_feature(feature, True) tiling_level = _enum_choice(args.tiling_optimization_level, trt.TilingOptimizationLevel) if tiling_level is not None and hasattr(config, "tiling_optimization_level"): config.tiling_optimization_level = tiling_level hardware_level = _enum_choice(args.hardware_compatibility_level, trt.HardwareCompatibilityLevel) if hardware_level is not None and hasattr(config, "hardware_compatibility_level"): config.hardware_compatibility_level = hardware_level profiling_verbosity = _enum_choice(args.profiling_verbosity, trt.ProfilingVerbosity) if profiling_verbosity is not None: config.profiling_verbosity = profiling_verbosity if args.version_compatible: config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE) if args.exclude_lean_runtime: if not args.version_compatible: raise ValueError("--exclude_lean_runtime requires --version_compatible") config.set_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME) timing_cache_path = Path(args.timing_cache) if args.timing_cache else None if timing_cache_path is not None: timing_blob = timing_cache_path.read_bytes() if timing_cache_path.is_file() else b"" timing_cache = config.create_timing_cache(timing_blob) if not config.set_timing_cache(timing_cache, ignore_mismatch=False): raise RuntimeError(f"TensorRT rejected timing cache: {timing_cache_path}") print(f"onnx_path={onnx_path}") print(f"engine_path={engine_path}") print(f"precision={args.precision}") print(f"workspace_gb={args.workspace_gb}") print(f"builder_optimization_level={args.builder_optimization_level}") print(f"max_aux_streams={args.max_aux_streams}") print(f"strongly_typed={args.strongly_typed}") print(f"sparse_weights={args.sparse_weights}") print(f"prefer_precision_constraints={args.prefer_precision_constraints}") print(f"obey_precision_constraints={args.obey_precision_constraints}") print(f"direct_io={args.direct_io}") print(f"no_tf32={args.no_tf32}") print(f"tactic_sources={args.tactic_sources}") print(f"preview_features={args.preview_features}") print(f"tiling_optimization_level={args.tiling_optimization_level}") print(f"hardware_compatibility_level={args.hardware_compatibility_level}") print(f"profiling_verbosity={args.profiling_verbosity}") print(f"timing_cache={args.timing_cache}") print(f"version_compatible={args.version_compatible}") print(f"exclude_lean_runtime={args.exclude_lean_runtime}") print(f"num_inputs={network.num_inputs}") for idx in range(network.num_inputs): tensor = network.get_input(idx) print(f"input[{idx}] name={tensor.name} shape={tuple(tensor.shape)} dtype={tensor.dtype}") print(f"num_outputs={network.num_outputs}") for idx in range(network.num_outputs): tensor = network.get_output(idx) print(f"output[{idx}] name={tensor.name} shape={tuple(tensor.shape)} dtype={tensor.dtype}") serialized = builder.build_serialized_network(network, config) if serialized is None: raise RuntimeError("TensorRT build_serialized_network returned None") engine_path.write_bytes(bytes(serialized)) print(f"engine_bytes={engine_path.stat().st_size}") if timing_cache_path is not None: timing_cache_path.parent.mkdir(parents=True, exist_ok=True) timing_cache_path.write_bytes(bytes(config.get_timing_cache().serialize())) print(f"timing_cache_bytes={timing_cache_path.stat().st_size}") if __name__ == "__main__": main()