python/structural_validator.py

from __future__ import annotations

from collections import Counter, defaultdict
from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

from fragment_map import FragmentElement


@dataclass(frozen=True)
class MfmaShape:
    target: str = "gfx942"
    opcode: str = "v_mfma_f32_16x16x16f16"
    m: int = 16
    n: int = 16
    k: int = 16
    wave_size: int = 64

    @property
    def expected_elements(self) -> Dict[str, int]:
        return {
            "A": self.m * self.k,
            "B": self.k * self.n,
            "C": self.m * self.n,
            "D": self.m * self.n,
        }

    @property
    def operand_bounds(self) -> Dict[str, Tuple[int, int]]:
        return {
            "A": (self.m, self.k),
            "B": (self.k, self.n),
            "C": (self.m, self.n),
            "D": (self.m, self.n),
        }


@dataclass(frozen=True)
class ValidationIssue:
    severity: str # "error" or "warning"
    code: str
    message: str
    operand: Optional[str] = None
    lane: Optional[int] = None
    vgpr: Optional[int] = None
    coordinate: Optional[Tuple[int, int]] = None


class FragmentMapValidationError(ValueError):
    def __init__(self, issues: Sequence[ValidationIssue]) -> None:
        self.issues = tuple(issues)
        errors = [x for x in issues if x.severity == "error"]

        lines = [
            f"MFMA fragment-map validation failed with {len(errors)} error(s)"
        ]
        for issue in errors[:32]:
            where = []
            if issue.operand is not None:
                where.append(f"operand={issue.operand}")
            if issue.lane is not None:
                where.append(f"lane={issue.lane}")
            if issue.vgpr is not None:
                where.append(f"vgpr={issue.vgpr}")
            if issue.coordinate is not None:
                where.append(f"coord={issue.coordinate}")

            suffix = f" ({', '.join(where)})" if where else ""
            lines.append(f"[{issue.code}] {issue.message}{suffix}")

        if len(errors) > 32:
            lines.append(f"... {len(errors) - 32} additional error(s) omitted")

        super().__init__("\n".join(lines))


@dataclass(frozen=True)
class FragmentMapReport:
    shape: MfmaShape
    issues: Tuple[ValidationIssue, ...]
    element_counts: Mapping[str, int]
    unique_coordinate_counts: Mapping[str, int]
    per_lane_element_counts: Mapping[str, Mapping[int, int]]
    per_lane_vgpr_counts: Mapping[str, Mapping[int, int]]

    @property
    def errors(self) -> Tuple[ValidationIssue, ...]:
        return tuple(x for x in self.issues if x.severity == "error")

    @property
    def warnings(self) -> Tuple[ValidationIssue, ...]:
        return tuple(x for x in self.issues if x.severity == "warning")

    @property
    def valid(self) -> bool:
        return not self.errors

    def raise_if_invalid(self) -> None:
        if self.errors:
            raise FragmentMapValidationError(self.issues)


