File size: 7,505 Bytes
c699c4c 026da6c c699c4c 026da6c c699c4c 026da6c c699c4c 026da6c 6ffd3f8 c699c4c 6ffd3f8 c699c4c 6ffd3f8 c699c4c 6ffd3f8 c699c4c 6ffd3f8 c699c4c 6ffd3f8 c699c4c 026da6c c699c4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Served-like benchmark: calls the server's own ``predict()`` handler (no HTTP framing) on the
sample 1600x900 JPEG, after the server's own warm-up, and reports the medians / minima of the
response's ``timing_ms`` keys (preprocess = base64 + JPEG decode + resize; device_forward;
postprocess; total) -- the definitions GPU_COMPARISON.md's served-like loop uses.
Run from code/ after `source tools/chipenv.sh <chip> <venv>` (PYTHONPATH must start with code/):
SP_DISPATCH=auto|eth|worker (default auto = ETH, 1 CQ) SP_N_ITER=100 python bench_served.py
Optional request knobs: SP_REQ_NMS_RADIUS (4), SP_REQ_MAX_KP (1024), SP_REQ_THR (0.005).
"""
import base64
import hashlib
import json
import os
import statistics
import time
import torch
import ttnn
from models.reference.superpoint_reference import load_reference_model # noqa: E402
from models.server import app as A # noqa: E402
from models.tt.superpoint_ttnn import TtSuperPoint # noqa: E402
HERE = os.path.dirname(os.path.abspath(__file__))
def main():
n = int(os.environ.get("SP_N_ITER", "100"))
from models.tt.device_open import open_ttnn_device, resolve_dispatch
mode = resolve_dispatch(os.environ.get("SP_DISPATCH", "auto")) # auto: ETH + 1 CQ + 12x10
kw = dict(l1_small_size=A.L1_SMALL_SIZE, trace_region_size=32 << 20)
dev = open_ttnn_device(0, dispatch=mode, **kw)
g = dev.compute_with_storage_grid_size()
print(f"dispatch={mode} num_command_queues=1 compute_grid={g.x}x{g.y}")
close_device = ttnn.close_device
try:
tm = load_reference_model()
A._setup_logging()
torch.set_grad_enabled(False)
model = TtSuperPoint(tm, dev, input_height=A.INPUT_HEIGHT, input_width=A.INPUT_WIDTH, fused=True)
tt_in = model.allocate_input(batch_size=1)
A.STATE.update(
model=model, tt_in=tt_in, fused=True, cfg={"fused": True},
model_config={"border_removal_distance": int(tm.config.border_removal_distance)},
)
dummy = torch.zeros(1, 3, A.INPUT_HEIGHT, A.INPUT_WIDTH)
A._warmup_fused(model, tt_in, dummy, ttnn, dev)
A.STATE["ready"] = True
print("warmup_ms", A.STATE.get("warmup_ms"))
if model.device_resize and os.environ.get("SP_FIRST_USE"):
# first request at a source size without a precompiled resize variant (trace capture)
import io
import numpy as np
from PIL import Image
for wh in ((1000, 700), (1001, 701)):
buf = io.BytesIO()
Image.fromarray(np.random.default_rng(0).integers(0, 256, (wh[1], wh[0], 3), dtype=np.uint8)).save(buf, "PNG")
rq = A.PredictRequest(image=base64.b64encode(buf.getvalue()).decode())
t0 = time.perf_counter(); r1 = A.predict_dict(rq); t1 = time.perf_counter(); r2 = A.predict_dict(rq); t2 = time.perf_counter()
print(f"first use {wh}: device_forward {r1['timing_ms']['device_forward']} ms (second {r2['timing_ms']['device_forward']} ms)")
img = base64.b64encode(open(os.path.join(HERE, "sample_data/house_in_field_1080p.jpg"), "rb").read()).decode()
req = A.PredictRequest(
image=img,
nms_radius=int(os.environ.get("SP_REQ_NMS_RADIUS", "4")),
max_keypoints=int(os.environ.get("SP_REQ_MAX_KP", "1024")),
keypoint_threshold=float(os.environ.get("SP_REQ_THR", "0.005")),
)
for _ in range(10):
r = A.predict_dict(req)
# Exactness vs the previous request path (RGB resize + fp32 /255 -> run_fused -> host
# post-processing from the NMS / descriptor maps), same decoded frame.
im = A._decode_image(img)
from models.tt import postprocess as _post
res = model.run_fused(tt_in, A._preprocess(im))
if req.nms_radius == model.nms_radius_traced:
a = _post.postprocess_from_nms_map(res.nms_map, res.descriptors_nchw, keypoint_threshold=req.keypoint_threshold,
max_keypoints=req.max_keypoints, border_removal_distance=4, with_descriptors=True)[0]
else: # host fold + simple_nms(r) on the traced scores (the previous behaviour for r != 4)
res = model.run_fused(tt_in, A._preprocess(im), nms_radius=req.nms_radius)
a = _post.postprocess_keypoints(res.scores_nchw, res.descriptors_nchw, nms_radius=req.nms_radius,
keypoint_threshold=req.keypoint_threshold, max_keypoints=req.max_keypoints,
border_removal_distance=4, with_descriptors=True)[0]
b = A._infer_fused(model, tt_in, A._preprocess_r8(im), max_keypoints=req.max_keypoints,
keypoint_threshold=req.keypoint_threshold, nms_radius=req.nms_radius,
return_descriptors=True, border=4)
oa, ob = torch.argsort(a[1], descending=True, stable=True), torch.argsort(b[1], descending=True, stable=True)
same = a[0].shape == b[0].shape and torch.equal(a[0][oa], b[0][ob]) and torch.equal(a[1][oa], b[1][ob])
dd = float((a[2][oa] - b[2][ob]).abs().max()) if same and a[0].shape[0] else float("nan")
print(f"host-resize request path vs RGB/run_fused path: kp/scores identical={same} n={b[0].shape[0]} desc max|diff|={dd:.2e}")
if model.device_resize:
# rsz stage: device resize of the full-size R plane vs the host Pillow resize (round-1 server)
c = A._infer_fused(model, tt_in, A._r_plane(im), max_keypoints=req.max_keypoints,
keypoint_threshold=req.keypoint_threshold, nms_radius=req.nms_radius,
return_descriptors=True, border=4)
same_c = all(torch.equal(x, y) for x, y in zip(b[:3], c[:3]))
print(f"device-resize request path vs host-resize request path: kp/scores/desc bit-identical={same_c}")
r_host = A.predict_dict(req)
model.fused_stages = model.fused_stages - {"rsz"}
r_old = A.predict_dict(req)
model.fused_stages = model.fused_stages | {"rsz"}
keys = ("num_keypoints", "keypoints", "scores")
same_r = all(r_host[k] == r_old[k] for k in keys) and r_host["descriptors"]["data"] == r_old["descriptors"]["data"]
print(f"predict() response (device resize) == predict() response (host resize): {same_r}")
keys = ("preprocess", "device_forward", "postprocess", "total")
T = {k: [] for k in keys + ("handler_wall",)}
for _ in range(n):
t0 = time.perf_counter()
r = A.predict_dict(req)
T["handler_wall"].append((time.perf_counter() - t0) * 1e3)
for k in keys:
T[k].append(r["timing_ms"][k])
print(f"mode={mode} n={n} num_keypoints={r['num_keypoints']} serving_path={r.get('serving_path')} "
f"top3={r['keypoints'][:3]} scores={r['scores'][:3]}")
body = {k: r[k] for k in ("num_keypoints", "keypoints", "scores", "descriptors")}
print("response_sha256", hashlib.sha256(json.dumps(body, sort_keys=True).encode()).hexdigest())
for k, v in T.items():
print(f"{k:16s} median={statistics.median(v):8.3f} ms min={min(v):8.3f} ms")
model.release()
finally:
close_device(dev)
if __name__ == "__main__":
main()
|