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

README.md

@@ -22,7 +22,11 @@ The ONNX export in [onnx-community/BiRefNet_lite-ONNX](https://huggingface.co/on
 Too many storage buffers in shader. Current: 17, Max is 16
 ```
 
-This repo replaces those 4 Concat and 59 Split nodes with trees of smaller ones (at most 6 inputs per Concat, 5 outputs per Split, so every node uses ≤ 7 storage buffers). Nothing else changes: the outputs are **bit-for-bit identical** to the original on ONNX Runtime CPU, for both files.
+This repo replaces those 4 Concat and 59 Split nodes with trees of smaller ones (at most 6 inputs per Concat, 5 outputs per Split, so every node uses ≤ 7 storage buffers).
+
+It also replaces the 20 `Sum` nodes with chains of `Add` (same order). ONNX Runtime's WebGPU backend has no Sum kernel, so Sum ran on the CPU, inside WebAssembly's 4 GB heap; in the deformable convolutions it adds four 392 MiB tensors (784 MiB each after the casts to fp32 the CPU needs), which fails with `std::bad_alloc`.
+
+Nothing else changes: the outputs are **bit-for-bit identical** to the original on ONNX Runtime CPU, for both files.
 
 | File | Precision | Size | Needs |
 |---|---|---|---|

onnx/model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a9e7aebec91bad0cf3949015e8ddc379602db6a2b7a06918e81340f4e0ce3716
-size 224182130
+oid sha256:fda4d8be0d03f3e4ee2f61aa4a4f2b72e0175b237839e0afc37db0e4d9420a11
+size 224191490

onnx/model_fp16.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e41401436e7e9a074f0f06556f5296702d48c8a0a227fb8a6f9ce4c7d08166f3
-size 114715263
+oid sha256:fb9bdae2f16fefd542929fe6efa52a0f8333bab5d92f35fdfcfb57ad0bc95e0d
+size 114724609

rebuild.py

@@ -13,6 +13,11 @@ values (concatenation and splitting are associative along one axis):
   - 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>
 """
 
@@ -128,11 +133,31 @@ class Rewriter:
         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}
+        stats = {"Concat": 0, "Split": 0, "Sum": 0}
         new_nodes: list[onnx.NodeProto] = []
         for node in self.graph.node:
-            if buffers(node) > MAX_BUFFERS and node.op_type == "Concat":
+            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":