ONNX
onnxruntime
onnx-mlir
quantization
fp32
ONNX_Models / scripts /mlir_graph_common.py
purejomo's picture
Finalize public ONNX/ONNX-MLIR validation release
ed3aeeb
Raw
History Blame Contribute Delete
43.9 kB
#!/usr/bin/env python3
"""Shared, read-only MLIR graph extraction utilities for T85.
The parser intentionally consumes the textual artifacts already recorded by
``reports/conversion/ir_stage_coverage.csv``. It does not invoke an MLIR toolchain. The
current ONNX-MLIR printer puts each operation on one physical line, including
very large dense constants; the scanner therefore recognizes constants from a
small prefix and never tokenizes their payload.
"""
from __future__ import annotations
import csv
import hashlib
import html
import json
import math
import os
import re
import tempfile
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable, Iterator
PARSER_SCHEMA_VERSION = "T85_MLIR_GRAPH_V2"
VARIANTS = ("fp32", "public_quantized")
AFFINE_PAIR_IDS = {"LM04", "SG06", "SG07", "SG08", "SP08", "VC03", "VC04"}
SSA_RE = re.compile(r"%[-A-Za-z0-9_.$]+(?:#\d+)?")
SYMBOL_RE = re.compile(r"@[-A-Za-z0-9_.$]+")
OP_PREFIX_RE = re.compile(
r"^\s*(?:(?P<lhs>%[-A-Za-z0-9_.$]+(?:\s*:\s*\d+)?"
r"(?:\s*,\s*%[-A-Za-z0-9_.$]+(?:\s*:\s*\d+)?)*)\s*=\s*)?"
r"(?P<quoted>\"(?P<quoted_name>[A-Za-z_][A-Za-z0-9_.$-]*)\")"
r"|^\s*(?:(?P<lhs_bare>%[-A-Za-z0-9_.$]+(?:\s*:\s*\d+)?"
r"(?:\s*,\s*%[-A-Za-z0-9_.$]+(?:\s*:\s*\d+)?)*)\s*=\s*)?"
r"(?P<bare_name>[A-Za-z_][A-Za-z0-9_.$-]*)"
)
FUNC_RE = re.compile(r"\b(?:func\.func|llvm\.func)\s+@(?P<name>[-A-Za-z0-9_.$]+)")
BLOCK_RE = re.compile(r"^\s*\^(?P<name>[-A-Za-z0-9_.$]+)(?:\((?P<args>.*)\))?\s*:")
ONNX_NODE_NAME_RE = re.compile(r'onnx_node_name\s*=\s*"((?:[^"\\]|\\.)*)"')
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def canonical_json_sha256(value: Any) -> str:
payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
def repo_path(path: Path, root: Path) -> str:
resolved = path.resolve()
try:
return str(resolved.relative_to(root.resolve()))
except ValueError:
return str(resolved)
def resolve_coverage_path(value: str, root: Path) -> Path:
"""Resolve source paths recorded before the current repository layout.
The coverage matrix remains immutable provenance. A missing absolute path
is relocated only by a recognized repository anchor; arbitrary basenames
are never searched.
"""
candidate = Path(value)
if candidate.is_file():
return candidate.resolve()
parts = candidate.parts
for anchor in ("models", "reports", "logs", "environment", "configs"):
if anchor in parts:
relocated = root.joinpath(*parts[parts.index(anchor) :]).resolve()
if relocated.is_file():
return relocated
if not candidate.is_absolute():
relocated = (root / candidate).resolve()
if relocated.is_file():
return relocated
raise FileNotFoundError(f"coverage artifact cannot be relocated: {value}")
def load_csv(path: Path) -> list[dict[str, str]]:
with path.open(newline="", encoding="utf-8") as handle:
return list(csv.DictReader(handle))
def atomic_text(path: Path, value: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
handle.write(value)
temporary = Path(handle.name)
os.replace(temporary, path)
def atomic_json(path: Path, value: Any) -> None:
atomic_text(path, json.dumps(value, indent=2, ensure_ascii=False, sort_keys=True) + "\n")
def atomic_csv(path: Path, rows: list[dict[str, Any]], fields: list[str]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile("w", encoding="utf-8", newline="", dir=path.parent, delete=False) as handle:
writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
writer.writeheader()
writer.writerows(rows)
temporary = Path(handle.name)
os.replace(temporary, path)
def file_record(path: Path, root: Path) -> dict[str, Any]:
return {"path": repo_path(path, root), "bytes": path.stat().st_size, "sha256": sha256(path)}
def split_top_level(value: str, delimiter: str = ",") -> list[str]:
parts: list[str] = []
start = 0
depths = {"(": 0, "[": 0, "{": 0, "<": 0}
closing = {")": "(", "]": "[", "}": "{", ">": "<"}
quoted = False
escaped = False
for index, char in enumerate(value):
if quoted:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
quoted = False
continue
if char == '"':
quoted = True
elif char in depths:
depths[char] += 1
elif char in closing:
key = closing[char]
depths[key] = max(0, depths[key] - 1)
elif char == delimiter and not any(depths.values()):
parts.append(value[start:index].strip())
start = index + 1
parts.append(value[start:].strip())
return [part for part in parts if part]
def expand_lhs(lhs: str | None) -> list[str]:
if not lhs:
return []
results: list[str] = []
for item in split_top_level(lhs):
item = item.strip()
match = re.fullmatch(r"(%[-A-Za-z0-9_.$]+)\s*:\s*(\d+)", item)
if match:
results.extend(f"{match.group(1)}#{index}" for index in range(int(match.group(2))))
else:
results.append(item)
return results
def operation_prefix(text: str) -> tuple[str | None, str] | None:
prefix = text[:4096]
match = OP_PREFIX_RE.match(prefix)
if not match:
return None
lhs = match.group("lhs") or match.group("lhs_bare")
name = match.group("quoted_name") or match.group("bare_name")
if name in {"module", "attributes"} or name.startswith("#"):
return None
return lhs, name
def _balanced_for_statement(text: str) -> bool:
"""Return whether parentheses/square brackets are balanced outside strings.
Angle brackets are deliberately ignored because the ``->`` token would
otherwise look like an unmatched close. Region-opening braces complete an
operation header and are handled by the scope parser.
"""
paren = square = curly = angle = 0
quoted = escaped = False
index = 0
while index < len(text):
char = text[index]
if not quoted and char == "/" and index + 1 < len(text) and text[index + 1] == "/":
break
if quoted:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
quoted = False
elif char == '"':
quoted = True
elif char == "(":
paren += 1
elif char == ")":
paren -= 1
elif char == "[":
square += 1
elif char == "]":
square -= 1
elif char == "{":
curly += 1
elif char == "}":
curly -= 1
elif char == "<":
angle += 1
elif char == ">" and not (index > 0 and text[index - 1] == "-"):
angle -= 1
index += 1
region_header = text.rstrip().endswith("{") and curly == 1
return (
paren <= 0
and square <= 0
and angle <= 0
and (curly <= 0 or region_header)
and not quoted
)
def _is_dense_constant(prefix: str) -> bool:
return "dense<" in prefix and any(
token in prefix
for token in ("onnx.Constant", "krnl.global", "arith.constant", "llvm.mlir.global")
)
@dataclass
class Statement:
text: str
start_line: int
end_line: int
source_sha256: str
dense_payload_omitted: bool = False
def iter_operation_statements(path: Path) -> Iterator[tuple[str, Statement]]:
"""Yield ``(kind, statement)`` events without materializing dense payloads.
Kinds are ``line`` for structural syntax and ``operation`` for operation
statements. Multiline operation syntax is joined until operand delimiters
close. Existing giant dense constants are represented by prefix+tail only.
"""
pending: list[str] = []
pending_hash = hashlib.sha256()
start_line = 0
with path.open(encoding="utf-8", errors="replace") as handle:
for line_number, raw in enumerate(handle, 1):
if pending:
pending_hash.update(raw.encode("utf-8"))
pending.append(raw)
combined = "".join(pending)
if _balanced_for_statement(combined):
yield "operation", Statement(
combined,
start_line,
line_number,
pending_hash.hexdigest(),
)
pending = []
pending_hash = hashlib.sha256()
continue
stripped = raw.strip()
prefix = raw[:4096]
if not stripped or stripped.startswith("//"):
continue
# Function/block/module declarations are structural syntax rather
# than executable operations. They happen to match the generic
# bare-operation token grammar, so classify them first.
if (
FUNC_RE.search(prefix)
or BLOCK_RE.match(prefix)
or stripped.startswith(("module ", "#", "}"))
):
yield "line", Statement(raw, line_number, line_number, hashlib.sha256(raw.encode()).hexdigest())
continue
if operation_prefix(prefix) is None:
yield "line", Statement(raw, line_number, line_number, hashlib.sha256(raw.encode()).hexdigest())
continue
digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()
if _is_dense_constant(prefix):
compact = raw if len(raw) <= 16384 else raw[:8192] + " ... <DENSE_PAYLOAD_OMITTED> ... " + raw[-4096:]
yield "operation", Statement(compact, line_number, line_number, digest, True)
elif _balanced_for_statement(raw):
yield "operation", Statement(raw, line_number, line_number, digest)
else:
start_line = line_number
pending = [raw]
pending_hash.update(raw.encode("utf-8"))
if pending:
raise ValueError(f"unterminated operation starting at {path}:{start_line}")
def _extract_parenthesized(value: str, start: int) -> str:
depth = 0
quoted = escaped = False
for index in range(start, len(value)):
char = value[index]
if quoted:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
quoted = False
continue
if char == '"':
quoted = True
elif char == "(":
depth += 1
elif char == ")":
depth -= 1
if depth == 0:
return value[start + 1 : index]
return value[start + 1 :]
def _loop_region_args(text: str, op_name: str) -> list[str]:
if op_name not in {"affine.for", "scf.for", "scf.parallel", "affine.parallel"}:
return []
args: list[str] = []
induction = re.search(r"\b(?:affine|scf)\.(?:for|parallel)\s+(%[-A-Za-z0-9_.$]+)", text)
if induction:
args.append(induction.group(1))
match = re.search(r"\biter_args\s*\((.*?)\)\s*(?:->|\{)", text, re.S)
if match:
for item in split_top_level(match.group(1)):
name = re.match(r"\s*(%[-A-Za-z0-9_.$]+)\s*=", item)
if name:
args.append(name.group(1))
return args
def _extract_operands(text: str, lhs: str | None, op_name: str, dense: bool) -> list[str]:
if dense or op_name in {"onnx.Constant", "krnl.global", "arith.constant", "llvm.mlir.global"}:
return []
prefix_match = OP_PREFIX_RE.match(text)
if not prefix_match:
return []
token_end = prefix_match.end()
remainder = text[token_end:]
if prefix_match.group("quoted"):
open_index = remainder.find("(")
operand_area = _extract_parenthesized(remainder, open_index) if open_index >= 0 else ""
else:
# Type and attribute suffixes do not carry SSA operands. The first
# top-level `` : `` is the stable printer boundary for current IR.
operand_area = remainder
colon = operand_area.find(" : ")
if colon >= 0:
operand_area = operand_area[:colon]
operands = SSA_RE.findall(operand_area)
region_args = set(_loop_region_args(text, op_name))
lhs_names = set(expand_lhs(lhs))
return [name for name in operands if name not in region_args and name not in lhs_names]
def _element_type(container_type: str) -> str:
match = re.search(r"(?:tensor|memref|vector)<(.+)>", container_type)
if not match:
return container_type.strip()
body = split_top_level(match.group(1))[0]
dtype = re.search(r"(?:^|x)((?:u|s)?i\d+|f\d+|bf16|index)$", body.strip())
return dtype.group(1) if dtype else "UNKNOWN"
def _extract_result_types(text: str, result_count: int, op_name: str) -> list[str]:
if result_count == 0:
return []
tail = text[-8192:]
arrow = tail.rfind("->")
if arrow >= 0:
value = tail[arrow + 2 :].strip().rstrip("{").strip()
else:
# In bare custom syntax ``: type`` is often an operand annotation
# (for example affine.load and arith.cmpi), not the result type. Only
# constant/global printers have an unambiguous trailing result type.
if op_name in {
"onnx.Constant", "krnl.global", "arith.constant", "llvm.mlir.global",
"memref.alloc", "memref.alloca", "memref.get_global",
}:
colon = tail.rfind(" : ")
if colon >= 0:
value = tail[colon + 3 :].strip().rstrip("{").strip()
elif op_name == "arith.constant" and re.search(r"\b(?:true|false)\b", tail):
value = "i1"
else:
value = "UNKNOWN"
elif op_name in {"affine.load", "memref.load"}:
colon = tail.rfind(" : ")
container = tail[colon + 3 :].strip() if colon >= 0 else "UNKNOWN"
value = _element_type(container)
elif op_name in {"arith.cmpi", "arith.cmpf"}:
value = "i1"
elif op_name in {"affine.apply", "affine.min", "affine.max", "memref.dim"}:
value = "index"
elif op_name.startswith(("arith.", "math.")):
if " to " in tail:
value = tail.rsplit(" to ", 1)[1].strip().rstrip("{").strip()
else:
colon = tail.rfind(" : ")
value = tail[colon + 3 :].strip().rstrip("{").strip() if colon >= 0 else "UNKNOWN"
elif op_name in {"memref.cast", "memref.reinterpret_cast", "memref.subview"} and " to " in tail:
value = tail.rsplit(" to ", 1)[1].strip().rstrip("{").strip()
elif op_name == "builtin.unrealized_conversion_cast" and " to " in tail:
value = tail.rsplit(" to ", 1)[1].strip().rstrip("{").strip()
else:
value = "UNKNOWN"
if value.startswith("(") and value.endswith(")"):
types = split_top_level(value[1:-1])
else:
types = [value]
if len(types) < result_count:
types.extend(["UNKNOWN"] * (result_count - len(types)))
return types[:result_count]
def _extract_operand_types(text: str, op_name: str, operand_count: int) -> list[str]:
if operand_count == 0:
return []
tail = text[-16384:]
generic = re.search(r"\)\s*:\s*\((.*)\)\s*->", tail, re.S)
if generic:
values = split_top_level(generic.group(1))
else:
colon = tail.rfind(" : ")
value = tail[colon + 3 :].strip().rstrip("{").strip() if colon >= 0 else "UNKNOWN"
if " to " in value:
value = value.split(" to ", 1)[0].strip()
values = [value] * operand_count
if len(values) < operand_count:
values.extend(["UNKNOWN"] * (operand_count - len(values)))
return values[:operand_count]
def type_facts(type_value: str) -> tuple[str, str, str]:
"""Return ``(shape, dtype, bytes)`` from a textual MLIR type."""
value = type_value.strip()
if value in {"", "UNKNOWN"}:
return "UNKNOWN", "UNKNOWN", "UNKNOWN"
if value == "none":
return "NOT_APPLICABLE", "none", "NOT_APPLICABLE"
container = re.fullmatch(r"(?:tensor|memref|vector)<(.+)>", value)
if container:
body = split_top_level(container.group(1))[0]
dtype_match = re.search(r"(?:^|x)((?:u|s)?i\d+|f\d+|bf16|index)$", body)
if not dtype_match:
return "UNKNOWN", "UNKNOWN", "UNKNOWN"
dtype = dtype_match.group(1)
prefix = body[: dtype_match.start(1)].rstrip("x")
dims = prefix.split("x") if prefix else []
else:
dtype_match = re.fullmatch(r"(?:u|s)?i\d+|f\d+|bf16|index", value)
if not dtype_match:
return "UNKNOWN", "UNKNOWN", "UNKNOWN"
dtype = value
dims = []
shape = "[" + ",".join(dims) + "]"
bits_match = re.search(r"(\d+)$", dtype)
if dtype == "index" or not bits_match or any(not dim.isdigit() for dim in dims):
return shape, dtype, "UNKNOWN"
element_bytes = max(1, (int(bits_match.group(1)) + 7) // 8)
elements = 1
for dim in dims:
elements *= int(dim)
return shape, dtype, str(elements * element_bytes)
def quantization_role(op_name: str) -> str:
lower = op_name.lower()
if "dequantize" in lower:
return "DEQUANTIZE"
if "requant" in lower:
return "REQUANTIZE"
if "dynamicquantize" in lower:
return "DYNAMIC_QUANTIZE"
if "quantizelinear" in lower or lower.endswith(".quantize"):
return "QUANTIZE"
if any(token in lower for token in ("qlinear", "matmulinteger", "convinteger")):
return "QUANTIZED_OPERATOR"
return "NONE"
def _symbol_references(text: str) -> list[str]:
# Remove quoted strings so model/node names containing '@' do not become
# symbol references. MLIR symbol uses such as ``func = @main_graph`` stay.
without_strings = re.sub(r'"(?:[^"\\]|\\.)*"', '""', text)
return SYMBOL_RE.findall(without_strings)
@dataclass
class Scope:
kind: str
name: str
function: str
region_path: str
block_id: str
parent_op_node_id: str = ""
region_args: list[str] = field(default_factory=list)
@dataclass
class ParsedGraph:
source: Path
operations: list[dict[str, Any]]
ssa_edges: list[dict[str, Any]]
relations: list[dict[str, Any]]
definitions: list[dict[str, Any]]
diagnostics: list[dict[str, Any]]
functions: list[str]
block_count: int
unresolved_use_count: int
duplicate_definition_count: int
producer_after_consumer_count: int
def parse_mlir(path: Path, graph_id: str) -> ParsedGraph:
operations: list[dict[str, Any]] = []
relations: list[dict[str, Any]] = []
definitions: list[dict[str, Any]] = []
diagnostics: list[dict[str, Any]] = []
scopes: list[Scope] = [Scope("module", "module", "module", "module", "module")]
function_names: list[str] = []
block_ids: set[tuple[str, str]] = set()
pending_region_scope: Scope | None = None
block_orders: Counter[tuple[str, str, str]] = Counter()
node_sequence = 0
def current() -> Scope:
return scopes[-1]
for kind, statement in iter_operation_statements(path):
text_value = statement.text
stripped = text_value.strip()
if kind == "line":
func_match = FUNC_RE.search(text_value[:8192])
if func_match:
name = func_match.group("name")
function_names.append(name)
scope = Scope("function", name, name, f"{name}/region0", "entry")
scopes.append(scope)
block_ids.add((name, "entry"))
signature = text_value[: text_value.rfind("->") if "->" in text_value else len(text_value)]
for arg in SSA_RE.findall(signature):
definitions.append({
"ssa_value": arg,
"producer_kind": "FUNCTION_ARG",
"producer_node_id": "",
"function": name,
"region_path": scope.region_path,
"block_id": scope.block_id,
"source_line": statement.start_line,
})
continue
block_match = BLOCK_RE.match(text_value[:8192])
if block_match:
name = block_match.group("name")
base = current()
block = Scope("block", name, base.function, base.region_path, name, base.parent_op_node_id)
if scopes and scopes[-1].kind == "block":
scopes.pop()
scopes.append(block)
block_ids.add((block.function, block.block_id))
for arg in SSA_RE.findall(block_match.group("args") or ""):
definitions.append({
"ssa_value": arg,
"producer_kind": "BLOCK_ARG",
"producer_node_id": "",
"function": block.function,
"region_path": block.region_path,
"block_id": block.block_id,
"source_line": statement.start_line,
})
continue
if stripped.startswith("}"):
if len(scopes) > 1:
scopes.pop()
continue
if stripped.startswith(("module ", "#", "//")) or stripped in {"{", "}"}:
continue
diagnostics.append({"code": "UNPARSED_STRUCTURAL_LINE", "line": statement.start_line, "text": stripped[:240]})
continue
parsed_prefix = operation_prefix(text_value)
if parsed_prefix is None:
diagnostics.append({"code": "UNPARSED_OPERATION", "line": statement.start_line, "text": stripped[:240]})
continue
lhs, op_name = parsed_prefix
results = expand_lhs(lhs)
dense = statement.dense_payload_omitted or _is_dense_constant(text_value[:4096])
operands = _extract_operands(text_value, lhs, op_name, dense)
result_types = _extract_result_types(text_value, len(results), op_name)
operand_types = _extract_operand_types(text_value, op_name, len(operands))
facts = [type_facts(type_value) for type_value in result_types]
scope = current()
order_key = (scope.function, scope.region_path, scope.block_id)
block_order = block_orders[order_key]
block_orders[order_key] += 1
node_id = f"{graph_id}:op{node_sequence:06d}"
node_sequence += 1
onnx_name_match = ONNX_NODE_NAME_RE.search(text_value if len(text_value) < 200000 else text_value[:65536])
entry = {
"node_id": node_id,
"node_kind": "OPERATION",
"function": scope.function,
"region_path": scope.region_path,
"block_id": scope.block_id,
"block_order": block_order,
"static_order": len(operations),
"operation": "func.return" if op_name == "return" else op_name,
"dialect": ("func" if op_name == "return" else op_name.split(".", 1)[0]),
"results": results,
"result_types": result_types,
"result_shapes": [item[0] for item in facts],
"result_dtypes": [item[1] for item in facts],
"result_bytes": [item[2] for item in facts],
"operands": operands,
"operand_types": operand_types,
"symbol_references": _symbol_references(text_value),
"quantization_role": quantization_role(op_name),
"onnx_node_name": onnx_name_match.group(1) if onnx_name_match else "",
"source_start_line": statement.start_line,
"source_end_line": statement.end_line,
"source_statement_sha256": statement.source_sha256,
"dense_payload_omitted": dense,
"parent_op_node_id": scope.parent_op_node_id,
"opens_region": False,
}
operations.append(entry)
for result, type_value in zip(results, result_types):
definitions.append({
"ssa_value": result,
"producer_kind": "OPERATION",
"producer_node_id": node_id,
"function": scope.function,
"region_path": scope.region_path,
"block_id": scope.block_id,
"source_line": statement.start_line,
"type": type_value,
})
# Program order is represented separately from SSA data dependency.
prior = next(
(
candidate for candidate in reversed(operations[:-1])
if candidate["function"] == scope.function
and candidate["region_path"] == scope.region_path
and candidate["block_id"] == scope.block_id
),
None,
)
if prior:
relations.append({
"relation_type": "PROGRAM_ORDER",
"source_node_id": prior["node_id"],
"target_node_id": node_id,
"source_block_id": scope.block_id,
"target_block_id": scope.block_id,
"detail": "consecutive operations in textual block order",
})
if scope.parent_op_node_id:
prior_in_region = [
item for item in operations[:-1]
if item["parent_op_node_id"] == scope.parent_op_node_id
and item["region_path"] == scope.region_path
]
if not prior_in_region:
relations.append({
"relation_type": "REGION_CONTAINS",
"source_node_id": scope.parent_op_node_id,
"target_node_id": node_id,
"source_block_id": "",
"target_block_id": scope.block_id,
"detail": scope.region_path,
})
if op_name in {"affine.yield", "scf.yield"}:
relations.append({
"relation_type": "REGION_YIELD",
"source_node_id": node_id,
"target_node_id": scope.parent_op_node_id,
"source_block_id": scope.block_id,
"target_block_id": "",
"detail": "structured-region yield/back-edge to parent operation",
})
# Explicit block successors are absent in the current snapshot, but are
# recorded when printed by cf/llvm branch operations.
successors = re.findall(r"\^([-A-Za-z0-9_.$]+)", text_value[:65536])
for successor in successors:
relations.append({
"relation_type": "CFG_SUCCESSOR",
"source_node_id": node_id,
"target_node_id": f"BLOCK:{scope.function}:{successor}",
"source_block_id": scope.block_id,
"target_block_id": successor,
"detail": op_name,
})
# A trailing top-level opening brace introduces an operation region.
compact_tail = text_value.rstrip()
opens_region = compact_tail.endswith("{") and op_name not in {"onnx.EntryPoint"}
if opens_region:
entry["opens_region"] = True
region_index = sum(1 for item in scopes if item.parent_op_node_id == node_id)
region_args = _loop_region_args(text_value, op_name)
new_scope = Scope(
"region",
f"region{region_index}",
scope.function,
f"{scope.region_path}/{node_id.rsplit(':', 1)[-1]}.region{region_index}",
f"{node_id.rsplit(':', 1)[-1]}.region{region_index}.entry",
node_id,
region_args,
)
scopes.append(new_scope)
block_ids.add((new_scope.function, new_scope.block_id))
for arg in region_args:
definitions.append({
"ssa_value": arg,
"producer_kind": "BLOCK_ARG",
"producer_node_id": "",
"function": new_scope.function,
"region_path": new_scope.region_path,
"block_id": new_scope.block_id,
"source_line": statement.start_line,
})
# Structured affine/scf loops have implicit control flow even when the
# printer omits an explicit block label or terminator. Represent the
# iteration edge from the last direct child (normally affine.yield) back to
# the first direct child. This is separate from SSA and program order.
direct_children: dict[str, list[dict[str, Any]]] = defaultdict(list)
for operation in operations:
if operation["parent_op_node_id"]:
direct_children[operation["parent_op_node_id"]].append(operation)
for parent in operations:
if parent["operation"] not in {"affine.for", "scf.for", "affine.parallel", "scf.parallel"}:
continue
children = direct_children.get(parent["node_id"], [])
if not children:
diagnostics.append({
"code": "EMPTY_STRUCTURED_LOOP_REGION",
"line": parent["source_start_line"],
"node_id": parent["node_id"],
})
continue
relations.append({
"relation_type": "LOOP_BACKEDGE",
"source_node_id": children[-1]["node_id"],
"target_node_id": children[0]["node_id"],
"source_block_id": children[-1]["block_id"],
"target_block_id": children[0]["block_id"],
"detail": f"implicit next iteration of {parent['operation']} ({parent['node_id']})",
})
# Resolve definitions after parsing because module-level operations and
# nested-region values can be referenced outside their physical scope.
definitions_by_value: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
definitions_by_scope: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
for definition in definitions:
definitions_by_value[(definition["function"], definition["ssa_value"])].append(definition)
definitions_by_scope[(
definition["function"], definition["region_path"], definition["ssa_value"]
)].append(definition)
duplicate_count = sum(max(0, len(items) - 1) for items in definitions_by_scope.values())
node_by_id = {item["node_id"]: item for item in operations}
ssa_edges: list[dict[str, Any]] = []
unresolved = 0
producer_after = 0
for consumer in operations:
for operand_index, operand in enumerate(consumer["operands"]):
candidates = definitions_by_value.get((consumer["function"], operand), [])
# MLIR values defined in a parent region are visible in nested
# regions; sibling/child definitions are not. Prefer the nearest
# lexical ancestor to handle reused block-argument names.
consumer_region = consumer["region_path"]
lexical = [
item for item in candidates
if consumer_region == item["region_path"]
or consumer_region.startswith(item["region_path"] + "/")
]
if lexical:
max_depth = max(item["region_path"].count("/") for item in lexical)
candidates = [item for item in lexical if item["region_path"].count("/") == max_depth]
if not candidates and consumer["function"] != "module":
candidates = definitions_by_value.get(("module", operand), [])
if not candidates:
unresolved += 1
ssa_edges.append({
"edge_id": f"{graph_id}:ssa{len(ssa_edges):07d}",
"producer_kind": "UNRESOLVED",
"producer_node_id": "",
"producer_result": operand,
"consumer_node_id": consumer["node_id"],
"consumer_operand_index": operand_index,
"consumer_operand": operand,
"producer_static_order": "",
"consumer_static_order": consumer["static_order"],
"producer_before_consumer": "UNKNOWN",
})
continue
# Prefer the closest definition that is textually before the use;
# MLIR names are normally unique so this also exposes duplicates.
before = [
item for item in candidates
if item.get("producer_node_id", "") == ""
or node_by_id[item["producer_node_id"]]["static_order"] < consumer["static_order"]
]
definition = before[-1] if before else candidates[0]
producer_node = definition.get("producer_node_id", "")
producer_order: int | str = ""
is_before: bool | str = True
if producer_node:
producer_order = node_by_id[producer_node]["static_order"]
is_before = int(producer_order) < int(consumer["static_order"])
if not is_before:
producer_after += 1
ssa_edges.append({
"edge_id": f"{graph_id}:ssa{len(ssa_edges):07d}",
"producer_kind": definition["producer_kind"],
"producer_node_id": producer_node,
"producer_result": operand,
"consumer_node_id": consumer["node_id"],
"consumer_operand_index": operand_index,
"consumer_operand": operand,
"producer_static_order": producer_order,
"consumer_static_order": consumer["static_order"],
"producer_before_consumer": is_before,
})
return ParsedGraph(
source=path,
operations=operations,
ssa_edges=ssa_edges,
relations=relations,
definitions=definitions,
diagnostics=diagnostics,
functions=sorted(set(function_names)),
block_count=len(block_ids),
unresolved_use_count=unresolved,
duplicate_definition_count=duplicate_count,
producer_after_consumer_count=producer_after,
)
def graph_fingerprint(
source_path: Path,
source_sha256: str,
source_status: str,
stage: str,
inkscape_version: str,
runtime_order_status: str,
) -> str:
return canonical_json_sha256({
"parser_schema": PARSER_SCHEMA_VERSION,
"renderer_schema": "T85_COMPACT_ORDER_GRID_V2",
"implementation_sha256": sha256(Path(__file__)),
"source_path": str(source_path),
"source_sha256": source_sha256,
"source_bytes": source_path.stat().st_size,
"source_status": source_status,
"stage": stage,
"inkscape_version": inkscape_version,
"runtime_order_status": runtime_order_status,
"layout": {
"order": "STATIC_MLIR_PROGRAM_ORDER",
"ssa_edge": "solid-blue",
"program_order": "solid-gray",
"region": "dashed-purple",
"cfg": "dashed-red",
"module_metadata_rendered": False,
},
})
def _dialect_color(dialect: str) -> str:
return {
"onnx": "#dbeafe",
"func": "#dcfce7",
"affine": "#fef3c7",
"scf": "#fde68a",
"memref": "#ede9fe",
"arith": "#fae8ff",
"krnl": "#fee2e2",
"builtin": "#e2e8f0",
"llvm": "#fed7aa",
}.get(dialect, "#f1f5f9")
def render_execution_dependency_svg(
parsed: ParsedGraph,
*,
title: str,
graph_id: str,
stage: str,
runtime_order_status: str = "RUNTIME_ORDER_UNAVAILABLE",
) -> tuple[str, dict[str, Any]]:
"""Render a compact, zoomable operation-order/SSA graph.
Every visible rectangle is one operation inventory row. Module metadata
operations (for example ``onnx.EntryPoint``) remain in the CSV evidence but
are not part of the compute graph. Large lowered graphs use a wider grid so
that PNG dimensions remain bounded while the SVG preserves per-node titles.
"""
operations = [item for item in parsed.operations if item["function"] != "module"]
count = len(operations)
if count <= 80:
columns = 4
elif count <= 400:
columns = 8
elif count <= 1600:
columns = 16
elif count <= 8000:
columns = 32
else:
columns = 48
cell_width = 142
cell_height = 22
gap_x = 8
gap_y = 8
margin_x = 32
header_height = 118
rows = max(1, math.ceil(count / columns))
width = max(920, margin_x * 2 + columns * (cell_width + gap_x))
height = header_height + rows * (cell_height + gap_y) + 40
positions: dict[str, tuple[float, float]] = {}
for index, operation in enumerate(operations):
row, column = divmod(index, columns)
# Alternate direction per row to keep consecutive program-order edges
# short at row boundaries while preserving the numeric labels.
visual_column = column if row % 2 == 0 else columns - 1 - column
x = margin_x + visual_column * (cell_width + gap_x)
y = header_height + row * (cell_height + gap_y)
positions[operation["node_id"]] = (x, y)
elements: list[str] = [
'<?xml version="1.0" encoding="UTF-8"?>',
f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
"<defs>",
'<marker id="arrow-ssa" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="4" markerHeight="4" orient="auto-start-reverse"><path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/></marker>',
'<marker id="arrow-order" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="3" markerHeight="3" orient="auto-start-reverse"><path d="M 0 0 L 10 5 L 0 10 z" fill="#94a3b8"/></marker>',
'<marker id="arrow-control" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="4" markerHeight="4" orient="auto-start-reverse"><path d="M 0 0 L 10 5 L 0 10 z" fill="#dc2626"/></marker>',
"</defs>",
f'<rect width="{width}" height="{height}" fill="#ffffff"/>',
f'<text x="{margin_x}" y="30" font-family="sans-serif" font-size="20" font-weight="700">{html.escape(title)}</text>',
f'<text x="{margin_x}" y="54" font-family="sans-serif" font-size="12">stage={html.escape(stage)} · order=STATIC_MLIR_PROGRAM_ORDER · runtime={html.escape(runtime_order_status)}</text>',
f'<text x="{margin_x}" y="74" font-family="sans-serif" font-size="12">{count:,} compute operations · {len(parsed.ssa_edges):,} total SSA uses / operation-to-operation edges rendered below · {parsed.block_count:,} blocks</text>',
f'<text x="{margin_x}" y="94" font-family="sans-serif" font-size="11">blue=SSA def-use · gray=program order · purple=region · red=CFG/loop/yield · orange border=quantization-related</text>',
]
rendered_ssa = 0
for edge in parsed.ssa_edges:
source = positions.get(edge["producer_node_id"])
target = positions.get(edge["consumer_node_id"])
if not source or not target:
continue
sx, sy = source
tx, ty = target
elements.append(
f'<line class="edge ssa-edge" data-edge-id="{html.escape(str(edge["edge_id"]))}" '
f'x1="{sx + cell_width / 2:.1f}" y1="{sy + cell_height / 2:.1f}" '
f'x2="{tx + cell_width / 2:.1f}" y2="{ty + cell_height / 2:.1f}" '
'stroke="#2563eb" stroke-width="0.7" opacity="0.16" marker-end="url(#arrow-ssa)"/>'
)
rendered_ssa += 1
relation_counts: Counter[str] = Counter()
for index, relation in enumerate(parsed.relations):
source = positions.get(relation["source_node_id"])
target = positions.get(relation["target_node_id"])
if not source or not target:
continue
relation_type = relation["relation_type"]
relation_counts[relation_type] += 1
if relation_type == "PROGRAM_ORDER":
color, opacity, dash, marker = "#94a3b8", "0.22", "", "arrow-order"
elif relation_type == "REGION_CONTAINS":
color, opacity, dash, marker = "#7c3aed", "0.32", "4 3", "arrow-order"
else:
color, opacity, dash, marker = "#dc2626", "0.42", "5 3", "arrow-control"
sx, sy = source
tx, ty = target
dash_attr = f' stroke-dasharray="{dash}"' if dash else ""
elements.append(
f'<line class="edge {relation_type.lower().replace("_", "-")}" data-relation-index="{index}" '
f'x1="{sx + cell_width / 2:.1f}" y1="{sy + cell_height / 2:.1f}" '
f'x2="{tx + cell_width / 2:.1f}" y2="{ty + cell_height / 2:.1f}" '
f'stroke="{color}" stroke-width="0.65" opacity="{opacity}"{dash_attr} marker-end="url(#{marker})"/>'
)
for operation in operations:
x, y = positions[operation["node_id"]]
op_name = operation["operation"]
short = op_name if len(op_name) <= 18 else op_name[:16] + "…"
tooltip = (
f"order={operation['static_order']} | {op_name} | function={operation['function']} | "
f"block={operation['block_id']} | source={operation['source_start_line']}:{operation['source_end_line']}"
)
quantized = operation.get("quantization_role", "NONE") != "NONE"
stroke = "#ea580c" if quantized else "#475569"
stroke_width = "1.5" if quantized else "0.55"
quant_attr = html.escape(str(operation.get("quantization_role", "NONE")))
elements.extend([
f'<g class="node operation-node" data-node-id="{html.escape(operation["node_id"])}" '
f'data-static-order="{operation["static_order"]}" data-operation="{html.escape(op_name)}" data-quantization-role="{quant_attr}">',
f'<title>{html.escape(tooltip)}</title>',
f'<rect x="{x:.1f}" y="{y:.1f}" width="{cell_width}" height="{cell_height}" rx="3" '
f'fill="{_dialect_color(operation["dialect"])}" stroke="{stroke}" stroke-width="{stroke_width}"/>',
f'<text x="{x + 4:.1f}" y="{y + 14:.1f}" font-family="monospace" font-size="8" fill="#0f172a">'
f'{operation["static_order"]:05d} {html.escape(short)}</text>',
"</g>",
])
elements.append("</svg>")
metadata = {
"graph_id": graph_id,
"width": width,
"height": height,
"columns": columns,
"rendered_operation_nodes": count,
"excluded_module_metadata_operations": len(parsed.operations) - count,
"rendered_ssa_edges": rendered_ssa,
"quantization_related_operation_nodes": sum(
item.get("quantization_role", "NONE") != "NONE" for item in operations
),
"rendered_relation_edges": sum(relation_counts.values()),
"rendered_relation_counts": dict(sorted(relation_counts.items())),
"order_semantics": "STATIC_MLIR_PROGRAM_ORDER",
"runtime_order_status": runtime_order_status,
}
return "\n".join(elements) + "\n", metadata