Replace Sum with Add chains (no WebGPU Sum kernel: std::bad_alloc in wasm)

1528ba5 verified about 6 hours ago

7.94 kB

"""
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()

	"""
	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()