# Copyright (C) 2026 Embedl AB
# mypy: disable-error-code="no-untyped-call"

"""
FP8 ViT on NVIDIA Jetson AGX Thor
=================================

Thor's Blackwell GPU runs FP8 GEMMs, so a Vision Transformer whose attention
and MLP matmuls are quantized to FP8 beats the FP16 model on latency without
touching the rest of the network. This tutorial takes the HuggingFace
``google/vit-base-patch16-224`` classifier through the whole path:

1. ``transform()`` fuses it with the TensorRT patterns under the
   ``tensorrt_fp8`` backend.
2. ``quantize()`` inserts FP8 Q/DQ stubs and calibrates them on ImageNette.
3. ``convert_float_to_float16()`` turns everything that is not FP8 into
   float16 *in PyTorch*, so ``torch.onnx.export`` writes the final graph.
4. ``trtexec --stronglyTyped`` builds an engine that keeps every tensor at the
   precision the ONNX says, and the engine is profiled and scored on the
   device.

The FP16 baseline goes through exactly the same steps minus ``quantize()``, so
the two engines differ only in the Q/DQ nodes. There is no ONNX-level rewriting
anywhere: what ``torch.onnx.export`` produces is what TensorRT builds.

The script runs end to end wherever ``trtexec`` is found, so the simplest way
to use it is to run the whole file on the Thor: the PyTorch half (steps 1--3)
takes a few minutes there and writes two ONNX files under ``artifacts/``, and
step 4 builds, profiles and scores them in place. On a workstation without
TensorRT the script stops after the export.

Measured on a Jetson AGX Thor with TensorRT 10.13.3 (batch 1, GPU compute time
with CUDA graphs, ImageNette validation, 3925 images):

.. code-block:: text

    engine                          median ms    top-1
    fp16, weakly typed --fp16           1.399    86.0 %
    fp16, strongly typed                1.402    86.0 %
    fp8 + fp16, strongly typed          1.065    86.2 %

FP8 is 1.32x faster than the FP16 engine at the same accuracy.

Install dependencies:

.. code-block:: bash

    pip install "embedl-deploy[tensorrt]" torchvision transformers onnx onnxscript

On the Thor itself, the PyPI ``torch`` wheel runs on the GPU as is (``torch
2.14, cu130``), and the ``tensorrt`` Python module comes with JetPack -- create
the virtualenv with ``--system-site-packages`` to see it.
"""

# %%
# Setup: the model and ImageNette
# --------------------------------
#
# The HuggingFace classifier takes ``pixel_values`` and returns a
# ``logits``-carrying output; a thin wrapper gives it the plain
# ``tensor -> tensor`` signature every step below expects. Attention is set
# to the ``eager`` implementation so the fuser sees the individual matmuls
# and softmax rather than a fused SDPA call.

import copy
import itertools
import json
import re
import shutil
import subprocess
import tarfile
import urllib.request
from pathlib import Path

import torch
import torchvision
from torch import fx, nn
from torchvision import transforms
from transformers import (  # type: ignore[import-not-found]
    AutoModelForImageClassification,
)

MODEL_ID = "google/vit-base-patch16-224"
IMAGENETTE_URL = (
    "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz"
)
DATA_DIR = Path("artifacts/data")
IMAGENETTE_DIR = DATA_DIR / "imagenette2-320"
OUT_DIR = Path("artifacts/thor_fp8_vit")
OUT_DIR.mkdir(parents=True, exist_ok=True)

#: ImageNette's ten classes, as ImageNet-1k indices.
IMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701]
CALIBRATION_IMAGES = 128
BATCH_SIZE = 1

torch.manual_seed(0)


