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)