From f1a5e1674d759ebb4eba8d65d97463668f324805 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Sun, 6 Sep 2026 22:04:03 +0300 Subject: [PATCH 1/2] compilation: run a rule's independent commands concurrently execute() ran every command of every rule one after another, so a design's kernel objects cost the SUM of their compiles: 7.9 s of a 17.1 s encoder-MHA build for two kernels, 2.6 s once parallel. Only the rule knows whether its commands are independent, so it declares it; KernelCompilationRule does. --- iron/common/compilation/base.py | 53 +++++++++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index af06dd128..3ab204c84 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -41,6 +41,7 @@ import shutil import zlib import logging +import concurrent.futures import subprocess import importlib.util from dataclasses import dataclass, field @@ -103,14 +104,46 @@ def plan( return [(rule, commands)] + plan(rules, graph, _seen_unavailable=unavailable) -def execute(plan_steps: list[tuple[CompilationRule, list[CompilationCommand]]]) -> None: - for rule, commands in plan_steps: - logging.debug(f"Applying rule: {rule.__class__.__name__}") +def _execute_rule(rule: CompilationRule, commands: list[CompilationCommand]) -> None: + """Run one rule's commands, concurrently when the rule says they are independent. + + A rule that sets `commands_are_independent` emits one command per artifact in + its worklist, with no artifact depending on another's output -- kernel object + compiles are the case that matters, and they were costing the SUM of their + walls (7.9 s of a 17.1 s encoder-MHA build for two kernels). + + Bounded by cores: each kernel compile is a single-threaded clang peaking near + 205 MB of RSS. Set IRON_COMPILE_JOBS to override. + """ + if not getattr(rule, "commands_are_independent", False) or len(commands) < 2: for command in commands: logging.debug(f" Executing command: {command}") - success = command.run() - if not success: + if not command.run(): raise RuntimeError(f"Command failed: {command}") + return + + try: + jobs = int(os.environ.get("IRON_COMPILE_JOBS", "0")) + except ValueError: + jobs = 0 + if jobs <= 0: + jobs = os.cpu_count() or 1 + jobs = min(jobs, len(commands)) + logging.debug(f" Executing {len(commands)} independent commands on {jobs} threads") + + with concurrent.futures.ThreadPoolExecutor(max_workers=jobs) as pool: + results = list(pool.map(lambda c: (c, c.run()), commands)) + # Report the first failure only after every sibling has finished, so a + # failing command cannot leave half-written outputs behind a raised error. + for command, success in results: + if not success: + raise RuntimeError(f"Command failed: {command}") + + +def execute(plan_steps: list[tuple[CompilationRule, list[CompilationCommand]]]) -> None: + for rule, commands in plan_steps: + logging.debug(f"Applying rule: {rule.__class__.__name__}") + _execute_rule(rule, commands) def compile( @@ -474,6 +507,12 @@ def __repr__(self) -> str: class CompilationRule(ABC): """A compilation rule is applied to a artifact graph, producing compilation commands and a transformed artifact graph.""" + #: Set by a rule whose `compile()` emits one command per worklist artifact, + #: none of them consuming another's output. `execute` may then run them + #: concurrently. Default off: a rule that batches dependent steps into one + #: application must keep its order, and only the rule knows which it is. + commands_are_independent: bool = False + @abstractmethod def matches(self, artifact: CompilationArtifactGraph) -> bool: """Return true if this rule can be applied to any artifact in the artifact graph.""" @@ -759,6 +798,10 @@ def _find_working_tool(name, peano_dir, mlir_aie_dir): class KernelCompilationRule(CompilationRule): """Compile KernelObjectArtifacts using Peano (clang++) or xchesscc.""" + # One command per KernelObjectArtifact, each reading its own source and + # writing its own object. + commands_are_independent = True + def __init__(self, peano_dir, mlir_aie_dir, use_chess=False, *args, **kwargs): self.peano_dir = peano_dir self.mlir_aie_dir = mlir_aie_dir From c4806696661db267bc7c3fa2062e3ac6d48703a6 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Fri, 11 Sep 2026 19:46:03 +0300 Subject: [PATCH 2/2] compilation: trim the _execute_rule docstring per review --- iron/common/compilation/base.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 3ab204c84..3d0d5a20a 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -108,12 +108,8 @@ def _execute_rule(rule: CompilationRule, commands: list[CompilationCommand]) -> """Run one rule's commands, concurrently when the rule says they are independent. A rule that sets `commands_are_independent` emits one command per artifact in - its worklist, with no artifact depending on another's output -- kernel object - compiles are the case that matters, and they were costing the SUM of their - walls (7.9 s of a 17.1 s encoder-MHA build for two kernels). - - Bounded by cores: each kernel compile is a single-threaded clang peaking near - 205 MB of RSS. Set IRON_COMPILE_JOBS to override. + its worklist, none of them consuming another's output, so they run in + parallel. Bounded by cores. Set IRON_COMPILE_JOBS to override. """ if not getattr(rule, "commands_are_independent", False) or len(commands) < 2: for command in commands: