| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import tensorrt as trt |
| from huggingface_hub import snapshot_download |
|
|
| ROOT = Path(snapshot_download( |
| os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX'), |
| repo_type='model', token=os.environ.get('HF_TOKEN'), local_dir='/tmp/cf-custom', |
| allow_patterns=['onnx/custom-ops/*.onnx','plugins/tensorrt/companion_sparse_trt.py'], |
| )) |
| sys.path.insert(0, str(ROOT/'plugins/tensorrt')) |
| import companion_sparse_trt |
|
|
| LOGGER = trt.Logger(trt.Logger.INFO) |
|
|
| PROFILES = { |
| 'SparseSubdivide': {'feats': (32,8), 'coords': (32,4)}, |
| 'SparseDownsample': {'feats': (64,8), 'coords': (64,4)}, |
| 'SparseUpsample': {'feats': (8,8), 'target_coords': (64,4), 'inverse': (64,)}, |
| 'SparseWindowAttention': {'qkv': (64,3,4,16), 'coords': (64,4)}, |
| 'SparseConv3D': {'feats': (64,8), 'coords': (64,4), 'weight': (16,3,3,3,8), 'bias': (16,)}, |
| 'MeshTopologyExtract': { |
| 'verts_grid': (274625,3), 'sdf': (274625,), 'cube_idx': (262144,8), |
| 'beta': (262144,12), 'alpha': (262144,8), 'gamma': (262144,), 'colors_grid': (274625,6), |
| }, |
| } |
|
|
|
|
| def build(name: str): |
| onnx_path = ROOT/'onnx/custom-ops'/f'{name}.onnx' |
| builder = trt.Builder(LOGGER) |
| network = builder.create_network(0) |
| parser = trt.OnnxParser(network, LOGGER) |
| if not parser.parse(onnx_path.read_bytes()): |
| errs = [str(parser.get_error(i)) for i in range(parser.num_errors)] |
| raise RuntimeError(f'{name} parse failed:\n' + '\n'.join(errs)) |
| cfg = builder.create_builder_config() |
| cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30) |
| prof = builder.create_optimization_profile() |
| input_report = {} |
| dynamic = False |
| for i in range(network.num_inputs): |
| inp = network.get_input(i) |
| shape = PROFILES[name].get(inp.name) |
| if shape is None: |
| raise KeyError(f'{name}: missing profile for {inp.name}') |
| input_report[inp.name] = tuple(inp.shape) |
| if any(int(d) < 0 for d in inp.shape): |
| dynamic = True |
| prof.set_shape(inp.name, shape, shape, shape) |
| if dynamic: |
| cfg.add_optimization_profile(prof) |
| blob = builder.build_serialized_network(network, cfg) |
| if blob is None: |
| raise RuntimeError(f'{name} build_serialized_network returned None') |
| out = Path('/tmp/custom-engines')/f'{name}.plan' |
| out.parent.mkdir(parents=True, exist_ok=True) |
| out.write_bytes(bytes(blob)) |
| return { |
| 'op': name, |
| 'onnx_bytes': onnx_path.stat().st_size, |
| 'plan_bytes': out.stat().st_size, |
| 'inputs': input_report, |
| 'outputs': [network.get_output(i).name for i in range(network.num_outputs)], |
| } |
|
|
|
|
| if __name__ == '__main__': |
| report=[] |
| for name in PROFILES: |
| print(f'=== BUILD {name} ===', flush=True) |
| r=build(name); print('OK',json.dumps(r),flush=True); report.append(r) |
| Path('/tmp/custom-engines/report.json').write_text(json.dumps(report,indent=2)) |
|
|
