from __future__ import annotations
import argparse,json,time
from pathlib import Path
import torch
from runtime.trt_torch import TorchTensorRTEngine

def bench(engine,feeds,warm=3,iters=20):
    for _ in range(warm): engine.run(feeds)
    torch.cuda.synchronize(); times=[]
    for _ in range(iters):
        t=time.perf_counter(); engine.run(feeds); torch.cuda.synchronize(); times.append((time.perf_counter()-t)*1000)
    times.sort(); return {'mean_ms':sum(times)/len(times),'p50_ms':times[len(times)//2],'p95_ms':times[int(len(times)*.95)-1]}
def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--root',default='.'); a=ap.parse_args(); r=Path(a.root); out={}
    d=TorchTensorRTEngine(r/'engines/l4-sm89/dinov2.plan'); out['dinov2']=bench(d,{'pixel_values':torch.randn(1,3,518,518,device='cuda')})
    s=TorchTensorRTEngine(r/'engines/l4-sm89/dsine.plan'); out['dsine']=bench(s,{'image':torch.randn(1,3,544,544,device='cuda'),'intrins':torch.eye(3,device='cuda').unsqueeze(0)})
    out['peak_vram_gib']=torch.cuda.max_memory_allocated()/1024**3; print(json.dumps(out,indent=2))
if __name__=='__main__':main()