ilgeco
254 lines
9.8 KiB
Python
254 lines
9.8 KiB
Python
#!/usr/bin/env python3.13
|
|
|
|
import argparse
|
|
import re
|
|
from collections import Counter, defaultdict
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
|
|
OP_PATTERNS = {
|
|
"tensor.extract_slice": re.compile(r"\btensor\.extract_slice\b"),
|
|
"tensor.insert_slice": re.compile(r"\btensor\.insert_slice\b"),
|
|
"spat.channel_send": re.compile(r"\bspat\.channel_send\b"),
|
|
"spat.channel_receive": re.compile(r"\bspat\.channel_receive\b"),
|
|
"scf.for": re.compile(r"\bscf\.for\b"),
|
|
"tensor.empty": re.compile(r"\btensor\.empty\b"),
|
|
}
|
|
|
|
VALUE_RE = re.compile(r"^\s*(%[\w.$-]+)\s*=\s*(.+)$")
|
|
TYPE_RE = re.compile(r":\s*([^:]+?)\s*(?:to|into|$)")
|
|
CHANNEL_RE = re.compile(r"channel\s+(%c[-\w.$]+)")
|
|
FROM_TO_RE = re.compile(r"from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)")
|
|
EXTRACT_SLICE_RE = re.compile(
|
|
r"^\s*(%[\w.$-]+)\s*=\s*tensor\.extract_slice\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+to\s+(.*)$"
|
|
)
|
|
INSERT_SLICE_RE = re.compile(
|
|
r"^\s*(%[\w.$-]+)\s*=\s*tensor\.insert_slice\s+(%[\w.$-]+)\s+into\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+into\s+(.*)$"
|
|
)
|
|
CHANNEL_RECEIVE_RE = re.compile(
|
|
r"^\s*(%[\w.$-]+)\s*=\s*spat\.channel_receive\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
|
|
)
|
|
CHANNEL_SEND_RE = re.compile(
|
|
r"^\s*spat\.channel_send\s+(%[\w.$-]+)\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
|
|
)
|
|
CONST_INDEX_RE = re.compile(r"^\s*(%c[\w.$-]+)\s*=\s*arith\.constant\s+(-?\d+)\s*:\s*index\b")
|
|
|
|
|
|
@dataclass
|
|
class ChainGroup:
|
|
kind: str
|
|
signature: str
|
|
count: int = 0
|
|
first_line: int = 0
|
|
last_line: int = 0
|
|
fragment_type: str = ""
|
|
dest_type: str = ""
|
|
varying_dims: set[int] = field(default_factory=set)
|
|
rows: list[int | None] = field(default_factory=list)
|
|
channels: list[int | None] = field(default_factory=list)
|
|
sources: list[int | None] = field(default_factory=list)
|
|
targets: list[int | None] = field(default_factory=list)
|
|
|
|
def add(self,
|
|
line_no: int,
|
|
fragment_type: str,
|
|
dest_type: str,
|
|
offsets: list[str],
|
|
channel: int | None = None,
|
|
source: int | None = None,
|
|
target: int | None = None) -> None:
|
|
self.count += 1
|
|
if self.first_line == 0:
|
|
self.first_line = line_no
|
|
self.last_line = line_no
|
|
self.fragment_type = fragment_type or self.fragment_type
|
|
self.dest_type = dest_type or self.dest_type
|
|
numeric_offsets = []
|
|
for idx, offset in enumerate(offsets):
|
|
try:
|
|
numeric_offsets.append(int(offset))
|
|
except ValueError:
|
|
self.varying_dims.add(idx)
|
|
numeric_offsets.append(None)
|
|
if self.rows is not None:
|
|
row = numeric_offsets[2] if len(numeric_offsets) > 2 else None
|
|
self.rows.append(row)
|
|
if len(self.rows) >= 2 and self.rows[-1] != self.rows[-2]:
|
|
self.varying_dims.add(2)
|
|
self.channels.append(channel)
|
|
self.sources.append(source)
|
|
self.targets.append(target)
|
|
|
|
|
|
def parse_const_indices(lines: Iterable[str]) -> dict[str, int]:
|
|
constants: dict[str, int] = {}
|
|
for line in lines:
|
|
match = CONST_INDEX_RE.match(line)
|
|
if match:
|
|
constants[match.group(1)] = int(match.group(2))
|
|
return constants
|
|
|
|
|
|
def split_index_list(value: str) -> list[str]:
|
|
return [piece.strip() for piece in value.split(",") if piece.strip()]
|
|
|
|
|
|
def decode_const_index(token: str, constants: dict[str, int]) -> int | None:
|
|
token = token.strip()
|
|
if token in constants:
|
|
return constants[token]
|
|
try:
|
|
return int(token)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def sequence_kind(values: list[int | None]) -> str:
|
|
concrete = [value for value in values if value is not None]
|
|
if not concrete:
|
|
return "dynamic"
|
|
if len(concrete) == len(values) and all(b - a == 1 for a, b in zip(concrete, concrete[1:])):
|
|
return "consecutive"
|
|
if len(set(concrete)) == 1:
|
|
return "constant"
|
|
return "table"
|
|
|
|
|
|
def analyze_file(path: Path) -> tuple[Counter, dict[tuple[str, str], ChainGroup]]:
|
|
text = path.read_text()
|
|
lines = text.splitlines()
|
|
consts = parse_const_indices(lines)
|
|
counts = Counter()
|
|
groups: dict[tuple[str, str], ChainGroup] = {}
|
|
|
|
for line in lines:
|
|
for name, pattern in OP_PATTERNS.items():
|
|
if pattern.search(line):
|
|
counts[name] += 1
|
|
|
|
value_defs: dict[str, tuple[str, int, re.Match[str] | None]] = {}
|
|
for line_no, line in enumerate(lines, start=1):
|
|
if match := CHANNEL_RECEIVE_RE.match(line):
|
|
value_defs[match.group(1)] = ("receive", line_no, match)
|
|
elif match := EXTRACT_SLICE_RE.match(line):
|
|
value_defs[match.group(1)] = ("extract", line_no, match)
|
|
elif match := INSERT_SLICE_RE.match(line):
|
|
source = match.group(2)
|
|
producer = value_defs.get(source)
|
|
if not producer:
|
|
continue
|
|
|
|
offsets = split_index_list(match.group(4))
|
|
sizes = split_index_list(match.group(5))
|
|
strides = split_index_list(match.group(6))
|
|
dest = match.group(3)
|
|
dest_type = match.group(8).strip()
|
|
|
|
if producer[0] == "receive":
|
|
recv = producer[2]
|
|
assert recv is not None
|
|
channel = decode_const_index(recv.group(2), consts)
|
|
source_core = decode_const_index(recv.group(3), consts)
|
|
target_core = decode_const_index(recv.group(4), consts)
|
|
signature = f"recv_insert:{recv.group(5).strip()}->{dest_type}|sizes={','.join(sizes)}|strides={','.join(strides)}"
|
|
group = groups.setdefault(
|
|
("receive_to_insert", signature),
|
|
ChainGroup("receive_to_insert", signature),
|
|
)
|
|
group.add(match.start() and producer[1] or line_no,
|
|
recv.group(5).strip(),
|
|
dest_type,
|
|
offsets,
|
|
channel=channel,
|
|
source=source_core,
|
|
target=target_core)
|
|
elif producer[0] == "extract":
|
|
extract = producer[2]
|
|
assert extract is not None
|
|
extract_offsets = split_index_list(extract.group(3))
|
|
signature = (
|
|
f"extract_insert:{extract.group(6).strip()}->{dest_type}|"
|
|
f"extract_sizes={extract.group(4).strip()}|insert_sizes={','.join(sizes)}|"
|
|
f"src={extract.group(2)}"
|
|
)
|
|
group = groups.setdefault(
|
|
("extract_to_insert", signature),
|
|
ChainGroup("extract_to_insert", signature),
|
|
)
|
|
group.add(producer[1], extract.group(7).strip(), dest_type, offsets)
|
|
if extract_offsets and decode_const_index(extract_offsets[0], consts) is None:
|
|
group.varying_dims.add(0)
|
|
|
|
elif match := CHANNEL_SEND_RE.match(line):
|
|
source_value = match.group(1)
|
|
producer = value_defs.get(source_value)
|
|
if not producer or producer[0] != "extract":
|
|
continue
|
|
extract = producer[2]
|
|
assert extract is not None
|
|
signature = (
|
|
f"extract_send:{extract.group(6).strip()}->{match.group(5).strip()}|"
|
|
f"extract_sizes={extract.group(4).strip()}|src={extract.group(2)}"
|
|
)
|
|
group = groups.setdefault(
|
|
("extract_to_send", signature),
|
|
ChainGroup("extract_to_send", signature),
|
|
)
|
|
group.add(
|
|
producer[1],
|
|
extract.group(7).strip(),
|
|
match.group(5).strip(),
|
|
split_index_list(extract.group(3)),
|
|
channel=decode_const_index(match.group(2), consts),
|
|
source=decode_const_index(match.group(3), consts),
|
|
target=decode_const_index(match.group(4), consts),
|
|
)
|
|
|
|
return counts, groups
|
|
|
|
|
|
def print_report(path: Path, counts: Counter, groups: dict[tuple[str, str], ChainGroup], limit: int) -> None:
|
|
print(f"== {path} ==")
|
|
print("counts:")
|
|
for name in OP_PATTERNS:
|
|
print(f" {name}: {counts[name]}")
|
|
|
|
ranked = sorted(groups.values(), key=lambda group: (-group.count, group.first_line))
|
|
print("hot chains:")
|
|
for group in ranked[:limit]:
|
|
varying = ",".join(str(dim) for dim in sorted(group.varying_dims)) or "none"
|
|
print(f" - kind: {group.kind}")
|
|
print(f" lines: {group.first_line}-{group.last_line}")
|
|
print(f" fragments: {group.count}")
|
|
print(f" fragment_type: {group.fragment_type}")
|
|
print(f" dest_type: {group.dest_type}")
|
|
print(f" varying_dims: {varying}")
|
|
if group.rows:
|
|
print(f" row_sequence: {sequence_kind(group.rows)}")
|
|
if group.channels:
|
|
print(f" channel_ids: {sequence_kind(group.channels)}")
|
|
if group.sources:
|
|
print(f" source_ids: {sequence_kind(group.sources)}")
|
|
if group.targets:
|
|
print(f" target_ids: {sequence_kind(group.targets)}")
|
|
print(f" signature: {group.signature}")
|
|
print()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Analyze repeated Spatial/PIM tensor IR cardinality patterns.")
|
|
parser.add_argument("paths", nargs="+", help="MLIR files to analyze.")
|
|
parser.add_argument("--limit", type=int, default=12, help="Maximum number of hot chains to print per file.")
|
|
args = parser.parse_args()
|
|
|
|
for path_arg in args.paths:
|
|
path = Path(path_arg)
|
|
counts, groups = analyze_file(path)
|
|
print_report(path, counts, groups, args.limit)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|