class ImageClassifier(nn.Module):
    """``pixel_values -> logits`` view of a HuggingFace image classifier."""

    def __init__(self, model_id: str) -> None:
        """Load the pretrained classifier `model_id` with eager attention."""
        super().__init__()
        self.model = AutoModelForImageClassification.from_pretrained(
            model_id, attn_implementation="eager"
        )

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Return the ImageNet-1k logits for `pixel_values`."""
        logits: torch.Tensor = self.model(pixel_values=pixel_values).logits
        return logits


def download_imagenette() -> None:
    """Download and extract ImageNette if not already present."""
    if IMAGENETTE_DIR.exists():
        return
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    tgz_path = DATA_DIR / "imagenette2-320.tgz"
    print(f"Downloading ImageNette to {tgz_path} ...")
    urllib.request.urlretrieve(IMAGENETTE_URL, str(tgz_path))
    with tarfile.open(tgz_path) as tar:
        tar.extractall(DATA_DIR, filter="data")
    tgz_path.unlink()


download_imagenette()

# ViT-B/16 was trained with 0.5 mean/std and a straight resize to 224.
preprocess = transforms.Compose(
    [
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)


def remap(label: int) -> int:
    """Map an ImageNette class index to its ImageNet class index."""
    return IMAGENETTE_TO_IMAGENET[label]


train_set = torchvision.datasets.ImageFolder(
    str(IMAGENETTE_DIR / "train"), transform=preprocess, target_transform=remap
)
val_set = torchvision.datasets.ImageFolder(
    str(IMAGENETTE_DIR / "val"), transform=preprocess, target_transform=remap
)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE)

calibration_loader = torch.utils.data.DataLoader(
    train_set, batch_size=BATCH_SIZE, shuffle=True
)
calibration_batches = [
    images
    for images, _ in itertools.islice(calibration_loader, CALIBRATION_IMAGES)
]

example_input = torch.randn(1, 3, 224, 224)

# %%
# Step 1: fuse for TensorRT with the FP8 policy
# ----------------------------------------------
#
# ``tensorrt_fp8`` is the TensorRT backend with the FP8 precision policy
# measured on Thor: the same conversions and fusions as ``tensorrt``, plus
# the knowledge of which kernels take FP8 profitably and which stay in
# floating point. ``transform()`` rewrites the model into fused modules
# around those kernels -- the attention projections, the MLP linears -- and
# leaves everything else as it is.

from embedl_deploy import transform
from embedl_deploy.backend import set_backend

set_backend("tensorrt_fp8")

model = ImageClassifier(MODEL_ID).eval()
fused = transform(model, (example_input,)).model

with torch.no_grad():
    drift = (model(example_input) - fused(example_input)).abs().max().item()
print(f"max |logit drift| after fusion: {drift:.2e}")

# %%
# Step 2: quantize to FP8
# ------------------------
#
# Activations and weights both get E4M3 FP8 Q/DQ. Calibration runs the 128
# training images through the model to pick each stub's scale. Layers the
# FP8 policy rules out -- the three-channel patch-embedding stem, for one --
# are left in floating point on its own. ``freeze_weights=True`` bakes the
# weight scales in, which the export in step 3 requires.

from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    TensorQuantConfig,
    quantize,
)


def forward_loop(m: nn.Module) -> None:
    """Feed the calibration images through `m`."""
    with torch.no_grad():
        for images in calibration_batches:
            m(images)


fp8 = TensorQuantConfig(Precision.FP8)
# ``quantize()`` works in place, so the baseline keeps its own copy.
quantized = quantize(
    copy.deepcopy(fused),
    (example_input,),
    config=QuantConfig(activation=fp8, weight=fp8),
    forward_loop=forward_loop,
    freeze_weights=True,
)


def accuracy(m: nn.Module, limit: int | None = None) -> float:
    """Top-1 accuracy of `m` on the validation split (or its first `limit`)."""
    correct = total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            correct += int((m(images).argmax(dim=1) == labels).sum())
            total += labels.numel()
            if limit is not None and total >= limit:
                break
    return correct / total


# A quick sanity check in PyTorch: the FP8 fake-quantized model on the first
# 256 validation images. The full-split numbers come from the engines.
print(f"FP32 top-1 (256 images):     {accuracy(fused, 256):.1%}")
print(f"FP8 QDQ top-1 (256 images):  {accuracy(quantized, 256):.1%}")

# %%
# Step 3: float16 in PyTorch, then export
# ----------------------------------------
#
# ``convert_float_to_float16()`` returns a copy whose floating parameters and
# buffers are float16, keeps the quantizer scales usable, and pins the few
# operations that need float32 (interpolation, cumulative sums, top-k) with
# casts at their boundaries. The Q/DQ nodes are emitted during export with
# float16 scales, so the exported graph is float16 end to end with FP8 in
# the quantized matmuls. The same half model is also saved as a ``.pt2``
# through ``torch.export`` for inspection or a later re-export. The baseline
# is the same conversion applied to the fused, unquantized model.

from embedl_deploy.quantize import convert_float_to_float16


def export_onnx(m: fx.GraphModule, path: Path) -> Path:
    """Export the float16 copy of `m` to ONNX, and to ``.pt2`` alongside."""
    half = convert_float_to_float16(m)
    with torch.no_grad():
        exported = torch.export.export(half, (example_input.half(),))
        torch.export.save(exported, str(path.with_suffix(".pt2")))
        torch.onnx.export(
            half,
            (example_input.half(),),
            str(path),
            dynamo=True,
            opset_version=20,
            input_names=["input"],
            output_names=["logits"],
        )
    # The exporter keeps the weights in a ``<name>.data`` file next to the ONNX.
    size = sum(p.stat().st_size for p in path.parent.glob(f"{path.name}*"))
    print(f"wrote {path} ({size / 1e6:.0f} MB)")
    return path


fp16_onnx = export_onnx(fused, OUT_DIR / "vit_fp16.onnx")
fp8_onnx = export_onnx(quantized, OUT_DIR / "vit_fp8.onnx")

# %%
# Step 4: build, profile and score on the Thor
# ---------------------------------------------
#
# ``--stronglyTyped`` tells TensorRT to honour the dtypes in the ONNX exactly:
# float16 tensors stay float16, the ``QuantizeLinear`` outputs stay FP8, and
# nothing is re-decided by the autotuner. That is what makes the comparison
# clean -- both engines run precisely the graph the export wrote. The FP16
# engine is also built once weakly typed (``--fp16``), the way most people
# run a float model, for reference.

TRTEXEC = shutil.which("trtexec") or "/usr/src/tensorrt/bin/trtexec"
if not Path(TRTEXEC).exists():
    print(
        f"trtexec not found -- copy {OUT_DIR} to the Thor and run this "
        "script there to build, profile and score the engines."
    )
    raise SystemExit(0)

_COMPUTE_TIME = re.compile(
    r"GPU Compute Time: min = ([\d.]+) ms, max = ([\d.]+) ms, "
    r"mean = ([\d.]+) ms, median = ([\d.]+) ms"
)


def build_engine(onnx_path: Path, label: str, *flags: str) -> dict[str, float]:
    """Build `onnx_path` with `flags`, time it, return the latency summary."""
    engine = OUT_DIR / f"{label}.engine"
    log = OUT_DIR / f"trtexec_{label}.log"
    cmd = [
        TRTEXEC,
        f"--onnx={onnx_path}",
        *flags,
        f"--saveEngine={engine}",
        f"--timingCacheFile={OUT_DIR / 'timing.cache'}",
        "--useCudaGraph",
        "--noDataTransfers",
        "--useSpinWait",
        "--warmUp=1000",
        "--iterations=500",
        "--avgRuns=100",
    ]
    print(f"[{label}] {' '.join(cmd)}")
    with log.open("w") as fh:
        subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, check=True)
    match = _COMPUTE_TIME.search(log.read_text())
    if match is None:
        raise RuntimeError(f"no GPU compute time summary in {log}")
    minimum, maximum, mean, median = (float(v) for v in match.groups())
    print(f"[{label}] median {median:.3f} ms (mean {mean:.3f} ms)")
    return {
        "median_ms": median,
        "mean_ms": mean,
        "min_ms": minimum,
        "max_ms": maximum,
    }


latency = {
    "fp16 (weakly typed --fp16)": build_engine(
        fp16_onnx, "fp16_weak", "--fp16"
    ),
    "fp16 (strongly typed)": build_engine(
        fp16_onnx, "fp16_strong", "--stronglyTyped"
    ),
    "fp8 + fp16 (strongly typed)": build_engine(
        fp8_onnx, "fp8_strong", "--stronglyTyped"
    ),
}

# %%
# Accuracy is measured by running the engines over the whole ImageNette
# validation split (3925 images) through the TensorRT Python API. The inputs
# are float16 because the strongly typed engines' input tensor is.

import tensorrt as trt


class Engine:
    """Minimal single-input, single-output TensorRT engine runner."""

    def __init__(self, path: Path) -> None:
        """Deserialize the engine at `path` and allocate its output buffer."""
        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
        self.engine = runtime.deserialize_cuda_engine(path.read_bytes())
        self.context = self.engine.create_execution_context()
        self.input_name = self.engine.get_tensor_name(0)
        self.output_name = self.engine.get_tensor_name(1)
        self.stream = torch.cuda.Stream()
        self.output = torch.empty(
            tuple(self.engine.get_tensor_shape(self.output_name)),
            dtype=_torch_dtype(self.engine.get_tensor_dtype(self.output_name)),
            device="cuda",
        )
        self.input_dtype = _torch_dtype(
            self.engine.get_tensor_dtype(self.input_name)
        )

    def __call__(self, images: torch.Tensor) -> torch.Tensor:
        """Run one batch through the engine and return float32 logits."""
        images = images.to("cuda", self.input_dtype).contiguous()
        self.context.set_tensor_address(self.input_name, images.data_ptr())
        self.context.set_tensor_address(
            self.output_name, self.output.data_ptr()
        )
        self.context.execute_async_v3(self.stream.cuda_stream)
        self.stream.synchronize()
        return self.output.float().cpu()


def _torch_dtype(dtype: "trt.DataType") -> torch.dtype:
    """Map a TensorRT tensor dtype to the torch dtype of its buffer."""
    return {trt.float32: torch.float32, trt.float16: torch.float16}[dtype]


def engine_accuracy(label: str) -> float:
    """Top-1 accuracy of the engine saved under `label`."""
    engine = Engine(OUT_DIR / f"{label}.engine")
    correct = total = 0
    for images, labels in val_loader:
        correct += int((engine(images).argmax(dim=1) == labels).sum())
        total += labels.numel()
    return correct / total


top1 = {
    "fp16 (weakly typed --fp16)": engine_accuracy("fp16_weak"),
    "fp16 (strongly typed)": engine_accuracy("fp16_strong"),
    "fp8 + fp16 (strongly typed)": engine_accuracy("fp8_strong"),
}

# %%
# Results
# -------

print(f"\n{'engine':<32}{'median ms':>12}{'top-1':>10}")
for label, stats in latency.items():
    print(f"{label:<32}{stats['median_ms']:>12.3f}{top1[label]:>10.1%}")
speedup = (
    latency["fp16 (strongly typed)"]["median_ms"]
    / latency["fp8 + fp16 (strongly typed)"]["median_ms"]
)
print(f"\nFP8 speed-up over the FP16 engine: {speedup:.2f}x")

results = {
    "model": MODEL_ID,
    "tensorrt": trt.__version__,
    "latency": latency,
    "top1": top1,
    "latency_is": "GPU compute time, batch 1, CUDA graphs, H2D/D2H excluded",
    "accuracy_is": "ImageNette validation, 3925 images, ImageNet-1k labels",
}
(OUT_DIR / "results.json").write_text(json.dumps(results, indent=2))
print(f"wrote {OUT_DIR / 'results.json'}")