NiccoloN
92 lines
4.7 KiB
Python
92 lines
4.7 KiB
Python
"""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["lowered_candidates"] == 1
|
|
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()
|
|
report = json.loads((out / "reports/pipeline_periodic.json").read_text())
|
|
assert report["diagnostic_only"]
|
|
assert not report["selection_eligible"]
|
|
assert not (out / "pim").exists()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|