NiccoloN
88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
import re
|
|
import shlex
|
|
import subprocess
|
|
from pathlib import Path
|
|
from colorama import Fore, Style
|
|
from subprocess_utils import run_command_with_reporter
|
|
|
|
PIM_PASS_LABELS = (
|
|
("ONNXToSpatialPass", "ONNX to Spatial"),
|
|
("MergeComputeNodesPass", "Merge Compute Nodes"),
|
|
("SpatialToPimPass", "Spatial to PIM"),
|
|
("PimBufferizationPass", "Bufferize PIM"),
|
|
("HostConstantFoldingPass", "Fold Host Constants"),
|
|
("PimMemoryCoalescingPass", "Coalesce Local Memory"),
|
|
("PimLocalMemoryPlanningPass", "Plan Local Memory"),
|
|
("VerificationPass", "Verify PIM"),
|
|
("EmitPimCodePass", "Emit PIM Code"),
|
|
)
|
|
PIM_PASS_LABEL_BY_SUFFIX = dict(PIM_PASS_LABELS)
|
|
TIMING_LINE_RE = re.compile(r"^\s*([0-9]+\.[0-9]+)\s+\(\s*[0-9.]+%\)\s+(.+?)\s*$")
|
|
|
|
|
|
def _parse_pim_pass_timings(output_text):
|
|
pass_timings = {}
|
|
for line in output_text.splitlines():
|
|
match = TIMING_LINE_RE.match(line)
|
|
if not match:
|
|
continue
|
|
|
|
duration = float(match.group(1))
|
|
pass_name = match.group(2)
|
|
for suffix, label in PIM_PASS_LABEL_BY_SUFFIX.items():
|
|
if pass_name.endswith(suffix):
|
|
pass_timings[label] = pass_timings.get(label, 0.0) + duration
|
|
break
|
|
|
|
return pass_timings
|
|
|
|
|
|
def _format_command(cmd):
|
|
return shlex.join(str(arg) for arg in cmd)
|
|
|
|
|
|
def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
|
|
crossbar_size, crossbar_count, core_count=None,
|
|
pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False,
|
|
reporter=None, timeout_sec=None):
|
|
# Define the arguments, with the possibility to set crossbar size and count
|
|
args = [
|
|
network_path,
|
|
"-o",
|
|
output_base,
|
|
"--maccel=PIM",
|
|
"--EmitPimCodegen",
|
|
f"--crossbar-size={crossbar_size}",
|
|
f"--crossbar-count={crossbar_count}",
|
|
]
|
|
if core_count is not None:
|
|
args.append(f"--core-count={core_count}")
|
|
if pim_memory_report != "none":
|
|
args.append(f"--pim-memory-report={pim_memory_report}")
|
|
if raptor_extra_args:
|
|
args.extend(str(arg) for arg in raptor_extra_args)
|
|
if verbose:
|
|
args.append("--enable-timing")
|
|
|
|
cmd = [str(raptor_onnx_path)] + [str(arg) for arg in args]
|
|
if reporter is not None:
|
|
reporter.log(f" Raptor command: {_format_command(cmd)}")
|
|
else:
|
|
print(f"Raptor command: {_format_command(cmd)}")
|
|
|
|
try:
|
|
output_text = run_command_with_reporter(
|
|
cmd,
|
|
cwd=cwd,
|
|
reporter=reporter,
|
|
capture_output=True,
|
|
timeout_sec=timeout_sec,
|
|
)
|
|
if reporter is None:
|
|
print(Fore.GREEN + "Raptor execution successful" + Style.RESET_ALL)
|
|
return _parse_pim_pass_timings(output_text)
|
|
except subprocess.CalledProcessError:
|
|
if reporter is None:
|
|
print(Fore.RED + "Raptor execution failed" + Style.RESET_ALL)
|
|
raise
|