File size: 7,942 Bytes
20f6e53 1528ba5 20f6e53 1528ba5 20f6e53 1528ba5 20f6e53 1528ba5 20f6e53 | 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """
Rebuild BiRefNet-lite ONNX so it runs on any WebGPU device.
The official export has Concat nodes with up to 1024 inputs and Split nodes
with up to 32 outputs. ONNX Runtime's WebGPU backend binds every input and
output of a node as a storage buffer, and GPUs cap that per shader
(maxStorageBuffersPerShaderStage: 8 minimum, 10 or 16 on many devices), so
those nodes fail with "Too many storage buffers in shader".
This replaces them with trees of small nodes that compute exactly the same
values (concatenation and splitting are associative along one axis):
- Concat: at most CONCAT_FAN_IN inputs per node (+1 output)
- Split: at most SPLIT_FAN_OUT outputs per node (+ data + sizes inputs)
Every node then uses at most MAX_BUFFERS storage buffers.
It also replaces Sum with a chain of Add nodes: ONNX Runtime's WebGPU backend
has no Sum kernel, so Sum runs on the CPU in WebAssembly memory, and the
deformable convolutions' Sums of four 392 MiB tensors exhaust its 4 GB heap
("std::bad_alloc").
Usage: python rebuild.py <in.onnx> <out.onnx>
"""
import sys
import numpy as np
import onnx
from onnx import helper, numpy_helper
CONCAT_FAN_IN = 6
SPLIT_FAN_OUT = 5
MAX_BUFFERS = 7
def buffers(node: onnx.NodeProto) -> int:
return len([i for i in node.input if i]) + len([o for o in node.output if o])
class Rewriter:
def __init__(self, model: onnx.ModelProto):
self.model = model
self.graph = model.graph
self.counter = 0
self.initializers = {t.name: t for t in self.graph.initializer}
self.constants = {
n.output[0]: n for n in self.graph.node if n.op_type == "Constant" and len(n.output) == 1
}
self.names = {o for n in self.graph.node for o in n.output} | set(self.initializers)
self.names |= {i.name for i in self.graph.input}
def fresh(self, base: str) -> str:
while True:
self.counter += 1
name = f"{base}/wgpu_{self.counter}"
if name not in self.names:
self.names.add(name)
return name
def const_value(self, name: str) -> np.ndarray:
if name in self.initializers:
return numpy_helper.to_array(self.initializers[name])
if name in self.constants:
attr = next(a for a in self.constants[name].attribute if a.name == "value")
return numpy_helper.to_array(attr.t)
raise ValueError(f"split sizes '{name}' are not constant")
def int64_initializer(self, base: str, values: list[int]) -> str:
name = self.fresh(base + "/sizes")
tensor = numpy_helper.from_array(np.array(values, dtype=np.int64), name)
self.graph.initializer.append(tensor)
self.initializers[name] = tensor
return name
def concat_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
axis = next(a.i for a in node.attribute if a.name == "axis")
inputs = list(node.input)
made: list[onnx.NodeProto] = []
while len(inputs) > CONCAT_FAN_IN:
grouped = []
for start in range(0, len(inputs), CONCAT_FAN_IN):
chunk = inputs[start : start + CONCAT_FAN_IN]
if len(chunk) == 1:
grouped.append(chunk[0])
continue
out = self.fresh(node.name)
made.append(helper.make_node("Concat", chunk, [out], name=self.fresh(node.name), axis=axis))
grouped.append(out)
inputs = grouped
made.append(helper.make_node("Concat", inputs, list(node.output), name=node.name, axis=axis))
return made
def split_tree(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
axis = next((a.i for a in node.attribute if a.name == "axis"), 0)
if len(node.input) < 2 or not node.input[1]:
raise ValueError(f"{node.name}: Split without explicit sizes isn't handled")
sizes = [int(v) for v in self.const_value(node.input[1]).tolist()]
outputs = list(node.output)
assert len(sizes) == len(outputs)
made: list[onnx.NodeProto] = []
def split(data: str, parts: list[tuple[int, str]], base: str) -> None:
# parts: (size, final output name)
if len(parts) <= SPLIT_FAN_OUT:
made.append(
helper.make_node(
"Split",
[data, self.int64_initializer(base, [s for s, _ in parts])],
[o for _, o in parts],
name=self.fresh(base),
axis=axis,
)
)
return
# Groups of roughly equal count, each split again below.
per_group = -(-len(parts) // SPLIT_FAN_OUT)
groups = [parts[i : i + per_group] for i in range(0, len(parts), per_group)]
group_outputs = []
for group in groups:
group_outputs.append(group[0][1] if len(group) == 1 else self.fresh(base))
made.append(
helper.make_node(
"Split",
[data, self.int64_initializer(base, [sum(s for s, _ in g) for g in groups])],
group_outputs,
name=self.fresh(base),
axis=axis,
)
)
for group, out in zip(groups, group_outputs):
if len(group) > 1:
split(out, group, base)
split(node.input[0], list(zip(sizes, outputs)), node.name)
return made
def add_chain(self, node: onnx.NodeProto) -> list[onnx.NodeProto]:
# ((a + b) + c) + d: the same order Sum adds in, but Add has a WebGPU kernel.
inputs = [i for i in node.input if i]
if len(inputs) == 1:
return [helper.make_node("Identity", inputs, list(node.output), name=node.name)]
made: list[onnx.NodeProto] = []
acc = inputs[0]
for index, other in enumerate(inputs[1:], start=2):
out = node.output[0] if index == len(inputs) else self.fresh(node.name)
made.append(helper.make_node("Add", [acc, other], [out], name=self.fresh(node.name)))
acc = out
return made
def run(self) -> dict[str, int]:
stats = {"Concat": 0, "Split": 0, "Sum": 0}
new_nodes: list[onnx.NodeProto] = []
for node in self.graph.node:
if node.op_type == "Sum":
# ONNX Runtime's WebGPU backend has no Sum kernel, so Sum runs on the
# CPU, inside WebAssembly's 4 GB heap. In the deformable convolutions it
# adds four 392 MiB tensors (784 MiB each after the fp16 → fp32 casts
# the CPU needs): "std::bad_alloc".
new_nodes.extend(self.add_chain(node))
stats["Sum"] += 1
elif buffers(node) > MAX_BUFFERS and node.op_type == "Concat":
new_nodes.extend(self.concat_tree(node))
stats["Concat"] += 1
elif buffers(node) > MAX_BUFFERS and node.op_type == "Split":
new_nodes.extend(self.split_tree(node))
stats["Split"] += 1
else:
new_nodes.append(node)
del self.graph.node[:]
self.graph.node.extend(new_nodes)
return stats
def main() -> None:
source, target = sys.argv[1], sys.argv[2]
model = onnx.load(source)
stats = Rewriter(model).run()
too_wide = [(n.op_type, n.name, buffers(n)) for n in model.graph.node if buffers(n) > MAX_BUFFERS]
if too_wide:
raise SystemExit(f"still too wide: {too_wide[:10]}")
onnx.checker.check_model(model)
onnx.save(model, target)
widest = max(buffers(n) for n in model.graph.node)
print(f"{source} -> {target}: rewrote {stats}, nodes {len(model.graph.node)}, widest node {widest} buffers")
if __name__ == "__main__":
main()
|