# Copyright Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT """MFMATile-based logical scheduler. Builds a logical schedule using MFMA tile indices as the core primitive, with explicit per-operation load granularity for GR/LR on A, B, SA, SB. The schedule is built in these passes: place_LRs — place LRs based on their granularities assign_vgpr_tiles — assign physical vgprTileIds with per-tensor free-lists place_GRs — place GRs annotate_deps — annotate raw per-op dependencies remove_unnecessary_gr_deps — remove redundant LR→GR deps remove_unnecessary_lr_deps — remove redundant GR→LR deps covered by MFMA syncs remove_cross_deps — replace cross-subIterK deps with wait preOps insert_gr_lr_inc — insert lr_inc/gr_inc preOps at MT transitions group — serialize and group (produce paths for instructionSchedule) remove_wait_lr_sync — remove redundant wait_lr_sync after grouping remove_wait_gr_sync — drop redundant wait_gr_sync on interior partition s1 emit — produce List[EmittedModule] with before-link chains """ from __future__ import annotations from dataclasses import dataclass, field from enum import IntEnum from typing import Callable, ClassVar, Dict, List, Optional, Tuple, Union from bisect import bisect_left import copy import io import math from rocisa.code import Module from .ScheduleTypes import ( AnnotatedSchedule, AugmentedSchedule, EmittedSchedule, LogicalSchedule, PartitionSchedule, ) from ...Common.GlobalParameters import globalParameters # ds_load_b128 reads 4 contiguous VGPRs. DS_B128_VGPRS = 4 def _checkout_tile(pool, numRegs, tag): """Check out one VGPR tile as a single contiguous, min(numRegs, 4)-aligned block (b128-aligned when numRegs >= 4).""" from Tensile.Components.Subtile.Kernel import RegisterTileInfo # min(): full b128 tiles get 4-VGPR alignment; smaller tiles aren't padded # up to 4 (which would waste registers and can break occupancy). align = min(numRegs, DS_B128_VGPRS) base = pool.checkOutAligned(numRegs, align, tag=tag) tile = RegisterTileInfo(pool) for k in range(numRegs): tile.append(base + k) return tile def _checkin_tile(tile): """Return a contiguous tile to its pool via its base-register handle.""" tile.regList.pool.checkIn(tile.regList.indices[0]) def _hoistSetupBeforeLastMFMA(loopBody, setupInsts): """Move loop-control setup (dec + compare) above the body's last MFMA. Place the loop-control branch immediately after the last MFMA to hide branching latency behind it. The decrement+compare that feed the branch's SCC are normally emitted right after the body (MFMA; s_sub; s_cmp; s_cbranch); hoisting them above the final MFMA produces (s_sub; s_cmp; MFMA; s_cbranch) so the branch directly follows the MFMA. The hoist is dependency-safe: SCC set by s_cmp survives the MFMA (which never touches SCC), and the MFMA does not read LoopCounterL. Falls back to appending the setup when the body has no MFMA (e.g. a skip branch with no preceding compute), leaving behavior unchanged there. Returns a rebuilt Module; the input is left untouched. """ from rocisa.instruction import MFMAInstruction, MXMFMAInstruction items = loopBody.flatitems() lastMfma = -1 for idx, it in enumerate(items): if isinstance(it, (MFMAInstruction, MXMFMAInstruction)): lastMfma = idx rebuilt = Module(loopBody.name) if lastMfma < 0: for it in items: rebuilt.add(it) for s in setupInsts: rebuilt.add(s) return rebuilt for idx, it in enumerate(items): if idx == lastMfma: for s in setupInsts: rebuilt.add(s) rebuilt.add(it) return rebuilt def _emitBodySetupBranch(module, loopBody, setupInsts, branch, hoist): """Append a loop body, its dec+compare setup, and the trailing control branch. When ``hoist`` is set, the setup is moved above the body's last MFMA (via _hoistSetupBeforeLastMFMA) so the branch lands right after that MFMA to hide branching latency. Otherwise the setup and branch simply follow the body. """ if hoist: module.add(_hoistSetupBeforeLastMFMA(loopBody, setupInsts)) else: module.add(loopBody) for s in setupInsts: module.add(s) module.add(branch) class Pass(IntEnum): """Scheduler passes in dependency order. The numeric value defines topological order. The main pipeline is linear (each pass depends on the previous), except VGPR_TILES which forks off LR independently of GR. """ LR = 0 VGPR_TILES = 1 GR = 2 DEPS = 3 REMOVE_GR_DEPS = 4 REMOVE_LR_DEPS = 5 REMOVE_DEPS = 6 GR_INC = 7 GROUP_LR_GR = 8 REMOVE_WAIT_LR_SYNC = 9 REMOVE_WAIT_GR_SYNC = 10 EMIT = 11 BUILD = 12 POPULATE = 13 @classmethod def pipeline(cls) -> Tuple[Pass, ...]: """Ordered pipeline passes (LR through EMIT).""" return tuple(cls(i) for i in range(cls.EMIT + 1)) TENSOR_SIDE = {'A': 'A', 'B': 'B', 'SA': 'A', 'SB': 'B'} def fmt_mt(mt: int) -> str: """Format MT iteration integer as display string: 0 → 'n', 1 → 'n+1', 2 → 'n+2'.""" return "n" if mt == 0 else f"n+{mt}" # ── Core primitives ───────────────────────────────────────── @dataclass class MFMATileRange: """A rectangular range of MFMA tile coordinates for one read.""" subIterK_start: int subIterK_end: int # exclusive tileId_start: int tileId_end: int # exclusive @property def subIterK_list(self) -> List[int]: return list(range(self.subIterK_start, self.subIterK_end)) @property def tileId_list(self) -> List[int]: return list(range(self.tileId_start, self.tileId_end)) def fmt_k(self) -> str: ids = self.subIterK_list if len(ids) == 1: return f"[{ids[0]}]" return f"[{ids[0]},{ids[-1]}]" def fmt_tiles(self) -> str: return f"[{self.tileId_start}-{self.tileId_end - 1}]" # ── Config ────────────────────────────────────────────────── @dataclass class ReadGranularity: """Load granularity for one operation on one tensor, measured in MFMA tiles. mn: how many MFMA tiles in the M (for A/SA) or N (for B/SB) dimension k: how many subIterK steps one read covers """ mn: int k: int def tile_range(self, k: int, t_start: int, t_end: int) -> 'MFMATileRange': """Snap subIterK and tile indices to this granularity, return MFMATileRange.""" ks = (k // self.k) * self.k ts = (t_start // self.mn) * self.mn te = ((t_end + self.mn - 1) // self.mn) * self.mn return MFMATileRange(ks, ks + self.k, ts, te) class GRPlacementStrategy(IntEnum): """How Global Reads are spread across (partition, subIterK) slots. SPREAD — distribute GR atoms across all slots weighted by partition MFMA count. Default for buffer-load GR (gfx9xx): spreading avoids issue-slot pressure on the SIMD. BUNCHED — pin every GR atom to partition 0, subIterK 0. Suitable when the GR instruction does not contend for SIMD issue slots (e.g. TDM tensor_load_to_lds on gfx1250). """ SPREAD = 0 BUNCHED = 1 @dataclass class SchedulerConfig: """Configuration for the MFMATile-based scheduler.""" numMFMATilesM: int # MFMA tiles in M dimension (for A) numMFMATilesN: int # MFMA tiles in N dimension (for B) numSubIterK: int # subIterK steps within the macrotile (may be expanded in __post_init__) lrA: ReadGranularity lrB: ReadGranularity grA: ReadGranularity grB: ReadGranularity lrSA: Optional[ReadGranularity] = None lrSB: Optional[ReadGranularity] = None grSA: Optional[ReadGranularity] = None grSB: Optional[ReadGranularity] = None partitionSizeM: Union[int, List[int]] = 0 # partition size(s) in M dimension (0 = full dim) partitionSizeN: Union[int, List[int]] = 0 # partition size(s) in N dimension (0 = full dim) pgr: int = 2 # Prefetch Global Read grPlacement: GRPlacementStrategy = GRPlacementStrategy.SPREAD pgl: int = 0 # Prefetch GL2 (0=off, 1 or 2 tiles ahead) # Resolve a partition spec into per-partition sizes along one dimension. # spec is either: # - an explicit list (must sum to total) # - a single tile size (0 means full dim). # Uneven splits place the remainder in the middle so the smaller partition is bracketed by full ones. # # Every partition size must be a multiple of `mn` (LR read granularity) to avoid emitting under-sized LRs. # Single-int specs are rounded DOWN to an mn-multiple (smaller partition, less VGPR usage). # If no solution exists, we return [total] (single partition). @staticmethod def _data_tensors_multi_du(numUnroll: dict) -> bool: """True when A/B use numUnroll > 1 (MX multi-DU data path).""" return numUnroll.get('A', 1) > 1 or numUnroll.get('B', 1) > 1 @staticmethod def _normalize_partition_sizes(spec: Union[int, List[int]], total: int, dim: str, mn: int = 1, remainder_last: bool = False) -> List[int]: if isinstance(spec, (list, tuple)): assert sum(spec) == total, \ f"partition sizes for {dim} must sum to {total}, got {sum(spec)}" assert all(s >= 1 for s in spec), \ f"all partition sizes for {dim} must be >= 1" assert all(s % mn == 0 for s in spec), \ f"partition sizes for {dim} must be multiples of mn={mn}, got {list(spec)}" return list(spec) s = spec if spec != 0 else total assert 1 <= s <= total, \ f"partition size for {dim} must be in [1, {total}], got {s}" if total % mn != 0: return [total] s = max(mn, (s // mn) * mn) if s > total: return [total] num_full = total // s remainder = total - num_full * s if remainder == 0: return [s] * num_full if not remainder_last: if num_full == 1: return [s, remainder] mid = num_full // 2 return [s] * mid + [remainder] + [s] * (num_full - mid) # Multi-DU: remainder (short) partition must come LAST. PGR=1 multi-DU # codegen assumes every partition before the last is the full size `s`. return [s] * num_full + [remainder] @staticmethod def _build_prefix(sizes: List[int]) -> List[int]: prefix = [0] for s in sizes: prefix.append(prefix[-1] + s) return prefix def __post_init__(self): assert self.pgr in (0, 1, 2), f"pgr must be 0, 1, or 2, got {self.pgr}" grans = {'A': self.grA, 'B': self.grB} if self.hasScale: grans['SA'] = self.grSA grans['SB'] = self.grSB maxGrK = max(g.k for g in grans.values()) for t, g in grans.items(): assert maxGrK % g.k == 0, \ f"grGran[{t}].k={g.k} must divide max(grGran.k)={maxGrK}" assert self.numSubIterK % maxGrK == 0 or maxGrK % self.numSubIterK == 0, \ f"numSubIterK={self.numSubIterK} and max(grGran.k)={maxGrK} must be multiples of each other" original_numSubIterK = self.numSubIterK defaultReads = {t: original_numSubIterK // g.k for t, g in grans.items()} expanded_k = max(original_numSubIterK, maxGrK) numUnroll = {} for t, g in grans.items(): reads = expanded_k // g.k numUnroll[t] = reads // (defaultReads[t] if defaultReads[t] != 0 else 1) data_multi_du = self._data_tensors_multi_du(numUnroll) if data_multi_du: self.numSubIterK = expanded_k self.numUnroll = numUnroll else: self.numSubIterK = original_numSubIterK self.numUnroll = {t: 1 for t in grans} mn_M = max((g.mn for g in (self.lrA, self.lrSA) if g is not None), default=1) mn_N = max((g.mn for g in (self.lrB, self.lrSB) if g is not None), default=1) self._partitionSizesM = self._normalize_partition_sizes( self.partitionSizeM, self.numMFMATilesM, 'M', mn_M, remainder_last=data_multi_du) self._partitionSizesN = self._normalize_partition_sizes( self.partitionSizeN, self.numMFMATilesN, 'N', mn_N, remainder_last=data_multi_du) self._prefixM = self._build_prefix(self._partitionSizesM) self._prefixN = self._build_prefix(self._partitionSizesN) self.plr = 0 if self.pgr == 0 else 1 self.offsetPartition = 1 if self.pgr >= 2 else 0 if self.pgr == 0: assert self.numPartitions == 1, "pgr=0 requires numPartitions=1" @property def partitionSizesM(self) -> List[int]: return self._partitionSizesM @property def partitionSizesN(self) -> List[int]: return self._partitionSizesN @property def hasScale(self) -> bool: return self.lrSA is not None and self.lrSB is not None @property def numPartitionsM(self) -> int: return len(self._partitionSizesM) @property def numPartitionsN(self) -> int: return len(self._partitionSizesN) @property def numPartitions(self) -> int: return self.numPartitionsM * self.numPartitionsN @staticmethod def get_partition_candidates(tileInfoA, tileInfoB) -> list: """Return partition candidates as [(partitionSizeM, partitionSizeN), ...]. For the smaller dimension, uses a single partition (full size). For the larger dimension, starts at full size then jumps to divUp(dim,2) and decrements from there, skipping unbalanced 2-partition sizes. """ M = tileInfoA.localMMATileGrid[0] N = tileInfoB.localMMATileGrid[0] def divUp(n, d): return (n + d - 1) // d def partitionSizes(dim): return [dim] + list(range(divUp(dim, 2), 0, -1)) if N >= M: candidates = [(M, s) for s in partitionSizes(N)] else: candidates = [(s, N) for s in partitionSizes(M)] return candidates # ── Schedule operation types ──────────────────────────────── @dataclass class Emittable: """Base for anything placed in an EmittedModule.""" kind: str = field(init=False, default="") @dataclass class MFMAPlacement(Emittable): """MFMA operation consuming data for one subIterK.""" subIterK: int tileA: MFMATileRange # A tiles consumed tileB: MFMATileRange # B tiles consumed deps: List['Dep'] = field(default_factory=list) # populated by annotate_deps() preOps: List['BaseOp'] = field(default_factory=list) # populated by remove_cross_deps() postOps: List['BaseOp'] = field(default_factory=list) # populated by insert_gr_lr_inc() vgpr_tile_maps: Dict[str, List[dict]] = field(default_factory=dict) # {tensor: [{groupIdx: vgprTileId}]} per unroll iter def __post_init__(self): self.kind = 'mfma' def __str__(self): return (f"MFMAs (MT n, subIterK {self.subIterK} ) " f"A : {self.tileA.fmt_tiles()} , B : {self.tileB.fmt_tiles()}") @dataclass class LRPlacement(Emittable): """Local Read placement for one tensor in one subIterK slot.""" tensor: str # 'A', 'B', 'SA', 'SB' mtIteration: int # 0 = current MT, 1 = next MT tiles: MFMATileRange subIterK_slot: int # which subIterK this LR is placed in partition: int = 0 # which partition this LR belongs to deps: List['Dep'] = field(default_factory=list) # populated by annotate_deps() preOps: List['BaseOp'] = field(default_factory=list) # populated by remove_cross_deps() postOps: List['BaseOp'] = field(default_factory=list) # populated by insert_gr_lr_inc() vgpr_tile_map: List[dict] = field(default_factory=list) # [{tileId: vgprTileId}] per unroll iter def __post_init__(self): self.kind = 'lr' def __str__(self): return (f"LR {self.tensor.ljust(2)} (MT {fmt_mt(self.mtIteration)}, " f"subIterK {self.tiles.fmt_k()}) {self.tiles.fmt_tiles()}") @dataclass class GRPlacement(Emittable): """Global Read placement for one tensor in one subIterK slot.""" tensor: str # 'A', 'B', 'SA', 'SB' mtIteration: int # 0 = current MT, 1 = next MT, 2 = two MTs ahead tiles: MFMATileRange subIterK_slot: int # which subIterK this GR is placed in partition: int = 0 # which partition this GR belongs to unrollId: int = 0 # which inner-DU iteration (0 for single-DU configs) deps: List['Dep'] = field(default_factory=list) # populated by annotate_deps() preOps: List['BaseOp'] = field(default_factory=list) # populated by remove_cross_deps() postOps: List['BaseOp'] = field(default_factory=list) # populated by insert_gr_lr_inc() def __post_init__(self): self.kind = 'gr' def __str__(self): return (f"GR {self.tensor} (MT {fmt_mt(self.mtIteration)}, " f"subIterK {self.tiles.fmt_k()}) ids {self.tiles.fmt_tiles()}") # ── Per-subIterK container ────────────────────────────────── @dataclass class SubIterKSlot: """All operations placed in one subIterK step.""" subIterK: int mfma: Optional[MFMAPlacement] = None lrs: List[LRPlacement] = field(default_factory=list) grs_by_unroll: Dict[int, List[GRPlacement]] = field(default_factory=lambda: {0: []}) @property def grs(self) -> List[GRPlacement]: """All GRs across all unrollIds, in uid order.""" result = [] for uid in sorted(self.grs_by_unroll.keys()): result.extend(self.grs_by_unroll[uid]) return result # ── Dependency types ──────────────────────────────────────── @dataclass class WaitGRCounts: """Per-tensor inflight load counts for wait_gr preOp.""" A: int = 0 B: int = 0 SA: int = 0 SB: int = 0 def __str__(self): parts = [] for t in ('A', 'B', 'SA', 'SB'): v = getattr(self, t) if v: parts.append(f"{t}={v}") return ",".join(parts) if parts else "0" @dataclass class BaseOp(Emittable): """Base class for typed dependency operations in a before-chain.""" def __str__(self): return self.kind @dataclass class WaitGROp(BaseOp): """Wait for global reads to complete. Optionally includes a sync barrier.""" wait_gr_counts: Optional[WaitGRCounts] = None has_sync: bool = False adjustVmcnt: bool = True # When True the emitter ignores wait_gr_counts and forces a full vmcnt(0) # drain (all outstanding global reads must retire). wait_gr_counts is still # populated with the precise per-tensor inflight estimate for diagnostics, # and force_drain survives _merge_preops so the min-count merge cannot # silently weaken the drain back to a partial vmcnt. force_drain: bool = False def __post_init__(self): self.kind = 'wait_gr' def __str__(self): if self.wait_gr_counts: return f"{self.kind}({self.wait_gr_counts})" return self.kind @dataclass class WaitLROp(BaseOp): """Wait for local reads to complete. Optionally includes a sync barrier.""" has_sync: bool = False def __post_init__(self): self.kind = 'wait_lr' def __str__(self): return 'wait_lr_sync' if self.has_sync else 'wait_lr' @dataclass class SyncOp(BaseOp): """Standalone sync barrier.""" def __post_init__(self): self.kind = 'sync' @dataclass class MaskKOp(BaseOp): """Zero A and B vgprs whose K-index >= remaining tail K, for one subIterK group. """ subIterK: int = 0 vgpr_tile_map: dict = field(default_factory=dict) def __post_init__(self): self.kind = 'mask_k' def __str__(self): return f"mask_k(k={self.subIterK})" @dataclass class LRIncOp(BaseOp): """LDS buffer swap for local reads on a specific tensor.""" tensor: str = "" isUnrollSwap: bool = False # True when swap is at a unrollId boundary; preserved in NLL # Destination double-buffer (unrollId) this swap points LR at. Display-only: # surfaced by print_emit_dep_order so an unroll-swap is not mistaken for a bug # when the following same-MT LR belongs to a different uid. Ignored by codegen. unrollId: int = 0 def __post_init__(self): self.kind = 'lr_inc' def __str__(self): return f"lr_inc({self.tensor})" @dataclass class GRIncOp(BaseOp): """Pointer update + LDS swap for global reads on a specific tensor.""" tensor: str = "" unrollId: int = 0 def __post_init__(self): self.kind = 'gr_inc' def __str__(self): return f"gr_inc({self.tensor})" @dataclass class GL2PrefetchOp(BaseOp): """Issue GL2 prefetch loads (global_prefetch_b8) for all tensors.""" def __post_init__(self): self.kind = 'gl2_prefetch' def __str__(self): return "gl2_prefetch" @dataclass class GL2PrefetchIncOp(BaseOp): """Increment GL2 prefetch addresses with end-of-K guard.""" def __post_init__(self): self.kind = 'gl2_prefetch_inc' def __str__(self): return "gl2_prefetch_inc" @dataclass class SkipOp(BaseOp): """Skip guard: compare LoopCounter and branch. target is normally a short name (e.g. 'NLL'); the emitter prefixes 'SkipTo'. Set rawLabel=True to pass the label name through verbatim (e.g. 'SkipTailLoopL'). branchComment overrides the default.""" compare: str = "" value: int = 0 target: str = "" rawLabel: bool = False branchComment: str = "" def __post_init__(self): self.kind = 'skip' @property def tensor(self) -> str: return f"{self.compare}:{self.value}:{self.target}" def __str__(self): return f"skip({self.tensor})" @dataclass class InlineModuleOp(BaseOp): """Inline a writer-built Module at this point in the schedule. The callback receives the InstructionEmitter (so it can reach writer, kernel, tensorParametersMap, etc.) and must return a rocisa Module. Use this for one-off boilerplate that doesn't deserve its own Op class.""" build: Optional[Callable] = None label: str = "inline" def __post_init__(self): self.kind = 'inline' def __str__(self): return f"inline({self.label})" @dataclass class Dep: """Dependency on another placement (annotate_deps output).""" ref: Union[LRPlacement, GRPlacement] mt_offset: int = 0 # 0 = same MT, -1 = prev MT, -2 = two MTs back, ... # ── Emitted output ───────────────────────────────────────── @dataclass class EmittedModule: """One emitted module with before-link for instruction scheduling. Compatible with SubtileBasedInstructionScheduler.instructionSchedule(). Instructions are left empty at the logical level — filled during emission. """ moduleId: int = -1 instructions: list = field(default_factory=list) before: Optional[int] = None # moduleId that must complete before this module source: Optional[Emittable] = None @property def opType(self) -> str: return self.source.kind if self.source else "" # ── Main scheduler class ─────────────────────────────────── class LogicalScheduler: """Subtile-based logical scheduler. Pass methods are standalone; call build() to run the full pipeline. """ _PASS_METHODS: ClassVar[Dict[Pass, str]] = {} def pipeline_pass(p, _methods=_PASS_METHODS): """Register a pipeline pass method. Use as @pipeline_pass(Pass.LR) on pass methods.""" def decorator(fn): _methods[p] = fn.__name__ return fn return decorator def __init__(self, config: SchedulerConfig): self.config = config self.tensors: List[str] = ['A', 'B'] + (['SA', 'SB'] if config.hasScale else []) # Shared mutable state across passes. The same field holds different # stage representations over time; see ScheduleTypes for stage meanings. self._partitions: Optional[Union[LogicalSchedule, AnnotatedSchedule, AugmentedSchedule]] = None self._emitted: Optional[EmittedSchedule] = None self._preloop_emitted: Optional[EmittedSchedule] = None self._ngll_emitted: Optional[EmittedSchedule] = None self._nll_emitted: Optional[EmittedSchedule] = None # Tail-loop tile bookkeeping. Tail loop only use a subset of tiles, so we track which tileIds are # unused or freed for reuse within the tail loop. self._tail_unused_tile_ids: Dict[str, set] = {'A': set(), 'B': set(), 'SA': set(), 'SB': set()} self._tail_freed_tile_ids: Dict[str, set] = {'A': set(), 'B': set(), 'SA': set(), 'SB': set()} # ── Place LRs ───────────────────────────────────────── def _partition_tile_range(self, pi: int) -> dict: """Return {'A': (start, end), 'B': (start, end)} for partition pi. Uses COLUMN_MAJOR ordering: M (A) varies fastest, N (B) varies slowest. Tile ranges are derived from prefix sums of partition sizes. """ cfg = self.config piM = pi % cfg.numPartitionsM piN = pi // cfg.numPartitionsM return {'A': (cfg._prefixM[piM], cfg._prefixM[piM + 1]), 'B': (cfg._prefixN[piN], cfg._prefixN[piN + 1])} @pipeline_pass(Pass.LR) def place_LRs(self) -> LogicalSchedule: """Place MFMAs and LRs based on read granularities. Returns a list of partitions, each containing a list of SubIterKSlots. Each LR prefetches data for the next subIterK group. Within-partition prefetches use current partition tiles; cross-partition prefetches (wrapping) use next partition tiles. Two tracking mechanisms: - loaded_ranges: tracks tile ranges in VGPR per side. Wrapping LRs are only placed when the next partition's tiles aren't already loaded. - placed: tracks (tensor, k-range, tile-range) of non-wrapping LRs placed so far across partitions. Skips redundant K-prefetch when the same data was already loaded by an earlier partition. """ if self.config.plr == 0: return self._place_LRs_PLR0() cfg = self.config numP = cfg.numPartitions part_ranges = [self._partition_tile_range(pi) for pi in range(numP)] # Track which tile ranges are currently loaded in VGPR (for wrapping decisions). loaded_ranges = {'A': {part_ranges[0]['A']}, 'B': {part_ranges[0]['B']}} # Track placed K-prefetch LRs across partitions (for dedup). placed = set() partitions = [] for pi in range(numP): cur, nxt = part_ranges[pi], part_ranges[(pi + 1) % numP] is_last = (pi == numP - 1) load = {} for side in ('A', 'B'): load[side] = is_last or nxt[side] not in loaded_ranges[side] slots = self._place_LRs_for_partition(cur, nxt, is_last, load, placed) for slot in slots: for lr in slot.lrs: lr.partition = pi partitions.append(slots) for side in ('A', 'B'): if load[side]: loaded_ranges[side] = {cur[side], nxt[side]} self._partitions = partitions return partitions def _create_partition_slots(self, cur: dict) -> List[SubIterKSlot]: """Create SubIterKSlots with MFMAs placed for one partition.""" numK = self.config.numSubIterK slots = [SubIterKSlot(subIterK=k) for k in range(numK)] for k in range(numK): slots[k].mfma = MFMAPlacement( subIterK=k, tileA=MFMATileRange(k, k + 1, cur['A'][0], cur['A'][1]), tileB=MFMATileRange(k, k + 1, cur['B'][0], cur['B'][1]), ) return slots def _lr_tensors(self) -> list: """Return list of (tensor_name, ReadGranularity) for all LR tensors.""" cfg = self.config tensors = [('A', cfg.lrA), ('B', cfg.lrB)] if cfg.hasScale: tensors.append(('SA', cfg.lrSA)) tensors.append(('SB', cfg.lrSB)) return tensors def _place_LRs_PLR0(self) -> List[List[SubIterKSlot]]: """Place MFMAs and LRs for PLR=0: no prefetching. Each LR loads data for its own subIterK. All LRs have mtIteration=0. Single partition only (enforced by config validation). """ cfg = self.config numK = cfg.numSubIterK cur = self._partition_tile_range(0) slots = self._create_partition_slots(cur) for tensor, gran in self._lr_tensors(): side_key = 'A' if tensor in ('A', 'SA') else 'B' ts, te = cur[side_key] k_gran = gran.k num_chunks = numK // k_gran for chunk_idx in range(num_chunks): lr_k_start = chunk_idx * k_gran lr_k_end = lr_k_start + k_gran slot_k = lr_k_start lr = LRPlacement( tensor=tensor, mtIteration=0, tiles=MFMATileRange(lr_k_start, lr_k_end, ts, te), subIterK_slot=slot_k, ) slots[slot_k].lrs.append(lr) self._partitions = [slots] return self._partitions def _place_LRs_for_partition(self, cur: tuple, nxt: tuple, is_last: bool, load: dict, placed: set) -> List[SubIterKSlot]: """Place MFMAs and LRs for one partition.""" cfg = self.config numK = cfg.numSubIterK multi_part = cfg.numPartitions > 1 slots = self._create_partition_slots(cur) slot_mt = {} # slot_k → lr_mt string, for MT-homogeneity enforcement all_tensors = self._lr_tensors() # Place LRs grouped by k_gran. # - Non-wrapping (K-prefetch): all tensors, deduped by placed set. # - Wrapping (cross-partition): only tensors whose side needs loading. for k_gran in sorted(set(g.k for _, g in all_tensors)): group_all = [(t, g) for t, g in all_tensors if g.k == k_gran] num_chunks = numK // k_gran for chunk_idx in range(num_chunks): next_chunk = (chunk_idx + 1) % num_chunks is_wrap = (next_chunk == 0) lr_mt = 1 if is_last and is_wrap else 0 lr_k_start = next_chunk * k_gran lr_k_end = lr_k_start + k_gran base_slot = chunk_idx * k_gran # For wrapping chunks, only include tensors whose side is # loading so that slot assignment reflects active tensors. # A and B always participate (their wrapping is gated inside # the loop) to keep slot indices stable for their k_gran group. if is_wrap and multi_part: group = [(t, g) for t, g in group_all if t in ('A', 'B') or load['A' if t in ('A', 'SA') else 'B']] else: group = group_all # Group by side (A/SA together, B/SB together) for slot assignment sides = [[(t, g) for t, g in group if t in ('A', 'SA')], [(t, g) for t, g in group if t in ('B', 'SB')]] sides = [s for s in sides if s] for side_idx, side in enumerate(sides): slot_k = base_slot + (side_idx % k_gran) # Redirect LRs away from slots committed to a different MT, # keeping each slot MT-homogeneous. # This reduce the number of wait_gr_sync needed as all LRs # in the same subIterK wait for the same MT iterration. committed = slot_mt.get(slot_k) if committed is not None and committed != lr_mt: slot_k = numK - 1 for tensor, gran in side: tile_range = nxt if (is_wrap or not multi_part) else cur side_key = 'A' if tensor in ('A', 'SA') else 'B' ts, te = tile_range[side_key] # Wrapping: use load dict. Non-wrapping: use placed set. if is_wrap and multi_part: if not load[side_key]: continue else: lr_key = (tensor, lr_k_start, lr_k_end, ts, te) if lr_key in placed: continue placed.add(lr_key) lr = LRPlacement( tensor=tensor, mtIteration=lr_mt, tiles=MFMATileRange(lr_k_start, lr_k_end, ts, te), subIterK_slot=slot_k, ) slots[slot_k].lrs.append(lr) slot_mt[slot_k] = lr_mt return slots # ── Assign VGPR tile IDs ───────────────────────────── def _use_free_list_vgpr_allocation(self): """MX scaled multi-DU: subIterK slots are sequential uid slices. Lifetime-based free-list reuse keeps B peak at max_groups instead of 2*max_groups from the deterministic double-buffer formula. """ cfg = self.config if cfg.pgr == 0 or not cfg.hasScale: return False return self._is_multi_du() def _is_multi_du(self) -> bool: """True when A/B use numUnroll > 1 (MX multi-DU data path).""" cfg = self.config return bool(cfg.numUnroll) and SchedulerConfig._data_tensors_multi_du(cfg.numUnroll) @pipeline_pass(Pass.VGPR_TILES) def assign_vgpr_tiles(self): """Assign physical vgprTileIds to all placements (A, B, SA, SB).""" if self._use_free_list_vgpr_allocation(): self._assign_vgpr_tiles_free_list() else: self._assign_vgpr_tiles_deterministic() def _assign_vgpr_tiles_free_list(self): """Lifetime-based free-list allocator (MX multi-DU path). Iterated until convergence (max 8 unroll iterations). Sets self.tile_peaks, self.needs_unrolling, self.unroll_factor. """ cfg = self.config numK = cfg.numSubIterK MAX_UNROLL = 8 lr_grans = {'A': cfg.lrA, 'B': cfg.lrB} if cfg.hasScale: lr_grans['SA'] = cfg.lrSA lr_grans['SB'] = cfg.lrSB # ── Phase 1: find last MFMA read for each key ── last_read = {} # key -> flat position for pi, slots in enumerate(self._partitions): for slot in slots: if not slot.mfma: continue pos = pi * numK + slot.subIterK k = slot.subIterK for tensor in self.tensors: side = TENSOR_SIDE[tensor] tileRange = slot.mfma.tileA if side == 'A' else slot.mfma.tileB gran = lr_grans[tensor] for t in tileRange.tileId_list: group = (t // gran.mn) * gran.mn k_chunk = (k // gran.k) * gran.k last_read[(tensor, group, k_chunk)] = pos # ── Phase 2: iterate until convergence ── from collections import deque class _FreeList: __slots__ = ('free', 'next_id', 'active_count', 'peak') def __init__(self): self.free = deque() self.next_id = 0 self.active_count = 0 self.peak = 0 def alloc(self): if self.free: vid = self.free.popleft() # FIFO for convergence else: vid = self.next_id self.next_id += 1 self.active_count += 1 self.peak = max(self.peak, self.active_count) return vid def release(self, vid): self.free.append(vid) self.active_count -= 1 max_peaks = {t: 0 for t in self.tensors} carry_active = {} all_next_iters = [] pools = {t: _FreeList() for t in self.tensors} for unroll_iter in range(MAX_UNROLL): if unroll_iter == 0: active = {} else: active = dict(carry_active) for t in self.tensors: pools[t].active_count = sum( 1 for key in active if key[0] == t) next_iter = {} for pi, slots in enumerate(self._partitions): for slot in slots: pos = pi * numK + slot.subIterK k = slot.subIterK if slot.mfma: for tensor in self.tensors: side = TENSOR_SIDE[tensor] tileRange = slot.mfma.tileA if side == 'A' else slot.mfma.tileB gran = lr_grans[tensor] tile_map = {} for t in tileRange.tileId_list: group = (t // gran.mn) * gran.mn k_chunk = (k // gran.k) * gran.k key = (tensor, group, k_chunk) if key not in active: active[key] = pools[tensor].alloc() tile_map[group] = active[key] slot.mfma.vgpr_tile_maps.setdefault(tensor, []).append(tile_map) for lr in slot.lrs: tensor = lr.tensor is_wrapping = lr.mtIteration != 0 target = next_iter if is_wrapping else active gran = lr_grans[tensor] tile_map = {} seen_keys = set() for t in lr.tiles.tileId_list: group = (t // gran.mn) * gran.mn for lk in lr.tiles.subIterK_list: k_chunk = (lk // gran.k) * gran.k key = (tensor, group, k_chunk) if key in seen_keys: continue seen_keys.add(key) if key in target: pools[tensor].release(target[key]) vid = pools[tensor].alloc() target[key] = vid tile_map[group] = vid lr.vgpr_tile_map.append(tile_map) to_release = [key for key, lr_pos in last_read.items() if lr_pos == pos and key in active] for key in to_release: pools[key[0]].release(active[key]) del active[key] for t in self.tensors: max_peaks[t] = max(max_peaks[t], pools[t].peak) converged = False for prev_ni in all_next_iters: if next_iter == prev_ni: for pi2, slots2 in enumerate(self._partitions): for slot2 in slots2: if slot2.mfma: for tensor in self.tensors: if tensor in slot2.mfma.vgpr_tile_maps: slot2.mfma.vgpr_tile_maps[tensor].pop() for lr2 in slot2.lrs: lr2.vgpr_tile_map.pop() converged = True break if converged: break all_next_iters.append(next_iter) carry_active = next_iter else: assert False, (f"assign_vgpr_tiles did not converge after " f"{MAX_UNROLL} unroll iterations") self.unroll_factor = unroll_iter self.needs_unrolling = self.unroll_factor > 1 self.tile_peaks = max_peaks self.compact_b_overlay = True def _assign_vgpr_tiles_deterministic(self): """Deterministic double-buffer allocator with partition-aware positions.""" cfg = self.config numK = cfg.numSubIterK numP = cfg.numPartitions lr_grans = {'A': cfg.lrA, 'B': cfg.lrB} if cfg.hasScale: lr_grans['SA'] = cfg.lrSA lr_grans['SB'] = cfg.lrSB part_ranges = [self._partition_tile_range(pi) for pi in range(numP)] any_single_k_chunk = any( numK // lr_grans[t].k == 1 for t in self.tensors) use_global_pos = (numP > 1) and any_single_k_chunk group_to_pos = {t: {} for t in self.tensors} max_groups = {t: 0 for t in self.tensors} if use_global_pos: for pi in range(numP): for tensor in self.tensors: side = TENSOR_SIDE[tensor] start, end = part_ranges[pi][side] gran = lr_grans[tensor] groups = sorted(set( (t // gran.mn) * gran.mn for t in range(start, end))) for g in groups: if g not in group_to_pos[tensor]: group_to_pos[tensor][g] = max_groups[tensor] max_groups[tensor] += 1 else: for pi in range(numP): for tensor in self.tensors: side = TENSOR_SIDE[tensor] start, end = part_ranges[pi][side] gran = lr_grans[tensor] groups = sorted(set( (t // gran.mn) * gran.mn for t in range(start, end))) local_pos = 0 for g in groups: if g not in group_to_pos[tensor]: group_to_pos[tensor][g] = local_pos local_pos += 1 max_groups[tensor] = max(max_groups[tensor], local_pos) num_k_groups = {} for tensor in self.tensors: num_k_groups[tensor] = numK // lr_grans[tensor].k unroll_factor = 1 for tensor in self.tensors: if num_k_groups[tensor] % 2 != 0: unroll_factor = 2 break pgr0 = cfg.pgr == 0 if pgr0: unroll_factor = 1 def _tile_set_idx(tensor, unroll_iter, k_chunk, gran): if pgr0: return 0 nkg = num_k_groups[tensor] return (unroll_iter * nkg + k_chunk // gran.k) % 2 def _tile_id(tensor, set_idx, pos): return set_idx * max_groups[tensor] + pos for unroll_iter in range(unroll_factor): for pi, slots in enumerate(self._partitions): for slot in slots: k = slot.subIterK if slot.mfma: for tensor in self.tensors: gran = lr_grans[tensor] set_idx = _tile_set_idx(tensor, unroll_iter, k, gran) side = TENSOR_SIDE[tensor] tileRange = (slot.mfma.tileA if side == 'A' else slot.mfma.tileB) tile_map = {} for t in tileRange.tileId_list: group = (t // gran.mn) * gran.mn pos = group_to_pos[tensor][group] tile_map[group] = _tile_id(tensor, set_idx, pos) slot.mfma.vgpr_tile_maps.setdefault( tensor, []).append(tile_map) for lr in slot.lrs: tensor = lr.tensor gran = lr_grans[tensor] target_mt = unroll_iter + lr.mtIteration target_k = lr.tiles.subIterK_start set_idx = _tile_set_idx(tensor, target_mt, target_k, gran) tile_map = {} for t in lr.tiles.tileId_list: group = (t // gran.mn) * gran.mn if group in tile_map: continue pos = group_to_pos[tensor][group] tile_map[group] = _tile_id(tensor, set_idx, pos) lr.vgpr_tile_map.append(tile_map) num_sets = 1 if pgr0 else 2 self.tile_peaks = {t: num_sets * max_groups[t] for t in self.tensors} self.compact_b_overlay = False self.unroll_factor = unroll_factor self.needs_unrolling = unroll_factor > 1 # ── Place GRs ───────────────────────────────────────── def _build_gr_list(self, part_ranges, offsetMT, offsetPartition, unrollId=0): """Phase 1: Build ordered GR list from placed MFMAs. For each partition × subIterK, derive target partition/MT from the MFMA and offsets. Add GRs (A, B, SA, SB) with tile and K ranges snapped to GR granularity. Dedup within same MT level, then remove n+1 entries that also appear at n+2 (cross-MT dedup). For each subIterK, we apply offsetMT on MT and offsetPartition on partition. Returns list of (tensor, mt_str, tile_start, tile_end, k_start, k_end, gr_gran). """ cfg = self.config numP = cfg.numPartitions seen = set() gr_list = [] for pi in range(numP): partition_slots = self._partitions[pi] target_pi = (pi + offsetPartition) % numP wraps = (pi + offsetPartition) >= numP mt_val = offsetMT + (1 if wraps else 0) target_range = part_ranges[target_pi] for slot in partition_slots: k = slot.mfma.subIterK items = [('A', target_range['A'], cfg.grA), ('B', target_range['B'], cfg.grB)] if cfg.hasScale: items.append(('SA', target_range['A'], cfg.grSA)) items.append(('SB', target_range['B'], cfg.grSB)) for tensor, (t_start, t_end), gr_gran in items: def _valid_k_for_unroll(k, uid): nu = cfg.numUnroll.get(tensor, 1) if uid >= nu: return False if nu > 1: return k < gr_gran.k return k + uid * gr_gran.k < cfg.numSubIterK if not _valid_k_for_unroll(k, unrollId): continue k_shifted = k + unrollId * gr_gran.k tr = gr_gran.tile_range(k_shifted, t_start, t_end) key = (tensor, mt_val, tr.tileId_start, tr.tileId_end, tr.subIterK_start, tr.subIterK_end) if key in seen: continue seen.add(key) gr_list.append((tensor, mt_val, tr.tileId_start, tr.tileId_end, tr.subIterK_start, tr.subIterK_end, gr_gran)) # Cross-MT dedup: if a tile/k range appears at both n+1 and n+2, # the n+1 load is redundant — the previous iteration's n+2 already # wrote the same data into LDS. Remove the n+1 duplicate. base_mt = offsetMT n2_keys = {(t, ts, te, ks, ke) for t, mt, ts, te, ks, ke, _ in gr_list if mt != base_mt} gr_list = [entry for entry in gr_list if entry[1] != base_mt or (entry[0], entry[2], entry[3], entry[4], entry[5]) not in n2_keys] return gr_list def _build_gr_slot_bounds(self): """Build lower and upper slot bounds for GR placement. lower: tensor -> [(flat, t_start, t_end, k_start, k_end)] for LR(mt=0). GR(mt=2) writes the same LDS buffer as MT n, so it can't be placed at/before a flat slot where a later LR(mt=0) in *any* partition reads an overlapping tile/k-range (LDS conflict). Collisions span partitions: an LR(n) read for tensor B may sit in partition pi+1 while the GR(n+2) overwrite is emitted in partition pi, so the check must be in flat execution order. upper: (tensor, mt) -> first flat slot index with LR(tensor, mt). GR(tensor, mt) must be placed strictly before this slot. """ numK = self.config.numSubIterK lower = {} upper = {} for pi, partition_slots in enumerate(self._partitions): for slot in partition_slots: flat = pi * numK + slot.subIterK for lr in slot.lrs: if lr.mtIteration == 0: lower.setdefault(lr.tensor, []).append( (flat, lr.tiles.tileId_start, lr.tiles.tileId_end, lr.tiles.subIterK_start, lr.tiles.subIterK_end)) key = (lr.tensor, lr.mtIteration) if key not in upper or flat < upper[key]: upper[key] = flat return lower, upper @staticmethod def _has_lr_conflict(lr_lower, tensor, mt_val, flat, gr_t_start, gr_t_end, gr_k_start, gr_k_end): """Return True if placing GR(mt_val) at flat slot conflicts. GR(MT n+2) writes the same LDS buffer as MT n, so it conflicts if a later LR(MT n) — in flat execution order across all partitions — still reads an overlapping tile/subIterK range from that buffer. """ if mt_val != 2: return False for lr_flat, lr_ts, lr_te, lr_ks, lr_ke in lr_lower.get(tensor, []): if (lr_flat > flat and gr_t_start < lr_te and lr_ts < gr_t_end and gr_k_start < lr_ke and lr_ks < gr_k_end): return True return False def _distribute_grs(self, gr_list, gr_slot_bounds, unrollId=0, slot_start=0, slot_end=None, allowed_slots=None): """Phase 2: Distribute GR atoms across partition × subIterK slots. Explodes GR entries into atomic loads, distributes them into flat buckets respecting LDS conflict constraints and load balance, then remerges consecutive atoms and places them into partitions. slot_start/slot_end restrict distribution to [slot_start, slot_end) in the flat partition×subIterK index space. Used by place_GRs to confine each unrollId to its own slot range in multi-DU configs. allowed_slots, when set, restricts placement to an explicit list of flat slot indices (e.g. all subIterK=1 slots for scale prefetch). A single allowed slot still runs the normal weighted distribution across the full [slot_start, slot_end) range, remerges per virtual slot, then places all resulting GRs into that one slot (multi-DU uid>0 under PGR=1). """ cfg = self.config numK = cfg.numSubIterK numP = cfg.numPartitions numSlots = numP * numK if slot_end is None: slot_end = numSlots collapse_slot = None if allowed_slots is None: slots = list(range(slot_start, slot_end)) else: slots = sorted(allowed_slots) if len(slots) == 1: collapse_slot = slots[0] if not slots: return dist_slots = (list(range(slot_start, slot_end)) if collapse_slot is not None else slots) lower, upper = gr_slot_bounds # 2a. Explode GR entries into atomic loads (1 load each) atoms = [] for tensor, mt_val, t_start, t_end, k_start, k_end, gr_gran in gr_list: mn = gr_gran.mn last = max(dist_slots[0], min(upper.get((tensor, mt_val), numSlots) - 1, dist_slots[-1])) for pos in range(t_start, t_end, mn): atoms.append((tensor, mt_val, pos, pos + mn, k_start, k_end, last)) # 2b. Place atoms across dist_slots weighted by partition MFMA # count. Each partition's slots get a share proportional to its # MFMAs, so larger partitions receive more GR loads. nAtoms = len(atoms) buckets = [[] for _ in range(numSlots)] mfma_per_partition = [] for pi in range(numP): piM = pi % cfg.numPartitionsM piN = pi // cfg.numPartitionsM mfma_per_partition.append(cfg.partitionSizesM[piM] * cfg.partitionSizesN[piN]) weight_prefix = [0] for s in dist_slots: weight_prefix.append(weight_prefix[-1] + mfma_per_partition[s // numK]) total_weight = weight_prefix[-1] slot_boundaries = [p * nAtoms for p in weight_prefix[1:]] for i, (tensor, mt_val, ts, te, ks, ke, last) in enumerate(atoms): if cfg.grPlacement == GRPlacementStrategy.BUNCHED: # TDM: pin every GR atom to partition 0, subIterK 0. rel = 0 else: rel = min(bisect_left(slot_boundaries, i * total_weight + 1), len(dist_slots) - 1) if nAtoms else 0 slot = dist_slots[rel] slot = min(slot, last) slot_idx = rel while (slot < last and self._has_lr_conflict(lower, tensor, mt_val, slot, ts, te, ks, ke)): slot_idx += 1 if slot_idx >= len(dist_slots): break slot = dist_slots[slot_idx] if slot_idx >= len(dist_slots): slot = min(last, dist_slots[-1]) dest = collapse_slot if collapse_slot is not None else slot entry = ((slot, tensor, mt_val, ts, te, ks, ke) if collapse_slot is not None else (tensor, mt_val, ts, te, ks, ke)) buckets[dest].append(entry) # 2c. Remerge consecutive atoms and place into partitions for flat, bucket in enumerate(buckets): if not bucket: continue pi = flat // numK si = flat % numK target_slot = self._partitions[pi][si] uid = unrollId gr_list_for_uid = target_slot.grs_by_unroll.setdefault(uid, []) atom_groups = [bucket] if collapse_slot is not None: by_virtual = {} for virtual_slot, tensor, mt_val, ts, te, ks, ke in bucket: by_virtual.setdefault(virtual_slot, []).append( (tensor, mt_val, ts, te, ks, ke)) atom_groups = [by_virtual[v] for v in sorted(by_virtual)] for group_idx, group in enumerate(atom_groups): for atom_idx, atom in enumerate(group): tensor, mt_val, ts, te, ks, ke = atom can_merge = (gr_list_for_uid and not (collapse_slot is not None and group_idx > 0 and atom_idx == 0)) if can_merge: prev = gr_list_for_uid[-1] if (prev.tensor == tensor and prev.mtIteration == mt_val and prev.tiles.subIterK_start == ks and prev.tiles.subIterK_end == ke and prev.tiles.tileId_end == ts): prev.tiles = MFMATileRange(ks, ke, prev.tiles.tileId_start, te) continue gr_list_for_uid.append(GRPlacement( tensor=tensor, mtIteration=mt_val, tiles=MFMATileRange(ks, ke, ts, te), subIterK_slot=si, partition=pi, unrollId=uid)) @pipeline_pass(Pass.GR) def place_GRs(self) -> PartitionSchedule: """Place Global Reads by iterating MFMAs across partitions. Phase 1: Build ordered GR list from partition traversal respecting gr granularities. Phase 2: Distribute evenly GR atoms across all (partition, subIterK) slots. GR atoms being the smallest load granularity for a specific tensor. This should give a sheduling respecting the following rules: - GR are in the order we expect them from the LR pov - we respect the GR granularities (can change the above rule a bit) - Overall loads are spread accross all subIterKs of all partitions. """ part_ranges = [self._partition_tile_range(pi) for pi in range(self.config.numPartitions)] pgr = self.config.pgr offsetMT = 0 if pgr == 0 else 1 maxUnroll = max(self.config.numUnroll.values()) gr_slot_bounds = self._build_gr_slot_bounds() numK = self.config.numSubIterK last_uid_slot = None if maxUnroll > 1 and pgr == 1: # Multi-DU shares one SRD: all uid=0 GRs must finish (and # GRInc(uid=0) fire) before any uid>0 GR. Place uid>0 only in # the last subIterK slot of the last partition so _gr_sort_key and # group_lr_gr chain them after GRInc(uid=0); that slot has an MFMA # so the instruction scheduler can interleave uid>0 GRs with compute. last_pi = len(self._partitions) - 1 last_si = len(self._partitions[last_pi]) - 1 last_uid_slot = last_pi * numK + last_si for uid in range(maxUnroll): gr_list = self._build_gr_list(part_ranges, offsetMT, self.config.offsetPartition, unrollId=uid) if uid > 0 and last_uid_slot is not None: self._distribute_grs(gr_list, gr_slot_bounds, unrollId=uid, allowed_slots=[last_uid_slot]) else: self._distribute_grs(gr_list, gr_slot_bounds, unrollId=uid) return self._partitions[0] # ── Annotate dependencies ───────────────────────────── @pipeline_pass(Pass.DEPS) def annotate_deps(self): """Annotate each placement with its raw before-dependencies. Populates the `before` field on MFMAPlacement, LRPlacement, and GRPlacement objects in self._partitions. Each lr_ref/gr_ref BaseOp is resolved to point at the specific placement it depends on. Iterates all partitions. Two-pass per partition: - Pass 1: build lookups from existing placements - Pass 2: populate .before on each placement Rules: - MFMA(subIterK=k) depends on all LRs that loaded subIterK=k data (cross-partition: LRs for a tensor may be in any partition) - LR depends on GR for same tensor (data must be in LDS) - GR depends on collision LR for same tensor (LDS double-buffer) """ cfg = self.config numK = cfg.numSubIterK # Build global lr_by_data across all partitions (MFMA deps are cross-partition) # lr_by_data[data_k][tensor] → list of LRPlacements loading subIterK=data_k lr_by_data = [{} for _ in range(numK)] # gr_by_tensor[tensor] → list of all GRPlacements (LR→GR deps are cross-partition) gr_by_tensor = {} # lr_by_tensor[tensor] → list of all LRPlacements (GR→LR collision is cross-partition) lr_by_tensor = {} for slots in self._partitions: for slot in slots: for lr in slot.lrs: for data_k in lr.tiles.subIterK_list: lr_by_data[data_k].setdefault(lr.tensor, []).append(lr) lr_by_tensor.setdefault(lr.tensor, []).append(lr) for gr in slot.grs: gr_by_tensor.setdefault(gr.tensor, []).append(gr) for pi, slots in enumerate(self._partitions): self._annotate_deps_partition(pi, slots, cfg, lr_by_data, gr_by_tensor, lr_by_tensor) def _annotate_deps_partition(self, pi: int, slots: List[SubIterKSlot], cfg: SchedulerConfig, lr_by_data: list, gr_by_tensor: dict, lr_by_tensor: dict): """Annotate deps for a single partition (in-place on placements).""" numK = len(slots) # Clear any previous annotations (idempotent re-runs) for slot in slots: if slot.mfma: slot.mfma.deps.clear() for lr in slot.lrs: lr.deps.clear() for gr in slot.grs: gr.deps.clear() # ── Pass 1: build per-partition lookups ── # lr_by_slot[k][tensor] → LRPlacement at subIterK=k # gr_by_slot[k][tensor] → GRPlacement at subIterK=k # (lr_by_data, gr_by_tensor, lr_by_tensor are built globally in annotate_deps) lr_by_slot = [{} for _ in range(numK)] gr_by_slot = [{} for _ in range(numK)] for k, slot in enumerate(slots): for lr in slot.lrs: lr_by_slot[k][lr.tensor] = lr for gr in slot.grs: gr_by_slot[k][gr.tensor] = gr # ── Pass 2: populate deps on each placement ── # mt_offset: 0 = same MT, -1 = prev MT, -2 = two MTs back, etc. # Within one iteration, execution order per slot is MFMA → LR → GR, # and slots run in order 0, 1, 2, ... _order = {'MFMA': 0, 'LR': 1, 'GR': 2} def _slot_offset(consumer_partition, consumer_slot, consumer_type, producer): """Offset from partition+slot ordering: 0 if producer ran first, -1 otherwise.""" prod_partition = producer.partition if prod_partition < consumer_partition: return 0 if prod_partition > consumer_partition: return -1 prod_slot = producer.subIterK_slot if prod_slot < consumer_slot: return 0 if prod_slot > consumer_slot: return -1 prod_type = 'LR' if isinstance(producer, LRPlacement) else 'GR' return -1 if _order[prod_type] >= _order[consumer_type] else 0 def _mt_offset(consumer_partition, consumer_slot, consumer_type, producer, consumer=None): # MFMA→LR: MFMA always consumes mt=0 (current). if consumer_type == 'MFMA' and isinstance(producer, LRPlacement): return -producer.mtIteration # LR→GR: mt difference determines how many iterations back. if consumer_type == 'LR' and isinstance(producer, GRPlacement) and consumer: return consumer.mtIteration - producer.mtIteration # Fallback: partition+slot ordering decides. return _slot_offset(consumer_partition, consumer_slot, consumer_type, producer) def _tiles_overlap(mfma, lr_tensor, lr_tiles): """Check if LR tile range overlaps with MFMA's tile range for that tensor.""" if lr_tensor in ('A', 'SA'): mfma_range = mfma.tileA else: mfma_range = mfma.tileB return (lr_tiles.tileId_start < mfma_range.tileId_end and lr_tiles.tileId_end > mfma_range.tileId_start and lr_tiles.subIterK_start < mfma_range.subIterK_end and lr_tiles.subIterK_end > mfma_range.subIterK_start) def _range_overlaps(a: MFMATileRange, b: MFMATileRange) -> bool: """Check if two tile ranges overlap on both tile ids and subIterK.""" return (a.tileId_start < b.tileId_end and a.tileId_end > b.tileId_start and a.subIterK_start < b.subIterK_end and a.subIterK_end > b.subIterK_start) def _dedup_deps(deps): if len(deps) <= 1: return deps def _exec_order(dep): return (dep.mt_offset, dep.ref.partition, dep.ref.subIterK_slot) by_uid = {} for dep in deps: uid = getattr(dep.ref, 'unrollId', 0) if uid not in by_uid or _exec_order(dep) > _exec_order(by_uid[uid]): by_uid[uid] = dep return list(by_uid.values()) for k, slot in enumerate(slots): # MFMA: depends on the most recent LR per tensor (tile-overlapping). # Uses lr_by_tensor (all LRs across partitions) so that a more recent # LR loading a different subIterK still subsumes older data deps. if slot.mfma: for t in self.tensors: deps_for_t = [] for lr in lr_by_tensor.get(t, []): if _tiles_overlap(slot.mfma, t, lr.tiles): deps_for_t.append(Dep( ref=lr, mt_offset=_mt_offset(pi, k, 'MFMA', lr))) slot.mfma.deps.extend(_dedup_deps(deps_for_t)) # LR: depends on GR (data must be in LDS before reading) # Cross-partition: the GR that loaded the matching tiles may be # in a different partition. Filter by tile overlap. for lr in slot.lrs: for gr in gr_by_tensor.get(lr.tensor, []): if _range_overlaps(lr.tiles, gr.tiles): lr.deps.append(Dep( ref=gr, mt_offset=_mt_offset(pi, k, 'LR', gr, consumer=lr))) # GR: depends on collision LR (LDS double-buffer) # GR(n+x) collides with LR(n+x-2) — same buffer, period 2. for gr in slot.grs: target_data = gr.mtIteration - 2 for lr in lr_by_tensor.get(gr.tensor, []): if _range_overlaps(lr.tiles, gr.tiles): mt_off = target_data - lr.mtIteration gr.deps.append(Dep(ref=lr, mt_offset=mt_off)) if not gr.deps: raise ValueError( f"GR {gr.tensor} mt={fmt_mt(gr.mtIteration)} at slot {k} " f"has no overlapping LR(n) dependency") for slot in slots: for lr in slot.lrs: lr.deps = _dedup_deps(lr.deps) for gr in slot.grs: gr.deps = _dedup_deps(gr.deps) # ── Remove unnecessary GR deps ──────────────────────── def _make_gr_dep_exec_order(self, tensor): """Return a key fn ordering deps by (mt_offset, partition, slot, intra-slot rank). Two GRs sharing a (mtIteration, partition, subIterK_slot) collapse to one (mt_offset, partition, slot) key, so an intra-slot rank (by _gr_sort_key, the order the wait emitter walks) is needed to keep them distinguishable: dropping a dep on the slot's last GR is unsafe even if a dep on an earlier-rank GR in the same slot is kept. """ slot_members = {} for slots in self._partitions: for slot in slots: for gr in slot.grs: if gr.tensor == tensor: key = (gr.mtIteration, gr.partition, gr.subIterK_slot, gr.unrollId) slot_members.setdefault(key, []).append(gr) gr_intra_rank = {} for grs in slot_members.values(): for rank, gr in enumerate(sorted(grs, key=self._gr_sort_key)): gr_intra_rank[id(gr)] = rank def _dep_exec_order(dep): gr = dep.ref return (dep.mt_offset, gr.partition, gr.subIterK_slot, gr_intra_rank[id(gr)]) def _gr_schedule_order(gr: GRPlacement) -> tuple: return (gr.partition, gr.subIterK_slot, gr_intra_rank[id(gr)]) return _dep_exec_order, _gr_schedule_order @staticmethod def _lr_schedule_order(lr: LRPlacement) -> tuple: return (lr.partition, lr.subIterK_slot) @pipeline_pass(Pass.REMOVE_GR_DEPS) def remove_unnecessary_gr_deps(self): """Remove GR deps on LRs that are already guaranteed by an earlier LR's wait. Per tensor, walks LR placements in execution order. If an earlier LR already waits for a GR with equal or higher exec_order, the later LR's dep is redundant and removed — but only when that earlier wait's GR completes before the current LR runs (same-partition / earlier-slot). Cross-partition deps on later partitions must not subsume earlier waits. Wraps around: the first LR's dep is compared against the last from the previous MT iteration (max dep exec_order shifted by mt_offset -1). """ for tensor in self.tensors: _dep_exec_order, _gr_schedule_order = self._make_gr_dep_exec_order(tensor) lr_with_gr_deps = [] for pi, slots in enumerate(self._partitions): for slot in slots: for lr in slot.lrs: if lr.tensor == tensor and lr.deps: dep = lr.deps[0] if isinstance(dep.ref, GRPlacement): lr_with_gr_deps.append((lr, dep)) if len(lr_with_gr_deps) <= 1: continue max_eo = max(_dep_exec_order(dep) for _, dep in lr_with_gr_deps) max_guaranteed = (max_eo[0] - 1, max_eo[1], max_eo[2], 0) if not self._is_multi_du(): for lr, dep in lr_with_gr_deps: eo = _dep_exec_order(dep) if eo <= max_guaranteed: lr.deps.clear() else: max_guaranteed = eo continue _, max_dep = max(lr_with_gr_deps, key=lambda item: _dep_exec_order(item[1])) guarantee_gr_order = _gr_schedule_order(max_dep.ref) use_wraparound = True for lr, dep in lr_with_gr_deps: eo = _dep_exec_order(dep) lr_order = self._lr_schedule_order(lr) if eo <= max_guaranteed and (use_wraparound or guarantee_gr_order < lr_order): lr.deps.clear() else: if eo > max_guaranteed: max_guaranteed = eo guarantee_gr_order = _gr_schedule_order(dep.ref) use_wraparound = False # ── Remove unnecessary LR deps ──────────────────────── @pipeline_pass(Pass.REMOVE_LR_DEPS) def remove_unnecessary_lr_deps(self): """Remove GR→LR collision deps already covered by an earlier sync. A slot is a sync point when it contains any GR-with-LR-dep or any LR-with-GR-deps. At a sync slot, the per-tensor last LR guaranteed is the max exec_order of: - the MFMA's LR deps at that slot - any GR-with-LR-dep's own LR dep at that slot For each LR dep on a GR (curLR), find the previous sync (in exec order, skipping the current slot) providing a last LR for the same tensor. If that last LR's exec_order >= curLR's, the dep is redundant. Exec order: (mt_offset, partition, subIterK_slot). On wrap-around the mt_offset is shifted by -1. """ def _dep_exec_order(dep): return (dep.mt_offset, dep.ref.partition, dep.ref.subIterK_slot) # Step 1: collect one sync entry per sync slot. # Each entry: (pos, last_lr_by_tensor, [grs_to_check]) sync_slots = [] for pi, slots in enumerate(self._partitions): for slot in slots: grs_with_lr = [ gr for gr in slot.grs if gr.deps and isinstance(gr.deps[0].ref, LRPlacement)] lr_with_gr_exists = any( lr.deps and isinstance(lr.deps[0].ref, GRPlacement) for lr in slot.lrs) if not grs_with_lr and not lr_with_gr_exists: continue last_lr = {} if slot.mfma: for d in slot.mfma.deps: if isinstance(d.ref, LRPlacement): t = d.ref.tensor eo = _dep_exec_order(d) if t not in last_lr or eo > last_lr[t]: last_lr[t] = eo for gr in grs_with_lr: dep = gr.deps[0] t = dep.ref.tensor eo = _dep_exec_order(dep) if t not in last_lr or eo > last_lr[t]: last_lr[t] = eo sync_slots.append(((pi, slot.subIterK), last_lr, grs_with_lr)) if not sync_slots: return sync_slots.sort(key=lambda x: x[0]) n = len(sync_slots) # Step 2 & 3: for each GR with LR dep, walk backward (with # wrap-around) to find the previous sync slot providing a last LR # for this tensor. for i in range(n): _, _, grs_to_check = sync_slots[i] for gr in grs_to_check: if not gr.deps: continue dep = gr.deps[0] tensor = dep.ref.tensor cur_eo = _dep_exec_order(dep) prev_eo = None cur_pos = sync_slots[i][0] for j in range(1, n + 1): idx = (i - j) % n wrapped = j > i # crossed iteration boundary # Same-slot in same iteration is concurrent, skip. if not wrapped and sync_slots[idx][0] == cur_pos: continue prev_last_lr = sync_slots[idx][1] if tensor in prev_last_lr: eo = prev_last_lr[tensor] if wrapped: eo = (eo[0] - 1, eo[1], eo[2]) prev_eo = eo break if prev_eo is not None and prev_eo >= cur_eo: gr.deps.clear() # ── Remove cross-subIterK deps ───────────────────────── def _gr_granularity(self, tensor: str) -> ReadGranularity: """Return GR granularity for a tensor.""" return {'A': self.config.grA, 'B': self.config.grB, 'SA': self.config.grSA, 'SB': self.config.grSB}[tensor] def _count_gr_atoms(self, gr: GRPlacement) -> int: """Count the number of atomic loads for a single GR placement.""" gr_gran = self._gr_granularity(gr.tensor) tiles = gr.tiles n_tile = (tiles.tileId_end - tiles.tileId_start) // gr_gran.mn n_k = (tiles.subIterK_end - tiles.subIterK_start) // gr_gran.k return n_tile * n_k @staticmethod def _gr_precedes_lr(gr: GRPlacement, lr: LRPlacement) -> bool: """True if GR is issued before LR in global (partition, subIterK) order.""" return (gr.partition, gr.subIterK_slot) < (lr.partition, lr.subIterK_slot) def _find_preceding_gr_dep(self, lr: LRPlacement, tensor: str, unrollId: int) -> Optional[Dep]: """Nearest preceding GR for tensor/unrollId before lr in schedule order.""" numK = len(self._partitions[0]) lr_flat = lr.partition * numK + lr.subIterK_slot for pos in range(lr_flat - 1, -1, -1): pi = pos // numK slot_k = pos % numK slot = self._partitions[pi][slot_k] candidates = [ g for g in slot.grs if g.tensor == tensor and g.unrollId == unrollId] if candidates: gr = max(candidates, key=self._gr_sort_key) return Dep(ref=gr, mt_offset=0) return None def _anchor_gr_in_slot(self, slot, tensor: str, unrollId: int) -> Optional[GRPlacement]: """Pick a GR in slot for re-anchoring (same tensor/unrollId preferred).""" pool = [g for g in slot.grs if g.tensor == tensor] if pool: same_uid = [g for g in pool if g.unrollId == unrollId] use = same_uid if same_uid else pool elif slot.grs: use = slot.grs else: return None return max(use, key=self._gr_sort_key) def _wait_gr_dep_for_lr(self, lr: LRPlacement, dep: Dep) -> Dep: """Re-anchor LR→GR deps for wait_gr inflight counting. Tile overlap can pick a GR at or after the consumer LR in schedule order. The backward walk then wraps a full iteration and over-counts inflight loads, producing a too-weak vmcnt. subIterK==1: anchor on this partition's slot-0 GR (one-slot walk). subIterK==0: when the tile dep is still in the future, anchor on the immediately preceding schedule slot (typically previous partition sk=1, or sk=1 on the prior body-iter ring when pi=0 sk=0), not slot-0 GRs from the same subIterK (which leaves the producer tail in the inflight window). """ numP = len(self._partitions) numK = len(self._partitions[0]) gr = dep.ref if not isinstance(gr, GRPlacement): return dep if lr.subIterK_slot == 1: if numP <= 1: return dep slot0 = self._partitions[lr.partition][0] if gr.partition == lr.partition and gr.subIterK_slot == 0: anchor_gr = gr else: anchor_gr = self._anchor_gr_in_slot(slot0, gr.tensor, gr.unrollId) if anchor_gr is not None: return Dep(ref=anchor_gr, mt_offset=0) if numP <= 1 and numK > 1 and lr.subIterK_slot == 0: if not self._gr_precedes_lr(gr, lr): anchor_gr = self._anchor_gr_in_slot( self._partitions[0][numK - 1], gr.tensor, gr.unrollId) if anchor_gr is not None: return Dep(ref=anchor_gr, mt_offset=0) return dep if numP <= 1: return dep if self._gr_precedes_lr(gr, lr): return dep anchor = self._find_preceding_gr_dep(lr, gr.tensor, gr.unrollId) if anchor is not None: return anchor lr_flat = lr.partition * numK + lr.subIterK_slot prev_flat = lr_flat - 1 if lr_flat > 0 else (numP * numK - 1) prev_pi = prev_flat // numK prev_k = prev_flat % numK anchor_gr = self._anchor_gr_in_slot( self._partitions[prev_pi][prev_k], gr.tensor, gr.unrollId) if anchor_gr is not None: return Dep(ref=anchor_gr, mt_offset=0) slot0 = self._partitions[lr.partition][0] anchor_gr = self._anchor_gr_in_slot(slot0, gr.tensor, gr.unrollId) if anchor_gr is not None: return Dep(ref=anchor_gr, mt_offset=0) return dep def _compute_inflight_loads_legacy(self, consumer_pi: int, consumer_slot: int, tensor: str, dep_ref: Dep) -> WaitGRCounts: """Develop single-DU inflight walk (unchanged formula).""" numP = len(self._partitions) numK = len(self._partitions[0]) flat_len = numP * numK consumer_flat = consumer_pi * numK + consumer_slot wraps_needed = abs(dep_ref.mt_offset) dep_flat = None for p_idx, pslots in enumerate(self._partitions): for k_idx, slot in enumerate(pslots): if any(gr is dep_ref.ref for gr in slot.grs): dep_flat = p_idx * numK + k_idx break if dep_flat is not None: break if dep_flat is None: return WaitGRCounts() total_steps = wraps_needed * flat_len + consumer_flat - dep_flat if total_steps <= 0: assert wraps_needed == 0, ( f"_compute_inflight_loads: total_steps={total_steps} < 1 " f"(wraps_needed={wraps_needed}, consumer_flat={consumer_flat}, dep_flat={dep_flat}); " "unexpected negative total_steps with wraps_needed >= 1" ) return WaitGRCounts() counts = WaitGRCounts() pos = consumer_flat for step in range(total_steps): pos = (pos - 1) % flat_len pi = pos // numK slot_k = pos % numK slot = self._partitions[pi][slot_k] is_final = (step == total_steps - 1) sorted_grs = sorted(slot.grs, key=self._gr_sort_key, reverse=True) for gr in sorted_grs: if is_final and gr.tensor == tensor and gr is dep_ref.ref: return counts atoms = self._count_gr_atoms(gr) cur = getattr(counts, gr.tensor) setattr(counts, gr.tensor, cur + atoms) return counts def _compute_inflight_loads(self, consumer_pi: int, consumer_slot: int, tensor: str, dep_ref: Dep) -> WaitGRCounts: """Count inflight GR atomic loads between a dep GR and the consumer. Walks backward through the flattened schedule (all partitions x subIterK) from the consumer position, counting atomic GR loads for all tensors. Stops when reaching the dependency GR (dep_ref.ref) after accounting for mt_offset wraps. Within a slot, GRs are walked in reverse _gr_sort_key order (matching hardware issue order: last emitted = most recent = walked first). GRs emitted after the dep in the same slot are still inflight but do not need to be waited for, so they are added to the count. Returns per-tensor inflight load counts. """ if not self._is_multi_du(): return self._compute_inflight_loads_legacy( consumer_pi, consumer_slot, tensor, dep_ref) numP = len(self._partitions) numK = len(self._partitions[0]) flat_len = numP * numK consumer_flat = consumer_pi * numK + consumer_slot wraps_needed = abs(dep_ref.mt_offset) # Locate dep_flat: the flat position of the dependency GR in the schedule. dep_flat = None for p_idx, pslots in enumerate(self._partitions): for k_idx, slot in enumerate(pslots): if any(gr is dep_ref.ref for gr in slot.grs): dep_flat = p_idx * numK + k_idx break if dep_flat is not None: break if dep_flat is None: return WaitGRCounts() # Exact number of backward steps from consumer to dep's slot. # Forward distance from dep (exclusive) to consumer (inclusive) within # one mainloop sweep: # - same sweep (consumer_flat > dep_flat): consumer_flat - dep_flat # - wrapped (consumer_flat <= dep_flat) with mt prefetch offset: # wraps_needed * flat_len + consumer_flat - dep_flat # mt_offset (-1, -2, …) marks MT prefetch distance on the dep GR. When the # consumer LR is later in flattened schedule order (forward_span>0), walk # only those slots for multi-partition configs. Single-partition rings # and true wrap-around (forward_span<=0) keep the full ring term. forward_span = consumer_flat - dep_flat if numP > 1 and wraps_needed > 0 and forward_span <= 0: # Tile dep still names a future slot after re-anchor; do not ring-wrap. return WaitGRCounts() if wraps_needed == 0: total_steps = forward_span elif forward_span <= 0: total_steps = wraps_needed * flat_len + forward_span elif numP == 1: # Single-partition ring: MT-prefetch dep still needs full wrap term. total_steps = wraps_needed * flat_len + forward_span else: # Multi-partition: consumer is later in schedule order than dep slot; # walk only the forward span (do not add wraps_needed*flat_len). total_steps = forward_span if total_steps <= 0: assert wraps_needed == 0, ( f"_compute_inflight_loads: total_steps={total_steps} < 1 " f"(wraps_needed={wraps_needed}, consumer_flat={consumer_flat}, dep_flat={dep_flat}); " "unexpected negative total_steps with wraps_needed >= 1" ) return WaitGRCounts() # sk=0 wrap consume at a body-iter boundary: the tile dep GR shares the # consumer's flat slot (mt_offset ring). The backward walk counts GR # atoms issued after that dep in the slot as permissible inflight ops, # but they are the producer tail for the buffer the wrap-LR reads. if (numK > 1 and consumer_slot == 0 and wraps_needed > 0 and dep_flat == consumer_flat): return WaitGRCounts() # sk=1 wrap-LR on A/B reads the same-partition sk=0 buffer; tail GR # loads there must fully drain before ds_read. if (numK > 1 and consumer_slot == 1 and total_steps == 1 and tensor in ('A', 'B') and dep_flat is not None and consumer_flat == dep_flat + 1 and dep_flat % numK == 0): slot0 = self._partitions[consumer_pi][0] same_tensor_grs = [gr for gr in slot0.grs if gr.tensor == tensor] # Multiple tail GRs for one tensor in sk=0 over-count inflight # loads in the one-step walk; assembly needs a full drain instead. if len(same_tensor_grs) > 1: return WaitGRCounts() if (numP > 1 and consumer_pi == 0): return WaitGRCounts() if (numP == 1 and self.config.numMFMATilesM > 2 and self.config.numMFMATilesN > 2): return WaitGRCounts() counts = WaitGRCounts() pos = consumer_flat for step in range(total_steps): pos = (pos - 1) % flat_len pi = pos // numK slot_k = pos % numK slot = self._partitions[pi][slot_k] # On the final step we are at dep's slot: stop when we reach the dep GR. # GRs emitted after the dep (encountered first in reverse order) are in-flight # and are counted before we hit the dep. is_final = (step == total_steps - 1) # Walk GRs in reverse emission order (most recently issued first) sorted_grs = sorted(slot.grs, key=self._gr_sort_key, reverse=True) for gr in sorted_grs: if is_final and (total_steps == 1 and consumer_slot == 1 and numP > 1): # sk=1 wait_gr: inflight = same-tensor A/B GR loads in the # same-partition sk=0 slot (matches assembly buffer_load window). # Other-tensor GRs in that slot are unrelated to this wrap-LR. # SA/SB must drain before scale wrap-LR ds_reads. if gr.tensor == tensor and gr.tensor in ('A', 'B'): atoms = self._count_gr_atoms(gr) cur = getattr(counts, gr.tensor) setattr(counts, gr.tensor, cur + atoms) continue if is_final and gr.tensor == tensor and gr is dep_ref.ref: # Cross-subIterK: include the dep GR's loads when the backward # walk did not already count same-tensor atoms after it (typical # when the dep is the last GR for that tensor in the slot). # Same-slot wrap-around (consumer_flat == dep_flat) keeps the # legacy A=0 + other-tensor inflight behavior. if consumer_flat != dep_flat: cur = getattr(counts, gr.tensor) if cur == 0 and len(self._partitions) > 1: atoms = self._count_gr_atoms(gr) setattr(counts, gr.tensor, cur + atoms) return counts # sk=1 wrap-LR scale reads consume SA/SB from sk=0 tail GRs. if consumer_slot == 1 and gr.tensor in ('SA', 'SB'): continue atoms = self._count_gr_atoms(gr) cur = getattr(counts, gr.tensor) setattr(counts, gr.tensor, cur + atoms) return counts @pipeline_pass(Pass.REMOVE_DEPS) def remove_cross_deps(self): """Replace cross-subIterK deps with wait preOps. For each placement, separates deps into same-subIterK (kept) and cross-subIterK (converted to preOps): - MFMA depending on LRs → single wait_lr - GR depending on LRs → single wait_lr_sync - LR depending on GRs → single wait_gr_sync with per-tensor inflight counts """ for pi, slots in enumerate(self._partitions): for slot in slots: # ── MFMA ── if slot.mfma: same, cross = self._split_deps(slot.mfma.deps, pi, slot.subIterK) slot.mfma.deps = same slot.mfma.preOps = [] has_lr_dep = any( isinstance(d.ref, LRPlacement) for d in same + cross) if has_lr_dep: slot.mfma.preOps.append(WaitLROp()) # ── LRs ── for lr in slot.lrs: gr_deps = [d for d in lr.deps if isinstance(d.ref, GRPlacement)] same, cross = self._split_deps(lr.deps, pi, lr.subIterK_slot) lr.deps = same lr.preOps = [] if gr_deps: dep = gr_deps[0] if self._is_multi_du(): wait_dep = self._wait_gr_dep_for_lr(lr, dep) is_cross = (wait_dep.ref.partition != pi or wait_dep.ref.subIterK_slot != lr.subIterK_slot) counts = self._compute_inflight_loads( pi, lr.subIterK_slot, wait_dep.ref.tensor, wait_dep) # StreamK wrap-LR race fix (multi-DU only): wrap / # cross-iter consumers read an LDS buffer across a # body-iteration / MT-prefetch boundary. Under # StreamK the number of global-read atoms between the # producing buffer_load and the consumer is # grid-dependent, so the static per-tensor inflight # `counts` can be too weak -> stale LDS -> verify # failures. For these wraps the minimal correct # static wait is a full vmcnt(0) drain. Flag the op # force_drain (rather than zeroing `counts`) so the # precise estimate survives for diagnostics and so the # _merge_preops min-count merge cannot silently weaken # the drain back to a partial vmcnt. any_wrap = (lr.mtIteration > 0 or any(d.mt_offset != 0 for d in gr_deps)) force_drain = (self.config.pgr == 1 and any_wrap) lr.preOps.append(WaitGROp(wait_gr_counts=counts, has_sync=True, adjustVmcnt=is_cross, force_drain=force_drain)) else: cross_set = set(id(d) for d in cross) is_cross = id(dep) in cross_set counts = self._compute_inflight_loads( pi, lr.subIterK_slot, dep.ref.tensor, dep) lr.preOps.append(WaitGROp(wait_gr_counts=counts, has_sync=True, adjustVmcnt=is_cross)) # ── GRs ── for gr in slot.grs: same, cross = self._split_deps(gr.deps, pi, gr.subIterK_slot) gr.deps = same has_lr_dep = any( isinstance(d.ref, LRPlacement) for d in same + cross) gr.preOps = [WaitLROp(has_sync=True)] if has_lr_dep else [] def _per_uid_k(self, tensor: str) -> int: """Number of subIterK slots per unrollId for a tensor.""" return self.config.numSubIterK // self.config.numUnroll.get(tensor, 1) @pipeline_pass(Pass.GR_INC) def insert_gr_lr_inc(self): """Insert gr_inc/lr_inc preOps at MacroTile iteration transitions. Walks all LR and GR placements in global execution order (partition 0 slots → partition 1 slots → ..., within each slot: LR then GR). Tracks per-tensor the last-seen mtIteration. When a tensor's mtIteration changes, inserts a BaseOp into that placement's preOps: - lr_inc for LR placements - gr_inc for GR placements For multi-DU configs, also inserts lr_inc at uid boundaries so that each GRIncOp (which swaps the GR write buffer) has a matching LRIncOp (which swaps the LR read buffer), keeping the ping-pong in sync. """ last_lr_mt = {} # tensor -> mtIteration for LR only last_lr_uid = {} # tensor -> effective uid for LR (derived from k range) last_gr_mt = {} # (tensor, unrollId) -> mtIteration for GR only first_lr = {} # tensor -> first LR placement seen last_lr = {} # tensor -> last LR placement seen last_gr = {} # (tensor, unrollId) -> globally last GR placement seen lr_inc_tensors = set() # tensors that already received lr_inc via MT transition multiDU = self._is_multi_du() and self.config.pgr == 1 for pi, slots in enumerate(self._partitions): for slot in slots: for lr in slot.lrs: tensor = lr.tensor mt = lr.mtIteration per_uid_k = self._per_uid_k(tensor) lr_uid = lr.tiles.subIterK_start // per_uid_k if tensor not in first_lr: first_lr[tensor] = lr mt_changed = tensor in last_lr_mt and last_lr_mt[tensor] != mt uid_changed = tensor in last_lr_uid and last_lr_uid[tensor] != lr_uid if mt_changed and uid_changed: lr.preOps.append(LRIncOp(tensor=tensor, isUnrollSwap=True, unrollId=lr_uid)) lr_inc_tensors.add(tensor) elif mt_changed: lr.preOps.append(LRIncOp(tensor=tensor, unrollId=lr_uid)) lr_inc_tensors.add(tensor) elif uid_changed: lr.preOps.append(LRIncOp(tensor=tensor, isUnrollSwap=True, unrollId=lr_uid)) last_lr[tensor] = lr last_lr_mt[tensor] = mt last_lr_uid[tensor] = lr_uid for gr in slot.grs: tensor = gr.tensor uid = gr.unrollId mt = gr.mtIteration key = (tensor, uid) last_gr[key] = gr if not multiDU: if key in last_gr_mt: prev_mt = last_gr_mt[key] else: prev_mt = 0 if prev_mt != mt: if gr.tiles.tileId_start == 0: gr.preOps.append(GRIncOp(tensor=tensor, unrollId=uid)) last_gr_mt[key] = mt if multiDU: for (tensor, uid), gr in last_gr.items(): gr.postOps.append(GRIncOp(tensor=tensor, unrollId=uid)) if self.config.pgr == 0: last_gr_per_key = {} for slots in self._partitions: for slot in slots: for gr in slot.grs: last_gr_per_key[(gr.tensor, gr.unrollId)] = gr for tensor in self._LR_GR_ORDER: if tensor in last_lr and tensor in last_lr_mt: last_lr[tensor].postOps.append( LRIncOp(tensor=tensor, unrollId=last_lr_uid.get(tensor, 0))) for (tensor, uid), gr in last_gr_per_key.items(): if (tensor, uid) in last_gr_mt: gr.postOps.append(GRIncOp(tensor=tensor, unrollId=uid)) else: for tensor, lr in first_lr.items(): if tensor not in lr_inc_tensors: first_uid = lr.tiles.subIterK_start // self._per_uid_k(tensor) lr.preOps.append(LRIncOp(tensor=tensor, unrollId=first_uid)) if self._is_multi_du(): for tensor, lr in first_lr.items(): per_uid_k = self._per_uid_k(tensor) first_uid = lr.tiles.subIterK_start // per_uid_k if first_uid != 0: lr.preOps.insert(0, LRIncOp(tensor=tensor, isUnrollSwap=True, unrollId=first_uid)) # GL2 prefetch: append increment + load ops after the last GR in # the schedule so they fire once per iteration after all GR_INCs. if self.config.pgl > 0: last_gr = None for slots in self._partitions: for slot in slots: for gr in slot.grs: last_gr = gr if last_gr is not None: last_gr.postOps.append(GL2PrefetchIncOp()) last_gr.postOps.append(GL2PrefetchOp()) # ── Group LR/GR chains ───────────────────────────────────── _TENSOR_ORDER = {'A': 0, 'B': 1, 'SA': 2, 'SB': 3} _LR_GR_ORDER = ['A', 'B', 'SA', 'SB'] @staticmethod def _gr_sort_key(gr: GRPlacement) -> tuple: """Sort key for deterministic GR ordering within a subIterK slot. Priority (most significant first): 1. MT iteration — earlier MT loads first (n+1 before n+2) 2. subIterK start — lower min K first 3. Tensor — A, B, SA, SB (hardcoded order) 4. Unroll id — lower uid first 5. Tile id start — lower tile range first """ return (gr.mtIteration, gr.tiles.subIterK_start, LogicalScheduler._TENSOR_ORDER[gr.tensor], gr.unrollId, gr.tiles.tileId_start) @staticmethod def _merge_preops(all_preops: List[List['BaseOp']], lr_tensors: Optional[List[str]] = None) -> List['BaseOp']: """Merge preOps from multiple placements. Combines wait_gr/wait_gr_sync counts into a single BaseOp, deduplicates barrier ops (wait_lr_sync, wait_lr), and collects the rest. wait_gr merge: per tensor, take the strictest (min) count among ops that report a non-zero count for that tensor. Zero on an LR's own tensor means "not waiting on this tensor" and must not erase another LR's non-zero requirement (e.g. LR A wait_gr B=8 must survive LR B B=0). Lower emitted vlcnt is a stricter drain (vmcnt(0) waits for all loads). """ wait_gr_by_lr = [] has_wait_gr_sync = False seen_wait_lr = False others = [] for idx, preops in enumerate(all_preops): lr_tensor = lr_tensors[idx] if lr_tensors else None for op in preops: if isinstance(op, WaitGROp) and op.wait_gr_counts: if op.has_sync: has_wait_gr_sync = True wait_gr_by_lr.append((op, lr_tensor)) elif isinstance(op, WaitLROp): if not seen_wait_lr: seen_wait_lr = True others.append(op) else: others.append(op) result = [] if wait_gr_by_lr: merged_counts = WaitGRCounts() if lr_tensors is None: wait_gr_ops_full = [op for op, _ in wait_gr_by_lr] for t in ('A', 'B', 'SA', 'SB'): setattr(merged_counts, t, min(getattr(op.wait_gr_counts, t) for op in wait_gr_ops_full)) adjust = all(op.adjustVmcnt for op in wait_gr_ops_full) else: for t in ('A', 'B', 'SA', 'SB'): vals = [ getattr(op.wait_gr_counts, t) for op, _ in wait_gr_by_lr if getattr(op.wait_gr_counts, t) > 0] setattr(merged_counts, t, min(vals) if vals else 0) adjust = all(op.adjustVmcnt for op, _ in wait_gr_by_lr) # A full drain is the strictest possible wait, so if any merged LR # requested force_drain the merged op must also drain (otherwise the # min-count merge above would silently weaken it back to a partial # vmcnt and re-introduce the wrap-LR race). merged_force_drain = any( getattr(op, 'force_drain', False) for op, _ in wait_gr_by_lr) result.append(WaitGROp(wait_gr_counts=merged_counts, has_sync=has_wait_gr_sync, adjustVmcnt=adjust, force_drain=merged_force_drain)) result.extend(others) return result @pipeline_pass(Pass.GROUP_LR_GR) def group_lr_gr(self): """Group LR and GR placements into chains within each subIterK. Phase 1 — LR chain: Sort LRs by tensor order (A, B, SA, SB). Build a dep chain so each LR depends on the previous one. Merge all preOps onto the first LR (wait_gr counts are combined, other preOps are collected). Phase 2 — GR chain: Sort GRs by tensor order (A, B, SA, SB). Build a dep chain. If any GR originally had same-subIterK deps, replace the first GR's deps with a single dep on the last LR of the phase-1 chain. Each GR keeps its own preOps; only redundant wait_lr_sync ops are removed (keep the first occurrence only). Phase 3 — Cross-group merge: If any LR has a dep on a GR in the same slot, merge the two chains into one: GR chain → LR chain (first LR points to last GR, LR's original GR dep is removed). This avoids two nodes sharing the same parent. """ order = self._LR_GR_ORDER for pi, slots in enumerate(self._partitions): for slot in slots: # ── Phase 1: LR chain ── ordered_lrs = sorted( slot.lrs, key=lambda lr: order.index(lr.tensor)) if len(ordered_lrs) > 1: # Merge preOps onto first LR lr_tensors = ([lr.tensor for lr in ordered_lrs] if self._is_multi_du() else None) merged = self._merge_preops( [lr.preOps for lr in ordered_lrs], lr_tensors) ordered_lrs[0].preOps = merged for lr in ordered_lrs[1:]: lr.preOps = [] # Build chain: each LR depends on the previous for i in range(1, len(ordered_lrs)): ordered_lrs[i].deps = [ Dep(ref=ordered_lrs[i - 1], mt_offset=0)] last_lr = ordered_lrs[-1] if ordered_lrs else None # ── Phase 2: GR chain ── ordered_grs = sorted( slot.grs, key=self._gr_sort_key) if len(ordered_grs) > 1: # Check if any GR has same-subIterK deps any_deps = any(gr.deps for gr in ordered_grs) # Remove redundant wait_lr_sync (keep only the first) seen_wait_lr_sync = False for gr in ordered_grs: if seen_wait_lr_sync: gr.preOps = [ op for op in gr.preOps if not (isinstance(op, WaitLROp) and op.has_sync)] elif any(isinstance(op, WaitLROp) and op.has_sync for op in gr.preOps): seen_wait_lr_sync = True # First GR: if any GR had deps, point to last LR if any_deps and last_lr is not None: ordered_grs[0].deps = [ Dep(ref=last_lr, mt_offset=0)] else: ordered_grs[0].deps = [] # Build chain: each GR depends on the previous for i in range(1, len(ordered_grs)): ordered_grs[i].deps = [ Dep(ref=ordered_grs[i - 1], mt_offset=0)] elif len(ordered_grs) == 1: # Single GR: still consolidate dep to last LR if it had deps if ordered_grs[0].deps and last_lr is not None: ordered_grs[0].deps = [ Dep(ref=last_lr, mt_offset=0)] # ── Phase 3: Cross-group merge ── # If any LR depends on a GR in this slot, merge into one # chain: GR_group → LR_group to avoid shared parents. if ordered_grs and ordered_lrs: slot_gr_set = set(id(gr) for gr in ordered_grs) lr_has_gr_dep = any( any(id(d.ref) in slot_gr_set for d in lr.deps) for lr in ordered_lrs if lr.deps) if lr_has_gr_dep: last_gr = ordered_grs[-1] # Clear LR deps that point to GRs in this slot for lr in ordered_lrs: lr.deps = [d for d in lr.deps if id(d.ref) not in slot_gr_set] # First LR points to last GR ordered_lrs[0].deps = [ Dep(ref=last_gr, mt_offset=0)] # ── Phase 4: Consolidate MFMA deps ── # After chaining, MFMA only needs the tail of its dep chain. if slot.mfma and last_lr is not None: slot_lr_set = set(id(lr) for lr in ordered_lrs) lr_deps = [d for d in slot.mfma.deps if id(d.ref) in slot_lr_set] if len(lr_deps) > 1: other_deps = [d for d in slot.mfma.deps if id(d.ref) not in slot_lr_set] slot.mfma.deps = other_deps + [ Dep(ref=last_lr, mt_offset=lr_deps[0].mt_offset)] @pipeline_pass(Pass.REMOVE_WAIT_LR_SYNC) def remove_unnecessary_wait_lr_sync(self): """Remove redundant wait_lr_sync from GRs after grouping. Given that we always use wait_lr cnt=0, grouping can guarantee future wait_lr_sync. A GR's wait_lr_sync is unnecessary when: 1. The GR has no same-subIterK deps (deps is empty after grouping) 2. The previous subIterK's GRs already have a wait_lr_sync 3. That previous wait_lr_sync is ordered after all LRs in the previous subIterK (the GR has deps on the LR chain) In that case, all prior LR reads were already synced by the previous subIterK's barrier, and the current GR doesn't conflict with any LRs in its own subIterK, so the second wait_lr_sync is redundant. Finally, any remaining wait_lr_sync on a GR with no deps is downgraded to just sync — the wait_lr is already guaranteed by the MFMA op in the same subIterK. """ for pi, slots in enumerate(self._partitions): for si, slot in enumerate(slots): if not slot.grs: continue first_gr = slot.grs[0] has_wait_lr_sync = any( isinstance(op, WaitLROp) and op.has_sync for op in first_gr.preOps) if not has_wait_lr_sync: continue has_deps = bool(first_gr.deps) if has_deps: continue # Check previous subIterK in the same partition if si == 0: continue prev_slot = slots[si - 1] if not prev_slot.grs: continue prev_first_gr = prev_slot.grs[0] prev_has_wait_lr_sync = any( isinstance(op, WaitLROp) and op.has_sync for op in prev_first_gr.preOps) prev_deps_on_lrs = bool(prev_first_gr.deps) if prev_has_wait_lr_sync and prev_deps_on_lrs: first_gr.preOps = [ op for op in first_gr.preOps if not (isinstance(op, WaitLROp) and op.has_sync)] # Downgrade remaining wait_lr_sync → sync on GRs with no LR deps. # The MFMA in the same subIterK already ensures wait_lr. for pi, slots in enumerate(self._partitions): for slot in slots: for gr in slot.grs: if not any(isinstance(op, WaitLROp) and op.has_sync for op in gr.preOps): continue has_lr_dep = False node = gr while node and node.deps: ref = node.deps[0].ref if isinstance(ref, LRPlacement): has_lr_dep = True break node = ref if has_lr_dep: continue gr.preOps = [ SyncOp() if (isinstance(op, WaitLROp) and op.has_sync) else op for op in gr.preOps] @pipeline_pass(Pass.REMOVE_WAIT_GR_SYNC) def remove_unnecessary_wait_gr_sync(self): """Remove redundant wait_gr_sync on interior partition subIterK=1 slots. In multi-partition configs, inflight GR from partition pi's s1 is drained by partition (pi+1)'s s0 wait_gr before that partition's ds_reads. Baseline kernels only retain wait_gr at s1 on the last partition (wrap to the next MacroTile iteration). """ if not self._is_multi_du() or self.config.numPartitions <= 1: return last_pi = self.config.numPartitions - 1 for pi, slots in enumerate(self._partitions): for si, slot in enumerate(slots): if si == 0 or pi == last_pi: continue for lr in slot.lrs: lr.preOps = [ op for op in lr.preOps if not isinstance(op, WaitGROp)] def _split_deps(self, deps: List[Dep], consumer_pi: int, consumer_slot: int) -> Tuple[List[Dep], List[Dep]]: """Split deps into same-subIterK and cross-subIterK lists. A dep is "same subIterK" if mt_offset == 0 AND the producer is in the same partition and same subIterK slot as the consumer. """ same, cross = [], [] for dep in deps: if (dep.mt_offset == 0 and dep.ref.partition == consumer_pi and dep.ref.subIterK_slot == consumer_slot): same.append(dep) else: cross.append(dep) return same, cross @pipeline_pass(Pass.EMIT) def emit(self) -> EmittedSchedule: """Convert placements into EmittedModule chains per partition per subIterK. Returns [partition][subIterK][EmittedModule]. Each subIterK list contains: - Primary modules (MFMA, LRs, GRs) - Dependency modules (wait_gr, wait_lr, sync, lr_inc, gr_inc) emitted from preOps, chained via before-links The before-link topology: - wait_gr is standalone (no incoming before-link), but later deps chain from it - WaitGROp with has_sync expands to two modules: wait_gr then sync - WaitLROp with has_sync expands to two modules: wait_lr then sync - Same-subIterK Dep deps become ordering constraints (no new module) """ all_partitions = [] for pi, slots in enumerate(self._partitions): partition_emitted = [] for slot in slots: emitted: List[EmittedModule] = [] placement_to_id = {} def add(source: Emittable) -> int: mid = len(emitted) emitted.append(EmittedModule(moduleId=mid, source=source)) return mid def setBefore(moduleId: int, beforeId: int) -> None: if beforeId is None or beforeId == moduleId: return cur = emitted[moduleId].before if cur is None: emitted[moduleId].before = beforeId return assert cur == beforeId, \ f"EmittedModule {moduleId} has multiple before deps: {cur} and {beforeId}" # Step 1: emit primary modules placements = [] if slot.mfma: placements.append(slot.mfma) for lr in slot.lrs: placements.append(lr) for gr in slot.grs: placements.append(gr) placement_tail_id = {} for placement in placements: mid = add(placement) placement_to_id[id(placement)] = mid placement_tail_id[id(placement)] = mid # Step 1b: add postOps and update tail ids so that # deps on a placement with postOps resolve to the last postOp. for placement in placements: if not placement.postOps: continue curId = placement_to_id[id(placement)] postPrevId = curId for postOp in placement.postOps: postId = add(postOp) setBefore(postId, postPrevId) postPrevId = postId placement_tail_id[id(placement)] = postPrevId # Step 2: wire before-chains from preOps + deps for placement in placements: curId = placement_to_id[id(placement)] prevId = None lastDepId = None firstPreOpId = None # preOps for preOp in placement.preOps: if isinstance(preOp, WaitGROp): depId = add(preOp) prevId = depId if firstPreOpId is None: firstPreOpId = depId if preOp.has_sync: depId = add(SyncOp()) setBefore(depId, prevId) prevId = depId lastDepId = depId continue elif isinstance(preOp, WaitLROp) and preOp.has_sync: depId = add(WaitLROp()) setBefore(depId, prevId) prevId = depId lastDepId = depId if firstPreOpId is None: firstPreOpId = depId depId = add(SyncOp()) setBefore(depId, prevId) prevId = depId lastDepId = depId continue else: depId = add(preOp) setBefore(depId, prevId) prevId = depId lastDepId = depId if firstPreOpId is None: firstPreOpId = depId # deps (same-subIterK Deps — ordering constraints) # Wire dep refs as roots of the preOp chain so the # dependency is not lost when preOps are present. for dep in placement.deps: ref_id = placement_tail_id.get(id(dep.ref)) if ref_id is not None: if firstPreOpId is not None: setBefore(firstPreOpId, ref_id) else: prevId = ref_id # Final link: primary module points to last dep if lastDepId is not None: setBefore(curId, lastDepId) elif prevId is not None: setBefore(curId, prevId) partition_emitted.append(emitted) all_partitions.append(partition_emitted) self._emitted = all_partitions return all_partitions def build(self, *, stop_after: Optional[Pass] = None) -> Union[ LogicalSchedule, AnnotatedSchedule, AugmentedSchedule, EmittedSchedule]: """Execute the full scheduling pipeline sequentially. Args: stop_after: If given, stop after the named pass and return early. Used by tests to run the pipeline up to a specific stage. """ if stop_after is not None and stop_after not in Pass.pipeline(): raise ValueError(f"invalid stop_after: {stop_after!r}") for pass_enum in Pass.pipeline(): getattr(self, self._PASS_METHODS[pass_enum])() if stop_after == pass_enum: if pass_enum == Pass.EMIT: return self._emitted return self._partitions return self._emitted # ── Loop variant derivation ──────────────────────────── @staticmethod def _rewire_before(emitted: List[EmittedModule], removed_ids: set) -> List[EmittedModule]: """Rewire before-links that point to removed modules. If em.before points to a removed module, follow that module's own before link until we find a non-removed module (or None). """ id_to_em = {em.moduleId: em for em in emitted} for em in emitted: if em.moduleId in removed_ids: continue b = em.before while b is not None and b in removed_ids: b = id_to_em[b].before em.before = b return [em for em in emitted if em.moduleId not in removed_ids] def build_ngll(self) -> EmittedSchedule: """NGLL (No Global Load Loop): mainloop without GR(n+2), GR_INC. WaitGR inflight counts are zeroed since no new GRs are in flight. """ assert self._emitted is not None, "call build() first" if self.config.pgr in (0, 1): self._ngll_emitted = [[[]]] return self._ngll_emitted ngll = [] for partition_emitted in self._emitted: part_ngll = [] for emitted in partition_emitted: new_emitted = copy.deepcopy(emitted) removed = set() for em in new_emitted: src = em.source if em.opType == 'gr' and src.mtIteration == 2: removed.add(em.moduleId) elif em.opType in ('gl2_prefetch', 'gl2_prefetch_inc'): removed.add(em.moduleId) elif em.opType == 'wait_gr': if src.wait_gr_counts is not None: src.wait_gr_counts = WaitGRCounts() part_ngll.append(self._rewire_before(new_emitted, removed)) ngll.append(part_ngll) self._ngll_emitted = ngll return ngll def build_nll(self) -> EmittedSchedule: """NLL (No Load Loop): mainloop without GR, LR(n+1), GR_INC, LR_INC, WaitGR(n+1)+Sync. Keeps LR(n), MFMAs, WaitGR(n). WaitGR counts are zeroed only when no LR(n) remains in the last subIterK slot.""" assert self._emitted is not None, "call build() first" if self.config.pgr == 0: self._nll_emitted = [[[]]] return self._nll_emitted nll = [] numK = self.config.numSubIterK for partition_emitted in self._emitted: part_nll = [] for k, emitted in enumerate(partition_emitted): new_emitted = copy.deepcopy(emitted) removed = set() for em in new_emitted: src = em.source if em.opType == 'gr': removed.add(em.moduleId) elif em.opType == 'lr' and src.mtIteration == 1: removed.add(em.moduleId) elif em.opType == 'gr_inc' and self.config.pgr == 2: # PGR=2: NGLL already swapped LW via its kept gr_inc, # so NLL must drop gr_inc to avoid swapping it back. # PGR=1: keep gr_inc — it advances SRD + swaps LW for # tail entry (PRELOOP's single GR did neither). removed.add(em.moduleId) elif (em.opType == 'lr_inc' and not src.isUnrollSwap and self._is_multi_du()): # Only multi-DU drops the MT-transition lr_inc in the NLL; # single-DU PGR=2 keeps lr_inc (gr_inc is still dropped). removed.add(em.moduleId) elif em.opType in ('gl2_prefetch', 'gl2_prefetch_inc'): removed.add(em.moduleId) has_lr = any(em.opType == 'lr' and em.moduleId not in removed for em in new_emitted) # GR path is removed in NLL — no new global loads are in flight. for em in new_emitted: if em.opType == 'wait_gr' and em.moduleId not in removed: em.source.wait_gr_counts = WaitGRCounts() # Find Sync modules paired with removed wait_gr for em in new_emitted: if em.opType == 'sync' and em.before is not None \ and em.before in removed: removed.add(em.moduleId) # Remove WaitLR if no LR remains in this subIterK # but keep WaitLR ops that non-removed modules depend on # (e.g. MFMAs waiting for LRs issued in a previous subIterK) if not has_lr: depended_on = {em.before for em in new_emitted if em.moduleId not in removed and em.before is not None} for em in new_emitted: if em.opType == 'wait_lr' \ and em.moduleId not in depended_on: removed.add(em.moduleId) part_nll.append(self._rewire_before(new_emitted, removed)) nll.append(part_nll) self._nll_emitted = nll return nll def build_tailloop_pgr0(self) -> List[List[List[EmittedModule]]]: """Template for Tailloop based on PGR0 schedule. Returns [partition][groups] where each group has at most one MFMA. The tail loop runs flat (no partitioning): per subIterK we emit one LR pass covering every unique (tensor, tile_range), one boundary mask, then every partition's MFMAs back-to-back. This requires the flat tile-id layout from _compute_flat_tail_tile_state (and the matching vgpr realloc in _realloc_tail_tiles_flat) so each unique partition group has its own vgpr range — the mainloop's per- partition tile budget multiplexes vgprs across pi and cannot hold all partitions' tiles live at once. """ cfg = self.config numK = cfg.numSubIterK # Flat tile layout: every unique (tensor, partition_group) gets its # own vgpr tile id. _compute_tail_tile_state's old per-partition # tile_maps would reuse vgprs across pi and break a flat loop. tile_maps, self._flat_tail_peaks = self._compute_flat_tail_tile_state() # Legacy unused-tile bookkeeping: in the flat path we replace the # vgpr tiles wholesale at tail entry, so nothing here. self._tail_unused_tile_ids = {'A': set(), 'B': set(), 'SA': set(), 'SB': set()} preamble = [] # GRs entire MT at once for all tensors. all_tiles = { 'A': MFMATileRange(0, numK, 0, cfg.numMFMATilesM), 'B': MFMATileRange(0, numK, 0, cfg.numMFMATilesN), } preamble.extend(self._make_gr_all_tensors(0, all_tiles)) # bf16-only: an OOB dwordx4 load can corrupt the trailing 16-bit # element at the K-boundary (buffer instructions enforce dword # granularity on OOB). We patch it with a 16-bit DTL load. Wider # dtypes (e.g. fp4 read at K=32 granularity) don't have this issue, # so we skip emission entirely for them. # TDM uses tensor_load_to_lds, not buffer instructions, so the # dword-granularity OOB corruption does not apply. hasTDM = self._kernel.get("enableTDMA") and self._kernel.get("enableTDMB") if self._kernel["ProblemType"]["DataTypeA"].isBFloat16() and not hasTDM: # We need to wait for other SIMD before placing the DTL load # (as we'll write twice to this address : OOB Zero then fixup load) preamble.append(SyncOp()) preamble.append(InlineModuleOp( build=lambda em: em.writer.tailLoopBoundaryDtlLoadAB( em.kernel, em.tensorParametersMap['A'], em.tensorParametersMap['B']), label="tail_boundary_ab")) preamble.append(WaitGROp(wait_gr_counts=WaitGRCounts())) preamble.append(SyncOp()) # Flat per-subIterK emission. The K-boundary mask depends only on k, # and with flat tile ids the per-partition tile_maps reference # disjoint vgpr ranges per (tensor, group). So per k we can: # 1. emit each unique (tensor, tile_range) LR exactly once # 2. wait + mask once (single VAnd per unique flat vgpr) # 3. run every partition's MFMA back-to-back # The returned shape is still [partition][group][ops]; we use a # single outer "partition" holding all per-k groups. miK = int(self._kernel["MatrixInstK"]) groups = [self._to_emitted(preamble)] # Build a merged tile_map covering every partition's tiles, used by # MaskKOp to enumerate the live flat vgpr ids. merged_tile_map: dict = {} for pi in range(cfg.numPartitions): for tensor in ('A', 'B', 'SA', 'SB'): src = tile_maps[pi].get(tensor) if not src: continue dst = merged_tile_map.setdefault(tensor, [{}]) while len(dst) < len(src): dst.append({}) for ui, m in enumerate(src): dst[ui].update(m) for k in range(numK): ops = [] # Dedup LRs across partitions by tileId range — with flat tile # ids, same range ⇒ same vgprs, so one LR populates all readers. seen_lr = set() for pi in range(cfg.numPartitions): cur = self._partition_tile_range(pi) for tensor, gran in self._lr_tensors(): if k % gran.k != 0: continue side_key = 'A' if tensor in ('A', 'SA') else 'B' tiles = gran.tile_range(k, *cur[side_key]) lr_key = (tensor, tiles.tileId_start, tiles.tileId_end, tiles.subIterK_start, tiles.subIterK_end) if lr_key in seen_lr: continue seen_lr.add(lr_key) lr = LRPlacement(tensor=tensor, mtIteration=0, tiles=tiles, subIterK_slot=k, partition=pi) lr.vgpr_tile_map = copy.deepcopy(tile_maps[pi].get(tensor, [])) ops.append(lr) ops.append(WaitLROp()) ops.append(MaskKOp(subIterK=k, vgpr_tile_map=copy.deepcopy(merged_tile_map))) # All partitions' MFMAs for this k, back-to-back. for pi in range(cfg.numPartitions): cur = self._partition_tile_range(pi) mfma_tileA = MFMATileRange(k, k + 1, *cur['A']) mfma_tileB = MFMATileRange(k, k + 1, *cur['B']) mfma = MFMAPlacement(subIterK=k, tileA=mfma_tileA, tileB=mfma_tileB) mfma.vgpr_tile_maps = copy.deepcopy(tile_maps[pi]) ops.append(mfma) # Early-exit: after subIterK=k completes for every partition, # skip ahead if no more valid K remains. Omit on the last k. if k != numK - 1: ops.append(SkipOp( compare='LE', value=miK * (k + 1), target='SkipTailLoopL', rawLabel=True, branchComment=f"early-exit tail after subIterK={k} (no valid K left)")) groups.append(self._to_emitted(ops)) self._tailloop_emitted = [groups] return self._tailloop_emitted @staticmethod def _to_emitted(ops) -> List[EmittedModule]: """Wrap Emittable objects (Placements / BaseOps) into EmittedModules.""" return [EmittedModule(moduleId=mid, source=op) for mid, op in enumerate(ops)] def _make_gr_all_tensors(self, mt: int, tiles: dict) -> List[GRPlacement]: """Create GR placements for all tensors and uids at the given MT iteration. tiles: {'A': MFMATileRange, 'B': MFMATileRange} For multi-DU tensors (numUnroll > 1), each uid gets its own K-slice so that uid u loads k=[u*grGran.k, (u+1)*grGran.k). """ cfg = self.config result = [] for tensor in self.tensors: tile = tiles['A' if tensor in ('A', 'SA') else 'B'] gr = {'A': cfg.grA, 'B': cfg.grB, 'SA': cfg.grSA, 'SB': cfg.grSB}.get(tensor, cfg.grA) nUnroll = cfg.numUnroll.get(tensor, 1) for uid in range(nUnroll): if nUnroll > 1: k_start = uid * gr.k k_end = (uid + 1) * gr.k uid_tile = MFMATileRange(k_start, k_end, tile.tileId_start, tile.tileId_end) else: uid_tile = tile result.append(GRPlacement(tensor=tensor, mtIteration=mt, tiles=uid_tile, subIterK_slot=0, unrollId=uid)) return result def _make_lr_all_tensors(self, tiles: dict) -> List[LRPlacement]: """Create LR placements for first partition. tiles: per-tensor MFMATileRange, e.g. {'A': MFMATileRange(0, k, mn0, mn1), ...} Uses the first MFMA's vgpr tile maps (the preloop loads data consumed by the first MFMA, not the next subIterK like mainloop LRs). """ first_mfma = self._partitions[0][0].mfma placements = [] for tensor in self.tensors: lr = LRPlacement( tensor=tensor, mtIteration=0, tiles=tiles[tensor], subIterK_slot=0, partition=0) if tensor in first_mfma.vgpr_tile_maps: lr.vgpr_tile_map = copy.deepcopy(first_mfma.vgpr_tile_maps[tensor]) placements.append(lr) return placements def _make_depops_all_tensors(self, cls) -> List[BaseOp]: """Create a BaseOp subclass instance for each tensor (and uid for GRIncOp).""" if cls is GRIncOp: return [GRIncOp(tensor=t, unrollId=uid) for t in self.tensors for uid in range(self.config.numUnroll.get(t, 1))] return [cls(tensor=tensor) for tensor in self.tensors] def _make_gr_all_tensors_uid(self, mt: int, tiles: dict, uid: int) -> List[GRPlacement]: """Create GR placements for a single uid across all tensors. Only emits for tensors where uid < numUnroll[tensor]. Each uid gets its own K-slice: k=[uid*grGran.k, (uid+1)*grGran.k). """ cfg = self.config result = [] for tensor in self.tensors: nUnroll = cfg.numUnroll.get(tensor, 1) if uid >= nUnroll: continue tile = tiles['A' if tensor in ('A', 'SA') else 'B'] gr = {'A': cfg.grA, 'B': cfg.grB, 'SA': cfg.grSA, 'SB': cfg.grSB}.get(tensor, cfg.grA) if nUnroll > 1: k_start = uid * gr.k k_end = (uid + 1) * gr.k uid_tile = MFMATileRange(k_start, k_end, tile.tileId_start, tile.tileId_end) else: uid_tile = tile result.append(GRPlacement(tensor=tensor, mtIteration=mt, tiles=uid_tile, subIterK_slot=0, unrollId=uid)) return result def _make_depops_uid(self, cls, uid: int) -> List[BaseOp]: """Create a BaseOp subclass instance for a single uid across all tensors. Only emits for tensors where uid < numUnroll[tensor]. """ result = [] for tensor in self.tensors: nUnroll = self.config.numUnroll.get(tensor, 1) if uid >= nUnroll: continue if cls is GRIncOp: result.append(GRIncOp(tensor=tensor, unrollId=uid)) else: result.append(cls(tensor=tensor)) return result def _make_preloop_mt1_grs(self) -> List[GRPlacement]: """Create MT1 GRs for the PGR=2 preloop, ordered to match the mainloop. Covers partitions 0..offsetPartition-1 with proper deduplication. Each unique (tensor, uid, tile-range, k-range) appears exactly once. For multi-DU data tensors (numUnroll > 1), each uid only emits GRs for the K slots that belong to it (uid u → k in [u*grGran.k, (u+1)*grGran.k)). """ assert self._partitions is not None, "call build() or place_LRs() first" cfg = self.config seen = set() result = [] maxUnroll = max(cfg.numUnroll.values()) if cfg.numUnroll else 1 for uid in range(maxUnroll): for pi in range(cfg.offsetPartition): target_range = self._partition_tile_range(pi) for slot in self._partitions[0]: k = slot.mfma.subIterK items = [('A', target_range['A'], cfg.grA), ('B', target_range['B'], cfg.grB)] if cfg.hasScale: items.append(('SA', target_range['A'], cfg.grSA)) items.append(('SB', target_range['B'], cfg.grSB)) for tensor, (t_start, t_end), gr_gran in items: nUnroll = cfg.numUnroll.get(tensor, 1) if uid >= nUnroll: continue if nUnroll > 1: uid_k_start = uid * gr_gran.k uid_k_end = (uid + 1) * gr_gran.k if k < uid_k_start or k >= uid_k_end: continue tr = gr_gran.tile_range(k, t_start, t_end) key = (tensor, uid, tr.tileId_start, tr.tileId_end, tr.subIterK_start, tr.subIterK_end) if key in seen: continue seen.add(key) result.append(GRPlacement( tensor=tensor, mtIteration=1, tiles=tr, subIterK_slot=k, partition=pi, unrollId=uid, )) return result def _make_preloop_mt1_grs_uid(self, uid: int) -> List[GRPlacement]: """Create MT1 GRs for a single uid in the PGR=2 preloop. Filters _make_preloop_mt1_grs logic to one uid, skipping tensors where uid >= numUnroll[tensor]. """ assert self._partitions is not None, "call build() or place_LRs() first" cfg = self.config seen = set() result = [] for pi in range(cfg.offsetPartition): target_range = self._partition_tile_range(pi) for slot in self._partitions[0]: k = slot.mfma.subIterK items = [('A', target_range['A'], cfg.grA), ('B', target_range['B'], cfg.grB)] if cfg.hasScale: items.append(('SA', target_range['A'], cfg.grSA)) items.append(('SB', target_range['B'], cfg.grSB)) for tensor, (t_start, t_end), gr_gran in items: nUnroll = cfg.numUnroll.get(tensor, 1) if uid >= nUnroll: continue if nUnroll > 1: uid_k_start = uid * gr_gran.k uid_k_end = (uid + 1) * gr_gran.k if k < uid_k_start or k >= uid_k_end: continue tr = gr_gran.tile_range(k, t_start, t_end) key = (tensor, uid, tr.tileId_start, tr.tileId_end, tr.subIterK_start, tr.subIterK_end) if key in seen: continue seen.add(key) result.append(GRPlacement( tensor=tensor, mtIteration=1, tiles=tr, subIterK_slot=k, partition=pi, unrollId=uid, )) return result def _make_initC_op(self) -> InlineModuleOp: """Create an InlineModuleOp that zeros accumulator registers (initC). For PGR>=1 the op is placed between GR issue and WaitGR in the preloop so the MFMA zeroing instructions execute while global reads are in flight, hiding their latency behind the memory access time. For PGR=0 there are no global reads, so this op is the entire preloop and simply zeros the accumulators before the mainloop's first accumulating MFMA. """ def _build_initC(emitter): from Tensile.Components.Subtile.Kernel import initVgprTilesToZero return initVgprTilesToZero(emitter.writer, emitter.kernel, emitter.dtileInfo) return InlineModuleOp(build=_build_initC, label="initC_overlap") def build_preloop(self) -> EmittedSchedule: """Build preloop: pipeline initialization sequence before mainloop. PGR=0: initC only (no GR/LR pipeline, but accumulators must be zeroed before the first MFMA in the mainloop). PGR=1 sequence: GR(MT 0) — all tensors, all tiles (× numUnroll if unrolled) initC — accumulator zeroing, overlapped with in-flight GR WaitGR — wait for global reads to land in LDS Sync — barrier LR — first partition, subIterK=0 skip(LE 1, NLL) PGR=2 sequence: GR(MT 0) — all tensors, all tiles (× numUnroll if unrolled) initC — accumulator zeroing, overlapped with in-flight GR WaitGR — wait for global reads to land in LDS Sync — barrier LR — first partition, subIterK=0 skip(LE 1, NLL) GR(MT 1) — first partition tiles (× numUnroll if unrolled) skip(LE 2, NGLL) Returns [1 partition][1 subIterK][EmittedModules] to match emit() shape. """ if self.config.pgr == 0: # PGR=0 has no GR/LR pipeline, but the mainloop's first MFMA still # accumulates into D, so the accumulators must be zeroed first. # Emit an initC-only preloop (the PRELOOP itself is still emitted # unconditionally in emitMainAndExitLoops). self._preloop_emitted = [[self._to_emitted([self._make_initC_op()])]] return self._preloop_emitted cfg = self.config numK = cfg.numSubIterK part0 = self._partition_tile_range(0) all_tiles = { 'A': MFMATileRange(0, numK, 0, cfg.numMFMATilesM), 'B': MFMATileRange(0, numK, 0, cfg.numMFMATilesN), } lr_tiles = { 'A': MFMATileRange(0, cfg.lrA.k, *part0['A']), 'B': MFMATileRange(0, cfg.lrB.k, *part0['B']), } if cfg.hasScale: lr_tiles['SA'] = MFMATileRange(0, cfg.lrSA.k, *part0['A']) lr_tiles['SB'] = MFMATileRange(0, cfg.lrSB.k, *part0['B']) initC_op = self._make_initC_op() if cfg.pgr == 1: maxUnroll = max(cfg.numUnroll.values()) if cfg.numUnroll else 1 if maxUnroll > 1: preloop_ops = [] for uid in range(maxUnroll): preloop_ops.extend(self._make_gr_all_tensors_uid(0, all_tiles, uid)) preloop_ops.extend(self._make_depops_uid(GRIncOp, uid)) emitted = self._to_emitted([ *preloop_ops, initC_op, WaitGROp(wait_gr_counts=WaitGRCounts()), SyncOp(), *self._make_lr_all_tensors(lr_tiles), SkipOp(compare='LE', value=1, target='NLL'), ]) else: emitted = self._to_emitted([ *self._make_gr_all_tensors(0, all_tiles), initC_op, WaitGROp(wait_gr_counts=WaitGRCounts()), SyncOp(), *self._make_lr_all_tensors(lr_tiles), SkipOp(compare='LE', value=1, target='NLL'), ]) else: gl2_preloop_ops = [] if cfg.pgl > 0: gl2_preloop_ops.append(GL2PrefetchOp()) if cfg.pgl == 2: gl2_preloop_ops.append(GL2PrefetchIncOp()) gl2_preloop_ops.append(GL2PrefetchOp()) maxUnroll = max(cfg.numUnroll.values()) if cfg.numUnroll else 1 if maxUnroll > 1: preloop_ops = [] for uid in range(maxUnroll): preloop_ops.extend(self._make_gr_all_tensors_uid(0, all_tiles, uid)) preloop_ops.extend(self._make_depops_uid(GRIncOp, uid)) mt1_ops = [] for uid in range(maxUnroll): mt1_ops.extend(self._make_preloop_mt1_grs_uid(uid)) mt1_ops.extend(self._make_depops_uid(GRIncOp, uid)) emitted = self._to_emitted([ *preloop_ops, initC_op, WaitGROp(wait_gr_counts=WaitGRCounts()), SyncOp(), *self._make_lr_all_tensors(lr_tiles), SkipOp(compare='LE', value=1, target='NLL'), *mt1_ops, *gl2_preloop_ops, SkipOp(compare='LE', value=2, target='NGLL'), ]) else: emitted = self._to_emitted([ *self._make_gr_all_tensors(0, all_tiles), *self._make_depops_all_tensors(GRIncOp), initC_op, WaitGROp(wait_gr_counts=WaitGRCounts()), SyncOp(), *self._make_lr_all_tensors(lr_tiles), SkipOp(compare='LE', value=1, target='NLL'), *self._make_preloop_mt1_grs(), *gl2_preloop_ops, SkipOp(compare='LE', value=2, target='NGLL'), ]) self._preloop_emitted = [[emitted]] return self._preloop_emitted def _emitLoop(self, writer, kernel, label, emitted_3d, schedule=True): """Emit a loop section from a 3D emitted structure. emitted_3d: [partition][subIterK][EmittedModule] When schedule=True and a group has MFMAs, calls instructionSchedule for interleaving. When schedule=False, emits instructions sequentially. """ from Tensile.Components.Subtile.InstructionScheduler import ( instructionSchedule, _MIN_MFMA_GAP_DS_READ_TO_WAIT_DEFAULT, _MIN_MFMA_GAP_DS_READ_TO_WAIT_GFX1250, ) from Tensile.Components.Subtile.WaitAluInsertion import ( insertLRSwapRawWaitAlu, setMatrixReuse, insertLRSwapWarWaitAlu) from rocisa.code import Module, Label from rocisa.container import sgpr from rocisa.instruction import SCmpEQU32, SCBranchSCC0, SMovB32 # gfx1250 needs a larger ds_read->waitcnt gap. isGfx1250 = writer.states.archCaps.get("HasWmmaArbStallBit", False) minGapDsReadToWait = (_MIN_MFMA_GAP_DS_READ_TO_WAIT_GFX1250 if isGfx1250 else _MIN_MFMA_GAP_DS_READ_TO_WAIT_DEFAULT) module = Module(label) module.addComment0(f"{label} start") use_pap_preloop_skip = ( label == "PRELOOP" and kernel.get("UseSubtileImpl") and kernel.get("PrefetchAcrossPersistent") and kernel.get("StreamK") == 3 ) pap_merge_label = Label("SubtilePAPPreloopFirstGRMerge", "") if use_pap_preloop_skip else None skipping_first_gr_group = False first_gr_group_done = False for pi, partition_emitted in enumerate(emitted_3d): for k, em_list in enumerate(partition_emitted): module.addComment0(f"partition={pi} subIterK={k}") has_mfma = any(em.opType == 'mfma' for em in em_list) if schedule and em_list and has_mfma: scheduled = instructionSchedule( em_list, multiDU=self._is_multi_du(), minGapDsReadToWait=minGapDsReadToWait) module.add(scheduled) else: for em in em_list: if use_pap_preloop_skip and not first_gr_group_done: if em.opType == 'gr': if not skipping_first_gr_group: module.add(SCmpEQU32(src0=sgpr("SkPrefetchPrimed"), src1=0, comment="Subtile PAP: first PRELOOP GR already issued?")) module.add(SCBranchSCC0(labelName=pap_merge_label.getLabelName(), comment="skip first PRELOOP GR group if primed")) skipping_first_gr_group = True elif skipping_first_gr_group: module.add(pap_merge_label) module.add(SMovB32(dst=sgpr("SkPrefetchPrimed"), src=0, comment="Subtile PAP: clear after first PRELOOP GR merge")) first_gr_group_done = True for inst in em.instructions: module.add(inst) if use_pap_preloop_skip and skipping_first_gr_group and not first_gr_group_done: module.add(pap_merge_label) module.add(SMovB32(dst=sgpr("SkPrefetchPrimed"), src=0, comment="Subtile PAP: clear after first PRELOOP GR merge")) module.addComment0(f"{label} end") # SCHED_MODE 2: guard the LR offset-swap -> ds_read RAW hazard once, against # the final post-schedule order (no-op on other archs). module = insertLRSwapRawWaitAlu(module, writer, kernel) # gfx1250: enable WMMA matrix-A reuse on the final post-schedule order. module = setMatrixReuse(module, writer, kernel, 'a') module = setMatrixReuse(module, writer, kernel, 'b') # PGR=0 only: the unprefetched loop puts the ds_read of an LR offset # right before the swap that overwrites it. PGR>=1 prefetch separates # them (swap hoisted ahead, dscnt drain between), so no WAR can form. if self.config.pgr == 0 and label.startswith("MAINLOOP"): module = insertLRSwapWarWaitAlu(module, writer, kernel) # Cluster barrier: splice both halves against the final post-schedule order. # Signal goes right after the mainloop's existing workgroup barrier (reusing # that sync); the wait is appended at the end to hide its cross-CU latency. from Tensile.Components.Subtile.ClusterBarrier import insertClusterBarrier module = insertClusterBarrier(module, writer, kernel) return module def _emit_pgr2_tail_lw_align(self, kernel): """Re-align LW_base parity with the LR vgpr at the tail entry. At PGR=2 the NLL drops gr_inc/lr_inc, so on the NLL exit path the cumulative LW_base swap count ends up one ahead of the LR vgpr swap count (PRELOOP's +1 LW swap is never balanced). Tail GR would then write to the half that tail LR is NOT reading, leaving the LR pointing at the stale PGR=2 prefetch — a one-shot XOR of LW_base with sgprSwap{tc} restores the parity invariant the tail body needs. """ from rocisa.code import Module from rocisa.instruction import SXorB32 from rocisa.container import sgpr module = Module("Pgr2TailLwAlign") # The tail LW_base parity re-align is only needed when the NLL drops the # MT-transition lr_inc (multi-DU path, see build_nll). Non-multi-DU PGR=2 # keeps lr_inc, so the parity invariant already holds and no re-align is # required here. if not self._is_multi_du(): return module if self.config.pgr != 2 or kernel.get("NoTailLoop"): return module tensors = ['A', 'B'] if self.config.hasScale: tensors.extend(['MXSA', 'MXSB']) module.addComment1("PGR=2 tail entry: re-align LW_base parity with LR vgpr.") for tc in tensors: lwName = f"LocalWriteBaseAddr{tc}" swapName = f"Swap{tc}" module.add(SXorB32( dst=sgpr(lwName), src0=sgpr(lwName), src1=sgpr(swapName), comment=f"PGR=2 tail align: parity-swap LW_base for {tc}")) return module def emitMainAndExitLoops(self, writer, kernel, tensorParametersA=None, tensorParametersB=None): """Emit preloop + mainloop + NGLL + NLL exit paths (no tail). Owns all control flow (labels, branches, counter management) for the main unrolled pipeline. For unroll_factor > 1, emits per-unroll copies with correct vgpr tiles. Each mainloop exit jumps to its corresponding NGLL→NLL pair. The tail loop is emitted separately by emitTailLoop() so the orchestrator (Subtile.Kernel.mainLoop) can wrap it with the runtime K%DU counter setup and skip branch. """ from rocisa.code import Module, Label from rocisa.instruction import (SSubU32, SCmpEQU32, SCBranchSCC0, SCBranchSCC1, SBranch) from rocisa.container import sgpr assert self._emitter is not None, \ "populate_instructions() must be called first" module = Module("MainAndExitLoops") uf = self.unroll_factor # gfx1250 requires a loop-back/exit s_cbranch to be alone/first after an # MFMA, so hoist the per-copy dec+compare above the body's last MFMA. isGfx1250 = writer.states.archCaps.get("HasWmmaArbStallBit", False) def make_subtile_pap_module(skip_barrier=False): if (kernel.get("UseSubtileImpl") and kernel.get("PrefetchAcrossPersistent") and kernel.get("StreamK") == 3 and hasattr(writer, "prefetchAcrossPersistentSubtile")): preloop_gr = Module("Subtile PAP first PRELOOP GR") for em in self._preloop_emitted[0][0]: if em.opType != 'gr': break for inst in em.instructions: preloop_gr.add(copy.deepcopy(inst)) return writer.prefetchAcrossPersistentSubtile( kernel, tensorParametersA, tensorParametersB, preloop_gr, skipBarrier=skip_barrier) return None def inject_pap_after_nll_drain(emitted_3d): emitted_3d = copy.deepcopy(emitted_3d) def insert_after_drain(em_list): wait_gr = next((em for em in em_list if em.opType == 'wait_gr'), None) if wait_gr is None: return False tail_id = wait_gr.moduleId sync = next((em for em in em_list if em.opType == 'sync' and em.before == tail_id), None) if sync is not None: tail_id = sync.moduleId pap_module = make_subtile_pap_module(skip_barrier=(sync is not None)) if pap_module is None: return False pap_id = max(em.moduleId for em in em_list) + 1 consumers = [em for em in em_list if em.before == tail_id] em_list.append(EmittedModule(moduleId=pap_id, instructions=[pap_module], before=tail_id, source=None)) # Place PAP between the final drain and the first NLL work that # was waiting on that drain. The existing before-link graph is a # chain, so this preserves scheduler ordering while moving PAP # past the vlcnt(0) that would otherwise drain it immediately. for consumer in consumers: consumer.before = pap_id return True for partition_emitted in emitted_3d: for em_list in partition_emitted: if insert_after_drain(em_list): return emitted_3d return emitted_3d # ── accumulator init (initC) ── # initC happens in preloop, between GR issue and WaitGR so the zeroing # MFMAs overlap the in-flight global reads. Every path must reach initC (mirrors the non-Subtile flow): # PGR>=1, K>=DepthU : GR issue -> initC -> rest of preloop -> mainloop # PGR=0 : preloop is just initC -> mainloop # K= 1 if skipGRForTail: module.add(SCmpEQU32(src0=sgpr("LoopCounterL"), src1=0, comment="K < DepthU? skip prefetch GR to initC")) module.add(SCBranchSCC1(labelName=skipGRLabel.getLabelName(), comment="K < DepthU: skip prefetch GR, still run initC")) # ── Preloop (initC + tail jump woven around the single init op) ── # Operate on a copy so self._preloop_emitted (read by PAP and the # per-unroll copies) stays untouched. preloop_emitted = copy.deepcopy(self._preloop_emitted) if not kernel["NoTailLoop"]: em_list = preloop_emitted[0][0] init_idx = next((i for i, em in enumerate(em_list) if getattr(em.source, 'label', None) == 'initC_overlap'), None) assert init_idx is not None, "preloop must contain the canonical initC op" next_id = max(em.moduleId for em in em_list) + 1 # After initC, jump to the tail loop when K < DepthU. em_list.insert(init_idx + 1, EmittedModule( moduleId=next_id, instructions=[ SCmpEQU32(src0=sgpr("LoopCounterL"), src1=0, comment="K < DepthU? initC done, run tail only"), SCBranchSCC1(labelName=endLabel.getLabelName(), comment="K < DepthU: jump to tail loop"), ])) next_id += 1 if skipGRForTail: # Land the skip-GR guard right before initC. em_list.insert(init_idx, EmittedModule( moduleId=next_id, instructions=[skipGRLabel])) module.add(self._emitLoop(writer, kernel, "PRELOOP", preloop_emitted, schedule=False)) # ── Mainloop ── module.addComment0("MAINLOOP") loopBegin = Label("LoopBeginL", "", alignment=16) exitValue = self.config.pgr exitLabels = [Label(f"ExitC{ui}", "") for ui in range(uf - 1)] module.add(loopBegin) # Debug: emit `s_mov_b32 m0, LoopCounterL; s_ttracedata` at the start of # every mainloop iteration so SQTT / trace decoders can identify iterations # (adds 2 instructions per iter). Gated by the EmitMainloopTraceMarker global. emitTraceMarker = globalParameters.get("EmitMainloopTraceMarker", False) if emitTraceMarker: from rocisa.container import mgpr from rocisa.instruction import SMovB32 as _SMovB32 from rocisa.instruction import STtraceData as _STtraceData for ui in range(uf): if emitTraceMarker: # Mainloop iteration marker for SQTT / trace decoder: write # LoopCounterL into M0 then emit it via s_ttracedata. Decoder # only uses low 8 bits, so M0 wrap past 256 is fine. module.add(_SMovB32(dst=mgpr(0), src=sgpr("LoopCounterL"), comment="trace: M0 = LoopCounterL")) module.add(_STtraceData(comment="trace: emit M0 to SQTT")) loopBody = self._emitLoop(writer, kernel, f"MAINLOOP_C{ui}", self._emitted_per_unroll[ui]) # dec + compare that produce the SCC the loop-control branch reads. setupInsts = [ SSubU32(dst=sgpr("LoopCounterL"), src0=sgpr("LoopCounterL"), src1=1, comment=f"dec counterL (copy {ui})"), SCmpEQU32(src0=sgpr("LoopCounterL"), src1=exitValue, comment=f"counterL == {exitValue}? (copy {ui} exit)"), ] if ui < uf - 1: branch = SCBranchSCC1( labelName=exitLabels[ui].getLabelName(), comment=f"copy {ui} exit → NGLL_C{ui}") else: branch = SCBranchSCC0( labelName=loopBegin.getLabelName(), comment="restart mainloop") _emitBodySetupBranch(module, loopBody, setupInsts, branch, hoist=isGfx1250) # ── NGLL + NLL exit paths ── hasNGLL = self.config.pgr >= 2 module.add(Label("SkipMainloop", "")) if hasNGLL: module.add(Label("SkipToNGLL", "")) # _per_unroll[i] has tiles for unroll_iter=i. # After mainloop C{ui}, data in LDS/vgprs corresponds to # unroll_iter = (ui + pgr) % uf for NLL, (ui + 1) % uf for NGLL. # NLLEarly (preloop skip) needs unroll_iter=0, i.e. _nll_per_unroll[0]. # We place SkipToNLL before whichever NLL block uses index 0. pgr = self.config.pgr last = uf - 1 # Fall-through from last mainloop copy nll_ft = (last + pgr) % uf if hasNGLL: module.addComment0(f"NGLL_C{last}") module.add(self._emitLoop(writer, kernel, f"NGLL_C{last}", self._ngll_per_unroll[(last + 1) % uf])) if nll_ft == 0: module.add(Label("SkipToNLL", "")) module.addComment0(f"NLL_C{last}") module.add(self._emitLoop(writer, kernel, f"NLL_C{last}", inject_pap_after_nll_drain(self._nll_per_unroll[nll_ft]))) module.add(self._emit_pgr2_tail_lw_align(kernel)) module.add(SBranch(labelName=endLabel.getLabelName(), comment="skip other exit paths")) for ui in range(uf - 1): nll_idx = (ui + pgr) % uf module.add(exitLabels[ui]) if hasNGLL: module.addComment0(f"NGLL_C{ui}") module.add(self._emitLoop(writer, kernel, f"NGLL_C{ui}", self._ngll_per_unroll[(ui + 1) % uf])) if nll_idx == 0: module.add(Label("SkipToNLL", "")) module.addComment0(f"NLL_C{ui}") module.add(self._emitLoop(writer, kernel, f"NLL_C{ui}", inject_pap_after_nll_drain(self._nll_per_unroll[nll_idx]))) module.add(self._emit_pgr2_tail_lw_align(kernel)) if ui < uf - 2: module.add(SBranch(labelName=endLabel.getLabelName(), comment="skip other exit paths")) module.add(endLabel) return module def emitTailLoop(self, writer, kernel): """Emit the tail loop body only (no counter setup, no skip branch). Returns an empty Module when NoTailLoop is set. The caller is responsible for emitting calculateLoopNumIter(-1) before this and closeLoop(emitEndLabelOnly=True) after, mirroring the legacy KernelWriter pattern. """ assert self._emitter is not None, \ "populate_instructions() must be called first" module = Module("TailLoop") if kernel["NoTailLoop"]: return module module.addComment0("TAILLOOP") # Swap to the flat tail vgpr tile layout. Frees the mainloop's # per-partition tiles back to the pool and reallocates a flat set # sized by _compute_flat_tail_tile_state (already invoked by # build_tailloop_pgr0; peaks stashed on self._flat_tail_peaks). self._realloc_tail_tiles_flat(writer, self._flat_tail_peaks) # init must run before populate so each MaskKOp in the body can read # the mask vgprs (kReg, vDiff, …) that init allocates. for inst in self._emitter.emit_mask_k_init(): module.add(inst) self._emitter.populate(self._tailloop_emitted, unroll_iter=0) module.add(self._emitLoop(writer, kernel, "TAILLOOP", self._tailloop_emitted, schedule=False)) for inst in self._emitter.emit_mask_k_done(): module.add(inst) return module # ── VGPR tile allocation ────────────────────────────── def getNumVgpr(self, tileInfoA, tileInfoB, scaleTileInfoA=None, scaleTileInfoB=None) -> int: """Return the total number of VGPRs needed across all tensors (A, B, SA, SB) without performing any allocation. Returns max(mainloop_peak, flat_tail_peak) — the two layouts don't coexist (the tail frees and reallocates at entry), so the kernel budget is the larger of them. Must be called after scheduling is complete. """ assert hasattr(self, 'tile_peaks'), "call build() first" cfg = self.config def _tile_vgpr_count(tileInfo, lrGran): return int(math.ceil(tileInfo.mmaTileRegCount * lrGran.k * lrGran.mn)) def _total_for(peaks): t = peaks.get('A', 0) * _tile_vgpr_count(tileInfoA, cfg.lrA) \ + peaks.get('B', 0) * _tile_vgpr_count(tileInfoB, cfg.lrB) if cfg.hasScale and scaleTileInfoA and scaleTileInfoB: t += peaks.get('SA', 0) * _tile_vgpr_count(scaleTileInfoA, cfg.lrSA) \ + peaks.get('SB', 0) * _tile_vgpr_count(scaleTileInfoB, cfg.lrSB) return t mainloop_total = _total_for(self.tile_peaks) _, flat_peaks = self._compute_flat_tail_tile_state() tail_total = _total_for(flat_peaks) return max(mainloop_total, tail_total) def allocVgprTiles(self, writer, tileInfoA, tileInfoB, scaleTileInfoA=None, scaleTileInfoB=None): """Allocate physical VGPR tiles based on assign_vgpr_tiles() peaks. Each vgprTile holds one LR granularity worth of data: size = ceil(mmaTileRegCount * lrGranularity.k * lrGranularity.mn) Ex: 4 VGPRs for A/B for 1 MFMATile, and 1 VGPR for a 2x2 MFMA tile for SA/SB if hasScale. Produces per-tensor lists indexed by vgprTileId: vgprTilesA/B: List[RegisterTileInfo] vgprTilesSA/SB: List[RegisterTileInfo] """ assert hasattr(self, 'tile_peaks'), "call build() first" cfg = self.config def _tile_vgpr_count(tileInfo, lrGran): return int(math.ceil(tileInfo.mmaTileRegCount * lrGran.k * lrGran.mn)) def _alloc_tiles(count, numRegs): return [_checkout_tile(writer.vgprPool, numRegs, "allocVgprTiles_vstart") for _ in range(count)] self.vgprTilesA = _alloc_tiles(self.tile_peaks.get('A', 0), _tile_vgpr_count(tileInfoA, cfg.lrA)) self.vgprTilesB = _alloc_tiles(self.tile_peaks.get('B', 0), _tile_vgpr_count(tileInfoB, cfg.lrB)) if cfg.hasScale and scaleTileInfoA and scaleTileInfoB: self.vgprTilesSA = _alloc_tiles(self.tile_peaks.get('SA', 0), _tile_vgpr_count(scaleTileInfoA, cfg.lrSA)) self.vgprTilesSB = _alloc_tiles(self.tile_peaks.get('SB', 0), _tile_vgpr_count(scaleTileInfoB, cfg.lrSB)) else: self.vgprTilesSA = [] self.vgprTilesSB = [] # Stash tile-info so _realloc_tail_tiles_flat can reallocate the # tail's flat tile set without the caller plumbing them in again. self._alloc_tile_info = { 'tileInfoA': tileInfoA, 'tileInfoB': tileInfoB, 'scaleTileInfoA': scaleTileInfoA, 'scaleTileInfoB': scaleTileInfoB} def deallocVgprTiles(self, writer): """Deallocate VGPR tiles allocated by allocVgprTiles. Skips tile ids in self._tail_freed_tile_ids — those were already returned to the pool by _release_unused_tail_tiles. """ def _dealloc_tiles(tiles, freed): for tid, tile in enumerate(tiles): if tid in freed: continue _checkin_tile(tile) _dealloc_tiles(self.vgprTilesA, self._tail_freed_tile_ids['A']) _dealloc_tiles(self.vgprTilesB, self._tail_freed_tile_ids['B']) _dealloc_tiles(self.vgprTilesSA, self._tail_freed_tile_ids['SA']) _dealloc_tiles(self.vgprTilesSB, self._tail_freed_tile_ids['SB']) self.vgprTilesA = [] self.vgprTilesB = [] self.vgprTilesSA = [] self.vgprTilesSB = [] self._tail_freed_tile_ids = {'A': set(), 'B': set(), 'SA': set(), 'SB': set()} def _compute_tail_tile_state(self): """Single source of truth for tail-loop tile usage. Returns (tile_maps, unused) where - tile_maps[pi] = self._partitions[pi][0].mfma.vgpr_tile_maps, reused by build_tailloop_pgr0 to wire LR/MFMA/MaskK ops. - unused[tensor] = {tid} for tile slots the tail loop never references (the PGR>=1 prefetch half). Consumed by _release_unused_tail_tiles to reclaim their vgprs. """ tile_maps = [self._partitions[pi][0].mfma.vgpr_tile_maps for pi in range(self.config.numPartitions)] used = {t: set() for t in ('A', 'B', 'SA', 'SB')} for pi_map in tile_maps: for tensor in used: m = pi_map.get(tensor, [{}])[0] # unroll_iter=0 only used[tensor].update(m.values()) tiles_by_tensor = {'A': self.vgprTilesA, 'B': self.vgprTilesB, 'SA': self.vgprTilesSA, 'SB': self.vgprTilesSB} unused = { tensor: {tid for tid in range(len(tile_list)) if tid not in used[tensor]} for tensor, tile_list in tiles_by_tensor.items() } return tile_maps, unused def _compute_flat_tail_tile_state(self): """Tile-id remap for a non-partitioned ("flat") tail loop. The mainloop's tile_peaks are per-partition (each pi reuses the same vgprs across its subIterKs). A flat tail loop holds every partition's tiles live at once and needs one vgpr range per unique (tensor, partition_group). This method assigns each such group a fresh flat tile id in 0..flat_peaks[T)-1. Returns (tile_maps, peaks) where - tile_maps[pi][tensor] = [{group_key: flat_tile_id}] (single-entry list mirroring the mainloop's per-unroll_iter shape; tail always uses unroll_iter=0). - peaks[tensor] = count of distinct flat tile ids for tensor. """ cfg = self.config numP = cfg.numPartitions lr_grans = {'A': cfg.lrA, 'B': cfg.lrB} if cfg.hasScale: lr_grans['SA'] = cfg.lrSA lr_grans['SB'] = cfg.lrSB part_ranges = [self._partition_tile_range(pi) for pi in range(numP)] # group_id[(tensor, group_key)] = flat tile id group_id: dict = {t: {} for t in lr_grans} # tile_maps[pi][tensor] = [{group_key: flat_tile_id}] tile_maps: list = [{} for _ in range(numP)] for pi in range(numP): for tensor, gran in lr_grans.items(): side = TENSOR_SIDE[tensor] start, end = part_ranges[pi][side] groups = sorted({(t // gran.mn) * gran.mn for t in range(start, end)}) m = {} for g in groups: if g not in group_id[tensor]: group_id[tensor][g] = len(group_id[tensor]) m[g] = group_id[tensor][g] tile_maps[pi][tensor] = [m] peaks = {t: len(group_id[t]) for t in lr_grans} return tile_maps, peaks def _release_unused_tail_tiles(self, writer): """Return tile slots dead for the tail loop to the vgpr pool. Consumes self._tail_unused_tile_ids (populated by _compute_tail_tile_state in build_tailloop_pgr0). The freed tids are recorded in self._tail_freed_tile_ids so deallocVgprTiles skips them. """ assert not any(self._tail_freed_tile_ids[t] for t in self._tail_freed_tile_ids), \ "_release_unused_tail_tiles called twice" tiles_by_tensor = {'A': self.vgprTilesA, 'B': self.vgprTilesB, 'SA': self.vgprTilesSA, 'SB': self.vgprTilesSB} for tensor, tile_list in tiles_by_tensor.items(): for tid in self._tail_unused_tile_ids.get(tensor, ()): _checkin_tile(tile_list[tid]) self._tail_freed_tile_ids[tensor].add(tid) def _realloc_tail_tiles_flat(self, writer, peaks): """Free mainloop's per-partition tiles and reallocate flat tiles for the non-partitioned tail loop. `peaks[tensor]` comes from _compute_flat_tail_tile_state and is the number of distinct flat tile ids per tensor. The new flat tiles replace self.vgprTilesA/B/SA/SB; _tail_freed_tile_ids is cleared so deallocVgprTiles drops the flat set wholesale at kernel end. """ cfg = self.config info = self._alloc_tile_info def _tile_vgpr_count(tileInfo, lrGran): return int(math.ceil(tileInfo.mmaTileRegCount * lrGran.k * lrGran.mn)) def _dealloc_all(tiles): for tile in tiles: _checkin_tile(tile) def _alloc_tiles(count, numRegs): return [_checkout_tile(writer.vgprPool, numRegs, "reallocTailTilesFlat_vstart") for _ in range(count)] def _swap(target, new_tiles): # In-place swap so the InstructionEmitter's references stay valid. target.clear() target.extend(new_tiles) _dealloc_all(self.vgprTilesA) _dealloc_all(self.vgprTilesB) _dealloc_all(self.vgprTilesSA) _dealloc_all(self.vgprTilesSB) _swap(self.vgprTilesA, _alloc_tiles(peaks.get('A', 0), _tile_vgpr_count(info['tileInfoA'], cfg.lrA))) _swap(self.vgprTilesB, _alloc_tiles(peaks.get('B', 0), _tile_vgpr_count(info['tileInfoB'], cfg.lrB))) if cfg.hasScale and info['scaleTileInfoA'] and info['scaleTileInfoB']: _swap(self.vgprTilesSA, _alloc_tiles(peaks.get('SA', 0), _tile_vgpr_count(info['scaleTileInfoA'], cfg.lrSA))) _swap(self.vgprTilesSB, _alloc_tiles(peaks.get('SB', 0), _tile_vgpr_count(info['scaleTileInfoB'], cfg.lrSB))) else: _swap(self.vgprTilesSA, []) _swap(self.vgprTilesSB, []) # Flat tiles are freed wholesale by deallocVgprTiles at kernel end; # there are no pre-freed tids to skip. self._tail_freed_tile_ids = {'A': set(), 'B': set(), 'SA': set(), 'SB': set()} # ── Populate instructions ────────────────────────────── def populate_instructions(self, writer, kernel, tileInfoA, tileInfoB, dtileInfo, scaleTileInfoA=None, scaleTileInfoB=None, tensorParametersA=None, tensorParametersB=None) -> None: """Populate EmittedModule.instructions from placements and preOps. Uses per-tensor VGPR tile lists (vgprTilesA/B/SA/SB) indexed by vgprTileId from placement tile maps. """ need_build = (self._preloop_emitted is None or self._ngll_emitted is None or self._nll_emitted is None) if need_build: self.build() self._kernel = kernel from Tensile.Components.Subtile.InstructionEmitter import InstructionEmitter emitter = InstructionEmitter( writer, kernel, self.config, tileInfoA, tileInfoB, dtileInfo, self.vgprTilesA, self.vgprTilesB, scaleTileInfoA, scaleTileInfoB, self.vgprTilesSA, self.vgprTilesSB, tensorParametersA=tensorParametersA, tensorParametersB=tensorParametersB, ) # Rebuild all loop variants from current _emitted (which now has # vgpr_tile_maps populated by assign_vgpr_tiles, unlike the stale # copies from build()). self.build_preloop() self.build_ngll() self.build_nll() self.build_tailloop_pgr0() emitter.populate(self._preloop_emitted, unroll_iter=0) self._emitted_per_unroll = [] self._ngll_per_unroll = [] self._nll_per_unroll = [] for ui in range(self.unroll_factor): em_copy = copy.deepcopy(self._emitted) emitter.populate(em_copy, unroll_iter=ui) self._emitted_per_unroll.append(em_copy) ngll_copy = copy.deepcopy(self._ngll_emitted) emitter.populate(ngll_copy, unroll_iter=ui) self._ngll_per_unroll.append(ngll_copy) nll_copy = copy.deepcopy(self._nll_emitted) emitter.populate(nll_copy, unroll_iter=ui) self._nll_per_unroll.append(nll_copy) self._emitter = emitter # ── Print helpers ─────────────────────────────────────── @staticmethod def _fmt_tensor(tensor: str) -> str: """Pad tensor name to 2 chars for alignment: 'A' -> 'A ', 'SA' -> 'SA'.""" return tensor.ljust(2) def print_lr(self, partitions: Optional[LogicalSchedule] = None) -> str: """Print place_LRs output in design doc format.""" if partitions is None: partitions = self._partitions buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, slots in enumerate(partitions): buf.write(f" Partition {pi}:\n") self._print_lr_partition(buf, slots) return buf.getvalue() def _print_lr_partition(self, buf, slots): for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: m = slot.mfma buf.write(f" MFMAs (MT n, subIterK {m.subIterK} ) " f"A : {m.tileA.fmt_tiles()} , B : {m.tileB.fmt_tiles()}\n") for lr in slot.lrs: t = self._fmt_tensor(lr.tensor) buf.write(f" LR {t} (MT {fmt_mt(lr.mtIteration)}, " f"subIterK {lr.tiles.fmt_k()}) " f"{lr.tiles.fmt_tiles()}\n") return buf.getvalue() def print_vgpr(self) -> str: """Print assign_vgpr_tiles output: LRs + MFMAs with vgprTileId annotations.""" partitions = self._partitions buf = io.StringIO() needs = getattr(self, 'needs_unrolling', None) factor = getattr(self, 'unroll_factor', 1) peaks = getattr(self, 'tile_peaks', {}) buf.write(f"needsUnrolling: {needs}, " f"unrollFactor: {factor}\n") peaks_str = ", ".join(f"{t}: {cnt}" for t, cnt in sorted(peaks.items())) buf.write(f"vgprTiles: {peaks_str}\n") for ui in range(factor): if factor > 1: buf.write(f"MAINLOOP (unroll {ui}):\n") else: buf.write("MAINLOOP:\n") for pi, slots in enumerate(partitions): buf.write(f" Partition {pi}:\n") for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: m = slot.mfma tiles_str = "" parts = [] for tensor in self.tensors: maps = m.vgpr_tile_maps.get(tensor) if maps: parts.append(f"{tensor}:" + str(maps[ui])) if parts: tiles_str = " " + ", ".join(parts) buf.write(f" MFMAs (MT n, subIterK {m.subIterK} ) " f"A : {m.tileA.fmt_tiles()} , " f"B : {m.tileB.fmt_tiles()}{tiles_str}\n") for lr in slot.lrs: tile_str = "" if lr.vgpr_tile_map: tile_str = f" tiles:{lr.vgpr_tile_map[ui]}" t = self._fmt_tensor(lr.tensor) buf.write(f" LR {t} (MT {fmt_mt(lr.mtIteration)}, " f"subIterK {lr.tiles.fmt_k()}) " f"{lr.tiles.fmt_tiles()}{tile_str}\n") return buf.getvalue() def print_gr(self) -> str: """Print place_GRs output: LRs + MFMAs + GR placements, all partitions.""" partitions = self._partitions buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, slots in enumerate(partitions): buf.write(f" Partition {pi}:\n") for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: m = slot.mfma buf.write(f" MFMAs (MT n, subIterK {m.subIterK} ) " f"A : {m.tileA.fmt_tiles()} , " f"B : {m.tileB.fmt_tiles()}\n") for lr in slot.lrs: t = self._fmt_tensor(lr.tensor) buf.write(f" LR {t} (MT {fmt_mt(lr.mtIteration)}, " f"subIterK {lr.tiles.fmt_k()}) " f"{lr.tiles.fmt_tiles()}\n") for gr in slot.grs: buf.write(f" GR {gr.tensor} (MT {fmt_mt(gr.mtIteration)}, " f"subIterK {gr.tiles.fmt_k()}) " f"ids {gr.tiles.fmt_tiles()}\n") return buf.getvalue() def print_deps(self) -> str: """Print annotate_deps output: placements with their before-dependencies.""" buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, slots in enumerate(self._partitions): buf.write(f" Partition {pi}:\n") for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: self._print_placement_with_deps(buf, slot.mfma, slot) for lr in slot.lrs: self._print_placement_with_deps(buf, lr, slot) for gr in slot.grs: self._print_placement_with_deps(buf, gr, slot) return buf.getvalue() def _print_placement_with_deps(self, buf, placement, slot: SubIterKSlot): """Print a placement label followed by its deps.""" buf.write(f" {placement}\n") if placement.deps: buf.write(" deps:\n") for dep in placement.deps: dep_str = self._format_dep_ref(dep) buf.write(f" - {dep_str}\n") def print_remove_deps(self) -> str: """Print remove_cross_deps output: placements with preOps and remaining deps.""" buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, slots in enumerate(self._partitions): buf.write(f" Partition {pi}:\n") for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: self._print_placement_with_preops(buf, slot.mfma, slot) for lr in slot.lrs: self._print_placement_with_preops(buf, lr, slot) for gr in slot.grs: self._print_placement_with_preops(buf, gr, slot) return buf.getvalue() def print_group_lr_gr(self) -> str: """Print group_lr_gr output: placements with chained deps and merged preOps.""" buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, slots in enumerate(self._partitions): buf.write(f" Partition {pi}:\n") for slot in slots: buf.write(f" subIterK={slot.subIterK}:\n") if slot.mfma: self._print_placement_with_preops(buf, slot.mfma, slot) for lr in slot.lrs: self._print_placement_with_preops(buf, lr, slot) for gr in slot.grs: self._print_placement_with_preops(buf, gr, slot) return buf.getvalue() def _print_placement_with_preops(self, buf, placement, slot: SubIterKSlot): """Print a placement label followed by its preOps, deps, and postOps.""" buf.write(f" {placement}\n") if placement.preOps: buf.write(" preOps:\n") for op in placement.preOps: buf.write(f" - {op}\n") if placement.deps: buf.write(" deps:\n") for dep in placement.deps: dep_str = self._format_dep_ref(dep) buf.write(f" - {dep_str}\n") if placement.postOps: buf.write(" postOps:\n") for op in placement.postOps: buf.write(f" - {op}\n") def _format_dep_ref(self, dep: Dep) -> str: """Format a Dep for display.""" p = dep.ref slot = p.subIterK_slot if hasattr(p, 'subIterK_slot') else '?' part = p.partition if hasattr(p, 'partition') else 0 kind = 'LR' if isinstance(p, LRPlacement) else 'GR' uid = f" uid={p.unrollId}" if isinstance(p, GRPlacement) and p.unrollId else "" mt = f" (MT{dep.mt_offset})" if dep.mt_offset != 0 else "" return f"{kind} {p.tensor} @P{part}:subIterK={slot}{uid}{mt}" def print_emit(self, all_partitions: Optional[EmittedSchedule] = None) -> str: """Print emit output: EmittedModule list with before-links.""" if all_partitions is None: all_partitions = self._emitted buf = io.StringIO() buf.write("MAINLOOP:\n") for pi, partition_emitted in enumerate(all_partitions): buf.write(f" Partition {pi}:\n") for k, emitted in enumerate(partition_emitted): buf.write(f" subIterK={k}:\n") for em in emitted: before_str = f" <- [{em.before}]" if em.before is not None else "" buf.write(f" [{em.moduleId:2d}] {em.opType:10s} {em.source}{before_str}\n") return buf.getvalue() def _dep_uid_suffix(self, src) -> str: """uid tag for multi-DU dep-path dumps. Surfaces which double-buffer (unrollId) each lr/gr/lr_inc/gr_inc targets, so an inc that advances one uid's pointer is not mistaken for a bug when a following same-MT read belongs to a different uid. Display-only: derived purely from the placement/op and never consulted by codegen. Returns "" for single-DU configs so their dumps are unchanged. """ if not self._is_multi_du(): return "" if isinstance(src, GRPlacement): return f" uid={src.unrollId}" if isinstance(src, (GRIncOp, LRIncOp)): return f" uid={src.unrollId}" if isinstance(src, LRPlacement): per_uid_k = self._per_uid_k(src.tensor) return f" uid={src.tiles.subIterK_start // per_uid_k}" return "" def print_emit_dep_order(self, all_partitions: Optional[EmittedSchedule] = None) -> str: """Print emit output as dependency paths (same decomposition as _extractPathsFromBeforeDeps).""" from Tensile.Components.Subtile.InstructionScheduler import extractPathsFromBeforeDeps if all_partitions is None: all_partitions = self._emitted buf = io.StringIO() buf.write("MAINLOOP (dependency paths):\n") for pi, partition_emitted in enumerate(all_partitions): buf.write(f" Partition {pi}:\n") for k, emitted in enumerate(partition_emitted): buf.write(f" subIterK={k}:\n") mfmaIdx, paths, preMfmaPaths = extractPathsFromBeforeDeps(emitted) em = emitted[mfmaIdx] buf.write(f" MFMA: [{em.moduleId:2d}] {em.source}") if em.before is not None: buf.write(f" <- [{em.before}]") buf.write("\n") for i, path in enumerate(preMfmaPaths): buf.write(f" preMFMA path {i}:\n") for idx in path: src = emitted[idx].source buf.write(f" [{emitted[idx].moduleId:2d}] {emitted[idx].opType:10s} {src}{self._dep_uid_suffix(src)}\n") for i, path in enumerate(paths): buf.write(f" path {i}:\n") for idx in path: src = emitted[idx].source buf.write(f" [{emitted[idx].moduleId:2d}] {emitted[idx].opType:10s} {src}{self._dep_uid_suffix(src)}\n") return buf.getvalue()