File size: 6,236 Bytes
d73aac1 | 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 | import json
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
import tensorrt as trt
from huggingface_hub import snapshot_download
REPO=os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX')
ROOT=Path(snapshot_download(REPO,repo_type='model',token=os.environ.get('HF_TOKEN'),local_dir='/tmp/cf-val',allow_patterns=['engines/l4-sm89/custom-ops/*.plan','plugins/tensorrt/companion_sparse_trt.py','plugins/reference/ops.py']))
sys.path.insert(0,str(ROOT/'plugins/tensorrt'));sys.path.insert(0,str(ROOT/'plugins/reference'))
import companion_sparse_trt # noqa
import ops as ref
LOGGER=trt.Logger(trt.Logger.ERROR)
DT={trt.float32:torch.float32,trt.float16:torch.float16,trt.int32:torch.int32,trt.int64:torch.int64,trt.bool:torch.bool}
if hasattr(trt,'bfloat16'): DT[trt.bfloat16]=torch.bfloat16
class OutputAllocator(trt.IOutputAllocator):
def __init__(self,dtype):
trt.IOutputAllocator.__init__(self);self.dtype=dtype;self.tensor=None;self.shape=None;self.nbytes=0
def reallocate_output(self,tensor_name,memory,size,alignment):
self.nbytes=int(size); n=max(1,(int(size)+torch.tensor([],dtype=self.dtype).element_size()-1)//torch.tensor([],dtype=self.dtype).element_size())
self.tensor=torch.empty(n,device='cuda',dtype=self.dtype);return int(self.tensor.data_ptr())
def reallocate_output_async(self,tensor_name,memory,size,alignment,stream):
return self.reallocate_output(tensor_name,memory,size,alignment)
def notify_shape(self,tensor_name,dims):
self.shape=tuple(int(x) for x in dims)
class Runner:
def __init__(self,name):
blob=(ROOT/'engines/l4-sm89/custom-ops'/f'{name}.plan').read_bytes();self.rt=trt.Runtime(LOGGER);self.eng=self.rt.deserialize_cuda_engine(blob);assert self.eng;self.ctx=self.eng.create_execution_context();self.name=name
def run(self,feeds):
keep=[]
for n,x in feeds.items():
dt=DT[self.eng.get_tensor_dtype(n)];x=x.to(dtype=dt).contiguous();feeds[n]=x;keep.append(x)
if self.eng.get_tensor_mode(n)==trt.TensorIOMode.INPUT:
self.ctx.set_input_shape(n,tuple(x.shape));self.ctx.set_tensor_address(n,int(x.data_ptr()))
outs={};allocs={}
for i in range(self.eng.num_io_tensors):
n=self.eng.get_tensor_name(i)
if self.eng.get_tensor_mode(n)!=trt.TensorIOMode.OUTPUT:continue
dt=DT[self.eng.get_tensor_dtype(n)];shape=tuple(int(x) for x in self.ctx.get_tensor_shape(n))
if any(x<0 for x in shape):
a=OutputAllocator(dt);allocs[n]=a;self.ctx.set_output_allocator(n,a)
else:
y=torch.empty(shape if shape else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);self.ctx.set_tensor_address(n,int(y.data_ptr()))
ok=self.ctx.execute_async_v3(stream_handle=torch.cuda.current_stream().cuda_stream);assert ok,f'{self.name} execute failed';torch.cuda.synchronize()
for n,a in allocs.items():
if a.tensor is None:raise RuntimeError(f'{self.name}/{n}: allocator not called')
if a.shape is None:raise RuntimeError(f'{self.name}/{n}: shape notification missing')
num=int(np.prod(a.shape)) if a.shape else 1
outs[n]=a.tensor[:num].view(a.shape)
return outs
def diff(a,b):
a=a.float();b=b.float();d=(a-b).abs();return {'max_abs':float(d.max()) if d.numel() else 0.0,'mean_abs':float(d.mean()) if d.numel() else 0.0,'cos':float(torch.nn.functional.cosine_similarity(a.reshape(1,-1),b.reshape(1,-1)).item()) if d.numel() else 1.0}
def grid4():
xyz=torch.cartesian_prod(torch.arange(4),torch.arange(4),torch.arange(4)).to(torch.int32);return torch.cat([torch.zeros((64,1),dtype=torch.int32),xyz],1).cuda()
report={}
torch.manual_seed(123)
# Subdivide
f=torch.randn(32,8,device='cuda',dtype=torch.float16);c=grid4()[:32].contiguous();r0,r1=ref.sparse_subdivide(f,c);o=Runner('SparseSubdivide').run({'feats':f,'coords':c});report['SparseSubdivide']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))}
# Downsample
f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();r0,r1,ri=ref.sparse_downsample(f,c,(2,2,2));o=Runner('SparseDownsample').run({'feats':f,'coords':c});m=int(o['count'].item());report['SparseDownsample']={'count':m,'ref_count':int(r0.shape[0]),'feats':diff(o['out_feats'][:m],r0),'coords_equal':bool(torch.equal(o['out_coords'][:m],r1)),'inverse_equal':bool(torch.equal(o['inverse'],ri))}
# Upsample
f=torch.randn(8,8,device='cuda',dtype=torch.float16);tc=grid4();inv=torch.arange(64,device='cuda',dtype=torch.int32)%8;r0,r1=ref.sparse_upsample(f,tc,inv);o=Runner('SparseUpsample').run({'feats':f,'target_coords':tc,'inverse':inv});report['SparseUpsample']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))}
# Window attention
qkv=torch.randn(64,3,4,16,device='cuda',dtype=torch.float16);c=grid4();r=ref.sparse_window_attention(qkv,c,8,(0,0,0));o=Runner('SparseWindowAttention').run({'qkv':qkv,'coords':c});report['SparseWindowAttention']=diff(o['out'],r)
# Conv: create actual spconv reference and feed its native weight layout.
import spconv.pytorch as spconv
f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();mod=spconv.SubMConv3d(8,16,3,bias=True,algo=spconv.ConvAlgo.Native).cuda().half().eval();st=spconv.SparseConvTensor(f,c,[64,64,64],1);ry=mod(st);w=mod.weight.detach().contiguous();b=mod.bias.detach().contiguous();print('SPCONV_WEIGHT_SHAPE',tuple(w.shape),flush=True)
try:
o=Runner('SparseConv3D').run({'feats':f,'coords':c,'weight':w,'bias':b});m=int(o['count'].item());report['SparseConv3D']={'count':m,'ref_count':int(ry.features.shape[0]),'feats':diff(o['out_feats'][:m],ry.features),'coords_equal':bool(torch.equal(o['out_coords'][:m],ry.indices.to(torch.int32)))}
except Exception as e:
report['SparseConv3D']={'error':repr(e),'weight_shape':tuple(w.shape)}
# Latency quick pass for successful fixed-shape ops
for name in ['SparseSubdivide','SparseUpsample','SparseWindowAttention']:
pass
print('REPORT',json.dumps(report,indent=2),flush=True)
Path('/tmp/custom_validation.json').write_text(json.dumps(report,indent=2))
|