Host-resident TensorRT shape-tensor inputs incur a device round trip and a per-call stream sync

"""Benchmark for host shape-tensor inputs to a TensorRT engine.

A TensorRT *shape tensor* input is read from host memory. ``setup_input_tensors`` in
``core/runtime/execute_engine.cpp`` nonetheless requires every input to be on the device
and copies a shape input back to the host with ``.cpu()``. When the caller already holds
the value on the host, that is an H2D move followed by a D2H copy, and the D2H copy
synchronizes the CUDA stream: the host blocks until the work already queued on the stream
has finished.

This script builds a one-engine model whose engine takes a TensorRT shape-tensor input and
reports two things:

  Part 1 -- per-call host overhead on an idle stream, with the shape input placed on the
            device and on the host. This is the cost the copy adds to every call.
  Part 2 -- the isolated cost of the exact copy the patch changes, measured on a stream that
            has a backlog of work queued on it, to show that the copy waits for that work.

On a runtime that copies a shape input to the host unconditionally, a host-resident shape
input is the slower of the two in Part 1 (an extra H2D then D2H). On a runtime that accepts a
host shape input as is, it is the faster, and Part 2's host path does not wait for the backlog.

Run under ``PYTHONPATH`` pointing at the torch_tensorrt build to compare.
"""

import gc
import statistics
import time
from typing import List, Tuple

import torch
import torch.nn as nn
import tensorrt as trt
import torch_tensorrt
from torch_tensorrt.dynamo.runtime import TorchTensorRTModule

SHAPE_VAL = 64
ITERS = 300
WARMUP = 30
BACKLOG_MATMULS = 20
BACKLOG_SIZE = 4096


class ShapeInputModel(nn.Module):
    """``arange(0, n, 1)`` makes ``n`` a TensorRT shape-tensor input binding."""

    def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
        a = torch.arange(0, n, 1, device=x.device).to(torch.float32)
        return (x * 2.0 + a).sum(dim=0, keepdim=True)


def build_engine_module() -> TorchTensorRTModule:
    x = torch.randn(SHAPE_VAL, device="cuda")
    n = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cuda")
    dim = torch.export.Dim("d", min=2, max=256)
    ep = torch.export.export(
        ShapeInputModel().eval().cuda(), (x, n), dynamic_shapes={"x": {0: dim}, "n": None}
    )
    gm = torch_tensorrt.dynamo.compile(
        ep, inputs=[x, n], min_block_size=1, pass_through_build_failures=True,
        use_python_runtime=False, assume_dynamic_shape_support=True,
    )
    m = next(mm for _, mm in gm.named_children() if isinstance(mm, TorchTensorRTModule))
    _assert_shape_binding(m)
    return m


def _assert_shape_binding(m: TorchTensorRTModule) -> None:
    rt = trt.Runtime(trt.Logger(trt.Logger.ERROR))
    ce = rt.deserialize_cuda_engine(m.serialized_engine)
    shape_ins = [
        ce.get_tensor_name(i)
        for i in range(ce.num_io_tensors)
        if ce.get_tensor_mode(ce.get_tensor_name(i)) == trt.TensorIOMode.INPUT
        and ce.is_shape_inference_io(ce.get_tensor_name(i))
    ]
    assert shape_ins, "engine has no shape-tensor input binding"


def ordered_inputs(m: TorchTensorRTModule, x: torch.Tensor, device: str) -> List[torch.Tensor]:
    args = {"x": x, "_local_scalar_dense": torch.tensor(SHAPE_VAL, dtype=torch.int64, device=device)}
    return [args[b] for b in m.input_binding_names]


def per_call_overhead(m: TorchTensorRTModule, x: torch.Tensor, device: str) -> Tuple[float, float]:
    inp = ordered_inputs(m, x, device)
    for _ in range(WARMUP):
        out = torch.ops.tensorrt.execute_engine(inp, m.engine)
    torch.cuda.synchronize()
    ts = []
    gc.disable()  # the call allocates its output; periodic GC otherwise skews the tail
    for _ in range(ITERS):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        torch.ops.tensorrt.execute_engine(inp, m.engine)
        ts.append((time.perf_counter() - t0) * 1e6)
    gc.enable()
    return min(ts), float(out[0])


def copy_cost_under_backlog() -> Tuple[float, float]:
    """Times the two shape-input copies the patch chooses between, on a busy stream."""
    a = torch.randn(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    b = torch.randn(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    c = torch.empty(BACKLOG_SIZE, BACKLOG_SIZE, device="cuda")
    dev_in = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cuda")
    host_in = torch.tensor(SHAPE_VAL, dtype=torch.int64, device="cpu")

    def queue_backlog() -> None:
        for _ in range(BACKLOG_MATMULS):
            torch.matmul(a, b, out=c)

    def bench(fn) -> float:
        for _ in range(5):
            fn()
        torch.cuda.synchronize()
        ts = []
        for _ in range(50):
            queue_backlog()
            t0 = time.perf_counter()
            fn()
            ts.append((time.perf_counter() - t0) * 1e3)
            torch.cuda.synchronize()
        return statistics.median(ts)

    dev_ms = bench(lambda: dev_in.clone().contiguous().cpu().to(torch.int64))
    host_ms = bench(lambda: host_in.contiguous().to(torch.int64))
    return dev_ms, host_ms


def main() -> None:
    print(f"torch {torch.__version__} | torch_tensorrt {torch_tensorrt.__file__}\n")
    m = build_engine_module()

    x = torch.randn(SHAPE_VAL, device="cuda")
    print(f"Part 1: per-call host overhead, idle stream (min of {ITERS} calls)")
    print(f"{'shape input on':>16} | {'host time':>11} | {'output':>10}")
    print("-" * 44)
    vals = []
    for device in ("cuda", "cpu"):
        mn, val = per_call_overhead(m, x, device)
        vals.append(val)
        print(f"{device:>16} | {mn:>8.1f} us | {val:>10.3f}")
    assert abs(vals[0] - vals[1]) < 1e-3, f"device and host inputs disagree: {vals}"
    print("  outputs match across placements")

    print(f"\nPart 2: cost of the shape-input copy on a stream with "
          f"{BACKLOG_MATMULS}x{BACKLOG_SIZE}^2 matmuls queued")
    dev_ms, host_ms = copy_cost_under_backlog()
    print(f"  device input  .clone().cpu().to(int64) : {dev_ms:7.2f} ms  (waits for the backlog)")
    print(f"  host input    .contiguous().to(int64)  : {host_ms:7.2f} ms  (does not)")


if __name__ == "__main__":
    main()