Upload twill/smt_joint.py with huggingface_hub

twill/smt_joint.py

@@ -0,0 +1,426 @@
+"""
+Phase 2: Joint SWP + WS via SMT Solver (Z3).
+
+Based on Section 4 of the paper (Figures 4, 5, and 6).
+
+Takes the initial modulo schedule M and initiation interval I from Phase 1,
+then formulates constraints over a straight-line program Q derived from M.
+
+Constraint groups:
+    1. Modulo Scheduling Constraints (Figure 4):
+       - Uniqueness, Consistency, Completion, Dependence, Capacity
+    2. Memory Allocation Constraints (Figure 5):
+       - Memory Capacity, Init, LiveProp-1, LiveProp-2, DeadProp-1, DeadProp-2
+    3. Warp Assignment Constraints (Figure 6):
+       - WarpUniqueness, VariableLatency, WarpCapacity, WarpSync
+
+Uses Z3 SMT solver (QFLIA theory - quantifier-free linear integer arithmetic).
+"""
+
+import z3
+import numpy as np
+from typing import Dict, List, Optional, Tuple, Set
+from twill.graph import DependenceGraph, Instruction, DependenceEdge, MachineDescription
+from twill.modulo_scheduler import ModuloScheduleResult
+
+
+class WarpAssignment:
+    """Result of warp assignment.
+    
+    Attributes:
+        assignment: Dict mapping instruction name -> warp index
+        warp_names: Optional names for warps (e.g., "producer", "consumer_A", "consumer_B")
+    """
+    def __init__(self, assignment: Dict[str, int], num_warps: int):
+        self.assignment = assignment
+        self.num_warps = num_warps
+        self.warp_names: Dict[int, str] = {}
+
+    def warp_of(self, instr_name: str) -> int:
+        return self.assignment[instr_name]
+
+    def instructions_on_warp(self, warp_idx: int) -> List[str]:
+        return [name for name, w in self.assignment.items() if w == warp_idx]
+
+    def __repr__(self):
+        warp_groups = {}
+        for name, w in self.assignment.items():
+            warp_groups.setdefault(w, []).append(name)
+        lines = []
+        for w in sorted(warp_groups):
+            label = self.warp_names.get(w, f"Warp {w}")
+            lines.append(f"  {label}: {warp_groups[w]}")
+        return "WarpAssignment(\n" + "\n".join(lines) + "\n)"
+
+
+class JointSWPWSResult:
+    """Result of the joint SWP + WS optimization.
+    
+    Attributes:
+        schedule: Dict mapping instruction name -> clock cycle in Q*
+        initiation_interval: I
+        length: L (schedule length)
+        warp_assignment: WarpAssignment
+        num_copies: ceil(L/I)
+    """
+    def __init__(
+        self,
+        schedule: Dict[str, int],
+        I: int,
+        L: int,
+        warp_assignment: WarpAssignment,
+    ):
+        self.schedule = schedule
+        self.initiation_interval = I
+        self.length = L
+        self.warp_assignment = warp_assignment
+
+    @property
+    def I(self) -> int:
+        return self.initiation_interval
+
+    @property
+    def num_copies(self) -> int:
+        return int(np.ceil(self.length / self.I))
+
+    def __repr__(self):
+        return (
+            f"JointSWPWSResult(\n"
+            f"  I={self.I}, L={self.length}, copies={self.num_copies}\n"
+            f"  schedule={self.schedule}\n"
+            f"  {self.warp_assignment}\n"
+            f")"
+        )
+
+
+def swp_and_ws(
+    graph: DependenceGraph,
+    initial_schedule: ModuloScheduleResult,
+    I: int,
+    L: int,
+    enable_memory_constraints: bool = True,
+    enable_warp_constraints: bool = True,
+    timeout_ms: int = 120000,
+    verbose: bool = False,
+) -> Optional[JointSWPWSResult]:
+    """Solve the joint SWP + WS problem using Z3.
+    
+    This is the core of Twill's Phase 2.
+    
+    Args:
+        graph: Loop dependence graph
+        initial_schedule: Phase 1 result (initial M and I)
+        I: Initiation interval (from Phase 1)
+        L: Schedule length to try
+        enable_memory_constraints: Whether to include memory capacity constraints
+        enable_warp_constraints: Whether to include warp assignment constraints
+        timeout_ms: Z3 timeout in milliseconds
+        verbose: Print constraint statistics
+    
+    Returns:
+        JointSWPWSResult if satisfiable, None otherwise
+    """
+    machine = graph.machine
+    V = graph.V
+    E = graph.E
+    n_copies = int(np.ceil(L / I))
+    # T = total time window for the straight-line program Q
+    # Must accommodate all overlapping copies: copy i starts at offset i*I
+    # The last copy (n_copies-1) starts at (n_copies-1)*I and needs L time steps
+    # So T = L + (n_copies - 1) * I
+    T = L + (n_copies - 1) * I
+
+    if verbose:
+        print(f"  SMT: I={I}, L={L}, T={T}, copies={n_copies}")
+        print(f"  SMT: |V|={len(V)}, |E|={len(E)}")
+
+    # Create Z3 solver
+    solver = z3.Solver()
+    solver.set("timeout", timeout_ms)
+
+    # ============================================================
+    # Boolean variables: op[v, i, t] 
+    # True if instruction v, iteration copy i, is scheduled at clock cycle t
+    # ============================================================
+    op = {}
+    for v in V:
+        for i in range(n_copies):
+            for t in range(T):
+                op[v.name, i, t] = z3.Bool(f"op_{v.name}_{i}_{t}")
+
+    # ============================================================
+    # Figure 4: Modulo Scheduling Constraints
+    # ============================================================
+
+    # (Uniqueness): ∀v,i: Σ_t op[v,i,t] = 1
+    for v in V:
+        for i in range(n_copies):
+            # Exactly one time slot: use PbEq (pseudo-boolean exactly-k)
+            solver.add(z3.PbEq([(op[v.name, i, t], 1) for t in range(T)], 1))
+
+    # (Consistency): ∀v, i∈[1,n_copies), t: op[v,0,t] => op[v,i,t+i*I]
+    for v in V:
+        for i in range(1, n_copies):
+            for t in range(T):
+                t_shifted = t + i * I
+                if 0 <= t_shifted < T:
+                    solver.add(z3.Implies(op[v.name, 0, t], op[v.name, i, t_shifted]))
+                else:
+                    # If shifted time is out of bounds, op[v,0,t] must be false
+                    # (because consistency requires scheduling at t_shifted which doesn't exist)
+                    # Actually: if t+i*I >= T, then this copy i can't start from t in iteration 0
+                    # Only add if the shift is valid
+                    if t_shifted >= T:
+                        solver.add(z3.Implies(op[v.name, 0, t], z3.BoolVal(False)))
+
+    # (Completion): ∀v,i,t: t + cycles(v) > T => ¬op[v,i,t]
+    for v in V:
+        for i in range(n_copies):
+            for t in range(T):
+                if t + v.cycles > T:
+                    solver.add(z3.Not(op[v.name, i, t]))
+
+    # (Dependence): ∀i,t, (u,v,d,δ)∈E, t'∈[0,t+d): op[u,i,t] => ¬op[v,i+δ,t']
+    for edge in E:
+        u_name = edge.src
+        v_name = edge.dst
+        d = edge.delay
+        delta = edge.iteration_delay
+        
+        for i in range(n_copies):
+            j = i + delta  # target iteration copy
+            if j >= n_copies:
+                continue  # skip if target iteration is beyond our window
+            
+            for t in range(T):
+                # If u is scheduled at t, then v cannot be at any t' < t + d
+                implications = []
+                for t_prime in range(min(t + d, T)):
+                    implications.append(z3.Not(op[v_name, j, t_prime]))
+                
+                if implications:
+                    solver.add(z3.Implies(op[u_name, i, t], z3.And(implications)))
+
+    # (Capacity): ∀t,f: Σ_{v,i,c∈[0,cycles(v))} op[v,i,t-c]·RRT[v][f,c] ≤ cap(f)
+    for t in range(T):
+        for f_idx, f_name in enumerate(machine.functional_units):
+            cap = machine.capacity(f_name)
+            if cap <= 0:
+                continue
+            
+            terms = []
+            for v in V:
+                for i in range(n_copies):
+                    for c in range(v.cycles):
+                        usage = int(v.rrt[c, f_idx])
+                        if usage > 0 and 0 <= t - c < T:
+                            terms.append((op[v.name, i, t - c], usage))
+            
+            if terms:
+                solver.add(z3.PbLe(terms, cap))
+
+    # ============================================================
+    # Figure 5: Memory Allocation Constraints (optional)
+    # ============================================================
+    live = {}
+    if enable_memory_constraints and machine.memory_spaces:
+        # live[v, i, t]: result of v (iteration i) is live at time t
+        for v in V:
+            if not v.memory_footprint:
+                continue
+            for i in range(n_copies):
+                for t in range(T + 1):  # include T for Init constraint
+                    live[v.name, i, t] = z3.Bool(f"live_{v.name}_{i}_{t}")
+
+        for v in V:
+            if not v.memory_footprint:
+                continue
+            for i in range(n_copies):
+                # (Init): loop-carried results of last copy are live at T
+                has_loop_carried_out = graph.has_loop_carried_output(v.name)
+                if i == n_copies - 1:
+                    if has_loop_carried_out:
+                        solver.add(live[v.name, i, T])
+                    else:
+                        solver.add(z3.Not(live[v.name, i, T]))
+
+                # (LiveProp-1): (live[v,i,t] ∧ op[v,i,t]) => ¬live[v,i,t-1]
+                for t in range(1, T + 1):
+                    if (v.name, i, t) in live and (v.name, i, t - 1) in live:
+                        if t < T:  # op only defined for t < T
+                            solver.add(z3.Implies(
+                                z3.And(live[v.name, i, t], op[v.name, i, t]),
+                                z3.Not(live[v.name, i, t - 1])
+                            ))
+
+                # (LiveProp-2): (live[v,i,t] ∧ ¬op[v,i,t]) => live[v,i,t-1]
+                for t in range(1, T + 1):
+                    if (v.name, i, t) in live and (v.name, i, t - 1) in live:
+                        if t < T:
+                            solver.add(z3.Implies(
+                                z3.And(live[v.name, i, t], z3.Not(op[v.name, i, t])),
+                                live[v.name, i, t - 1]
+                            ))
+
+                # (DeadProp): propagation of deadness
+                for t in range(1, T + 1):
+                    if (v.name, i, t) not in live or (v.name, i, t - 1) not in live:
+                        continue
+                    
+                    # Collect all consumers of v at time t
+                    consumer_ops = []
+                    for edge in graph.outgoing_edges(v.name):
+                        j = i + edge.iteration_delay
+                        if j < n_copies and t < T:
+                            consumer_ops.append(op[edge.dst, j, t])
+                    
+                    if consumer_ops:
+                        # (DeadProp-1): (¬live[v,i,t] ∧ ∨consumers_active) => live[v,i,t-1]
+                        solver.add(z3.Implies(
+                            z3.And(z3.Not(live[v.name, i, t]), z3.Or(consumer_ops)),
+                            live[v.name, i, t - 1]
+                        ))
+                        # (DeadProp-2): (¬live[v,i,t] ∧ ∧¬consumers_active) => ¬live[v,i,t-1]
+                        solver.add(z3.Implies(
+                            z3.And(z3.Not(live[v.name, i, t]),
+                                   *[z3.Not(c) for c in consumer_ops]),
+                            z3.Not(live[v.name, i, t - 1])
+                        ))
+                    else:
+                        # No consumers at this time -> propagate deadness
+                        solver.add(z3.Implies(
+                            z3.Not(live[v.name, i, t]),
+                            z3.Not(live[v.name, i, t - 1])
+                        ))
+
+        # (Memory Capacity): ∀t,m: Σ_{v,i} live[v,i,t]·footprint(v,m) ≤ capacity(m)
+        for t in range(T):
+            for mem_name, mem_cap in machine.memory_spaces.items():
+                terms = []
+                for v in V:
+                    if mem_name in v.memory_footprint and v.memory_footprint[mem_name] > 0:
+                        for i in range(n_copies):
+                            if (v.name, i, t) in live:
+                                terms.append(
+                                    (live[v.name, i, t], v.memory_footprint[mem_name])
+                                )
+                if terms:
+                    solver.add(z3.PbLe(terms, mem_cap))
+
+    # ============================================================
+    # Figure 6: Warp Assignment Constraints (optional)
+    # ============================================================
+    opw = {}
+    if enable_warp_constraints:
+        num_warps = machine.num_warps
+        W_vl = machine.variable_latency_warp
+
+        # opw[v, w]: instruction v is assigned to warp w
+        for v in V:
+            for w in range(num_warps):
+                opw[v.name, w] = z3.Bool(f"opw_{v.name}_{w}")
+
+        # (WarpUniqueness): ∀v: Σ_w opw[v,w] = 1
+        for v in V:
+            solver.add(z3.PbEq([(opw[v.name, w], 1) for w in range(num_warps)], 1))
+
+        # (VariableLatency): ∀v: variable_latency(v) <=> opw[v, W_vl]
+        # Only enforce if there ARE variable-latency ops (otherwise W_vl is unused
+        # and all warps should be available for compute)
+        has_var_lat_ops = any(v.variable_latency for v in V)
+        for v in V:
+            if v.variable_latency:
+                # Variable-latency ops MUST go to W_vl
+                solver.add(opw[v.name, W_vl])
+            elif has_var_lat_ops:
+                # Non-variable-latency ops must NOT go to W_vl (reserved for var-lat)
+                solver.add(z3.Not(opw[v.name, W_vl]))
+
+        # (WarpCapacity): Per-warp resource constraints
+        # For each time slot t ∈ [0, I), warp w, functional unit f:
+        # Σ_{v,c} opw[v,w] · op[v,0,t_shift] · RRT[v][f,c] ≤ per_warp_cap(f)
+        # 
+        # Simplified: for each warp, the total resource usage across
+        # its assigned instructions in a single I-window must fit
+        # This is a tighter constraint that prevents resource conflicts within a warp
+        for w in range(num_warps):
+            for f_idx, f_name in enumerate(machine.functional_units):
+                cap = machine.capacity(f_name)
+                if cap <= 0:
+                    continue
+                
+                # Total usage across one initiation interval per warp
+                terms = []
+                for v in V:
+                    total_fu_use = int(v.rrt[:, f_idx].sum())
+                    if total_fu_use > 0:
+                        terms.append((opw[v.name, w], total_fu_use))
+                
+                if terms:
+                    # Per-warp capacity: at most cap per I cycles
+                    solver.add(z3.PbLe(terms, cap * I))
+
+        # (Cross-warp synchronization): When two ops are on different warps
+        # and have a dependence, a barrier is needed (mbarrier on Hopper/Blackwell).
+        # This is modeled implicitly: the dependence constraints already enforce
+        # timing, and the WS assignment determines which barriers are needed.
+        # Twill emits barrier annotations in code generation.
+
+    # ============================================================
+    # Solve
+    # ============================================================
+    if verbose:
+        stats = solver.statistics()
+        print(f"  SMT: Constraints added. Solving...")
+
+    result = solver.check()
+
+    if result != z3.sat:
+        if verbose:
+            print(f"  SMT: {result}")
+        return None
+
+    model = solver.model()
+
+    # Extract schedule M*
+    new_schedule = {}
+    for v in V:
+        for t in range(T):
+            if z3.is_true(model.evaluate(op[v.name, 0, t])):
+                new_schedule[v.name] = t
+                break
+        else:
+            # Should not happen if constraints are correct
+            raise RuntimeError(f"No time slot found for instruction {v.name} in iteration 0")
+
+    # Extract warp assignment A*
+    warp_assign = {}
+    if enable_warp_constraints:
+        for v in V:
+            for w in range(machine.num_warps):
+                if z3.is_true(model.evaluate(opw[v.name, w])):
+                    warp_assign[v.name] = w
+                    break
+            else:
+                raise RuntimeError(f"No warp found for instruction {v.name}")
+    else:
+        # Default: all on warp 0
+        for v in V:
+            warp_assign[v.name] = 0
+
+    wa = WarpAssignment(warp_assign, machine.num_warps)
+
+    # Label warps
+    if enable_warp_constraints:
+        vl_warp = machine.variable_latency_warp
+        wa.warp_names[vl_warp] = f"Warp {vl_warp} (variable-latency/producer)"
+        for w in range(machine.num_warps):
+            if w != vl_warp and w not in wa.warp_names:
+                wa.warp_names[w] = f"Warp {w} (compute/consumer)"
+
+    return JointSWPWSResult(
+        schedule=new_schedule,
+        I=I,
+        L=L,
+        warp_assignment=wa,
+    )