Improved Examples

examples/benchmark.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
+#
+# MIT License
+"""
+File showcasing how to benchmark TensorRT engines.
+
+Demonstrates :func:`trtutils.benchmark_engine` for a single engine and
+:func:`trtutils.benchmark_engines` for side-by-side comparison. Builds an
+FP32 and FP16 variant of YOLOv8n and reports latency statistics for each.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from trtutils import benchmark_engine, benchmark_engines, build_engine, set_log_level
+from trtutils.download import download
+
+
+def main() -> None:
+    onnx_path = Path("/tmp/yolov8n.onnx")  # noqa: S108
+    fp32_engine = Path("/tmp/yolov8n_fp32.engine")  # noqa: S108
+    fp16_engine = Path("/tmp/yolov8n_fp16.engine")  # noqa: S108
+
+    if not onnx_path.exists():
+        print("Downloading yolov8n ONNX model...")
+        download("yolov8n", onnx_path, imgsz=640, simplify=True)
+
+    shapes = [("images", (1, 3, 640, 640))]
+    if not fp32_engine.exists():
+        print("Building FP32 engine...")
+        build_engine(onnx_path, fp32_engine, shapes=shapes)
+    if not fp16_engine.exists():
+        print("Building FP16 engine...")
+        build_engine(onnx_path, fp16_engine, fp16=True, shapes=shapes)
+
+    print("\nSingle-engine benchmarks:")
+    for label, path in [("FP32", fp32_engine), ("FP16", fp16_engine)]:
+        result = benchmark_engine(path, iterations=200, warmup_iterations=20)
+        m = result.latency
+        print(
+            f"  {label}: mean={m.mean:.3f} ms  median={m.median:.3f} ms  "
+            f"min={m.min:.3f} ms  max={m.max:.3f} ms"
+        )
+
+    print("\nbenchmark_engines (serial):")
+    serial = benchmark_engines([fp32_engine, fp16_engine], iterations=200, warmup_iterations=20)
+    for label, result in zip(["FP32", "FP16"], serial):
+        print(f"  {label}: mean={result.latency.mean:.3f} ms")
+
+    print("\nbenchmark_engines (parallel, both engines run in lockstep):")
+    parallel = benchmark_engines(
+        [fp32_engine, fp16_engine],
+        iterations=200,
+        warmup_iterations=20,
+        parallel=True,
+    )
+    print(f"  combined: mean={parallel[0].latency.mean:.3f} ms")
+
+
+if __name__ == "__main__":
+    set_log_level("ERROR")
+    main()

examples/build.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
+#
+# MIT License
+"""
+File showcasing how to build a TensorRT engine from an ONNX model.
+
+Demonstrates :func:`trtutils.builder.read_onnx` for peeking at the network
+before building, then :func:`trtutils.build_engine` to produce a serialized
+engine file. Bootstraps its own ONNX via :func:`trtutils.download.download`.
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+from trtutils import build_engine, set_log_level
+from trtutils.builder import read_onnx
+from trtutils.download import download
+
+
+def main() -> None:
+    onnx_path = Path("/tmp/yolov8n.onnx")  # noqa: S108
+    engine_path = Path("/tmp/yolov8n.engine")  # noqa: S108
+
+    if not onnx_path.exists():
+        print("Downloading yolov8n ONNX model...")
+        download("yolov8n", onnx_path, imgsz=640, simplify=True)
+
+    # peek at the parsed network before we build
+    network, _builder, _config, _parser = read_onnx(onnx_path)
+    print(f"ONNX network: {network.num_layers} layers")
+    for i in range(network.num_inputs):
+        t = network.get_input(i)
+        print(f"  input {t.name}: shape={tuple(t.shape)}, dtype={t.dtype}")
+    for i in range(network.num_outputs):
+        t = network.get_output(i)
+        print(f"  output {t.name}: shape={tuple(t.shape)}, dtype={t.dtype}")
+    # release the parser-side handles before building
+    del network, _builder, _config, _parser
+
+    if engine_path.exists():
+        engine_path.unlink()
+
+    t0 = time.perf_counter()
+    build_engine(
+        onnx_path,
+        engine_path,
+        fp16=True,
+        shapes=[("images", (1, 3, 640, 640))],
+    )
+    t1 = time.perf_counter()
+
+    size_mb = engine_path.stat().st_size / (1024 * 1024)
+    print(f"Built FP16 engine in {t1 - t0:.2f} s -> {engine_path} ({size_mb:.2f} MB)")
+
+
+if __name__ == "__main__":
+    set_log_level("ERROR")
+    main()

