better throughput scheduling with steady state cost model (WIP)

2026-10-23 11:27:06 +02:00
parent 48e9fb3795
commit fd493ae6e8
51 changed files with 3581 additions and 247 deletions
@@ -0,0 +1,89 @@
+"""Intermediate emission must stop before unsupported backend transformations."""
+import argparse
+import json
+import os
+from pathlib import Path
+import subprocess
+import tempfile
+
+import onnx
+from onnx import TensorProto, helper
+
+REPO = Path(__file__).resolve().parents[2]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--compiler", required=True)
+    args = parser.parse_args()
+    with tempfile.TemporaryDirectory(prefix="raptor-emission-") as temporary:
+        root = Path(temporary)
+
+        def model(name, elements):
+            tensors = [helper.make_tensor_value_info(n, TensorProto.FLOAT, [1, elements])
+                       for n in ("x", "y", "z")]
+            graph = helper.make_graph([helper.make_node("Add", ["x", "y"], ["z"])],
+                                      name, tensors[:2], tensors[2:])
+            value = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
+            value.ir_version = 10
+            path = root / f"{name}.onnx"
+            onnx.save(value, path)
+            return path
+
+        small = model("small", 16)
+        # This static type is legal in Spatial. PIM's signed byte-size operand cannot encode it.
+        large = model("large", 536870912)
+        for source, targets in ((small, ("EmitSpatial", "EmitPim", "EmitPimBufferized", "EmitPimCodegen")),
+                                (large, ("EmitSpatial", "EmitPim"))):
+            for target in targets:
+                out = root / source.stem / target
+                out.mkdir(parents=True)
+                result = subprocess.run([args.compiler, str(source), "-o", str(out / "model"),
+                    "--maccel=PIM", f"--{target}", "--pipeline=2", "--core-count=168",
+                    "--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading", "--pim-dump-ir",
+                    f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
+                    capture_output=True, text=True, timeout=30, env=dict(os.environ, OMP_NUM_THREADS="1"))
+                output = result.stdout + result.stderr
+                if source == large and target == "EmitPim":
+                    assert result.returncode != 0 and "host input copy byte size is outside representable range" in output, output
+                    continue
+                assert result.returncode == 0, output
+                if target == "EmitPimCodegen":
+                    assert (out / "pim/config.json").exists()
+                    report = json.loads((out / "reports/pipeline_selection.json").read_text())
+                    assert report["generated_candidates"] == 1
+                    assert report["single_candidate_analysis_skipped"]
+                    assert report["analyzed_candidates"] == report["lowered_candidates"] == 0
+                    assert report["unavailable_estimates"] == 0
+                    assert not report["reused_lowered_candidate"]
+                    assert "no cost comparison required" in report["reason"]
+                else:
+                    assert "intermediate emission retains the original schedule" in output
+                    assert not (out / "reports/pipeline_selection.json").exists()
+                    assert not (out / "reports/pipeline_periodic.json").exists()
+                    assert not (out / "pim").exists()
+                    name = {"EmitSpatial": "spatial4_scheduled", "EmitPim": "pim0",
+                            "EmitPimBufferized": "pim1_buff"}[target]
+                    ir = (out / f"dialects/{name}.mlir").read_text()
+                    if target == "EmitSpatial":
+                        assert "spat.scheduled_compute" in ir and "pim.core" not in ir
+                        assert not (out / "dialects/pim0.mlir").exists()
+                    elif target == "EmitPim":
+                        assert "pim.core" in ir and "tensor<" in ir
+                        assert not (out / "dialects/pim1_buff.mlir").exists()
+                    else:
+                        assert "pim.core" in ir and "memref<" in ir
+        out = root / "explicit-diagnostic"
+        out.mkdir()
+        result = subprocess.run([args.compiler, str(small), "-o", str(out / "model"),
+            "--maccel=PIM", "--EmitSpatial", "--pipeline=2", "--core-count=168",
+            "--crossbar-count=96", "--crossbar-size=128", "--mlir-disable-threading",
+            "--pim-periodic-diagnostics",
+            f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-a/throughput_config_1000ms.json"],
+            capture_output=True, text=True, timeout=30)
+        assert result.returncode == 0, result.stdout + result.stderr
+        assert (out / "reports/pipeline_periodic.json").exists()
+
+
+if __name__ == "__main__":
+    main()