Upload verify_tftrt_implicit_remote.py with huggingface_hub

verify_tftrt_implicit_remote.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import ctypes
+import hashlib
+import json
+import shutil
+import tempfile
+import urllib.request
+from pathlib import Path
+
+import tensorrt as trt
+
+
+BASE = "https://huggingface.co/hacnho/tensorrt-efficientnms-tftrt-implicit-bypass-poc/resolve/main"
+FILES = {
+    "control": "control.engine",
+    "neg_score": "neg_score.engine",
+}
+
+
+def sha256_file(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        while True:
+            chunk = f.read(1024 * 1024)
+            if not chunk:
+                break
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def load_cudart():
+    cudart = ctypes.CDLL("libcudart.so")
+    cuda_malloc = cudart.cudaMalloc
+    cuda_malloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
+    cuda_malloc.restype = ctypes.c_int
+    cuda_free = cudart.cudaFree
+    cuda_free.argtypes = [ctypes.c_void_p]
+    cuda_free.restype = ctypes.c_int
+    cuda_memcpy = cudart.cudaMemcpy
+    cuda_memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+    cuda_memcpy.restype = ctypes.c_int
+    cuda_memset = cudart.cudaMemset
+    cuda_memset.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
+    cuda_memset.restype = ctypes.c_int
+    return {
+        "malloc": cuda_malloc,
+        "free": cuda_free,
+        "memcpy": cuda_memcpy,
+        "memset": cuda_memset,
+        "h2d": 1,
+        "d2h": 2,
+    }
+
+
+def upload_floats(cuda: dict, ptr: ctypes.c_void_p, values: list[float]) -> None:
+    arr = (ctypes.c_float * len(values))(*values)
+    rc = cuda["memcpy"](ptr, ctypes.cast(arr, ctypes.c_void_p), ctypes.sizeof(arr), cuda["h2d"])
+    if rc != 0:
+        raise RuntimeError(f"cudaMemcpy H2D failed rc={rc}")
+
+
+def presets():
+    return {
+        "all_negative": {
+            "boxes": [0.0, 0.0, 1.0, 1.0] * 4,
+            "scores": [-0.2, -0.2, -0.2, -0.2],
+        },
+        "mixed_scores": {
+            "boxes": [0.0, 0.0, 1.0, 1.0] * 4,
+            "scores": [0.6, 0.4, 0.2, -0.1],
+        },
+        "all_zero": {
+            "boxes": [0.0, 0.0, 1.0, 1.0] * 4,
+            "scores": [0.0, 0.0, 0.0, 0.0],
+        },
+    }
+
+
+def run_engine(engine_path: Path) -> dict[str, object]:
+    cuda = load_cudart()
+    logger = trt.Logger(trt.Logger.ERROR)
+    trt.init_libnvinfer_plugins(logger, "")
+    runtime = trt.Runtime(logger)
+    blob = engine_path.read_bytes()
+    engine = runtime.deserialize_cuda_engine(blob)
+
+    result: dict[str, object] = {
+        "engine_path": str(engine_path),
+        "engine_sha256": sha256_file(engine_path),
+        "presets": {},
+    }
+
+    for preset_name, preset in presets().items():
+        ctx = engine.create_execution_context()
+        ptrs: dict[str, tuple[ctypes.c_void_p, int, str]] = {}
+        try:
+            for i in range(engine.num_io_tensors):
+                tensor_name = engine.get_tensor_name(i)
+                shape = engine.get_tensor_shape(tensor_name)
+                count = 1
+                for dim in shape:
+                    count *= dim
+                dtype = str(engine.get_tensor_dtype(tensor_name))
+                nbytes = count * 4
+                ptr = ctypes.c_void_p()
+                assert cuda["malloc"](ctypes.byref(ptr), nbytes) == 0
+                assert cuda["memset"](ptr, 0, nbytes) == 0
+                assert ctx.set_tensor_address(tensor_name, ptr.value)
+                ptrs[tensor_name] = (ptr, nbytes, dtype)
+                if tensor_name in preset:
+                    upload_floats(cuda, ptr, preset[tensor_name])
+
+            infer = ctx.infer_shapes()
+            exec_ok = ctx.execute_async_v3(0)
+            outputs = {}
+            output_order = []
+            for i in range(engine.num_io_tensors):
+                tensor_name = engine.get_tensor_name(i)
+                if "OUTPUT" not in str(engine.get_tensor_mode(tensor_name)):
+                    continue
+                output_order.append(tensor_name)
+                ptr, nbytes, dtype = ptrs[tensor_name]
+                if "INT32" in dtype:
+                    host = (ctypes.c_int32 * min(max(nbytes // 4, 1), 16))()
+                    rc = cuda["memcpy"](ctypes.byref(host), ptr, min(nbytes, 64), cuda["d2h"])
+                    outputs[tensor_name] = {"copy_rc": rc, "values": list(host)}
+                else:
+                    host = (ctypes.c_float * min(max(nbytes // 4, 1), 16))()
+                    rc = cuda["memcpy"](ctypes.byref(host), ptr, min(nbytes, 64), cuda["d2h"])
+                    outputs[tensor_name] = {"copy_rc": rc, "values": [float(x) for x in host]}
+
+            num_detections = None
+            score_values = []
+            if output_order:
+                first = outputs.get(output_order[0], {})
+                if first.get("values"):
+                    num_detections = first["values"]
+            if len(output_order) >= 3:
+                score_values = outputs.get(output_order[2], {}).get("values", [])
+            result["presets"][preset_name] = {
+                "infer_shapes": infer,
+                "execute_ok": exec_ok,
+                "num_detections": num_detections,
+                "score_values": score_values,
+                "outputs": outputs,
+            }
+        finally:
+            for ptr, _, _ in ptrs.values():
+                if ptr.value:
+                    cuda["free"](ptr)
+    return result
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--local-dir", type=Path, help="reuse files from a local directory instead of downloading from HF")
+    args = ap.parse_args()
+
+    td = None
+    try:
+        if args.local_dir:
+            local = {label: args.local_dir / name for label, name in FILES.items()}
+        else:
+            td = Path(tempfile.mkdtemp(prefix="trt_efficientnms_tftrt_implicit_remote_"))
+            local = {}
+            for label, name in FILES.items():
+                dst = td / name
+                urllib.request.urlretrieve(f"{BASE}/{name}", dst)
+                local[label] = dst
+
+        payload = {
+            "control": run_engine(local["control"]),
+            "neg_score": run_engine(local["neg_score"]),
+        }
+        print(json.dumps(payload, indent=2, ensure_ascii=False))
+    finally:
+        if td is not None:
+            shutil.rmtree(td, ignore_errors=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())