derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1393 lines
45 KiB
Python
Executable File
1393 lines
45 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# Copyright (C) 2025-2026 Advanced Micro Devices, Inc.
|
|
#
|
|
# Validate s_wait_* insertion in STIR (StinkyTofu IR) files.
|
|
#
|
|
# Checks that register (VGPR/SGPR/acc) and LDS (memtoken) data dependencies
|
|
# are correctly satisfied by s_wait_dscnt / s_wait_loadcnt / s_wait_tensorcnt
|
|
# instructions already present in the input STIR.
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import re
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Counter kinds (mirrors WaitDataflow::CounterKind)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# Enum member order mirrors C++ CounterKind (CK_DS, CK_Buffer, CK_KM,
|
|
# CK_Tensor) so iteration order matches the reference dataflow.
|
|
class CK(Enum):
|
|
DS = "DS"
|
|
BUFFER = "Buffer"
|
|
KM = "KM"
|
|
TENSOR = "Tensor"
|
|
|
|
|
|
COUNTER_WAIT_OPS: Dict[CK, str] = {
|
|
CK.DS: "s_wait_dscnt",
|
|
CK.BUFFER: "s_wait_loadcnt",
|
|
CK.KM: "s_wait_kmcnt",
|
|
CK.TENSOR: "s_wait_tensorcnt",
|
|
}
|
|
|
|
WAIT_MOD_FIELDS: Dict[str, Tuple[CK, str]] = {
|
|
"s_wait_dscnt": (CK.DS, "dlcnt"),
|
|
"s_wait_loadcnt": (CK.BUFFER, "vlcnt"),
|
|
"s_wait_kmcnt": (CK.KM, "kmcnt"),
|
|
"s_wait_tensorcnt": (CK.TENSOR, "tlcnt"),
|
|
}
|
|
|
|
K_UNUSED = -1
|
|
K_MAX_IN_FLIGHT = 64
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lexer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TokKind(Enum):
|
|
EOF = auto()
|
|
NEWLINE = auto()
|
|
IDENT = auto()
|
|
INT = auto()
|
|
FLOAT = auto()
|
|
STRING = auto()
|
|
LPAREN = auto()
|
|
RPAREN = auto()
|
|
LBRACE = auto()
|
|
RBRACE = auto()
|
|
LBRACK = auto()
|
|
RBRACK = auto()
|
|
COMMA = auto()
|
|
EQ = auto()
|
|
COLON = auto()
|
|
TRUE = auto()
|
|
FALSE = auto()
|
|
AT = auto()
|
|
CARET = auto()
|
|
|
|
|
|
@dataclass
|
|
class Token:
|
|
kind: TokKind
|
|
value: str = ""
|
|
line: int = 1
|
|
col: int = 1
|
|
|
|
|
|
class Lexer:
|
|
def __init__(self, text: str) -> None:
|
|
self.text = text
|
|
self.pos = 0
|
|
self.line = 1
|
|
self.col = 1
|
|
|
|
def _peek(self, n: int = 0) -> str:
|
|
i = self.pos + n
|
|
return self.text[i] if i < len(self.text) else ""
|
|
|
|
def _advance(self, n: int = 1) -> None:
|
|
for _ in range(n):
|
|
if self.pos < len(self.text):
|
|
if self.text[self.pos] == "\n":
|
|
self.line += 1
|
|
self.col = 1
|
|
else:
|
|
self.col += 1
|
|
self.pos += 1
|
|
|
|
def _skip_ws(self) -> None:
|
|
while self.pos < len(self.text):
|
|
c = self.text[self.pos]
|
|
if c in " \t\r\n":
|
|
self._advance()
|
|
elif c == "#":
|
|
while self.pos < len(self.text) and self.text[self.pos] != "\n":
|
|
self._advance()
|
|
elif c == "/" and self._peek(1) == "/":
|
|
while self.pos < len(self.text) and self.text[self.pos] != "\n":
|
|
self._advance()
|
|
elif c == "/" and self._peek(1) == "*":
|
|
self._advance(2)
|
|
while self.pos < len(self.text):
|
|
if self.text[self.pos] == "*" and self._peek(1) == "/":
|
|
self._advance(2)
|
|
break
|
|
self._advance()
|
|
else:
|
|
break
|
|
|
|
def next_token(self) -> Token:
|
|
self._skip_ws()
|
|
line, col = self.line, self.col
|
|
if self.pos >= len(self.text):
|
|
return Token(TokKind.EOF, "", line, col)
|
|
|
|
c = self.text[self.pos]
|
|
|
|
singles = {
|
|
"(": TokKind.LPAREN,
|
|
")": TokKind.RPAREN,
|
|
"{": TokKind.LBRACE,
|
|
"}": TokKind.RBRACE,
|
|
"[": TokKind.LBRACK,
|
|
"]": TokKind.RBRACK,
|
|
",": TokKind.COMMA,
|
|
"=": TokKind.EQ,
|
|
":": TokKind.COLON,
|
|
"@": TokKind.AT,
|
|
"^": TokKind.CARET,
|
|
}
|
|
if c in singles:
|
|
self._advance()
|
|
return Token(singles[c], c, line, col)
|
|
|
|
if c == '"':
|
|
self._advance()
|
|
start = self.pos
|
|
while self.pos < len(self.text) and self.text[self.pos] != '"':
|
|
if self.text[self.pos] == "\\":
|
|
self._advance()
|
|
self._advance()
|
|
val = self.text[start : self.pos]
|
|
self._advance()
|
|
return Token(TokKind.STRING, val, line, col)
|
|
|
|
if c.isdigit() or (c == "-" and self._peek(1).isdigit()):
|
|
start = self.pos
|
|
if c == "-":
|
|
self._advance()
|
|
while self.pos < len(self.text) and (
|
|
self.text[self.pos].isdigit()
|
|
or self.text[self.pos] in "abcdefABCDEF"
|
|
or (
|
|
self.text[self.pos] == "x"
|
|
and "x" in self.text[start : self.pos + 1]
|
|
)
|
|
):
|
|
self._advance()
|
|
# Fractional part of a decimal float (e.g. 1.000000). Only consume
|
|
# the dot when it is followed by a digit so we don't swallow other
|
|
# punctuation, and never for hex literals.
|
|
if (
|
|
self._peek() == "."
|
|
and self._peek(1).isdigit()
|
|
and "x" not in self.text[start : self.pos]
|
|
):
|
|
self._advance()
|
|
while self.pos < len(self.text) and self.text[self.pos].isdigit():
|
|
self._advance()
|
|
raw = self.text[start : self.pos]
|
|
if "." in raw:
|
|
return Token(TokKind.FLOAT, raw, line, col)
|
|
return Token(TokKind.INT, raw, line, col)
|
|
|
|
if c.isalpha() or c == "_" or c == ".":
|
|
start = self.pos
|
|
while self.pos < len(self.text) and (
|
|
self.text[self.pos].isalnum() or self.text[self.pos] in "._+"
|
|
):
|
|
self._advance()
|
|
val = self.text[start : self.pos]
|
|
if val == "true":
|
|
return Token(TokKind.TRUE, val, line, col)
|
|
if val == "false":
|
|
return Token(TokKind.FALSE, val, line, col)
|
|
return Token(TokKind.IDENT, val, line, col)
|
|
|
|
raise SyntaxError(f"Unexpected character {c!r} at line {line}, col {col}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AST nodes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RegUnit:
|
|
cls: str
|
|
index: int
|
|
|
|
def __str__(self) -> str:
|
|
return f"{self.cls}{self.index}"
|
|
|
|
|
|
@dataclass
|
|
class Instruction:
|
|
opcode: str
|
|
dests: List[RegUnit] = field(default_factory=list)
|
|
srcs: List[Any] = field(default_factory=list) # RegUnit | int | str
|
|
attrs: Dict[str, Any] = field(default_factory=dict)
|
|
line: int = 0
|
|
block: str = ""
|
|
uid: int = 0
|
|
|
|
def memtokens(self) -> List[int]:
|
|
mod = self.attrs.get("mod.memtoken")
|
|
if mod and isinstance(mod, dict):
|
|
tokens = mod.get("tokens", [])
|
|
if isinstance(tokens, list):
|
|
return [int(t) for t in tokens]
|
|
return []
|
|
|
|
def label(self) -> str:
|
|
return f"{self.opcode} @ line {self.line}"
|
|
|
|
|
|
@dataclass
|
|
class BasicBlock:
|
|
name: str
|
|
instructions: List[Instruction] = field(default_factory=list)
|
|
successors: List[str] = field(default_factory=list)
|
|
line: int = 0
|
|
|
|
|
|
@dataclass
|
|
class Function:
|
|
name: str
|
|
blocks: Dict[str, BasicBlock] = field(default_factory=dict)
|
|
entry: str = "entry"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parser helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Parser:
|
|
_uid = 0
|
|
|
|
def __init__(self, text: str) -> None:
|
|
self.lexer = Lexer(text)
|
|
self.tok = self.lexer.next_token()
|
|
self._peek_tok: Optional[Token] = None
|
|
|
|
def _peek(self) -> Token:
|
|
return self.tok
|
|
|
|
def _eat(self, kind: Optional[TokKind] = None) -> Token:
|
|
t = self.tok
|
|
if kind is not None and t.kind != kind:
|
|
raise SyntaxError(
|
|
f"Expected {kind}, got {t.kind} ({t.value!r}) at line {t.line}"
|
|
)
|
|
self.tok = self.lexer.next_token()
|
|
return t
|
|
|
|
def _match(self, kind: TokKind) -> bool:
|
|
if self.tok.kind == kind:
|
|
self._eat()
|
|
return True
|
|
return False
|
|
|
|
def parse(self) -> List[Function]:
|
|
funcs: List[Function] = []
|
|
while self.tok.kind != TokKind.EOF:
|
|
if self.tok.kind == TokKind.IDENT and self.tok.value in ("st", "st.func"):
|
|
funcs.append(self._parse_func())
|
|
elif self.tok.kind in (TokKind.IDENT, TokKind.CARET):
|
|
funcs.append(self._parse_flat())
|
|
else:
|
|
self._eat()
|
|
return funcs
|
|
|
|
def _parse_func(self) -> Function:
|
|
hdr = self._eat(TokKind.IDENT).value
|
|
if hdr == "st":
|
|
self._eat(TokKind.IDENT) # func
|
|
self._eat(TokKind.AT)
|
|
name = self._eat(TokKind.IDENT).value
|
|
self._eat(TokKind.LPAREN)
|
|
self._eat(TokKind.RPAREN)
|
|
self._eat(TokKind.LBRACE)
|
|
func = Function(name=name)
|
|
while self.tok.kind != TokKind.RBRACE and self.tok.kind != TokKind.EOF:
|
|
bb = self._parse_block()
|
|
func.blocks[bb.name] = bb
|
|
self._eat(TokKind.RBRACE)
|
|
if func.blocks:
|
|
func.entry = next(iter(func.blocks))
|
|
return func
|
|
|
|
def _parse_flat(self) -> Function:
|
|
func = Function(name="flat")
|
|
bb = BasicBlock(name="entry")
|
|
while self.tok.kind != TokKind.EOF:
|
|
if self.tok.kind == TokKind.CARET:
|
|
bb = self._parse_block_header()
|
|
continue
|
|
if self.tok.kind == TokKind.IDENT and self._peek_label():
|
|
self._eat(TokKind.IDENT)
|
|
self._eat(TokKind.COLON)
|
|
continue
|
|
if self.tok.kind == TokKind.STRING or (
|
|
self.tok.kind == TokKind.IDENT and self._peek_assign()
|
|
):
|
|
bb.instructions.append(self._parse_instruction(bb.name))
|
|
continue
|
|
break
|
|
func.blocks[bb.name] = bb
|
|
func.entry = bb.name
|
|
return func
|
|
|
|
def _peek_label(self) -> bool:
|
|
if self.tok.kind != TokKind.IDENT:
|
|
return False
|
|
saved = self.tok
|
|
self.tok = self.lexer.next_token()
|
|
is_label = self.tok.kind == TokKind.COLON
|
|
self.tok = saved
|
|
return is_label
|
|
|
|
def _save_lexer_state(self) -> Tuple[Token, int, int, int]:
|
|
return (self.tok, self.lexer.pos, self.lexer.line, self.lexer.col)
|
|
|
|
def _restore_lexer_state(self, state: Tuple[Token, int, int, int]) -> None:
|
|
self.tok, self.lexer.pos, self.lexer.line, self.lexer.col = state
|
|
|
|
def _peek_assign(self) -> bool:
|
|
saved = self._save_lexer_state()
|
|
try:
|
|
while self.tok.kind not in (TokKind.EOF, TokKind.EQ, TokKind.STRING):
|
|
self._eat()
|
|
return self.tok.kind == TokKind.EQ
|
|
finally:
|
|
self._restore_lexer_state(saved)
|
|
|
|
def _parse_block(self) -> BasicBlock:
|
|
if self.tok.kind == TokKind.CARET:
|
|
return self._parse_block_header()
|
|
raise SyntaxError(f"Expected block label at line {self.tok.line}")
|
|
|
|
def _parse_block_header(self) -> BasicBlock:
|
|
self._eat(TokKind.CARET)
|
|
name = self._eat(TokKind.IDENT).value
|
|
self._eat(TokKind.COLON)
|
|
bb = BasicBlock(name=name, line=self.tok.line)
|
|
while self.tok.kind not in (TokKind.CARET, TokKind.RBRACE, TokKind.EOF):
|
|
if self.tok.kind == TokKind.IDENT and self.tok.value == "Successors":
|
|
self._parse_successors(bb)
|
|
elif self.tok.kind == TokKind.STRING or (
|
|
self.tok.kind == TokKind.IDENT and self._peek_assign()
|
|
):
|
|
bb.instructions.append(self._parse_instruction(bb.name))
|
|
else:
|
|
self._eat()
|
|
return bb
|
|
|
|
def _parse_successors(self, bb: BasicBlock) -> None:
|
|
self._eat(TokKind.IDENT)
|
|
self._eat(TokKind.COLON)
|
|
while True:
|
|
if self.tok.kind == TokKind.CARET:
|
|
self._eat(TokKind.CARET)
|
|
bb.successors.append(self._eat(TokKind.IDENT).value)
|
|
else:
|
|
break
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
else:
|
|
break
|
|
|
|
def _parse_instruction(self, block: str) -> Instruction:
|
|
dests: List[RegUnit] = []
|
|
if self.tok.kind == TokKind.IDENT:
|
|
dests = self._parse_reg_list()
|
|
self._eat(TokKind.EQ)
|
|
|
|
opcode_tok = self._eat(TokKind.STRING)
|
|
opcode = opcode_tok.value
|
|
if opcode.startswith("st."):
|
|
opcode = opcode[3:]
|
|
|
|
self._eat(TokKind.LPAREN)
|
|
srcs: List[Any] = []
|
|
if self.tok.kind != TokKind.RPAREN:
|
|
srcs = self._parse_operand_list()
|
|
self._eat(TokKind.RPAREN)
|
|
|
|
attrs: Dict[str, Any] = {}
|
|
if self.tok.kind == TokKind.LBRACE:
|
|
attrs = self._parse_attr_block()
|
|
|
|
Parser._uid += 1
|
|
return Instruction(
|
|
opcode=opcode,
|
|
dests=dests,
|
|
srcs=srcs,
|
|
attrs=attrs,
|
|
line=opcode_tok.line,
|
|
block=block,
|
|
uid=Parser._uid,
|
|
)
|
|
|
|
def _parse_reg_list(self) -> List[RegUnit]:
|
|
regs: List[RegUnit] = []
|
|
while True:
|
|
regs.extend(self._parse_register())
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
else:
|
|
break
|
|
return regs
|
|
|
|
def _parse_operand_list(self) -> List[Any]:
|
|
ops: List[Any] = []
|
|
while True:
|
|
ops.append(self._parse_operand())
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
else:
|
|
break
|
|
return ops
|
|
|
|
def _parse_operand(self) -> Any:
|
|
if self.tok.kind in (TokKind.INT, TokKind.FLOAT):
|
|
raw = self._eat().value
|
|
try:
|
|
if raw.startswith("0x") or raw.startswith("-0x"):
|
|
return int(raw, 16)
|
|
if "." in raw:
|
|
return float(raw)
|
|
return int(raw)
|
|
except ValueError:
|
|
return raw
|
|
if self.tok.kind == TokKind.STRING:
|
|
return self._eat(TokKind.STRING).value
|
|
if self.tok.kind == TokKind.TRUE:
|
|
self._eat()
|
|
return True
|
|
if self.tok.kind == TokKind.FALSE:
|
|
self._eat()
|
|
return False
|
|
if self.tok.kind == TokKind.IDENT:
|
|
if self.tok.value == "hwreg":
|
|
return self._parse_hwreg()
|
|
regs = self._parse_register()
|
|
if len(regs) == 1:
|
|
return regs[0]
|
|
return regs
|
|
raise SyntaxError(f"Unexpected operand at line {self.tok.line}")
|
|
|
|
def _parse_hwreg(self) -> str:
|
|
self._eat(TokKind.IDENT)
|
|
self._eat(TokKind.LPAREN)
|
|
parts = []
|
|
while self.tok.kind != TokKind.RPAREN:
|
|
parts.append(self._eat().value)
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
self._eat(TokKind.RPAREN)
|
|
return "hwreg(" + ",".join(parts) + ")"
|
|
|
|
def _parse_register(self) -> List[RegUnit]:
|
|
name = self._eat(TokKind.IDENT).value
|
|
cls = _normalize_reg_class(name)
|
|
m = re.match(r"^([A-Za-z_]+)(\d+)$", name)
|
|
if m and self.tok.kind != TokKind.LBRACK:
|
|
return [RegUnit(_normalize_reg_class(m.group(1)), int(m.group(2)))]
|
|
if self.tok.kind == TokKind.LBRACK:
|
|
self._eat(TokKind.LBRACK)
|
|
lo = int(self._eat(TokKind.INT).value)
|
|
if self.tok.kind == TokKind.COLON:
|
|
self._eat(TokKind.COLON)
|
|
hi = int(self._eat(TokKind.INT).value)
|
|
self._eat(TokKind.RBRACK)
|
|
return [RegUnit(cls, i) for i in range(lo, hi + 1)]
|
|
self._eat(TokKind.RBRACK)
|
|
return [RegUnit(cls, lo)]
|
|
|
|
return [RegUnit(cls, 0)]
|
|
|
|
def _parse_attr_block(self) -> Dict[str, Any]:
|
|
self._eat(TokKind.LBRACE)
|
|
attrs: Dict[str, Any] = {}
|
|
while self.tok.kind != TokKind.RBRACE:
|
|
key = self._eat(TokKind.IDENT).value
|
|
self._eat(TokKind.EQ)
|
|
attrs[key] = self._parse_attr_value()
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
self._eat(TokKind.RBRACE)
|
|
return attrs
|
|
|
|
def _parse_attr_value(self) -> Any:
|
|
if self.tok.kind == TokKind.LBRACE:
|
|
self._eat(TokKind.LBRACE)
|
|
d: Dict[str, Any] = {}
|
|
while self.tok.kind != TokKind.RBRACE:
|
|
k = self._eat(TokKind.IDENT).value
|
|
self._eat(TokKind.EQ)
|
|
d[k] = self._parse_attr_value()
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
self._eat(TokKind.RBRACE)
|
|
return d
|
|
if self.tok.kind == TokKind.LBRACK:
|
|
self._eat(TokKind.LBRACK)
|
|
arr: List[Any] = []
|
|
while self.tok.kind != TokKind.RBRACK:
|
|
arr.append(self._parse_attr_value())
|
|
if self.tok.kind == TokKind.COMMA:
|
|
self._eat(TokKind.COMMA)
|
|
self._eat(TokKind.RBRACK)
|
|
return arr
|
|
if self.tok.kind == TokKind.STRING:
|
|
return self._eat(TokKind.STRING).value
|
|
if self.tok.kind == TokKind.INT:
|
|
return int(self._eat(TokKind.INT).value)
|
|
if self.tok.kind == TokKind.FLOAT:
|
|
return float(self._eat(TokKind.FLOAT).value)
|
|
if self.tok.kind == TokKind.TRUE:
|
|
self._eat()
|
|
return True
|
|
if self.tok.kind == TokKind.FALSE:
|
|
self._eat()
|
|
return False
|
|
if self.tok.kind == TokKind.IDENT:
|
|
return self._eat(TokKind.IDENT).value
|
|
raise SyntaxError(f"Bad attribute value at line {self.tok.line}")
|
|
|
|
|
|
def _normalize_reg_class(name: str) -> str:
|
|
lower = name.lower()
|
|
if lower in ("v", "s", "acc", "a", "agpr", "lds", "scc", "vcc", "exec", "m"):
|
|
if lower in ("a", "agpr"):
|
|
return "acc"
|
|
if lower == "lds":
|
|
return "LDS"
|
|
if lower == "scc":
|
|
return "SCC"
|
|
return lower
|
|
if name in ("SCC", "LDS", "AGPR"):
|
|
return "acc" if name == "AGPR" else name
|
|
if name.startswith("vcc"):
|
|
return name
|
|
if name.startswith("exec"):
|
|
return name
|
|
return name
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Instruction classification
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# Counter classification mirrors WaitDataflow.cpp's defaultCounterPolicy:
|
|
# CK_DS : ds_read / ds_write / ds_atomic -> s_wait_dscnt
|
|
# CK_Buffer : vector buffer/global/flat load+store -> s_wait_loadcnt
|
|
# CK_KM : scalar SMEM loads (s_load / s_buffer_load)-> s_wait_kmcnt
|
|
# CK_Tensor : tensor_load_to_lds -> s_wait_tensorcnt
|
|
def classify_counter(inst: Instruction) -> Optional[CK]:
|
|
op = inst.opcode
|
|
if _is_ds_producer(op):
|
|
return CK.DS
|
|
if _is_km_producer(op):
|
|
return CK.KM
|
|
if _is_buffer_producer(op):
|
|
return CK.BUFFER
|
|
if op == "tensor_load_to_lds":
|
|
return CK.TENSOR
|
|
return None
|
|
|
|
|
|
def _is_ds_producer(op: str) -> bool:
|
|
return (
|
|
op.startswith("ds_read")
|
|
or op.startswith("ds_load")
|
|
or op.startswith("ds_write")
|
|
or op.startswith("ds_store")
|
|
or op.startswith("ds_atomic")
|
|
)
|
|
|
|
|
|
def _is_km_producer(op: str) -> bool:
|
|
# Scalar memory loads complete on the kmcnt counter.
|
|
return op.startswith("s_load") or op.startswith("s_buffer_load")
|
|
|
|
|
|
def _is_buffer_producer(op: str) -> bool:
|
|
prefixes = (
|
|
"buffer_load",
|
|
"buffer_store",
|
|
"global_load",
|
|
"global_store",
|
|
"flat_load",
|
|
"flat_store",
|
|
)
|
|
return any(op.startswith(p) for p in prefixes)
|
|
|
|
|
|
def is_barrier(inst: Instruction) -> bool:
|
|
return inst.opcode.startswith("s_barrier")
|
|
|
|
|
|
def is_ds_read(inst: Instruction) -> bool:
|
|
return inst.opcode.startswith("ds_read") or inst.opcode.startswith("ds_load")
|
|
|
|
|
|
def is_ds_write(inst: Instruction) -> bool:
|
|
return inst.opcode.startswith("ds_write") or inst.opcode.startswith("ds_store")
|
|
|
|
|
|
def is_ds_atomic(inst: Instruction) -> bool:
|
|
return inst.opcode.startswith("ds_atomic")
|
|
|
|
|
|
def is_tensor_anchor(inst: Instruction) -> bool:
|
|
return (
|
|
is_barrier(inst) or is_ds_read(inst) or is_ds_write(inst) or is_ds_atomic(inst)
|
|
)
|
|
|
|
|
|
def is_lds_writer_anchor(inst: Instruction) -> bool:
|
|
return inst.opcode == "tensor_load_to_lds" or is_ds_write(inst)
|
|
|
|
|
|
def is_wait_inst(inst: Instruction) -> bool:
|
|
return inst.opcode in WAIT_MOD_FIELDS
|
|
|
|
|
|
def is_branch(inst: Instruction) -> bool:
|
|
return inst.opcode.startswith("s_cbranch") or inst.opcode.startswith("s_branch")
|
|
|
|
|
|
def get_wait_counts(inst: Instruction) -> Dict[CK, int]:
|
|
"""Return counter -> keep-count for each wait instruction present."""
|
|
result: Dict[CK, int] = {}
|
|
if inst.opcode not in WAIT_MOD_FIELDS:
|
|
return result
|
|
ck, field_name = WAIT_MOD_FIELDS[inst.opcode]
|
|
mod_key = "mod.swaitcnt" if ck != CK.TENSOR else "mod.swaittensorcnt"
|
|
mod = inst.attrs.get(mod_key)
|
|
if isinstance(mod, dict) and field_name in mod:
|
|
val = int(mod[field_name])
|
|
if val != K_UNUSED:
|
|
result[ck] = val
|
|
return result
|
|
for src in inst.srcs:
|
|
if isinstance(src, int):
|
|
result[ck] = src
|
|
break
|
|
return result
|
|
|
|
|
|
def has_token_overlap(a: List[int], b: List[int]) -> bool:
|
|
return bool(set(a) & set(b))
|
|
|
|
|
|
def is_on_same_pipeline(a: Instruction, b: Instruction) -> bool:
|
|
return classify_counter(a) == CK.DS and classify_counter(b) == CK.DS
|
|
|
|
|
|
def raw_needs_wait_tensor(inst: Instruction, num_waves: int) -> bool:
|
|
return is_barrier(inst) or num_waves == 1
|
|
|
|
|
|
def raw_needs_wait(inst: Instruction, ck: CK, num_waves: int) -> bool:
|
|
if ck == CK.TENSOR:
|
|
return raw_needs_wait_tensor(inst, num_waves)
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Simulation state
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class WaitReport:
|
|
inst: Instruction
|
|
counter: CK
|
|
keep: int
|
|
drained: List[Instruction] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class Violation:
|
|
kind: str # MISSING | ERROR
|
|
consumer: Instruction
|
|
producer: Instruction
|
|
counter: CK
|
|
message: str
|
|
detail: str = ""
|
|
|
|
|
|
@dataclass
|
|
class PerPredQueue:
|
|
"""One queue of in-flight memops on a counter, tagged by the CFG
|
|
predecessor it was seeded from (None == in-block / synthetic).
|
|
|
|
Mirrors C++ waitcnt::PerPredQueue. For an op at index I in a queue of
|
|
size N, the wait value to drain it is N - I - 1 (count_from == N - I).
|
|
"""
|
|
|
|
pred: Optional[str] = None
|
|
ops: List[Instruction] = field(default_factory=list)
|
|
|
|
def count_from(self, op: Instruction) -> int:
|
|
for i, o in enumerate(self.ops):
|
|
if o.uid == op.uid:
|
|
return len(self.ops) - i
|
|
return 0
|
|
|
|
|
|
@dataclass
|
|
class CounterState:
|
|
# One LIST of PerPredQueue per counter (kept per-pred, never collapsed
|
|
# into a single union), mirroring DataflowState::queues. Per-pred queues
|
|
# let a consumer at a join compute the strictest required wait as the
|
|
# min over paths instead of an over-deep union.
|
|
queues: Dict[CK, List[PerPredQueue]] = field(
|
|
default_factory=lambda: {ck: [] for ck in CK}
|
|
)
|
|
reg_defs: Dict[RegUnit, Instruction] = field(default_factory=dict)
|
|
num_waves: int = 0
|
|
|
|
def copy(self) -> "CounterState":
|
|
return CounterState(
|
|
queues={
|
|
ck: [PerPredQueue(q.pred, list(q.ops)) for q in qs]
|
|
for ck, qs in self.queues.items()
|
|
},
|
|
reg_defs=dict(self.reg_defs),
|
|
num_waves=self.num_waves,
|
|
)
|
|
|
|
def seed_from_pred(self, pred: str, other: "CounterState") -> None:
|
|
"""Seed one PerPredQueue per counter from a predecessor's exit
|
|
queues, retagged as via-`pred`. Identical (pred, ops) queues are
|
|
deduplicated (mirrors mergeFromPredecessors). reg_defs are merged
|
|
first-writer-wins as an SSA approximation."""
|
|
for ck in CK:
|
|
for predQ in other.queues[ck]:
|
|
uids = [o.uid for o in predQ.ops]
|
|
dup = any(
|
|
q.pred == pred and [o.uid for o in q.ops] == uids
|
|
for q in self.queues[ck]
|
|
)
|
|
if not dup:
|
|
self.queues[ck].append(PerPredQueue(pred=pred, ops=list(predQ.ops)))
|
|
for reg, prod in other.reg_defs.items():
|
|
if reg not in self.reg_defs:
|
|
self.reg_defs[reg] = prod
|
|
|
|
def iter_ops(self, ck: CK) -> Iterator[Instruction]:
|
|
"""Yield each distinct in-flight op across all per-pred queues."""
|
|
seen: Set[int] = set()
|
|
for q in self.queues[ck]:
|
|
for op in q.ops:
|
|
if op.uid not in seen:
|
|
seen.add(op.uid)
|
|
yield op
|
|
|
|
def in_flight(self, ck: CK) -> bool:
|
|
return any(q.ops for q in self.queues[ck])
|
|
|
|
def count_from(self, ck: CK, producer: Instruction) -> int:
|
|
"""Strictest residual depth of `producer`: the MIN positive
|
|
count_from across per-pred queues (the closest-to-tail path bounds
|
|
the wait), or 0 if drained on every path."""
|
|
best = 0
|
|
for q in self.queues[ck]:
|
|
n = q.count_from(producer)
|
|
if n > 0 and (best == 0 or n < best):
|
|
best = n
|
|
return best
|
|
|
|
def trim_queue(self, ck: CK, keep: int) -> List[Instruction]:
|
|
"""Trim every per-pred queue to keep at most `keep` tail ops;
|
|
return the distinct ops drained (for reporting)."""
|
|
drained: List[Instruction] = []
|
|
seen: Set[int] = set()
|
|
|
|
def _record(op: Instruction) -> None:
|
|
if op.uid not in seen:
|
|
seen.add(op.uid)
|
|
drained.append(op)
|
|
|
|
for q in self.queues[ck]:
|
|
if keep <= 0:
|
|
for op in q.ops:
|
|
_record(op)
|
|
q.ops.clear()
|
|
elif len(q.ops) > keep:
|
|
n = len(q.ops) - keep
|
|
for op in q.ops[:n]:
|
|
_record(op)
|
|
del q.ops[:n]
|
|
return drained
|
|
|
|
def append_producer(self, inst: Instruction, ck: CK) -> None:
|
|
"""Append an in-block memop to every per-pred path (it is in flight
|
|
on all paths through this block). Each queue is capped at
|
|
K_MAX_IN_FLIGHT so an undrained counter cannot grow forever."""
|
|
qs = self.queues[ck]
|
|
if not qs:
|
|
qs.append(PerPredQueue())
|
|
for q in qs:
|
|
q.ops.append(inst)
|
|
while len(q.ops) > K_MAX_IN_FLIGHT:
|
|
q.ops.pop(0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dependency collection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _operand_regs(srcs: List[Any]) -> List[RegUnit]:
|
|
regs: List[RegUnit] = []
|
|
for s in srcs:
|
|
if isinstance(s, RegUnit):
|
|
regs.append(s)
|
|
elif isinstance(s, list):
|
|
regs.extend(r for r in s if isinstance(r, RegUnit))
|
|
return regs
|
|
|
|
|
|
def _lds_pseudo_tokens(regs: List[RegUnit]) -> List[int]:
|
|
return [r.index for r in regs if r.cls == "LDS"]
|
|
|
|
|
|
def collect_register_deps(
|
|
inst: Instruction, state: CounterState
|
|
) -> List[Tuple[Instruction, CK]]:
|
|
deps: List[Tuple[Instruction, CK]] = []
|
|
seen: Set[int] = set()
|
|
for reg in _operand_regs(inst.srcs):
|
|
prod = state.reg_defs.get(reg)
|
|
if prod is None:
|
|
continue
|
|
ck = classify_counter(prod)
|
|
if ck is None:
|
|
continue
|
|
if prod.uid in seen:
|
|
continue
|
|
if not raw_needs_wait(inst, ck, state.num_waves):
|
|
continue
|
|
if state.count_from(ck, prod) > 0:
|
|
deps.append((prod, ck))
|
|
seen.add(prod.uid)
|
|
return deps
|
|
|
|
|
|
def is_lds_reader(inst: Instruction) -> bool:
|
|
"""Instructions that consume LDS data (RAW on memtoken), not writers."""
|
|
if is_ds_read(inst):
|
|
return True
|
|
if is_lds_writer_anchor(inst):
|
|
return False
|
|
if _lds_pseudo_tokens(_operand_regs(inst.srcs)):
|
|
return True
|
|
if is_tensor_anchor(inst) and inst.memtokens():
|
|
return True
|
|
return False
|
|
|
|
|
|
def collect_lds_raw_deps(
|
|
inst: Instruction, state: CounterState
|
|
) -> List[Tuple[Instruction, CK]]:
|
|
if not is_lds_reader(inst):
|
|
return []
|
|
tokens = inst.memtokens() or _lds_pseudo_tokens(_operand_regs(inst.srcs))
|
|
if not tokens:
|
|
return []
|
|
deps: List[Tuple[Instruction, CK]] = []
|
|
seen: Set[int] = set()
|
|
for ck in CK:
|
|
if not raw_needs_wait(inst, ck, state.num_waves):
|
|
continue
|
|
for prod in state.iter_ops(ck):
|
|
pt = prod.memtokens()
|
|
if not pt:
|
|
continue
|
|
if not has_token_overlap(tokens, pt):
|
|
continue
|
|
# Same DS FIFO orders consecutive DS ops; only cross-pipeline
|
|
# memtoken RAW needs an explicit wait.
|
|
if is_on_same_pipeline(inst, prod):
|
|
continue
|
|
if prod.uid in seen:
|
|
continue
|
|
if state.count_from(ck, prod) > 0:
|
|
deps.append((prod, ck))
|
|
seen.add(prod.uid)
|
|
return deps
|
|
|
|
|
|
def collect_war_deps(
|
|
inst: Instruction, state: CounterState
|
|
) -> List[Tuple[Instruction, CK]]:
|
|
deps: List[Tuple[Instruction, CK]] = []
|
|
barrier_mode = is_barrier(inst)
|
|
if not is_lds_writer_anchor(inst) and not barrier_mode:
|
|
return deps
|
|
|
|
anchor_tokens = inst.memtokens()
|
|
if anchor_tokens is None:
|
|
anchor_tokens = []
|
|
|
|
seen: Set[int] = set()
|
|
for op in state.iter_ops(CK.DS):
|
|
if op.uid == inst.uid:
|
|
continue
|
|
if not barrier_mode and not (is_ds_read(op) or is_ds_atomic(op)):
|
|
continue
|
|
if is_on_same_pipeline(inst, op):
|
|
continue
|
|
op_tokens = op.memtokens()
|
|
overlap = (
|
|
op_tokens is None
|
|
or not anchor_tokens
|
|
or has_token_overlap(op_tokens, anchor_tokens)
|
|
)
|
|
if not overlap:
|
|
continue
|
|
if op.uid in seen:
|
|
continue
|
|
if state.count_from(CK.DS, op) > 0:
|
|
deps.append((op, CK.DS))
|
|
seen.add(op.uid)
|
|
return deps
|
|
|
|
|
|
def collect_conservative_deps(
|
|
inst: Instruction, state: CounterState
|
|
) -> List[Tuple[Instruction, CK]]:
|
|
deps: List[Tuple[Instruction, CK]] = []
|
|
if is_tensor_anchor(inst) and not inst.memtokens() and state.in_flight(CK.TENSOR):
|
|
for prod in state.iter_ops(CK.TENSOR):
|
|
if state.count_from(CK.TENSOR, prod) > 0:
|
|
deps.append((prod, CK.TENSOR))
|
|
if is_lds_writer_anchor(inst) and not inst.memtokens() and not is_ds_write(inst):
|
|
if state.in_flight(CK.DS):
|
|
for prod in state.iter_ops(CK.DS):
|
|
if state.count_from(CK.DS, prod) > 0:
|
|
deps.append((prod, CK.DS))
|
|
if is_barrier(inst) and state.in_flight(CK.DS):
|
|
barrier_untagged = not inst.memtokens()
|
|
for prod in state.iter_ops(CK.DS):
|
|
if barrier_untagged or not prod.memtokens():
|
|
if state.count_from(CK.DS, prod) > 0:
|
|
deps.append((prod, CK.DS))
|
|
if is_tensor_anchor(inst) and inst.memtokens():
|
|
for prod in state.iter_ops(CK.TENSOR):
|
|
if not prod.memtokens() and state.count_from(CK.TENSOR, prod) > 0:
|
|
deps.append((prod, CK.TENSOR))
|
|
return deps
|
|
|
|
|
|
def collect_all_deps(
|
|
inst: Instruction, state: CounterState
|
|
) -> List[Tuple[Instruction, CK]]:
|
|
if is_wait_inst(inst):
|
|
return []
|
|
all_deps: List[Tuple[Instruction, CK]] = []
|
|
seen: Set[Tuple[int, CK]] = set()
|
|
for collector in (
|
|
collect_register_deps,
|
|
collect_lds_raw_deps,
|
|
collect_war_deps,
|
|
collect_conservative_deps,
|
|
):
|
|
for prod, ck in collector(inst, state):
|
|
key = (prod.uid, ck)
|
|
if key not in seen:
|
|
all_deps.append((prod, ck))
|
|
seen.add(key)
|
|
return all_deps
|
|
|
|
|
|
def required_wait_value(
|
|
deps: List[Tuple[Instruction, CK]], state: CounterState
|
|
) -> Dict[CK, int]:
|
|
required: Dict[CK, int] = {ck: K_UNUSED for ck in CK}
|
|
for prod, ck in deps:
|
|
n = state.count_from(ck, prod)
|
|
if n <= 0:
|
|
continue
|
|
w = n - 1
|
|
if required[ck] == K_UNUSED or w < required[ck]:
|
|
required[ck] = w
|
|
return required
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CFG utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def infer_successors(func: Function) -> None:
|
|
block_names = list(func.blocks.keys())
|
|
for i, (name, bb) in enumerate(func.blocks.items()):
|
|
if bb.successors:
|
|
continue
|
|
for inst in reversed(bb.instructions):
|
|
if is_branch(inst):
|
|
for src in inst.srcs:
|
|
if isinstance(src, str) and not src.startswith("0"):
|
|
target = src.lstrip("^")
|
|
if target in func.blocks:
|
|
bb.successors.append(target)
|
|
break
|
|
if not bb.successors and i + 1 < len(block_names):
|
|
bb.successors.append(block_names[i + 1])
|
|
|
|
|
|
def compute_predecessors(func: Function) -> Dict[str, List[str]]:
|
|
preds: Dict[str, List[str]] = {name: [] for name in func.blocks}
|
|
for name, bb in func.blocks.items():
|
|
for succ in bb.successors:
|
|
if succ in preds:
|
|
preds[succ].append(name)
|
|
return preds
|
|
|
|
|
|
def compute_rpo(func: Function, preds: Dict[str, List[str]]) -> List[str]:
|
|
visited: Set[str] = set()
|
|
order: List[str] = []
|
|
|
|
def dfs(name: str) -> None:
|
|
if name in visited:
|
|
return
|
|
visited.add(name)
|
|
bb = func.blocks.get(name)
|
|
if bb:
|
|
for succ in bb.successors:
|
|
dfs(succ)
|
|
order.append(name)
|
|
|
|
dfs(func.entry)
|
|
for name in func.blocks:
|
|
if name not in visited:
|
|
dfs(name)
|
|
order.reverse()
|
|
return order
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Validator
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class ValidationResult:
|
|
wait_reports: List[WaitReport] = field(default_factory=list)
|
|
violations: List[Violation] = field(default_factory=list)
|
|
infos: List[str] = field(default_factory=list)
|
|
|
|
|
|
def _dump_violation_detail(
|
|
consumer: Instruction,
|
|
bb_name: str,
|
|
ck: CK,
|
|
undrained: List[Tuple[Instruction, CK]],
|
|
state: CounterState,
|
|
) -> str:
|
|
"""Render the live per-pred queue snapshot for counter `ck` at the
|
|
consumer, marking the undrained producer(s) and the wait each needs."""
|
|
wait_op = COUNTER_WAIT_OPS[ck]
|
|
bad_uids = {p.uid for p, c in undrained if c == ck}
|
|
|
|
lines: List[str] = []
|
|
lines.append(
|
|
f" Queue dump for {wait_op} at consumer {consumer.label()} in ^{bb_name}:"
|
|
)
|
|
|
|
# Required wait: strictest (smallest keep) that still drains every
|
|
# undrained producer on this counter, i.e. min over producers of
|
|
# (count_from - 1).
|
|
required = K_UNUSED
|
|
for prod, c in undrained:
|
|
if c != ck:
|
|
continue
|
|
n = state.count_from(ck, prod)
|
|
if n > 0:
|
|
w = n - 1
|
|
required = w if required == K_UNUSED else min(required, w)
|
|
req_str = "none" if required == K_UNUSED else str(required)
|
|
lines.append(
|
|
f" required (drains ALL {len(bad_uids)} undrained producer(s) "
|
|
f"this consumer uses): {wait_op} {req_str}"
|
|
)
|
|
|
|
queues = state.queues[ck]
|
|
lines.append(f" live queues ({ck.value}): {len(queues)} per-pred queue(s)")
|
|
if not queues:
|
|
lines.append(" (empty — nothing in flight)")
|
|
for qi, q in enumerate(queues):
|
|
pred = q.pred if q.pred is not None else "<in-block>"
|
|
lines.append(f" queue[{qi}] pred=^{pred} depth={len(q.ops)}")
|
|
size = len(q.ops)
|
|
for idx, op in enumerate(q.ops):
|
|
drain = size - idx - 1 # keep value that drains this op
|
|
mark = " <-- undrained producer" if op.uid in bad_uids else ""
|
|
lines.append(
|
|
f" idx {idx}: {op.opcode} @ line {op.line} "
|
|
f"(drain with {wait_op} {drain}){mark}"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
class WaitCntValidator:
|
|
def __init__(self, num_waves: int = 0) -> None:
|
|
self.num_waves = num_waves
|
|
self.result = ValidationResult()
|
|
# Counters that already have a reported violation; once a queue has
|
|
# its first violation we stop reporting further ones for it.
|
|
self._reported_counters: Set[CK] = set()
|
|
|
|
def validate_function(self, func: Function) -> ValidationResult:
|
|
infer_successors(func)
|
|
preds = compute_predecessors(func)
|
|
rpo = compute_rpo(func, preds)
|
|
|
|
exit_states: Dict[str, CounterState] = {}
|
|
max_iter = min(256, max(K_MAX_IN_FLIGHT + 8, 2 * len(func.blocks)))
|
|
|
|
for _ in range(max_iter):
|
|
changed = False
|
|
for bb_name in rpo:
|
|
bb = func.blocks[bb_name]
|
|
entry = self._merge_entry(preds, bb_name, exit_states)
|
|
exit_state = self._process_block(bb, entry, collect_reports=False)
|
|
prev = exit_states.get(bb_name)
|
|
if prev is None or not _states_equal(prev, exit_state):
|
|
exit_states[bb_name] = exit_state
|
|
changed = True
|
|
if not changed:
|
|
break
|
|
|
|
self._reported_counters.clear()
|
|
for bb_name in rpo:
|
|
bb = func.blocks[bb_name]
|
|
entry = self._merge_entry(preds, bb_name, exit_states)
|
|
self._process_block(bb, entry, collect_reports=True)
|
|
if len(self._reported_counters) == len(CK):
|
|
break
|
|
|
|
return self.result
|
|
|
|
def _merge_entry(
|
|
self,
|
|
preds: Dict[str, List[str]],
|
|
bb_name: str,
|
|
exit_states: Dict[str, CounterState],
|
|
) -> CounterState:
|
|
entry = CounterState(num_waves=self.num_waves)
|
|
for p in preds.get(bb_name, []):
|
|
if p in exit_states:
|
|
entry.seed_from_pred(p, exit_states[p])
|
|
return entry
|
|
|
|
def _process_block(
|
|
self, bb: BasicBlock, state: CounterState, *, collect_reports: bool
|
|
) -> CounterState:
|
|
state = state.copy()
|
|
state.num_waves = self.num_waves
|
|
|
|
for inst in bb.instructions:
|
|
if is_wait_inst(inst):
|
|
self._process_wait(inst, state, collect_reports=collect_reports)
|
|
continue
|
|
|
|
deps = collect_all_deps(inst, state)
|
|
undrained = [(p, ck) for p, ck in deps if state.count_from(ck, p) > 0]
|
|
if undrained and collect_reports:
|
|
for prod, ck in undrained:
|
|
# Report only the FIRST violation per counter/queue, then
|
|
# stop tracking that queue (one diagnostic per s_wait_* kind).
|
|
if ck in self._reported_counters:
|
|
continue
|
|
self._reported_counters.add(ck)
|
|
n = state.count_from(ck, prod)
|
|
w_needed = n - 1
|
|
self.result.violations.append(
|
|
Violation(
|
|
kind="MISSING",
|
|
consumer=inst,
|
|
producer=prod,
|
|
counter=ck,
|
|
message=(
|
|
f"Consumer {inst.label()} in ^{bb.name} uses async "
|
|
f"producer {prod.label()} ({COUNTER_WAIT_OPS[ck]}) "
|
|
f"still in-flight (queue depth={n}, needs wait "
|
|
f"{COUNTER_WAIT_OPS[ck]} {w_needed})"
|
|
),
|
|
detail=_dump_violation_detail(
|
|
inst, bb.name, ck, undrained, state
|
|
),
|
|
)
|
|
)
|
|
|
|
ck_prod = classify_counter(inst)
|
|
if ck_prod is not None:
|
|
state.append_producer(inst, ck_prod)
|
|
|
|
for reg in inst.dests:
|
|
state.reg_defs[reg] = inst
|
|
|
|
return state
|
|
|
|
def _process_wait(
|
|
self, inst: Instruction, state: CounterState, *, collect_reports: bool
|
|
) -> None:
|
|
counts = get_wait_counts(inst)
|
|
for ck, keep in counts.items():
|
|
drained = state.trim_queue(ck, keep)
|
|
if collect_reports:
|
|
self.result.wait_reports.append(
|
|
WaitReport(inst=inst, counter=ck, keep=keep, drained=drained)
|
|
)
|
|
if not drained:
|
|
self.result.infos.append(
|
|
f"INFO: {inst.label()} {COUNTER_WAIT_OPS[ck]} {keep} "
|
|
f"— no in-flight ops drained (possibly redundant)"
|
|
)
|
|
|
|
|
|
def _states_equal(a: CounterState, b: CounterState) -> bool:
|
|
for ck in CK:
|
|
qa, qb = a.queues[ck], b.queues[ck]
|
|
if len(qa) != len(qb):
|
|
return False
|
|
for x, y in zip(qa, qb):
|
|
if x.pred != y.pred:
|
|
return False
|
|
if [o.uid for o in x.ops] != [o.uid for o in y.ops]:
|
|
return False
|
|
if a.reg_defs.keys() != b.reg_defs.keys():
|
|
return False
|
|
for reg, prod in a.reg_defs.items():
|
|
if prod.uid != b.reg_defs[reg].uid:
|
|
return False
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reporting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def format_report(
|
|
func: Function, result: ValidationResult, verbose: bool = False
|
|
) -> str:
|
|
lines: List[str] = []
|
|
lines.append(f"=== Waitcnt Validation: @{func.name} ===")
|
|
lines.append("")
|
|
|
|
if verbose and result.wait_reports:
|
|
lines.append("--- Wait Instructions ---")
|
|
for wr in result.wait_reports:
|
|
op = COUNTER_WAIT_OPS[wr.counter]
|
|
lines.append(
|
|
f" {op} {wr.keep} @ line {wr.inst.line} " f"(block ^{wr.inst.block})"
|
|
)
|
|
if wr.drained:
|
|
lines.append(" Waiting for:")
|
|
for prod in wr.drained:
|
|
ck = classify_counter(prod)
|
|
ck_name = ck.value if ck else "?"
|
|
lines.append(
|
|
f" - [{ck_name}] {prod.opcode} @ line {prod.line} "
|
|
f"(block ^{prod.block})"
|
|
)
|
|
else:
|
|
lines.append(" Waiting for: (nothing in-flight)")
|
|
lines.append("")
|
|
|
|
if verbose and result.infos:
|
|
lines.append("--- Notes ---")
|
|
for info in result.infos:
|
|
lines.append(f" {info}")
|
|
lines.append("")
|
|
|
|
if result.violations:
|
|
lines.append("--- Violations ---")
|
|
for v in result.violations:
|
|
lines.append(f" {v.kind}: {v.message}")
|
|
if v.detail:
|
|
lines.append(v.detail)
|
|
lines.append("")
|
|
lines.append(f"RESULT: FAIL ({len(result.violations)} violation(s))")
|
|
else:
|
|
lines.append("RESULT: PASS")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def main(argv: Optional[List[str]] = None) -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Validate s_wait_* insertion in STIR files."
|
|
)
|
|
parser.add_argument("stir_file", help="Path to STIR file to validate")
|
|
parser.add_argument(
|
|
"--num-waves",
|
|
type=int,
|
|
default=0,
|
|
help="NumWaves tile config (0=multi-wave: tensor waits only at barriers)",
|
|
)
|
|
parser.add_argument(
|
|
"-v",
|
|
"--verbose",
|
|
action="store_true",
|
|
help="Print full report (wait instructions + notes) in addition to "
|
|
"violations; default prints only violations",
|
|
)
|
|
parser.add_argument(
|
|
"--quiet",
|
|
action="store_true",
|
|
help="Only print summary line",
|
|
)
|
|
args = parser.parse_args(argv)
|
|
|
|
try:
|
|
with open(args.stir_file, encoding="utf-8") as f:
|
|
source = f.read()
|
|
except OSError as e:
|
|
print(f"Error reading {args.stir_file}: {e}", file=sys.stderr)
|
|
return 2
|
|
|
|
try:
|
|
funcs = Parser(source).parse()
|
|
except SyntaxError as e:
|
|
print(f"Parse error: {e}", file=sys.stderr)
|
|
return 2
|
|
|
|
if not funcs:
|
|
print("No functions found in STIR file.", file=sys.stderr)
|
|
return 2
|
|
|
|
exit_code = 0
|
|
for func in funcs:
|
|
validator = WaitCntValidator(num_waves=args.num_waves)
|
|
result = validator.validate_function(func)
|
|
if args.quiet:
|
|
status = "PASS" if not result.violations else "FAIL"
|
|
print(f"@{func.name}: {status} ({len(result.violations)} violations)")
|
|
else:
|
|
print(format_report(func, result, verbose=args.verbose))
|
|
if result.violations:
|
|
exit_code = 1
|
|
|
|
return exit_code
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|