File size: 3,116 Bytes
5a6a851 | 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 | from __future__ import annotations
from pathlib import Path
import math
import numpy as np
import torch
class _Allocator:
def __new__(cls,*args,**kwargs):
import tensorrt as trt
class Impl(trt.IOutputAllocator):
def __init__(self,dtype):
trt.IOutputAllocator.__init__(self);self.dtype=dtype;self.tensor=None;self.shape=None
def reallocate_output(self,name,memory,size,alignment):
es=torch.empty((),dtype=self.dtype).element_size();n=max(1,(int(size)+es-1)//es)
self.tensor=torch.empty(n,device='cuda',dtype=self.dtype);return int(self.tensor.data_ptr())
def reallocate_output_async(self,name,memory,size,alignment,stream):return self.reallocate_output(name,memory,size,alignment)
def notify_shape(self,name,dims):self.shape=tuple(int(x) for x in dims)
return Impl(*args,**kwargs)
class TensorRTDDS:
"""TensorRT 11 zero-copy runner with IOutputAllocator for DDS plugin outputs."""
def __init__(self,path):
import tensorrt as trt
self.trt=trt;self.logger=trt.Logger(trt.Logger.ERROR);self.runtime=trt.Runtime(self.logger);self.engine=self.runtime.deserialize_cuda_engine(Path(path).read_bytes())
if self.engine is None:raise RuntimeError(f'Could not deserialize {path}')
self.context=self.engine.create_execution_context()
self.inputs=[];self.outputs=[]
for i in range(self.engine.num_io_tensors):
n=self.engine.get_tensor_name(i);(self.inputs if self.engine.get_tensor_mode(n)==trt.TensorIOMode.INPUT else self.outputs).append(n)
def _dtype(self,dt):
trt=self.trt;m={trt.float32:torch.float32,trt.float16:torch.float16,trt.int32:torch.int32,trt.int64:torch.int64,trt.bool:torch.bool}
if hasattr(trt,'bfloat16'):m[trt.bfloat16]=torch.bfloat16
return m[dt]
def run(self,feeds):
keep=[]
for n in self.inputs:
x=feeds[n];dt=self._dtype(self.engine.get_tensor_dtype(n));x=x.to(device='cuda',dtype=dt).contiguous();feeds[n]=x;keep.append(x);self.context.set_input_shape(n,tuple(x.shape));self.context.set_tensor_address(n,int(x.data_ptr()))
outs={};allocs={}
for n in self.outputs:
dt=self._dtype(self.engine.get_tensor_dtype(n));shape=tuple(int(q) for q in self.context.get_tensor_shape(n))
if any(q<0 for q in shape):
a=_Allocator(dt);allocs[n]=a;self.context.set_output_allocator(n,a)
else:
y=torch.empty(shape if shape else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);self.context.set_tensor_address(n,int(y.data_ptr()))
ok=self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
if not ok:raise RuntimeError('TensorRT execute_async_v3 failed')
torch.cuda.synchronize()
for n,a in allocs.items():
if a.tensor is None or a.shape is None:raise RuntimeError(f'DDS output {n} was not allocated/notified')
num=math.prod(a.shape) if a.shape else 1;outs[n]=a.tensor[:num].view(a.shape)
return outs
|