top-sort-asm/topsort_asm/deps.py

"""Build the dependence DAG of a straight-line block.

An edge ``u -> v`` means "u must execute before v". Three kinds exist for
every storage location (registers, the implicit flags register, memory):

* RAW (true / flow): v reads what u wrote.             ``add r1,..`` then ``mul ..,r1,..``
* WAR (anti):        v overwrites what u still reads.  ``mul ..,r1,..`` then ``load r1,..``
* WAW (output):      v overwrites what u wrote.        ``li r1,1`` then ``li r1,2``

Only RAW carries a value; WAR and WAW are *name* conflicts that disappear
when every write gets a fresh register (see :mod:`topsort_asm.rename`).

Every edge points forward in program order, so the graph is always acyclic
and the original order is always one valid topological order. The point of
topological sorting is that it is usually *not the only one*.
"""

from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple

from .isa import Instr

MEM = "mem"


@dataclass(frozen=True)
class Edge:
    src: int
    dst: int
    kind: str  # "RAW" | "WAR" | "WAW"
    via: str  # register name, "flags", or "mem"

    def __str__(self) -> str:
        return f"{self.src} -> {self.dst}  {self.kind} on {self.via}"


@dataclass
class DAG:
    instrs: List[Instr]
    edges: List[Edge]
    succ: Dict[int, Set[int]] = field(default_factory=dict)
    pred: Dict[int, Set[int]] = field(default_factory=dict)

    def __post_init__(self) -> None:
        n = len(self.instrs)
        self.succ = {i: set() for i in range(n)}
        self.pred = {i: set() for i in range(n)}
        for e in self.edges:
            self.succ[e.src].add(e.dst)
            self.pred[e.dst].add(e.src)

    @property
    def n(self) -> int:
        return len(self.instrs)

    def pairs(self) -> Set[Tuple[int, int]]:
        return {(e.src, e.dst) for e in self.edges}

    def edges_between(self, u: int, v: int) -> List[Edge]:
        return [e for e in self.edges if e.src == u and e.dst == v]

    def delay(self, u: int, v: int) -> int:
        """Cycles v must wait after u issues.

        A RAW edge waits for u's result; WAR/WAW edges only need v to issue
        after u, which a single-issue in-order pipeline gives for free.
        """
        if any(e.kind == "RAW" for e in self.edges_between(u, v)):
            return self.instrs[u].latency
        return 1


def may_alias(a: Tuple[str, Optional[int], int], b: Tuple[str, Optional[int], int]) -> bool:
    """Conservative alias test for two memory addresses.

    Each address is ``(base register, defining instruction of base, offset)``.
    The only thing we can *prove* is that ``[r+i]`` and ``[r+j]`` with i != j
    are disjoint when both read the same definition of ``r``. Different base
    registers, or the same base register redefined in between, may point
    anywhere -- so we must assume they alias.

    Within one block the redefinition case is already covered: the earlier
    access reads ``r`` before the redefinition (WAR) and the later one reads
    it after (RAW), so the pair is ordered through the register edges and the
    memory edge added here is implied. It matters once accesses are compared
    across blocks or after code motion, where those register edges are gone.
    """
    base_a, def_a, off_a = a
    base_b, def_b, off_b = b
    if base_a == base_b and def_a == def_b:
        return off_a == off_b
    return True


def build_dag(instrs: List[Instr]) -> DAG:
    edges: List[Edge] = []
    last_writer: Dict[str, int] = {}
    readers: Dict[str, List[int]] = defaultdict(list)  # readers since last write
    mem_ops: List[Tuple[int, str, Tuple[str, Optional[int], int]]] = []

    for j, ins in enumerate(instrs):
        # Register (and flags) dependences: only the most recent writer and
        # the readers since it matter; older ones are already ordered before
        # them, so edges to them would be implied transitively.
        for r in dict.fromkeys(ins.reads):
            if r in last_writer:
                edges.append(Edge(last_writer[r], j, "RAW", r))
        for w in ins.writes:
            if w in last_writer:
                edges.append(Edge(last_writer[w], j, "WAW", w))
            for i in readers[w]:
                edges.append(Edge(i, j, "WAR", w))

        # Memory dependences: a later access may alias any earlier one, so we
        # compare against all of them. Two loads never conflict.
        if ins.is_mem:
            addr = (ins.mem_base, last_writer.get(ins.mem_base), ins.imm)
            for i, kind_i, addr_i in mem_ops:
                if kind_i == "load" and ins.op == "load":
                    continue
                if not may_alias(addr_i, addr):
                    continue
                kind = {("store", "load"): "RAW",
                        ("load", "store"): "WAR",
                        ("store", "store"): "WAW"}[(kind_i, ins.op)]
                edges.append(Edge(i, j, kind, MEM))
            mem_ops.append((j, ins.op, addr))

        for r in dict.fromkeys(ins.reads):
            readers[r].append(j)
        for w in ins.writes:
            last_writer[w] = j
            readers[w] = []

    return DAG(instrs, edges)


def transitive_reduction(dag: DAG) -> Set[Tuple[int, int]]:
    """Edges not implied by a longer path -- the minimal drawing of the DAG."""
    n = dag.n
    reach: List[Set[int]] = [set() for _ in range(n)]
    for u in reversed(range(n)):  # program order is a topological order
        for v in dag.succ[u]:
            reach[u] |= {v} | reach[v]
    keep = set()
    for u, v in dag.pairs():
        if not any(v in reach[w] for w in dag.succ[u] if w != v):
            keep.add((u, v))
    return keep


def to_dot(dag: DAG, title: str = "deps") -> str:
    """Graphviz rendering; implied (transitively redundant) edges are omitted."""
    style = {"RAW": "solid", "WAR": "dashed", "WAW": "dotted"}
    lines = [f'digraph "{title}" {{', "  rankdir=TB;",
             '  node [shape=box, fontname="monospace"];']
    for ins in dag.instrs:
        lines.append(f'  i{ins.index} [label="{ins.index}: {ins.text()}"];')
    keep = transitive_reduction(dag)
    for u, v in sorted(keep):
        es = dag.edges_between(u, v)
        label = ", ".join(f"{e.kind} {e.via}" for e in es)
        kind = "RAW" if any(e.kind == "RAW" for e in es) else es[0].kind
        lines.append(f'  i{u} -> i{v} [label="{label}", style={style[kind]}];')
    lines.append("}")
    return "\n".join(lines)