def validate_fragment_map(
    fragments: Mapping[str, Sequence[FragmentElement]],
    *,
    shape: MfmaShape = MfmaShape(),
    strict_register_layout: bool = True,
    require_all_lanes_for_ab: bool = True,
    require_all_lanes_for_cd: bool = True,
    require_c_d_same_layout: bool = True,
) -> FragmentMapReport:
    """
    Validate an imported gfx942 v_mfma_f32_16x16x16f16 fragment map.

    The validator establishes structural facts:

      * A has exactly M*K unique coordinates in [0,M) x [0,K).
      * B has exactly K*N unique coordinates in [0,K) x [0,N).
      * C and D each have exactly M*N unique coordinates in [0,M) x [0,N).
      * All elements identify the correct operand and a lane in [0,wave_size).
      * A/B are packed FP16: each logical element has packed_half in {0,1}.
      * C/D are FP32: packed_half is None.
      * A/B each use exactly 4 FP16 elements per lane for a 16x16x16 tile.
      * C/D each use exactly 4 FP32 elements per lane for a 16x16 output tile.
      * Each lane's A/B halves form valid packed dwords:
            (lane, vgpr) -> exactly one low and one high half.
      * No lane maps two distinct C/D elements to the same accumulator VGPR.
      * C and D use the same lane/VGPR/coordinate ownership map.

    It does NOT claim that a given lane/VGPR/coordinate formula is the
    hardware's canonical MFMA formula. Compare that stronger claim against
    an ISA-calculator export before treating the map as opcode-authoritative.
    """
    issues: List[ValidationIssue] = []

    required_operands = ("A", "B", "C", "D")
    expected_elements = shape.expected_elements
    bounds = shape.operand_bounds

    normalized: Dict[str, List[FragmentElement]] = {}

    # ------------------------------------------------------------------
    # 1. Schema and element-level checks.
    # ------------------------------------------------------------------
    for operand in required_operands:
        if operand not in fragments:
            issues.append(ValidationIssue(
                severity="error",
                code="missing-operand",
                message=f"Fragment map is missing required operand {operand}",
                operand=operand,
            ))
            normalized[operand] = []
            continue

        elems = list(fragments[operand])
        normalized[operand] = elems

        if len(elems) != expected_elements[operand]:
            issues.append(ValidationIssue(
                severity="error",
                code="wrong-element-count",
                message=(
                    f"Expected {expected_elements[operand]} logical elements, "
                    f"found {len(elems)}"
                ),
                operand=operand,
            ))

        row_limit, col_limit = bounds[operand]

        for e in elems:
            if e.operand != operand:
                issues.append(ValidationIssue(
                    severity="error",
                    code="wrong-operand-tag",
                    message=(
                        f"Element appears in {operand} list but has "
                        f"operand tag {e.operand!r}"
                    ),
                    operand=operand,
                    lane=e.lane,
                    vgpr=e.vgpr,
                    coordinate=(e.logical_row, e.logical_col),
                ))

            if not (0 <= e.lane < shape.wave_size):
                issues.append(ValidationIssue(
                    severity="error",
                    code="lane-out-of-range",
                    message=f"Lane must be in [0, {shape.wave_size})",
                    operand=operand,
                    lane=e.lane,
                    vgpr=e.vgpr,
                    coordinate=(e.logical_row, e.logical_col),
                ))

            if e.vgpr < 0:
                issues.append(ValidationIssue(
                    severity="error",
                    code="negative-vgpr",
                    message="VGPR index must be non-negative",
                    operand=operand,
                    lane=e.lane,
                    vgpr=e.vgpr,
                    coordinate=(e.logical_row, e.logical_col),
                ))

            if not (0 <= e.logical_row < row_limit):
                issues.append(ValidationIssue(
                    severity="error",
                    code="row-out-of-range",
                    message=f"Row must be in [0, {row_limit})",
                    operand=operand,
                    lane=e.lane,
                    vgpr=e.vgpr,
                    coordinate=(e.logical_row, e.logical_col),
                ))

            if not (0 <= e.logical_col < col_limit):
                issues.append(ValidationIssue(
                    severity="error",
                    code="column-out-of-range",
                    message=f"Column must be in [0, {col_limit})",
                    operand=operand,
                    lane=e.lane,
                    vgpr=e.vgpr,
                    coordinate=(e.logical_row, e.logical_col),
                ))

            if operand in ("A", "B"):
                if e.packed_half not in (0, 1):
                    issues.append(ValidationIssue(
                        severity="error",
                        code="invalid-fp16-half",
                        message=(
                            "A/B entries must identify packed_half=0 (low) "
                            "or packed_half=1 (high)"
                        ),
                        operand=operand,
                        lane=e.lane,
                        vgpr=e.vgpr,
                        coordinate=(e.logical_row, e.logical_col),
                    ))
            else:
                if e.packed_half is not None:
                    issues.append(ValidationIssue(
                        severity="error",
                        code="invalid-fp32-packing",
                        message=(
                            "C/D entries are FP32 accumulator values and "
                            "must use packed_half=None"
                        ),
                        operand=operand,
                        lane=e.lane,
                        vgpr=e.vgpr,
                        coordinate=(e.logical_row, e.logical_col),
                    ))

    unexpected = sorted(set(fragments) - set(required_operands))
    for operand in unexpected:
        issues.append(ValidationIssue(
            severity="warning",
            code="unexpected-operand",
            message=f"Ignoring unexpected fragment-map operand {operand!r}",
            operand=operand,
        ))

    # ------------------------------------------------------------------
    # 2. Coordinate bijectivity: every logical matrix element must appear
    # exactly once across the wave.
    # ------------------------------------------------------------------
    unique_coordinate_counts: Dict[str, int] = {}

    for operand in required_operands:
        elems = normalized[operand]
        coord_to_entries: Dict[Tuple[int, int], List[FragmentElement]] = defaultdict(list)

        for e in elems:
            coord_to_entries[(e.logical_row, e.logical_col)].append(e)

        unique_coordinate_counts[operand] = len(coord_to_entries)

        row_limit, col_limit = bounds[operand]
        expected_coords = {
            (row, col)
            for row in range(row_limit)
            for col in range(col_limit)
        }

        actual_coords = set(coord_to_entries)
        missing = sorted(expected_coords - actual_coords)
        extra = sorted(actual_coords - expected_coords)

        if missing:
            issues.append(ValidationIssue(
                severity="error",
                code="missing-logical-coordinates",
                message=(
                    f"Map omits {len(missing)} logical coordinate(s); "
                    f"first few: {missing[:8]}"
                ),
                operand=operand,
            ))

        if extra:
            issues.append(ValidationIssue(
                severity="error",
                code="extra-logical-coordinates",
                message=(
                    f"Map contains {len(extra)} out-of-domain coordinate(s); "
                    f"first few: {extra[:8]}"
                ),
                operand=operand,
            ))

        for coordinate, entries in coord_to_entries.items():
            if len(entries) > 1:
                owners = [(e.lane, e.vgpr, e.packed_half) for e in entries]
                issues.append(ValidationIssue(
                    severity="error",
                    code="duplicate-logical-coordinate",
                    message=(
                        f"Logical matrix element has {len(entries)} owners: "
                        f"{owners}"
                    ),
                    operand=operand,
                    coordinate=coordinate,
                ))

    # ------------------------------------------------------------------
    # 3. Per-lane occupancy.
    #
    # For this exact tile:
    # A: 16*16 / 64 = 4 FP16 values per lane
    # B: 16*16 / 64 = 4 FP16 values per lane
    # C: 16*16 / 64 = 4 FP32 values per lane
    # D: 16*16 / 64 = 4 FP32 values per lane
    # ------------------------------------------------------------------
    expected_per_lane = {"A": 4, "B": 4, "C": 4, "D": 4}
    per_lane_element_counts: Dict[str, Dict[int, int]] = {}
    per_lane_vgpr_counts: Dict[str, Dict[int, int]] = {}

    for operand in required_operands:
        elems = normalized[operand]
        counts = Counter(e.lane for e in elems)
        per_lane_element_counts[operand] = {
            lane: counts.get(lane, 0)
            for lane in range(shape.wave_size)
        }

        distinct_vgprs_by_lane: Dict[int, set[int]] = defaultdict(set)
        for e in elems:
            if 0 <= e.lane < shape.wave_size:
                distinct_vgprs_by_lane[e.lane].add(e.vgpr)

        per_lane_vgpr_counts[operand] = {
            lane: len(distinct_vgprs_by_lane.get(lane, set()))
            for lane in range(shape.wave_size)
        }

        require_all = (
            operand in ("A", "B") and require_all_lanes_for_ab
        ) or (
            operand in ("C", "D") and require_all_lanes_for_cd
        )

        for lane in range(shape.wave_size):
            actual = counts.get(lane, 0)

            if require_all and actual != expected_per_lane[operand]:
                issues.append(ValidationIssue(
                    severity="error",
                    code="wrong-per-lane-element-count",
                    message=(
                        f"Expected {expected_per_lane[operand]} elements in "
                        f"this lane, found {actual}"
                    ),
                    operand=operand,
                    lane=lane,
                ))
            elif not require_all and actual not in (0, expected_per_lane[operand]):
                issues.append(ValidationIssue(
                    severity="error",
                    code="partial-lane-fragment",
                    message=(
                        f"Lane owns {actual} values; expected either 0 or "
                        f"{expected_per_lane[operand]}"
                    ),
                    operand=operand,
                    lane=lane,
                ))

    # ------------------------------------------------------------------
    # 4. Packed FP16 register validity for A and B.
    #
    # Every input VGPR dword represented in this model must contain exactly
    # a low and high FP16 value for the same lane. With 4 values/lane this
    # gives exactly two distinct source VGPR dwords per lane.
    # ------------------------------------------------------------------
    for operand in ("A", "B"):
        by_lane_vgpr: Dict[Tuple[int, int], List[FragmentElement]] = defaultdict(list)

        for e in normalized[operand]:
            if 0 <= e.lane < shape.wave_size:
                by_lane_vgpr[(e.lane, e.vgpr)].append(e)

        for lane in range(shape.wave_size):
            lane_regs = [
                vgpr
                for (entry_lane, vgpr) in by_lane_vgpr
                if entry_lane == lane
            ]

            if strict_register_layout and len(lane_regs) != 2:
                issues.append(ValidationIssue(
                    severity="error",
                    code="wrong-input-vgpr-count",
                    message=(
                        "Expected exactly 2 packed-FP16 source VGPR dwords "
                        "for this lane"
                    ),
                    operand=operand,
                    lane=lane,
                ))

        for (lane, vgpr), entries in by_lane_vgpr.items():
            half_counts = Counter(e.packed_half for e in entries)

            if len(entries) != 2:
                issues.append(ValidationIssue(
                    severity="error",
                    code="wrong-packed-vgpr-arity",
                    message=(
                        f"Packed FP16 source VGPR must own exactly 2 logical "
                        f"halves, found {len(entries)}"
                    ),
                    operand=operand,
                    lane=lane,
                    vgpr=vgpr,
                ))
                continue

            if half_counts.get(0, 0) != 1 or half_counts.get(1, 0) != 1:
                issues.append(ValidationIssue(
                    severity="error",
                    code="invalid-packed-half-pair",
                    message=(
                        "Packed FP16 source VGPR must contain exactly one "
                        "low half and one high half"
                    ),
                    operand=operand,
                    lane=lane,
                    vgpr=vgpr,
                ))

    # ------------------------------------------------------------------
    # 5. FP32 accumulator register validity for C and D.
    #
    # A lane owns 4 output values. In the conventional model, they occupy
    # four distinct accumulator-register positions. No two distinct
    # coordinates may alias one (lane, vgpr) location.
    # ------------------------------------------------------------------
    for operand in ("C", "D"):
        by_lane_vgpr: Dict[Tuple[int, int], List[FragmentElement]] = defaultdict(list)

        for e in normalized[operand]:
            if 0 <= e.lane < shape.wave_size:
                by_lane_vgpr[(e.lane, e.vgpr)].append(e)

        for lane in range(shape.wave_size):
            regs = {
                e.vgpr
                for e in normalized[operand]
                if e.lane == lane
            }

            if strict_register_layout and len(regs) != 4:
                issues.append(ValidationIssue(
                    severity="error",
                    code="wrong-accumulator-vgpr-count",
                    message=(
                        "Expected exactly 4 distinct FP32 accumulator VGPRs "
                        "for this lane"
                    ),
                    operand=operand,
                    lane=lane,
                ))

        for (lane, vgpr), entries in by_lane_vgpr.items():
            if len(entries) != 1:
                coords = [(e.logical_row, e.logical_col) for e in entries]
                issues.append(ValidationIssue(
                    severity="error",
                    code="accumulator-vgpr-alias",
                    message=(
                        f"One accumulator VGPR aliases {len(entries)} "
                        f"distinct FP32 values: {coords}"
                    ),
                    operand=operand,
                    lane=lane,
                    vgpr=vgpr,
                ))

    # ------------------------------------------------------------------
    # 6. C/D correspondence.
    #
    # An MFMA updates C into D with identical fragment ownership. The values
    # differ, but (lane, vgpr) -> (logical row, logical column) should match.
    # ------------------------------------------------------------------
    if require_c_d_same_layout:
        def accumulator_ownership(
            entries: Iterable[FragmentElement],
        ) -> Dict[Tuple[int, int], Tuple[int, int]]:
            result: Dict[Tuple[int, int], Tuple[int, int]] = {}

            for e in entries:
                key = (e.lane, e.vgpr)
                value = (e.logical_row, e.logical_col)

                if key not in result:
                    result[key] = value

            return result

        c_layout = accumulator_ownership(normalized["C"])
        d_layout = accumulator_ownership(normalized["D"])

        if c_layout != d_layout:
            c_keys = set(c_layout)
            d_keys = set(d_layout)

            missing_in_d = sorted(c_keys - d_keys)
            extra_in_d = sorted(d_keys - c_keys)
            changed = sorted(
                key for key in (c_keys & d_keys)
                if c_layout[key] != d_layout[key]
            )

            issues.append(ValidationIssue(
                severity="error",
                code="c-d-layout-mismatch",
                message=(
                    "C and D must have identical accumulator ownership; "
                    f"missing-in-D={missing_in_d[:8]}, "
                    f"extra-in-D={extra_in_d[:8]}, "
                    f"changed={[(key, c_layout[key], d_layout[key]) for key in changed[:8]]}"
                ),
            ))

    # ------------------------------------------------------------------
    # 7. Soft checks: source VGPR numbering may be local tuple offsets
    # rather than absolute hardware VGPR IDs. Emit warnings only.
    # ------------------------------------------------------------------
    for operand in ("A", "B", "C", "D"):
        used = sorted({e.vgpr for e in normalized[operand]})

        if not used:
            continue

        contiguous = used == list(range(used[0], used[-1] + 1))
        if not contiguous:
            issues.append(ValidationIssue(
                severity="warning",
                code="noncontiguous-vgpr-numbering",
                message=(
                    f"{operand} uses non-contiguous VGPR indices {used}; "
                    "this can be valid for an absolute register allocation, "
                    "but is unexpected for a compact local fragment tuple"
                ),
                operand=operand,
            ))

    report = FragmentMapReport(
        shape=shape,
        issues=tuple(issues),
        element_counts={
            operand: len(normalized[operand])
            for operand in required_operands
        },
        unique_coordinate_counts=unique_coordinate_counts,
        per_lane_element_counts=per_lane_element_counts,
        per_lane_vgpr_counts=per_lane_vgpr_counts,
    )

    report.raise_if_invalid()
    return report


# -----------------------------
# Example usage
# -----------------------------
if __name__ == "__main__":
    from fragment_map import generate_v_mfma_f32_16x16x16f16_fragments

    frags = generate_v_mfma_f32_16x16x16f16_fragments()

    report = validate_fragment_map(
        frags,
        shape=MfmaShape(
            target="gfx942",
            opcode="v_mfma_f32_16x16x16f16",
            m=16,
            n=16,
            k=16,
            wave_size=64,
        ),
    )

    print("Fragment map is structurally valid.")
    print(f"Element counts: {report.element_counts}")
    print(f"Unique coordinate counts: {report.unique_coordinate_counts}")