NiccoloN
113 lines
5.8 KiB
Python
113 lines
5.8 KiB
Python
"""Spatial selection lowers one winner and rejects incomplete comparisons."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
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-spatial-selection-") as temporary:
|
|
root = Path(temporary)
|
|
guard = root / "static-audit.so"
|
|
subprocess.run([os.environ.get("CC", "cc"), "-shared", "-fPIC",
|
|
str(REPO / "test/PIM/StaticSchedulingAudit.c"), "-o", str(guard)], check=True)
|
|
env = dict(os.environ, LD_PRELOAD=str(guard), OMP_NUM_THREADS="1")
|
|
iteration = 0
|
|
|
|
def compile(source, pipeline=2, emit="EmitPimCodegen", extra=(), succeeds=True):
|
|
nonlocal iteration
|
|
output = root / str(iteration)
|
|
iteration += 1
|
|
output.mkdir()
|
|
command = [args.compiler, str(source), "-o", str(output / "model"), "--maccel=PIM",
|
|
f"--{emit}", f"--pipeline={pipeline}", "--core-count=138", "--crossbar-count=128",
|
|
"--crossbar-size=128", "--mlir-disable-threading", "--pim-detect-communication-deadlock",
|
|
f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json"]
|
|
command.extend(extra)
|
|
result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=120)
|
|
assert (result.returncode == 0) == succeeds, result.stdout + result.stderr
|
|
path = output / "reports/pipeline_selection.json"
|
|
report = json.loads(path.read_text()) if path.exists() else None
|
|
program = {str(p.relative_to(output / "pim")): hashlib.sha256(p.read_bytes()).hexdigest()
|
|
for p in (output / "pim").rglob("*") if p.is_file()}
|
|
if "--pim-export-spatial-dataflow=all" in extra:
|
|
for stage in ("spatial3_scheduled_no_comm", "spatial4_scheduled"):
|
|
for suffix in ("nodes.csv", "edges.csv"):
|
|
assert (output / f"reports/{stage}.{suffix}").stat().st_size > 0
|
|
assert (output / "dialects/spatial4_scheduled.mlir").exists()
|
|
assert (output / "dialects/pim1_buff.mlir").exists()
|
|
shutil.rmtree(output)
|
|
return report, program, result.stdout + result.stderr
|
|
|
|
vgg = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
|
|
selected, program, _ = compile(vgg)
|
|
count = selected["generated_candidates"]
|
|
assert count > 1 and selected["strategy"] == "spatial-queue-aware"
|
|
assert selected["lowered_candidates"] == 1
|
|
assert len(selected["backend_attempts"]) == 1
|
|
assert selected["backend_attempts"][0]["candidate"] == selected["selection"]
|
|
assert not selected["prediction_error_bound_available"]
|
|
assert selected["selection_margin"] == 0.021
|
|
for candidate in selected["candidates"]:
|
|
assert candidate["cost_available"] and candidate["convergence_proved"]
|
|
assert candidate["lowering_supported"] and candidate["eligible"]
|
|
assert "request_order" not in candidate["spatial_cost"]
|
|
assert "without_local_copy_work" not in candidate["spatial_cost"]
|
|
diagnostic, diagnostic_program, _ = compile(vgg, extra=(
|
|
"--pim-export-spatial-dataflow=all", "--pim-dump-ir", "--pim-periodic-diagnostics"))
|
|
assert diagnostic_program == program and diagnostic["selection"] == selected["selection"]
|
|
|
|
# A single schedule needs no optional timing contract.
|
|
source = root / "softmax.onnx"
|
|
tensor = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 8])
|
|
output = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8])
|
|
model = helper.make_model(helper.make_graph([helper.make_node("Softmax", ["x"], ["y"], axis=1)],
|
|
"softmax", [tensor], [output]), opset_imports=[helper.make_opsetid("", 18)])
|
|
model.ir_version = 10
|
|
onnx.save(model, source)
|
|
selected, program, _ = compile(source)
|
|
assert selected["single_candidate_analysis_skipped"] and selected["lowered_candidates"] == 1
|
|
assert "spatial_cost" not in selected["candidates"][0]
|
|
|
|
# An unknown estimate cannot silently exclude any alternative.
|
|
model = onnx.load(vgg)
|
|
output = model.graph.output[0].name
|
|
internal = output + "_before_softmax"
|
|
for node in model.graph.node:
|
|
for index, value in enumerate(node.output):
|
|
if value == output:
|
|
node.output[index] = internal
|
|
for index, value in enumerate(node.input):
|
|
if value == output:
|
|
node.input[index] = internal
|
|
model.graph.node.append(helper.make_node("Softmax", [internal], [output], axis=-1))
|
|
source = root / "vgg-softmax.onnx"
|
|
onnx.save(model, source)
|
|
report, program, diagnostics = compile(source, succeeds=False)
|
|
assert report["generated_candidates"] > 1 and report["status"] == "incomplete-costs"
|
|
assert "cannot compare the complete candidate set" in diagnostics
|
|
assert "unsupported abstract operation" in diagnostics
|
|
assert not program and "backend_attempts" not in report
|
|
# Latency PEFT and earlier emission do not invoke candidate costing.
|
|
report, program, _ = compile(vgg, pipeline=1)
|
|
assert report is None and program
|
|
report, program, stderr = compile(vgg, emit="EmitSpatial")
|
|
assert report is None and not program
|
|
assert "intermediate emission retains the original schedule" in stderr
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|