| import os,sys |
| from pathlib import Path |
| import torch,tensorrt as trt |
| from huggingface_hub import snapshot_download |
| ROOT=Path(snapshot_download('patdev/Companion-Forge-L4-ONNX',repo_type='model',token=os.environ.get('HF_TOKEN'),local_dir='/tmp/cf-conv',allow_patterns=['engines/l4-sm89/custom-ops/SparseConv3D.plan','plugins/tensorrt/companion_sparse_trt.py'])) |
| sys.path.insert(0,str(ROOT/'plugins/tensorrt'));import companion_sparse_trt |
| LOGGER=trt.Logger(trt.Logger.VERBOSE);rt=trt.Runtime(LOGGER);eng=rt.deserialize_cuda_engine((ROOT/'engines/l4-sm89/custom-ops/SparseConv3D.plan').read_bytes());ctx=eng.create_execution_context() |
| class OA(trt.IOutputAllocator): |
| def __init__(self,dt):super().__init__();self.dt=dt;self.t=None;self.shape=None |
| def reallocate_output(self,n,mem,size,align): |
| es=torch.empty((),dtype=self.dt).element_size();self.t=torch.empty(max(1,(int(size)+es-1)//es),device='cuda',dtype=self.dt);print('ALLOC',n,size,align,self.t.numel(),flush=True);return int(self.t.data_ptr()) |
| def reallocate_output_async(self,n,mem,size,align,stream):return self.reallocate_output(n,mem,size,align) |
| def notify_shape(self,n,d):self.shape=tuple(d);print('SHAPE',n,self.shape,flush=True) |
| DT={trt.float16:torch.float16,trt.float32:torch.float32,trt.int32:torch.int32,trt.int64:torch.int64} |
| xyz=torch.cartesian_prod(torch.arange(4),torch.arange(4),torch.arange(4)).to(torch.int32);c=torch.cat([torch.zeros((64,1),dtype=torch.int32),xyz],1).cuda();f=torch.randn(64,8,device='cuda',dtype=torch.float16) |
| import spconv.pytorch as spconv |
| mod=spconv.SubMConv3d(8,16,3,bias=True,algo=spconv.ConvAlgo.Native).cuda().half().eval();w=mod.weight.detach().contiguous();b=mod.bias.detach().contiguous();print('WEIGHT',w.shape,w.dtype,flush=True) |
| feeds={'feats':f,'coords':c,'weight':w,'bias':b};keep=[];alloc={} |
| for n,x in feeds.items():ctx.set_input_shape(n,tuple(x.shape));ctx.set_tensor_address(n,int(x.data_ptr()));keep.append(x);print('INPUT',n,x.shape,x.dtype,flush=True) |
| for i in range(eng.num_io_tensors): |
| n=eng.get_tensor_name(i) |
| if eng.get_tensor_mode(n)!=trt.TensorIOMode.OUTPUT:continue |
| dt=DT[eng.get_tensor_dtype(n)];sh=tuple(ctx.get_tensor_shape(n));print('OUTPUT_DESC',n,sh,dt,flush=True) |
| if any(int(q)<0 for q in sh):a=OA(dt);alloc[n]=a;ctx.set_output_allocator(n,a) |
| else:y=torch.empty(sh if sh else (),device='cuda',dtype=dt);keep.append(y);ctx.set_tensor_address(n,int(y.data_ptr())) |
| print('EXEC',flush=True);ok=ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream);print('EXEC_RETURN',ok,flush=True);torch.cuda.synchronize();print('DONE',flush=True) |
|
|
