OverviewFilesCommitsIssuesDocsReleasesPeers

terrain-lookdev commits / 3856aad

optimisation: ONNX decoder shim — numerically perfect, and no faster
Scoped to the decoder because the profile says so: 59.1% of a tile is decoder_model, 36.0% base_model, and 95% is inside model forwards, so one model to shim gets most of the benefit. The repo ships an exporter and a published ONNX model but NO inference path -- WorldPipeline calls self.decoder_model(x, noise_labels, conditional_inputs) directly. onnx_decoder.py replaces that module's forward with an onnxruntime session, leaving the module object intact so .to()/.config still work. IO binding passes torch device pointers straight to ORT; a naive version copying GPU->CPU->ORT->CPU->GPU would swamp any kernel gain at 549 ms/call. The published graph has dynamic spatial dims, so it works at any decoder_tile_size. Upstream's own exporter is broken for this model anyway -- its conv2d padding shim turns MPConv's 1-tuple into [] and export dies. GATE: numerically equivalent. 512^2 max|diff| 1.320e-05 rel 7.32e-06 256^2 max|diff| 8.345e-05 rel 7.88e-06 SPEED: no gain. Alternating, 512^2: pytorch 250.0 ms onnx-cuda 247.8 ms 1.01x pytorch 260.2 ms onnx-cuda 266.0 ms 0.98x pytorch 276.0 ms onnx-cuda 289.4 ms 0.95x That is the third kernel-level lever to come back empty -- torch.compile was slower, reduce-overhead worse, ONNX flat -- while autocast fp16, which halves memory traffic rather than improving kernels, gave 1.5x. The consistent reading is that this decoder is BANDWIDTH bound on a 2070, not compute bound, and faster kernels have nothing to win. TensorRT is unproven: onnxruntime-gpu 1.29 needs CUDA 13 (we are on cu121) so 1.20.2 is pinned, and the TRT engine build exceeded a 10 minute budget and was abandoned. The shim supports it (TD_ONNX_EP=tensorrt) if anyone wants to spend the build time. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
3856aaddf608985e1c2d168344ac41d9db559e17
by plankatron <noreply@anthropic.com> · 2026-08-31T20:25:11.000Z

Changed files (2) +98 -0

__pycache__/onnx_decoder.cpython-312.pyc +0 -0
(no content changes — binary or rename)
onnx_decoder.py +98 -0
……@@ -0,0 +1,98 @@
1"""
2ONNX decoder shim for terrain-diffusion.
3
4WHY THE DECODER, AND ONLY THE DECODER
5
6Profiled, not guessed. For one 512^2 tile:
7
8 decoder_model 4.39 s 8 calls 549 ms 59.1% of total
9 base_model 2.67 s 6 calls 445 ms 36.0%
10 model forward 95.0%
11 python / tiling / copies 5.0%
12
1395% of the time is inside model forwards, so a faster runtime is aimed at the
14real cost rather than at overhead -- and the decoder alone is 59% of it. One
15model to shim instead of three, for most of the benefit.
16
17WHY A SHIM AT ALL
18
19The repo ships an ONNX *exporter* and a published ONNX model, but no inference
20path: WorldPipeline calls self.decoder_model(x, noise_labels, conditional_inputs)
21directly. This replaces that module's forward with an onnxruntime session while
22leaving the module object intact, so .to(), .config and everything else the
23pipeline touches still work.
24
25ZERO COPY OR IT IS POINTLESS
26
27The naive version moves the tensor GPU -> CPU -> ORT -> CPU -> GPU per call. At
28549 ms/call and 8 calls a tile that would swamp any kernel gain. IO binding
29passes torch's device pointers straight to ORT, so nothing crosses PCIe.
30
31The published graph has dynamic spatial dims (['batch', 5, 'height', 'width']),
32so it works at any decoder_tile_size rather than only the exported one.
33"""
34import os
35import numpy as np
36import torch
37
38
39def _providers(kind):
40 if kind == "tensorrt":
41 return [("TensorrtExecutionProvider",
42 {"trt_fp16_enable": True,
43 "trt_engine_cache_enable": True,
44 "trt_engine_cache_path": "/tmp/td_trt_cache"}),
45 "CUDAExecutionProvider", "CPUExecutionProvider"]
46 if kind == "cuda":
47 return ["CUDAExecutionProvider", "CPUExecutionProvider"]
48 return ["CPUExecutionProvider"]
49
50
51class OnnxDecoderRunner:
52 def __init__(self, path, kind="cuda", device=0):
53 import onnxruntime as ort
54 self.ort = ort
55 so = ort.SessionOptions()
56 so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
57 self.sess = ort.InferenceSession(path, so, providers=_providers(kind))
58 self.device = device
59 self.kind = kind
60 got = self.sess.get_providers()
61 print(f" ONNX decoder: {os.path.basename(path)} providers={got}")
62 if kind == "tensorrt" and "TensorrtExecutionProvider" not in got:
63 print(" WARNING: TensorRT provider not active; fell back.")
64
65 @torch.no_grad()
66 def __call__(self, x, noise_labels=None, conditional_inputs=None):
67 # The graph is fp32. Under autocast x may arrive as fp16, so normalise
68 # rather than silently feeding a mismatched dtype.
69 xt = x.detach().to(torch.float32).contiguous()
70 nt = noise_labels.detach().to(torch.float32).contiguous()
71 if not xt.is_cuda:
72 xt = xt.cuda()
73 if not nt.is_cuda:
74 nt = nt.cuda()
75 b, _c, h, w = xt.shape
76 out = torch.empty((b, 1, h, w), dtype=torch.float32, device=xt.device)
77
78 io = self.sess.io_binding()
79 io.bind_input("x", "cuda", self.device, np.float32, tuple(xt.shape), xt.data_ptr())
80 io.bind_input("noise_labels", "cuda", self.device, np.float32,
81 tuple(nt.shape), nt.data_ptr())
82 io.bind_output("output", "cuda", self.device, np.float32,
83 tuple(out.shape), out.data_ptr())
84 self.sess.run_with_iobinding(io)
85 return out
86
87
88def install(pipeline, path=None, kind=None):
89 """Swap the pipeline's decoder forward for an ONNX session. Returns True if
90 installed. Controlled by TD_ONNX_DECODER (path) and TD_ONNX_EP."""
91 path = path or os.environ.get("TD_ONNX_DECODER")
92 if not path or not os.path.exists(path):
93 return False
94 kind = kind or os.environ.get("TD_ONNX_EP", "cuda")
95 runner = OnnxDecoderRunner(path, kind=kind)
96 pipeline.decoder_model.forward = runner
97 pipeline._onnx_decoder = runner # keep a reference alive
98 return True