examples/builder/dla.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
+#
+# MIT License
+"""
+File showcasing how to build a TensorRT engine targeting the DLA.
+
+Demonstrates :func:`trtutils.builder.can_run_on_dla` to inspect which layers of
+an ONNX model are DLA-compatible, and :func:`trtutils.builder.build_dla_engine`
+to build a hybrid DLA/GPU engine. INT8 calibration is mandatory for DLA builds,
+so we feed a :class:`trtutils.builder.SyntheticBatcher`.
+
+Exits cleanly when the system has no DLA hardware.
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import numpy as np
+
+from trtutils import FLAGS, TRTEngine, set_log_level
+from trtutils.builder import SyntheticBatcher, build_dla_engine, can_run_on_dla
+from trtutils.download import download
+
+
+def main() -> None:
+    if not FLAGS.HAS_DLA:
+        print(f"Skipping: no DLA cores available (NUM_DLA_CORES={FLAGS.NUM_DLA_CORES}).")
+        return
+
+    onnx_path = Path("/tmp/yolov8n.onnx")  # noqa: S108
+    engine_path = Path("/tmp/yolov8n_dla.engine")  # noqa: S108
+
+    if not onnx_path.exists():
+        print("Downloading yolov8n ONNX model...")
+        download("yolov8n", onnx_path, imgsz=640, simplify=True)
+
+    full_dla, chunks = can_run_on_dla(onnx_path)
+    print(f"Fully DLA-compatible: {full_dla}")
+    print(f"Found {len(chunks)} layer chunks:")
+    for i, (layers, start, end, on_dla) in enumerate(chunks):
+        target = "DLA" if on_dla else "GPU"
+        print(f"  chunk {i}: layers [{start}-{end}] ({len(layers)} layers) -> {target}")
+
+    # DLA builds need INT8 calibration data; use synthetic data for the demo
+    batcher = SyntheticBatcher(
+        shape=(640, 640, 3),
+        dtype=np.float32,
+        batch_size=1,
+        num_batches=8,
+    )
+
+    if engine_path.exists():
+        engine_path.unlink()
+
+    t0 = time.perf_counter()
+    build_dla_engine(
+        onnx_path,
+        engine_path,
+        data_batcher=batcher,
+        dla_core=0,
+        shapes=[("images", (1, 3, 640, 640))],
+    )
+    t1 = time.perf_counter()
+
+    size_mb = engine_path.stat().st_size / (1024 * 1024)
+    print(f"Built DLA engine in {t1 - t0:.2f} s -> {engine_path} ({size_mb:.2f} MB)")
+
+    # confirm the engine is loadable; cuda_graph=False since DLA + graphs don't mix
+    engine = TRTEngine(engine_path, dla_core=0, warmup=True, cuda_graph=False)
+    engine.mock_execute()
+    print(f"Loaded {engine.name}, mock_execute OK")
+    del engine
+
+
+if __name__ == "__main__":
+    set_log_level("ERROR")
+    main()

examples/config.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
+#
+# MIT License
+"""
+File showcasing trtutils configuration, logging, and profiling toggles.
+
+Demonstrates :data:`trtutils.FLAGS`, :data:`trtutils.CONFIG`,
+:func:`trtutils.set_log_level`, the :data:`trtutils.NVTX` context manager,
+and :func:`trtutils.register_jit` / :data:`trtutils.JIT` for Numba JIT
+compilation. No model needed.
+"""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+
+from trtutils import (
+    CONFIG,
+    FLAGS,
+    JIT,
+    NVTX,
+    register_jit,
+    set_log_level,
+)
+
+
+@register_jit(fastmath=True)
+def sum_squares(arr: np.ndarray) -> float:
+    """Trivial numeric kernel — Numba JITs this when JIT is enabled."""
+    total = 0.0
+    for value in arr:
+        total += value * value
+    return float(total)
+
+
+def main() -> None:
+    print("FLAGS:")
+    for attr in sorted(
+        a for a in dir(FLAGS) if not a.startswith("_") and not callable(getattr(FLAGS, a))
+    ):
+        print(f"  {attr}: {getattr(FLAGS, attr)}")
+
+    print("\nCONFIG: loading TensorRT plugins (idempotent)...")
+    CONFIG.load_plugins()
+    print("CONFIG: plugins loaded.")
+
+    print("\nLog level demo — toggle between INFO and ERROR:")
+    set_log_level("INFO")
+    print("  log level set to INFO (TensorRT messages would print here)")
+    set_log_level("ERROR")
+    print("  log level set back to ERROR")
+
+    print("\nNVTX context manager — ranges are visible to Nsight Systems:")
+    with NVTX("example::demo"):
+        time.sleep(0.001)
+    print(f"  NVTX_ENABLED after context exit: {FLAGS.NVTX_ENABLED}")
+
+    print("\nJIT context manager — toggle Numba compilation around a hot loop:")
+    data = np.random.default_rng(0).standard_normal(100_000).astype(np.float32)
+
+    # Warm up either path so we measure steady-state cost
+    sum_squares(data)
+    t0 = time.perf_counter()
+    sum_squares(data)
+    no_jit_ms = (time.perf_counter() - t0) * 1000.0
+    print(f"  baseline (JIT={FLAGS.JIT}): {no_jit_ms:.3f} ms")
+
+    with JIT:
+        sum_squares(data)  # one warmup so the JIT compile cost is excluded
+        t0 = time.perf_counter()
+        sum_squares(data)
+        with_jit_ms = (time.perf_counter() - t0) * 1000.0
+        print(
+            f"  inside JIT block (JIT={FLAGS.JIT}, Numba={FLAGS.FOUND_NUMBA}): {with_jit_ms:.3f} ms"
+        )
+
+
+if __name__ == "__main__":
+    set_log_level("ERROR")
+    main()

examples/core.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com)
+#
+# MIT License
+"""
+File showcasing a high-level tour of the ``trtutils.core`` CUDA backend.
+
+Demonstrates :class:`trtutils.core.Device`, device introspection
+(:func:`trtutils.core.get_device`, :func:`trtutils.core.get_device_name`,
+:func:`trtutils.core.get_compute_capability`), stream lifecycle
+(:func:`trtutils.core.create_stream` / :func:`trtutils.core.destroy_stream`,
+:func:`trtutils.core.stream_synchronize`), explicit device memory
+(:func:`trtutils.core.cuda_malloc`, :func:`trtutils.core.cuda_free`,
+:func:`trtutils.core.memcpy_host_to_device` /
+:func:`trtutils.core.memcpy_device_to_host`), and :class:`trtutils.core.CUDAGraph`
+capture/replay wrapped around a synthetic stream sleep.
+
+The goal is to show that the ``core`` module exists and to give a guided
+overview of its building blocks; it is not a deep dive into any one piece.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from trtutils import set_log_level
+from trtutils.core import (
+    CUDAGraph,
+    Device,
+    create_stream,
+    cuda_free,
+    cuda_malloc,
+    destroy_stream,
+    get_compute_capability,
+    get_device,
+    get_device_count,
+    get_device_name,
+    get_num_dla_cores,
+    memcpy_device_to_host,
+    memcpy_host_to_device,
+    memcpy_host_to_device_async,
+    stream_synchronize,
+)
+
+
+def main() -> None:
+    print("Device info:")
+    print(f"  current device index: {get_device()}")
+    print(f"  device count:         {get_device_count()}")
+    print(f"  device name:          {get_device_name()}")
+    print(f"  compute capability:   {get_compute_capability()}")
+    print(f"  DLA cores:            {get_num_dla_cores()}")
+
+    # Device(idx) saves/restores the current device on exit; Device(None) is a no-op.
+    with Device(get_device()):
+        print(f"  inside Device guard:  {get_device()}")
+
+    # synchronous memcpy roundtrip
+    host = np.arange(8, dtype=np.float32)
+    nbytes = host.nbytes
+    device_ptr = cuda_malloc(nbytes)
+    memcpy_host_to_device(device_ptr, host)
+    roundtrip = np.zeros_like(host)
+    memcpy_device_to_host(roundtrip, device_ptr)
+    cuda_free(device_ptr)
+    print(f"\nSync memcpy roundtrip: {host.tolist()} -> {roundtrip.tolist()}")
+
+    # async memcpy through a stream
+    stream = create_stream()
+    device_ptr = cuda_malloc(nbytes)
+    memcpy_host_to_device_async(device_ptr, host, stream)
+    stream_synchronize(stream)
+    memcpy_device_to_host(roundtrip, device_ptr)
+    print(f"Async memcpy result:   {roundtrip.tolist()}")
+    cuda_free(device_ptr)
+
+    # CUDA graph capture/replay around an async memcpy
+    src = np.arange(16, dtype=np.float32)
+    dst = np.zeros_like(src)
+    device_ptr = cuda_malloc(src.nbytes)
+    graph = CUDAGraph(stream)
+    with graph:
+        memcpy_host_to_device_async(device_ptr, src, stream)
+    if graph.is_captured:
+        graph.launch()
+        stream_synchronize(stream)
+        memcpy_device_to_host(dst, device_ptr)
+        print(f"\nCUDA graph captured and replayed: dst[:4]={dst[:4].tolist()}")
+    else:
+        print("\nCUDA graph capture failed on this stream (skipping launch).")
+    graph.invalidate()
+    cuda_free(device_ptr)
+    destroy_stream(stream)
+
+
+if __name__ == "__main__":
+    set_log_level("ERROR")
+    main()