working on code

[bigint-presentation-code.git] / src / bigint_presentation_code / compiler_ir2.py

from collections import defaultdict
import enum
from abc import ABCMeta, abstractmethod
from enum import Enum, unique
from functools import lru_cache, total_ordering
from typing import (AbstractSet, Any, Callable, Generic, Iterable, Iterator,
                    Sequence, TypeVar, overload)
from weakref import WeakValueDictionary as _WeakVDict

from cached_property import cached_property
from nmutil.plain_data import fields, plain_data

from bigint_presentation_code.type_util import Self, assert_never, final, Literal
from bigint_presentation_code.util import BitSet, FBitSet, FMap, OFSet, OSet


@final
class Fn:
    def __init__(self):
        self.ops = []  # type: list[Op]
        self.__op_names = _WeakVDict()  # type: _WeakVDict[str, Op]
        self.__next_name_suffix = 2

    def _add_op_with_unused_name(self, op, name=""):
        # type: (Op, str) -> str
        if op.fn is not self:
            raise ValueError("can't add Op to wrong Fn")
        if hasattr(op, "name"):
            raise ValueError("Op already named")
        orig_name = name
        while True:
            if name != "" and name not in self.__op_names:
                self.__op_names[name] = op
                return name
            name = orig_name + str(self.__next_name_suffix)
            self.__next_name_suffix += 1

    def __repr__(self):
        # type: () -> str
        return "<Fn>"

    def append_op(self, op):
        # type: (Op) -> None
        if op.fn is not self:
            raise ValueError("can't add Op to wrong Fn")
        self.ops.append(op)

    def append_new_op(self, kind, input_vals=(), immediates=(), name="",
                      maxvl=1):
        # type: (OpKind, Iterable[SSAVal], Iterable[int], str, int) -> Op
        retval = Op(fn=self, properties=kind.instantiate(maxvl=maxvl),
                    input_vals=input_vals, immediates=immediates, name=name)
        self.append_op(retval)
        return retval

    def pre_ra_sim(self, state):
        # type: (PreRASimState) -> None
        for op in self.ops:
            op.pre_ra_sim(state)

    def pre_ra_insert_copies(self):
        # type: () -> None
        orig_ops = list(self.ops)
        copied_outputs = {}  # type: dict[SSAVal, SSAVal]
        self.ops.clear()
        for op in orig_ops:
            for i in range(len(op.input_vals)):
                inp = copied_outputs[op.input_vals[i]]
                if inp.ty.base_ty is BaseTy.I64:
                    maxvl = inp.ty.reg_len
                    if inp.ty.reg_len != 1:
                        setvl = self.append_new_op(
                            OpKind.SetVLI, immediates=[maxvl],
                            name=f"{op.name}.inp{i}.setvl")
                        vl = setvl.outputs[0]
                        mv = self.append_new_op(
                            OpKind.VecCopyToReg, input_vals=[inp, vl],
                            maxvl=maxvl, name=f"{op.name}.inp{i}.copy")
                    else:
                        mv = self.append_new_op(
                            OpKind.CopyToReg, input_vals=[inp],
                            name=f"{op.name}.inp{i}.copy")
                    op.input_vals[i] = mv.outputs[0]
                elif inp.ty.base_ty is BaseTy.CA \
                        or inp.ty.base_ty is BaseTy.VL_MAXVL:
                    # all copies would be no-ops, so we don't need to copy
                    op.input_vals[i] = inp
                else:
                    assert_never(inp.ty.base_ty)
            self.ops.append(op)
            for i, out in enumerate(op.outputs):
                if out.ty.base_ty is BaseTy.I64:
                    maxvl = out.ty.reg_len
                    if out.ty.reg_len != 1:
                        setvl = self.append_new_op(
                            OpKind.SetVLI, immediates=[maxvl],
                            name=f"{op.name}.out{i}.setvl")
                        vl = setvl.outputs[0]
                        mv = self.append_new_op(
                            OpKind.VecCopyFromReg, input_vals=[out, vl],
                            maxvl=maxvl, name=f"{op.name}.out{i}.copy")
                    else:
                        mv = self.append_new_op(
                            OpKind.CopyFromReg, input_vals=[out],
                            name=f"{op.name}.out{i}.copy")
                    copied_outputs[out] = mv.outputs[0]
                elif out.ty.base_ty is BaseTy.CA \
                        or out.ty.base_ty is BaseTy.VL_MAXVL:
                    # all copies would be no-ops, so we don't need to copy
                    copied_outputs[out] = out
                else:
                    assert_never(out.ty.base_ty)


@final
@unique
@total_ordering
class OpStage(Enum):
    value: Literal[0, 1]  # type: ignore

    def __new__(cls, value):
        # type: (int) -> OpStage
        value = int(value)
        if value not in (0, 1):
            raise ValueError("invalid value")
        retval = object.__new__(cls)
        retval._value_ = value
        return retval

    Early = 0
    """ early stage of Op execution, where all input reads occur.
    all output writes with `write_stage == Early` occur here too, and therefore
    conflict with input reads, telling the compiler that it that can't share
    that output's register with any inputs that the output isn't tied to.

    All outputs, even unused outputs, can't share registers with any other
    outputs, independent of `write_stage` settings.
    """
    Late = 1
    """ late stage of Op execution, where all output writes with
    `write_stage == Late` occur, and therefore don't conflict with input reads,
    telling the compiler that any inputs can safely use the same register as
    those outputs.

    All outputs, even unused outputs, can't share registers with any other
    outputs, independent of `write_stage` settings.
    """

    def __repr__(self):
        # type: () -> str
        return f"OpStage.{self._name_}"

    def __lt__(self, other):
        # type: (OpStage | object) -> bool
        if isinstance(other, OpStage):
            return self.value < other.value
        return NotImplemented


assert OpStage.Early < OpStage.Late, "early must be less than late"


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
@total_ordering
class ProgramPoint:
    __slots__ = "op_index", "stage"

    def __init__(self, op_index, stage):
        # type: (int, OpStage) -> None
        self.op_index = op_index
        self.stage = stage

    @property
    def int_value(self):
        # type: () -> int
        """ an integer representation of `self` such that it keeps ordering and
        successor/predecessor relations.
        """
        return self.op_index * 2 + self.stage.value

    @staticmethod
    def from_int_value(int_value):
        # type: (int) -> ProgramPoint
        op_index, stage = divmod(int_value, 2)
        return ProgramPoint(op_index=op_index, stage=OpStage(stage))

    def next(self, steps=1):
        # type: (int) -> ProgramPoint
        return ProgramPoint.from_int_value(self.int_value + steps)

    def prev(self, steps=1):
        # type: (int) -> ProgramPoint
        return self.next(steps=-steps)

    def __lt__(self, other):
        # type: (ProgramPoint | Any) -> bool
        if not isinstance(other, ProgramPoint):
            return NotImplemented
        if self.op_index != other.op_index:
            return self.op_index < other.op_index
        return self.stage < other.stage

    def __repr__(self):
        # type: () -> str
        return f"<ops[{self.op_index}]:{self.stage._name_}>"


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class ProgramRange(Sequence[ProgramPoint]):
    __slots__ = "start", "stop"

    def __init__(self, start, stop):
        # type: (ProgramPoint, ProgramPoint) -> None
        self.start = start
        self.stop = stop

    @cached_property
    def int_value_range(self):
        # type: () -> range
        return range(self.start.int_value, self.stop.int_value)

    @staticmethod
    def from_int_value_range(int_value_range):
        # type: (range) -> ProgramRange
        if int_value_range.step != 1:
            raise ValueError("int_value_range must have step == 1")
        return ProgramRange(
            start=ProgramPoint.from_int_value(int_value_range.start),
            stop=ProgramPoint.from_int_value(int_value_range.stop))

    @overload
    def __getitem__(self, __idx):
        # type: (int) -> ProgramPoint
        ...

    @overload
    def __getitem__(self, __idx):
        # type: (slice) -> ProgramRange
        ...

    def __getitem__(self, __idx):
        # type: (int | slice) -> ProgramPoint | ProgramRange
        v = range(self.start.int_value, self.stop.int_value)[__idx]
        if isinstance(v, int):
            return ProgramPoint.from_int_value(v)
        return ProgramRange.from_int_value_range(v)

    def __len__(self):
        # type: () -> int
        return len(self.int_value_range)

    def __iter__(self):
        # type: () -> Iterator[ProgramPoint]
        return map(ProgramPoint.from_int_value, self.int_value_range)

    def __repr__(self):
        # type: () -> str
        start = repr(self.start).lstrip("<").rstrip(">")
        stop = repr(self.stop).lstrip("<").rstrip(">")
        return f"<range:{start}..{stop}>"


@plain_data(frozen=True, eq=False)
@final
class FnAnalysis:
    __slots__ = ("fn", "uses", "op_indexes", "live_ranges", "live_at",
                 "def_program_ranges", "use_program_points",
                 "all_program_points")

    def __init__(self, fn):
        # type: (Fn) -> None
        self.fn = fn
        self.op_indexes = FMap((op, idx) for idx, op in enumerate(fn.ops))
        self.all_program_points = ProgramRange(
            start=ProgramPoint(op_index=0, stage=OpStage.Early),
            stop=ProgramPoint(op_index=len(fn.ops), stage=OpStage.Early))
        def_program_ranges = {}  # type: dict[SSAVal, ProgramRange]
        use_program_points = {}  # type: dict[SSAUse, ProgramPoint]
        uses = {}  # type: dict[SSAVal, OSet[SSAUse]]
        live_range_stops = {}  # type: dict[SSAVal, ProgramPoint]
        for op in fn.ops:
            for use in op.input_uses:
                uses[use.ssa_val].add(use)
                use_program_point = self.__get_use_program_point(use)
                use_program_points[use] = use_program_point
                live_range_stops[use.ssa_val] = max(
                    live_range_stops[use.ssa_val], use_program_point.next())
            for out in op.outputs:
                uses[out] = OSet()
                def_program_range = self.__get_def_program_range(out)
                def_program_ranges[out] = def_program_range
                live_range_stops[out] = def_program_range.stop
        self.uses = FMap((k, OFSet(v)) for k, v in uses.items())
        self.def_program_ranges = FMap(def_program_ranges)
        self.use_program_points = FMap(use_program_points)
        live_ranges = {}  # type: dict[SSAVal, ProgramRange]
        live_at = {i: OSet[SSAVal]() for i in self.all_program_points}
        for ssa_val in uses.keys():
            live_ranges[ssa_val] = live_range = ProgramRange(
                start=self.def_program_ranges[ssa_val].start,
                stop=live_range_stops[ssa_val])
            for program_point in live_range:
                live_at[program_point].add(ssa_val)
        self.live_ranges = FMap(live_ranges)
        self.live_at = FMap((k, OFSet(v)) for k, v in live_at.items())

    def __get_def_program_range(self, ssa_val):
        # type: (SSAVal) -> ProgramRange
        write_stage = ssa_val.defining_descriptor.write_stage
        start = ProgramPoint(
            op_index=self.op_indexes[ssa_val.op], stage=write_stage)
        # always include late stage of ssa_val.op, to ensure outputs always
        # overlap all other outputs.
        # stop is exclusive, so we need the next program point.
        stop = ProgramPoint(op_index=start.op_index, stage=OpStage.Late).next()
        return ProgramRange(start=start, stop=stop)

    def __get_use_program_point(self, ssa_use):
        # type: (SSAUse) -> ProgramPoint
        assert ssa_use.defining_descriptor.write_stage is OpStage.Early, \
            "assumed here, ensured by GenericOpProperties.__init__"
        return ProgramPoint(
            op_index=self.op_indexes[ssa_use.op], stage=OpStage.Early)

    def __eq__(self, other):
        # type: (FnAnalysis | Any) -> bool
        if isinstance(other, FnAnalysis):
            return self.fn == other.fn
        return NotImplemented

    def __hash__(self):
        # type: () -> int
        return hash(self.fn)


@unique
@final
class BaseTy(Enum):
    I64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def only_scalar(self):
        # type: () -> bool
        if self is BaseTy.I64:
            return False
        elif self is BaseTy.CA or self is BaseTy.VL_MAXVL:
            return True
        else:
            assert_never(self)

    @cached_property
    def max_reg_len(self):
        # type: () -> int
        if self is BaseTy.I64:
            return 128
        elif self is BaseTy.CA or self is BaseTy.VL_MAXVL:
            return 1
        else:
            assert_never(self)

    def __repr__(self):
        return "BaseTy." + self._name_


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class Ty:
    __slots__ = "base_ty", "reg_len"

    @staticmethod
    def validate(base_ty, reg_len):
        # type: (BaseTy, int) -> str | None
        """ return a string with the error if the combination is invalid,
        otherwise return None
        """
        if base_ty.only_scalar and reg_len != 1:
            return f"can't create a vector of an only-scalar type: {base_ty}"
        if reg_len < 1 or reg_len > base_ty.max_reg_len:
            return "reg_len out of range"
        return None

    def __init__(self, base_ty, reg_len):
        # type: (BaseTy, int) -> None
        msg = self.validate(base_ty=base_ty, reg_len=reg_len)
        if msg is not None:
            raise ValueError(msg)
        self.base_ty = base_ty
        self.reg_len = reg_len

    def __repr__(self):
        # type: () -> str
        if self.reg_len != 1:
            reg_len = f"*{self.reg_len}"
        else:
            reg_len = ""
        return f"<{self.base_ty._name_}{reg_len}>"


@unique
@final
class LocKind(Enum):
    GPR = enum.auto()
    StackI64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def base_ty(self):
        # type: () -> BaseTy
        if self is LocKind.GPR or self is LocKind.StackI64:
            return BaseTy.I64
        if self is LocKind.CA:
            return BaseTy.CA
        if self is LocKind.VL_MAXVL:
            return BaseTy.VL_MAXVL
        else:
            assert_never(self)

    @cached_property
    def loc_count(self):
        # type: () -> int
        if self is LocKind.StackI64:
            return 1024
        if self is LocKind.GPR or self is LocKind.CA \
                or self is LocKind.VL_MAXVL:
            return self.base_ty.max_reg_len
        else:
            assert_never(self)

    def __repr__(self):
        return "LocKind." + self._name_


@final
@unique
class LocSubKind(Enum):
    BASE_GPR = enum.auto()
    SV_EXTRA2_VGPR = enum.auto()
    SV_EXTRA2_SGPR = enum.auto()
    SV_EXTRA3_VGPR = enum.auto()
    SV_EXTRA3_SGPR = enum.auto()
    StackI64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def kind(self):
        # type: () -> LocKind
        # pyright fails typechecking when using `in` here:
        # reported: https://github.com/microsoft/pyright/issues/4102
        if self in (LocSubKind.BASE_GPR, LocSubKind.SV_EXTRA2_VGPR,
                    LocSubKind.SV_EXTRA2_SGPR, LocSubKind.SV_EXTRA3_VGPR,
                    LocSubKind.SV_EXTRA3_SGPR):
            return LocKind.GPR
        if self is LocSubKind.StackI64:
            return LocKind.StackI64
        if self is LocSubKind.CA:
            return LocKind.CA
        if self is LocSubKind.VL_MAXVL:
            return LocKind.VL_MAXVL
        assert_never(self)

    @property
    def base_ty(self):
        return self.kind.base_ty

    @lru_cache()
    def allocatable_locs(self, ty):
        # type: (Ty) -> LocSet
        if ty.base_ty != self.base_ty:
            raise ValueError("type mismatch")
        if self is LocSubKind.BASE_GPR:
            starts = range(32)
        elif self is LocSubKind.SV_EXTRA2_VGPR:
            starts = range(0, 128, 2)
        elif self is LocSubKind.SV_EXTRA2_SGPR:
            starts = range(64)
        elif self is LocSubKind.SV_EXTRA3_VGPR \
                or self is LocSubKind.SV_EXTRA3_SGPR:
            starts = range(128)
        elif self is LocSubKind.StackI64:
            starts = range(LocKind.StackI64.loc_count)
        elif self is LocSubKind.CA or self is LocSubKind.VL_MAXVL:
            return LocSet([Loc(kind=self.kind, start=0, reg_len=1)])
        else:
            assert_never(self)
        retval = []  # type: list[Loc]
        for start in starts:
            loc = Loc.try_make(kind=self.kind, start=start, reg_len=ty.reg_len)
            if loc is None:
                continue
            conflicts = False
            for special_loc in SPECIAL_GPRS:
                if loc.conflicts(special_loc):
                    conflicts = True
                    break
            if not conflicts:
                retval.append(loc)
        return LocSet(retval)

    def __repr__(self):
        return "LocSubKind." + self._name_


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericTy:
    __slots__ = "base_ty", "is_vec"

    def __init__(self, base_ty, is_vec):
        # type: (BaseTy, bool) -> None
        self.base_ty = base_ty
        if base_ty.only_scalar and is_vec:
            raise ValueError(f"base_ty={base_ty} requires is_vec=False")
        self.is_vec = is_vec

    def instantiate(self, maxvl):
        # type: (int) -> Ty
        # here's where subvl and elwid would be accounted for
        if self.is_vec:
            return Ty(self.base_ty, maxvl)
        return Ty(self.base_ty, 1)

    def can_instantiate_to(self, ty):
        # type: (Ty) -> bool
        if self.base_ty != ty.base_ty:
            return False
        if self.is_vec:
            return True
        return ty.reg_len == 1


@plain_data(frozen=True, unsafe_hash=True)
@final
class Loc:
    __slots__ = "kind", "start", "reg_len"

    @staticmethod
    def validate(kind, start, reg_len):
        # type: (LocKind, int, int) -> str | None
        msg = Ty.validate(base_ty=kind.base_ty, reg_len=reg_len)
        if msg is not None:
            return msg
        if reg_len > kind.loc_count:
            return "invalid reg_len"
        if start < 0 or start + reg_len > kind.loc_count:
            return "start not in valid range"
        return None

    @staticmethod
    def try_make(kind, start, reg_len):
        # type: (LocKind, int, int) -> Loc | None
        msg = Loc.validate(kind=kind, start=start, reg_len=reg_len)
        if msg is not None:
            return None
        return Loc(kind=kind, start=start, reg_len=reg_len)

    def __init__(self, kind, start, reg_len):
        # type: (LocKind, int, int) -> None
        msg = self.validate(kind=kind, start=start, reg_len=reg_len)
        if msg is not None:
            raise ValueError(msg)
        self.kind = kind
        self.reg_len = reg_len
        self.start = start

    def conflicts(self, other):
        # type: (Loc) -> bool
        return (self.kind == other.kind
                and self.start < other.stop and other.start < self.stop)

    @staticmethod
    def make_ty(kind, reg_len):
        # type: (LocKind, int) -> Ty
        return Ty(base_ty=kind.base_ty, reg_len=reg_len)

    @cached_property
    def ty(self):
        # type: () -> Ty
        return self.make_ty(kind=self.kind, reg_len=self.reg_len)

    @property
    def stop(self):
        # type: () -> int
        return self.start + self.reg_len

    def try_concat(self, *others):
        # type: (*Loc | None) -> Loc | None
        reg_len = self.reg_len
        stop = self.stop
        for other in others:
            if other is None or other.kind != self.kind:
                return None
            if stop != other.start:
                return None
            stop = other.stop
            reg_len += other.reg_len
        return Loc(kind=self.kind, start=self.start, reg_len=reg_len)


SPECIAL_GPRS = (
    Loc(kind=LocKind.GPR, start=0, reg_len=1),
    Loc(kind=LocKind.GPR, start=1, reg_len=1),
    Loc(kind=LocKind.GPR, start=2, reg_len=1),
    Loc(kind=LocKind.GPR, start=13, reg_len=1),
)


@plain_data(frozen=True, eq=False)
@final
class LocSet(AbstractSet[Loc]):
    __slots__ = "starts", "ty"

    def __init__(self, __locs=()):
        # type: (Iterable[Loc]) -> None
        if isinstance(__locs, LocSet):
            self.starts = __locs.starts  # type: FMap[LocKind, FBitSet]
            self.ty = __locs.ty  # type: Ty | None
            return
        starts = {i: BitSet() for i in LocKind}
        ty = None
        for loc in __locs:
            if ty is None:
                ty = loc.ty
            if ty != loc.ty:
                raise ValueError(f"conflicting types: {ty} != {loc.ty}")
            starts[loc.kind].add(loc.start)
        self.starts = FMap(
            (k, FBitSet(v)) for k, v in starts.items() if len(v) != 0)
        self.ty = ty

    @cached_property
    def stops(self):
        # type: () -> FMap[LocKind, FBitSet]
        if self.ty is None:
            return FMap()
        sh = self.ty.reg_len
        return FMap(
            (k, FBitSet(bits=v.bits << sh)) for k, v in self.starts.items())

    @property
    def kinds(self):
        # type: () -> AbstractSet[LocKind]
        return self.starts.keys()

    @property
    def reg_len(self):
        # type: () -> int | None
        if self.ty is None:
            return None
        return self.ty.reg_len

    @property
    def base_ty(self):
        # type: () -> BaseTy | None
        if self.ty is None:
            return None
        return self.ty.base_ty

    def concat(self, *others):
        # type: (*LocSet) -> LocSet
        if self.ty is None:
            return LocSet()
        base_ty = self.ty.base_ty
        reg_len = self.ty.reg_len
        starts = {k: BitSet(v) for k, v in self.starts.items()}
        for other in others:
            if other.ty is None:
                return LocSet()
            if other.ty.base_ty != base_ty:
                return LocSet()
            for kind, other_starts in other.starts.items():
                if kind not in starts:
                    continue
                starts[kind].bits &= other_starts.bits >> reg_len
                if starts[kind] == 0:
                    del starts[kind]
                    if len(starts) == 0:
                        return LocSet()
            reg_len += other.ty.reg_len

        def locs():
            # type: () -> Iterable[Loc]
            for kind, v in starts.items():
                for start in v:
                    loc = Loc.try_make(kind=kind, start=start, reg_len=reg_len)
                    if loc is not None:
                        yield loc
        return LocSet(locs())

    def __contains__(self, loc):
        # type: (Loc | Any) -> bool
        if not isinstance(loc, Loc) or loc.ty != self.ty:
            return False
        if loc.kind not in self.starts:
            return False
        return loc.start in self.starts[loc.kind]

    def __iter__(self):
        # type: () -> Iterator[Loc]
        if self.ty is None:
            return
        for kind, starts in self.starts.items():
            for start in starts:
                yield Loc(kind=kind, start=start, reg_len=self.ty.reg_len)

    @cached_property
    def __len(self):
        return sum((len(v) for v in self.starts.values()), 0)

    def __len__(self):
        return self.__len

    @cached_property
    def __hash(self):
        return super()._hash()

    def __hash__(self):
        return self.__hash

    @lru_cache(maxsize=None, typed=True)
    def max_conflicts_with(self, other):
        # type: (LocSet | Loc) -> int
        """the largest number of Locs in `self` that a single Loc
        from `other` can conflict with
        """
        if isinstance(other, LocSet):
            return max(self.max_conflicts_with(i) for i in other)
        else:
            return sum(other.conflicts(i) for i in self)


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericOperandDesc:
    """generic Op operand descriptor"""
    __slots__ = ("ty", "fixed_loc", "sub_kinds", "tied_input_index", "spread",
                 "write_stage")

    def __init__(
        self, ty,  # type: GenericTy
        sub_kinds,  # type: Iterable[LocSubKind]
        *,
        fixed_loc=None,  # type: Loc | None
        tied_input_index=None,  # type: int | None
        spread=False,  # type: bool
        write_stage=OpStage.Early,  # type: OpStage
    ):
        # type: (...) -> None
        self.ty = ty
        self.sub_kinds = OFSet(sub_kinds)
        if len(self.sub_kinds) == 0:
            raise ValueError("sub_kinds can't be empty")
        self.fixed_loc = fixed_loc
        if fixed_loc is not None:
            if tied_input_index is not None:
                raise ValueError("operand can't be both tied and fixed")
            if not ty.can_instantiate_to(fixed_loc.ty):
                raise ValueError(
                    f"fixed_loc has incompatible type for given generic "
                    f"type: fixed_loc={fixed_loc} generic ty={ty}")
            if len(self.sub_kinds) != 1:
                raise ValueError(
                    "multiple sub_kinds not allowed for fixed operand")
            for sub_kind in self.sub_kinds:
                if fixed_loc not in sub_kind.allocatable_locs(fixed_loc.ty):
                    raise ValueError(
                        f"fixed_loc not in given sub_kind: "
                        f"fixed_loc={fixed_loc} sub_kind={sub_kind}")
        for sub_kind in self.sub_kinds:
            if sub_kind.base_ty != ty.base_ty:
                raise ValueError(f"sub_kind is incompatible with type: "
                                 f"sub_kind={sub_kind} ty={ty}")
        if tied_input_index is not None and tied_input_index < 0:
            raise ValueError("invalid tied_input_index")
        self.tied_input_index = tied_input_index
        self.spread = spread
        if spread:
            if self.tied_input_index is not None:
                raise ValueError("operand can't be both spread and tied")
            if self.fixed_loc is not None:
                raise ValueError("operand can't be both spread and fixed")
            if self.ty.is_vec:
                raise ValueError("operand can't be both spread and vector")
        self.write_stage = write_stage

    def tied_to_input(self, tied_input_index):
        # type: (int) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds,
                                  tied_input_index=tied_input_index,
                                  write_stage=self.write_stage)

    def with_fixed_loc(self, fixed_loc):
        # type: (Loc) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds, fixed_loc=fixed_loc,
                                  write_stage=self.write_stage)

    def with_write_stage(self, write_stage):
        # type: (OpStage) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds,
                                  fixed_loc=self.fixed_loc,
                                  tied_input_index=self.tied_input_index,
                                  spread=self.spread,
                                  write_stage=write_stage)

    def instantiate(self, maxvl):
        # type: (int) -> Iterable[OperandDesc]
        # assumes all spread operands have ty.reg_len = 1
        rep_count = 1
        if self.spread:
            rep_count = maxvl
            maxvl = 1
        ty = self.ty.instantiate(maxvl=maxvl)

        def locs():
            # type: () -> Iterable[Loc]
            if self.fixed_loc is not None:
                if ty != self.fixed_loc.ty:
                    raise ValueError(
                        f"instantiation failed: type mismatch with fixed_loc: "
                        f"instantiated type: {ty} fixed_loc: {self.fixed_loc}")
                yield self.fixed_loc
                return
            for sub_kind in self.sub_kinds:
                yield from sub_kind.allocatable_locs(ty)
        loc_set_before_spread = LocSet(locs())
        for idx in range(rep_count):
            if not self.spread:
                idx = None
            yield OperandDesc(loc_set_before_spread=loc_set_before_spread,
                              tied_input_index=self.tied_input_index,
                              spread_index=idx, write_stage=self.write_stage)


@plain_data(frozen=True, unsafe_hash=True)
@final
class OperandDesc:
    """Op operand descriptor"""
    __slots__ = ("loc_set_before_spread", "tied_input_index", "spread_index",
                 "write_stage")

    def __init__(self, loc_set_before_spread, tied_input_index, spread_index,
                 write_stage):
        # type: (LocSet, int | None, int | None, OpStage) -> None
        if len(loc_set_before_spread) == 0:
            raise ValueError("loc_set_before_spread must not be empty")
        self.loc_set_before_spread = loc_set_before_spread
        self.tied_input_index = tied_input_index
        if self.tied_input_index is not None and self.spread_index is not None:
            raise ValueError("operand can't be both spread and tied")
        self.spread_index = spread_index
        self.write_stage = write_stage

    @cached_property
    def ty_before_spread(self):
        # type: () -> Ty
        ty = self.loc_set_before_spread.ty
        assert ty is not None, (
            "__init__ checked that the LocSet isn't empty, "
            "non-empty LocSets should always have ty set")
        return ty

    @cached_property
    def ty(self):
        """ Ty after any spread is applied """
        if self.spread_index is not None:
            # assumes all spread operands have ty.reg_len = 1
            return Ty(base_ty=self.ty_before_spread.base_ty, reg_len=1)
        return self.ty_before_spread

    @property
    def reg_offset_in_unspread(self):
        """ the number of reg-sized slots in the unspread Loc before self's Loc

        e.g. if the unspread Loc containing self is:
        `Loc(kind=LocKind.GPR, start=8, reg_len=4)`
        and self's Loc is `Loc(kind=LocKind.GPR, start=10, reg_len=1)`
        then reg_offset_into_unspread == 2 == 10 - 8
        """
        if self.spread_index is None:
            return 0
        return self.spread_index * self.ty.reg_len


OD_BASE_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.BASE_GPR])
OD_EXTRA3_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.SV_EXTRA3_SGPR])
OD_EXTRA3_VGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=True),
    sub_kinds=[LocSubKind.SV_EXTRA3_VGPR])
OD_EXTRA2_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.SV_EXTRA2_SGPR])
OD_EXTRA2_VGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=True),
    sub_kinds=[LocSubKind.SV_EXTRA2_VGPR])
OD_CA = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.CA, is_vec=False),
    sub_kinds=[LocSubKind.CA])
OD_VL = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.VL_MAXVL, is_vec=False),
    sub_kinds=[LocSubKind.VL_MAXVL])


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericOpProperties:
    __slots__ = ("demo_asm", "inputs", "outputs", "immediates",
                 "is_copy", "is_load_immediate", "has_side_effects")

    def __init__(
        self, demo_asm,  # type: str
        inputs,  # type: Iterable[GenericOperandDesc]
        outputs,  # type: Iterable[GenericOperandDesc]
        immediates=(),  # type: Iterable[range]
        is_copy=False,  # type: bool
        is_load_immediate=False,  # type: bool
        has_side_effects=False,  # type: bool
    ):
        # type: (...) -> None
        self.demo_asm = demo_asm  # type: str
        self.inputs = tuple(inputs)  # type: tuple[GenericOperandDesc, ...]
        for inp in self.inputs:
            if inp.tied_input_index is not None:
                raise ValueError(
                    f"tied_input_index is not allowed on inputs: {inp}")
            if inp.write_stage is not OpStage.Early:
                raise ValueError(
                    f"write_stage is not allowed on inputs: {inp}")
        self.outputs = tuple(outputs)  # type: tuple[GenericOperandDesc, ...]
        fixed_locs = []  # type: list[tuple[Loc, int]]
        for idx, out in enumerate(self.outputs):
            if out.tied_input_index is not None:
                if out.tied_input_index >= len(self.inputs):
                    raise ValueError(f"tied_input_index out of range: {out}")
                tied_inp = self.inputs[out.tied_input_index]
                if tied_inp.tied_to_input(out.tied_input_index) != out:
                    raise ValueError(f"output can't be tied to non-equivalent "
                                     f"input: {out} tied to {tied_inp}")
            if out.fixed_loc is not None:
                for other_fixed_loc, other_idx in fixed_locs:
                    if not other_fixed_loc.conflicts(out.fixed_loc):
                        continue
                    raise ValueError(
                        f"conflicting fixed_locs: outputs[{idx}] and "
                        f"outputs[{other_idx}]: {out.fixed_loc} conflicts "
                        f"with {other_fixed_loc}")
                fixed_locs.append((out.fixed_loc, idx))
        self.immediates = tuple(immediates)  # type: tuple[range, ...]
        self.is_copy = is_copy  # type: bool
        self.is_load_immediate = is_load_immediate  # type: bool
        self.has_side_effects = has_side_effects  # type: bool


@plain_data(frozen=True, unsafe_hash=True)
@final
class OpProperties:
    __slots__ = "kind", "inputs", "outputs", "maxvl"

    def __init__(self, kind, maxvl):
        # type: (OpKind, int) -> None
        self.kind = kind  # type: OpKind
        inputs = []  # type: list[OperandDesc]
        for inp in self.generic.inputs:
            inputs.extend(inp.instantiate(maxvl=maxvl))
        self.inputs = tuple(inputs)  # type: tuple[OperandDesc, ...]
        outputs = []  # type: list[OperandDesc]
        for out in self.generic.outputs:
            outputs.extend(out.instantiate(maxvl=maxvl))
        self.outputs = tuple(outputs)  # type: tuple[OperandDesc, ...]
        self.maxvl = maxvl  # type: int

    @property
    def generic(self):
        # type: () -> GenericOpProperties
        return self.kind.properties

    @property
    def immediates(self):
        # type: () -> tuple[range, ...]
        return self.generic.immediates

    @property
    def demo_asm(self):
        # type: () -> str
        return self.generic.demo_asm

    @property
    def is_copy(self):
        # type: () -> bool
        return self.generic.is_copy

    @property
    def is_load_immediate(self):
        # type: () -> bool
        return self.generic.is_load_immediate

    @property
    def has_side_effects(self):
        # type: () -> bool
        return self.generic.has_side_effects


IMM_S16 = range(-1 << 15, 1 << 15)

_PRE_RA_SIM_FN = Callable[["Op", "PreRASimState"], None]
_PRE_RA_SIM_FN2 = Callable[[], _PRE_RA_SIM_FN]
_PRE_RA_SIMS = {}  # type: dict[GenericOpProperties | Any, _PRE_RA_SIM_FN2]


@unique
@final
class OpKind(Enum):
    def __init__(self, properties):
        # type: (GenericOpProperties) -> None
        super().__init__()
        self.__properties = properties

    @property
    def properties(self):
        # type: () -> GenericOpProperties
        return self.__properties

    def instantiate(self, maxvl):
        # type: (int) -> OpProperties
        return OpProperties(self, maxvl=maxvl)

    def __repr__(self):
        # type: () -> str
        return "OpKind." + self._name_

    @cached_property
    def pre_ra_sim(self):
        # type: () -> _PRE_RA_SIM_FN
        return _PRE_RA_SIMS[self.properties]()

    @staticmethod
    def __clearca_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = False,
    ClearCA = GenericOpProperties(
        demo_asm="addic 0, 0, 0",
        inputs=[],
        outputs=[OD_CA.with_write_stage(OpStage.Late)],
    )
    _PRE_RA_SIMS[ClearCA] = lambda: OpKind.__clearca_pre_ra_sim

    @staticmethod
    def __setca_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = True,
    SetCA = GenericOpProperties(
        demo_asm="subfc 0, 0, 0",
        inputs=[],
        outputs=[OD_CA.with_write_stage(OpStage.Late)],
    )
    _PRE_RA_SIMS[SetCA] = lambda: OpKind.__setca_pre_ra_sim

    @staticmethod
    def __svadde_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RA = state.ssa_vals[op.input_vals[0]]
        RB = state.ssa_vals[op.input_vals[1]]
        carry, = state.ssa_vals[op.input_vals[2]]
        VL, = state.ssa_vals[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = RA[i] + RB[i] + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = (v >> GPR_SIZE_IN_BITS) != 0
        state.ssa_vals[op.outputs[0]] = tuple(RT)
        state.ssa_vals[op.outputs[1]] = carry,
    SvAddE = GenericOpProperties(
        demo_asm="sv.adde *RT, *RA, *RB",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_VGPR, OD_CA, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_CA],
    )
    _PRE_RA_SIMS[SvAddE] = lambda: OpKind.__svadde_pre_ra_sim

    @staticmethod
    def __svsubfe_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RA = state.ssa_vals[op.input_vals[0]]
        RB = state.ssa_vals[op.input_vals[1]]
        carry, = state.ssa_vals[op.input_vals[2]]
        VL, = state.ssa_vals[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = (~RA[i] & GPR_VALUE_MASK) + RB[i] + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = (v >> GPR_SIZE_IN_BITS) != 0
        state.ssa_vals[op.outputs[0]] = tuple(RT)
        state.ssa_vals[op.outputs[1]] = carry,
    SvSubFE = GenericOpProperties(
        demo_asm="sv.subfe *RT, *RA, *RB",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_VGPR, OD_CA, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_CA],
    )
    _PRE_RA_SIMS[SvSubFE] = lambda: OpKind.__svsubfe_pre_ra_sim

    @staticmethod
    def __svmaddedu_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RA = state.ssa_vals[op.input_vals[0]]
        RB, = state.ssa_vals[op.input_vals[1]]
        carry, = state.ssa_vals[op.input_vals[2]]
        VL, = state.ssa_vals[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = RA[i] * RB + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = v >> GPR_SIZE_IN_BITS
        state.ssa_vals[op.outputs[0]] = tuple(RT)
        state.ssa_vals[op.outputs[1]] = carry,
    SvMAddEDU = GenericOpProperties(
        demo_asm="sv.maddedu *RT, *RA, RB, RC",
        inputs=[OD_EXTRA2_VGPR, OD_EXTRA2_SGPR, OD_EXTRA2_SGPR, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_EXTRA2_SGPR.tied_to_input(2)],
    )
    _PRE_RA_SIMS[SvMAddEDU] = lambda: OpKind.__svmaddedu_pre_ra_sim

    @staticmethod
    def __setvli_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = op.immediates[0],
    SetVLI = GenericOpProperties(
        demo_asm="setvl 0, 0, imm, 0, 1, 1",
        inputs=(),
        outputs=[OD_VL.with_write_stage(OpStage.Late)],
        immediates=[range(1, 65)],
        is_load_immediate=True,
    )
    _PRE_RA_SIMS[SetVLI] = lambda: OpKind.__setvli_pre_ra_sim

    @staticmethod
    def __svli_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        VL, = state.ssa_vals[op.input_vals[0]]
        imm = op.immediates[0] & GPR_VALUE_MASK
        state.ssa_vals[op.outputs[0]] = (imm,) * VL
    SvLI = GenericOpProperties(
        demo_asm="sv.addi *RT, 0, imm",
        inputs=[OD_VL],
        outputs=[OD_EXTRA3_VGPR],
        immediates=[IMM_S16],
        is_load_immediate=True,
    )
    _PRE_RA_SIMS[SvLI] = lambda: OpKind.__svli_pre_ra_sim

    @staticmethod
    def __li_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        imm = op.immediates[0] & GPR_VALUE_MASK
        state.ssa_vals[op.outputs[0]] = imm,
    LI = GenericOpProperties(
        demo_asm="addi RT, 0, imm",
        inputs=(),
        outputs=[OD_BASE_SGPR.with_write_stage(OpStage.Late)],
        immediates=[IMM_S16],
        is_load_immediate=True,
    )
    _PRE_RA_SIMS[LI] = lambda: OpKind.__li_pre_ra_sim

    @staticmethod
    def __veccopytoreg_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = state.ssa_vals[op.input_vals[0]]
    VecCopyToReg = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=True),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR, LocSubKind.StackI64],
        ), OD_VL],
        outputs=[OD_EXTRA3_VGPR.with_write_stage(OpStage.Late)],
        is_copy=True,
    )
    _PRE_RA_SIMS[VecCopyToReg] = lambda: OpKind.__veccopytoreg_pre_ra_sim

    @staticmethod
    def __veccopyfromreg_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = state.ssa_vals[op.input_vals[0]]
    VecCopyFromReg = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[OD_EXTRA3_VGPR, OD_VL],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=True),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR, LocSubKind.StackI64],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _PRE_RA_SIMS[VecCopyFromReg] = lambda: OpKind.__veccopyfromreg_pre_ra_sim

    @staticmethod
    def __copytoreg_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = state.ssa_vals[op.input_vals[0]]
    CopyToReg = GenericOpProperties(
        demo_asm="mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR,
                       LocSubKind.StackI64],
        )],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _PRE_RA_SIMS[CopyToReg] = lambda: OpKind.__copytoreg_pre_ra_sim

    @staticmethod
    def __copyfromreg_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = state.ssa_vals[op.input_vals[0]]
    CopyFromReg = GenericOpProperties(
        demo_asm="mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR],
        )],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR,
                       LocSubKind.StackI64],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _PRE_RA_SIMS[CopyFromReg] = lambda: OpKind.__copyfromreg_pre_ra_sim

    @staticmethod
    def __concat_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        state.ssa_vals[op.outputs[0]] = tuple(
            state.ssa_vals[i][0] for i in op.input_vals[:-1])
    Concat = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR],
            spread=True,
        ), OD_VL],
        outputs=[OD_EXTRA3_VGPR.with_write_stage(OpStage.Late)],
        is_copy=True,
    )
    _PRE_RA_SIMS[Concat] = lambda: OpKind.__concat_pre_ra_sim

    @staticmethod
    def __spread_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        for idx, inp in enumerate(state.ssa_vals[op.input_vals[0]]):
            state.ssa_vals[op.outputs[idx]] = inp,
    Spread = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[OD_EXTRA3_VGPR, OD_VL],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR],
            spread=True,
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _PRE_RA_SIMS[Spread] = lambda: OpKind.__spread_pre_ra_sim

    @staticmethod
    def __svld_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RA, = state.ssa_vals[op.input_vals[0]]
        VL, = state.ssa_vals[op.input_vals[1]]
        addr = RA + op.immediates[0]
        RT = []  # type: list[int]
        for i in range(VL):
            v = state.load(addr + GPR_SIZE_IN_BYTES * i)
            RT.append(v & GPR_VALUE_MASK)
        state.ssa_vals[op.outputs[0]] = tuple(RT)
    SvLd = GenericOpProperties(
        demo_asm="sv.ld *RT, imm(RA)",
        inputs=[OD_EXTRA3_SGPR, OD_VL],
        outputs=[OD_EXTRA3_VGPR],
        immediates=[IMM_S16],
    )
    _PRE_RA_SIMS[SvLd] = lambda: OpKind.__svld_pre_ra_sim

    @staticmethod
    def __ld_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RA, = state.ssa_vals[op.input_vals[0]]
        addr = RA + op.immediates[0]
        v = state.load(addr)
        state.ssa_vals[op.outputs[0]] = v & GPR_VALUE_MASK,
    Ld = GenericOpProperties(
        demo_asm="ld RT, imm(RA)",
        inputs=[OD_BASE_SGPR],
        outputs=[OD_BASE_SGPR.with_write_stage(OpStage.Late)],
        immediates=[IMM_S16],
    )
    _PRE_RA_SIMS[Ld] = lambda: OpKind.__ld_pre_ra_sim

    @staticmethod
    def __svstd_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RS = state.ssa_vals[op.input_vals[0]]
        RA, = state.ssa_vals[op.input_vals[1]]
        VL, = state.ssa_vals[op.input_vals[2]]
        addr = RA + op.immediates[0]
        for i in range(VL):
            state.store(addr + GPR_SIZE_IN_BYTES * i, value=RS[i])
    SvStd = GenericOpProperties(
        demo_asm="sv.std *RS, imm(RA)",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_SGPR, OD_VL],
        outputs=[],
        immediates=[IMM_S16],
        has_side_effects=True,
    )
    _PRE_RA_SIMS[SvStd] = lambda: OpKind.__svstd_pre_ra_sim

    @staticmethod
    def __std_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        RS, = state.ssa_vals[op.input_vals[0]]
        RA, = state.ssa_vals[op.input_vals[1]]
        addr = RA + op.immediates[0]
        state.store(addr, value=RS)
    Std = GenericOpProperties(
        demo_asm="std RT, imm(RA)",
        inputs=[OD_BASE_SGPR, OD_BASE_SGPR],
        outputs=[],
        immediates=[IMM_S16],
        has_side_effects=True,
    )
    _PRE_RA_SIMS[Std] = lambda: OpKind.__std_pre_ra_sim

    @staticmethod
    def __funcargr3_pre_ra_sim(op, state):
        # type: (Op, PreRASimState) -> None
        pass  # return value set before simulation
    FuncArgR3 = GenericOpProperties(
        demo_asm="",
        inputs=[],
        outputs=[OD_BASE_SGPR.with_fixed_loc(
            Loc(kind=LocKind.GPR, start=3, reg_len=1))],
    )
    _PRE_RA_SIMS[FuncArgR3] = lambda: OpKind.__funcargr3_pre_ra_sim


@plain_data(frozen=True, unsafe_hash=True, repr=False)
class SSAValOrUse(metaclass=ABCMeta):
    __slots__ = "op", "operand_idx"

    def __init__(self, op, operand_idx):
        # type: (Op, int) -> None
        super().__init__()
        self.op = op
        if operand_idx < 0 or operand_idx >= len(self.descriptor_array):
            raise ValueError("invalid operand_idx")
        self.operand_idx = operand_idx

    @abstractmethod
    def __repr__(self):
        # type: () -> str
        ...

    @property
    @abstractmethod
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        ...

    @cached_property
    def defining_descriptor(self):
        # type: () -> OperandDesc
        return self.descriptor_array[self.operand_idx]

    @cached_property
    def ty(self):
        # type: () -> Ty
        return self.defining_descriptor.ty

    @cached_property
    def ty_before_spread(self):
        # type: () -> Ty
        return self.defining_descriptor.ty_before_spread

    @property
    def base_ty(self):
        # type: () -> BaseTy
        return self.ty_before_spread.base_ty

    @property
    def reg_offset_in_unspread(self):
        """ the number of reg-sized slots in the unspread Loc before self's Loc

        e.g. if the unspread Loc containing self is:
        `Loc(kind=LocKind.GPR, start=8, reg_len=4)`
        and self's Loc is `Loc(kind=LocKind.GPR, start=10, reg_len=1)`
        then reg_offset_into_unspread == 2 == 10 - 8
        """
        return self.defining_descriptor.reg_offset_in_unspread

    @property
    def unspread_start_idx(self):
        # type: () -> int
        return self.operand_idx - (self.defining_descriptor.spread_index or 0)

    @property
    def unspread_start(self):
        # type: () -> Self
        return self.__class__(op=self.op, operand_idx=self.unspread_start_idx)


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class SSAVal(SSAValOrUse):
    __slots__ = ()

    def __repr__(self):
        # type: () -> str
        return f"<{self.op.name}.outputs[{self.operand_idx}]: {self.ty}>"

    @cached_property
    def def_loc_set_before_spread(self):
        # type: () -> LocSet
        return self.defining_descriptor.loc_set_before_spread

    @cached_property
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.outputs

    @cached_property
    def tied_input(self):
        # type: () -> None | SSAUse
        if self.defining_descriptor.tied_input_index is None:
            return None
        return SSAUse(op=self.op,
                      operand_idx=self.defining_descriptor.tied_input_index)

    @property
    def write_stage(self):
        # type: () -> OpStage
        return self.defining_descriptor.write_stage


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class SSAUse(SSAValOrUse):
    __slots__ = ()

    @cached_property
    def use_loc_set_before_spread(self):
        # type: () -> LocSet
        return self.defining_descriptor.loc_set_before_spread

    @cached_property
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.inputs

    def __repr__(self):
        # type: () -> str
        return f"<{self.op.name}.input_uses[{self.operand_idx}]: {self.ty}>"

    @property
    def ssa_val(self):
        # type: () -> SSAVal
        return self.op.input_vals[self.operand_idx]

    @ssa_val.setter
    def ssa_val(self, ssa_val):
        # type: (SSAVal) -> None
        self.op.input_vals[self.operand_idx] = ssa_val


_T = TypeVar("_T")
_Desc = TypeVar("_Desc")


class OpInputSeq(Sequence[_T], Generic[_T, _Desc]):
    @abstractmethod
    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, _T | Any, _Desc) -> None
        raise NotImplementedError

    @final
    def _verify_write(self, idx, item):
        # type: (int | Any, _T | Any) -> int
        if not isinstance(idx, int):
            if isinstance(idx, slice):
                raise TypeError(
                    f"can't write to slice of {self.__class__.__name__}")
            raise TypeError(f"can't write with index {idx!r}")
        # normalize idx, raising IndexError if it is out of range
        idx = range(len(self.descriptors))[idx]
        desc = self.descriptors[idx]
        self._verify_write_with_desc(idx, item, desc)
        return idx

    def _on_set(self, idx, new_item, old_item):
        # type: (int, _T, _T | None) -> None
        pass

    @abstractmethod
    def _get_descriptors(self):
        # type: () -> tuple[_Desc, ...]
        raise NotImplementedError

    @cached_property
    @final
    def descriptors(self):
        # type: () -> tuple[_Desc, ...]
        return self._get_descriptors()

    @property
    @final
    def op(self):
        return self.__op

    def __init__(self, items, op):
        # type: (Iterable[_T], Op) -> None
        super().__init__()
        self.__op = op
        self.__items = []  # type: list[_T]
        for idx, item in enumerate(items):
            if idx >= len(self.descriptors):
                raise ValueError("too many items")
            _ = self._verify_write(idx, item)
            self.__items.append(item)
        if len(self.__items) < len(self.descriptors):
            raise ValueError("not enough items")

    @final
    def __iter__(self):
        # type: () -> Iterator[_T]
        yield from self.__items

    @overload
    def __getitem__(self, idx):
        # type: (int) -> _T
        ...

    @overload
    def __getitem__(self, idx):
        # type: (slice) -> list[_T]
        ...

    @final
    def __getitem__(self, idx):
        # type: (int | slice) -> _T | list[_T]
        return self.__items[idx]

    @final
    def __setitem__(self, idx, item):
        # type: (int, _T) -> None
        idx = self._verify_write(idx, item)
        self.__items[idx] = item

    @final
    def __len__(self):
        # type: () -> int
        return len(self.__items)

    def __repr__(self):
        # type: () -> str
        return f"{self.__class__.__name__}({self.__items}, op=...)"


@final
class OpInputVals(OpInputSeq[SSAVal, OperandDesc]):
    def _get_descriptors(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.inputs

    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, SSAVal | Any, OperandDesc) -> None
        if not isinstance(item, SSAVal):
            raise TypeError("expected value of type SSAVal")
        if item.ty != desc.ty:
            raise ValueError(f"assigned item's type {item.ty!r} doesn't match "
                             f"corresponding input's type {desc.ty!r}")

    def _on_set(self, idx, new_item, old_item):
        # type: (int, SSAVal, SSAVal | None) -> None
        SSAUses._on_op_input_set(self, idx, new_item, old_item)  # type: ignore

    def __init__(self, items, op):
        # type: (Iterable[SSAVal], Op) -> None
        if hasattr(op, "inputs"):
            raise ValueError("Op.inputs already set")
        super().__init__(items, op)


@final
class OpImmediates(OpInputSeq[int, range]):
    def _get_descriptors(self):
        # type: () -> tuple[range, ...]
        return self.op.properties.immediates

    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, int | Any, range) -> None
        if not isinstance(item, int):
            raise TypeError("expected value of type int")
        if item not in desc:
            raise ValueError(f"immediate value {item!r} not in {desc!r}")

    def __init__(self, items, op):
        # type: (Iterable[int], Op) -> None
        if hasattr(op, "immediates"):
            raise ValueError("Op.immediates already set")
        super().__init__(items, op)


@plain_data(frozen=True, eq=False, repr=False)
@final
class Op:
    __slots__ = ("fn", "properties", "input_vals", "input_uses", "immediates",
                 "outputs", "name")

    def __init__(self, fn, properties, input_vals, immediates, name=""):
        # type: (Fn, OpProperties, Iterable[SSAVal], Iterable[int], str) -> None
        self.fn = fn
        self.properties = properties
        self.input_vals = OpInputVals(input_vals, op=self)
        inputs_len = len(self.properties.inputs)
        self.input_uses = tuple(SSAUse(self, i) for i in range(inputs_len))
        self.immediates = OpImmediates(immediates, op=self)
        outputs_len = len(self.properties.outputs)
        self.outputs = tuple(SSAVal(self, i) for i in range(outputs_len))
        self.name = fn._add_op_with_unused_name(self, name)  # type: ignore

    @property
    def kind(self):
        # type: () -> OpKind
        return self.properties.kind

    def __eq__(self, other):
        # type: (Op | Any) -> bool
        if isinstance(other, Op):
            return self is other
        return NotImplemented

    def __hash__(self):
        # type: () -> int
        return object.__hash__(self)

    def __repr__(self):
        # type: () -> str
        field_vals = []  # type: list[str]
        for name in fields(self):
            if name == "properties":
                name = "kind"
            elif name == "fn":
                continue
            try:
                value = getattr(self, name)
            except AttributeError:
                field_vals.append(f"{name}=<not set>")
                continue
            if isinstance(value, OpInputSeq):
                value = list(value)  # type: ignore
            field_vals.append(f"{name}={value!r}")
        field_vals_str = ", ".join(field_vals)
        return f"Op({field_vals_str})"

    def pre_ra_sim(self, state):
        # type: (PreRASimState) -> None
        for inp in self.input_vals:
            if inp not in state.ssa_vals:
                raise ValueError(f"SSAVal {inp} not yet assigned when "
                                 f"running {self}")
            if len(state.ssa_vals[inp]) != inp.ty.reg_len:
                raise ValueError(
                    f"value of SSAVal {inp} has wrong number of elements: "
                    f"expected {inp.ty.reg_len} found "
                    f"{len(state.ssa_vals[inp])}: {state.ssa_vals[inp]!r}")
        for out in self.outputs:
            if out in state.ssa_vals:
                if self.kind is OpKind.FuncArgR3:
                    continue
                raise ValueError(f"SSAVal {out} already assigned before "
                                 f"running {self}")
        self.kind.pre_ra_sim(self, state)
        for out in self.outputs:
            if out not in state.ssa_vals:
                raise ValueError(f"running {self} failed to assign to {out}")
            if len(state.ssa_vals[out]) != out.ty.reg_len:
                raise ValueError(
                    f"value of SSAVal {out} has wrong number of elements: "
                    f"expected {out.ty.reg_len} found "
                    f"{len(state.ssa_vals[out])}: {state.ssa_vals[out]!r}")


GPR_SIZE_IN_BYTES = 8
BITS_IN_BYTE = 8
GPR_SIZE_IN_BITS = GPR_SIZE_IN_BYTES * BITS_IN_BYTE
GPR_VALUE_MASK = (1 << GPR_SIZE_IN_BITS) - 1


@plain_data(frozen=True, repr=False)
@final
class PreRASimState:
    __slots__ = "ssa_vals", "memory"

    def __init__(self, ssa_vals, memory):
        # type: (dict[SSAVal, tuple[int, ...]], dict[int, int]) -> None
        self.ssa_vals = ssa_vals  # type: dict[SSAVal, tuple[int, ...]]
        self.memory = memory  # type: dict[int, int]

    def load_byte(self, addr):
        # type: (int) -> int
        addr &= GPR_VALUE_MASK
        return self.memory.get(addr, 0) & 0xFF

    def store_byte(self, addr, value):
        # type: (int, int) -> None
        addr &= GPR_VALUE_MASK
        value &= 0xFF
        self.memory[addr] = value

    def load(self, addr, size_in_bytes=GPR_SIZE_IN_BYTES, signed=False):
        # type: (int, int, bool) -> int
        if addr % size_in_bytes != 0:
            raise ValueError(f"address not aligned: {hex(addr)} "
                             f"required alignment: {size_in_bytes}")
        retval = 0
        for i in range(size_in_bytes):
            retval |= self.load_byte(addr + i) << i * BITS_IN_BYTE
        if signed and retval >> (size_in_bytes * BITS_IN_BYTE - 1) != 0:
            retval -= 1 << size_in_bytes * BITS_IN_BYTE
        return retval

    def store(self, addr, value, size_in_bytes=GPR_SIZE_IN_BYTES):
        # type: (int, int, int) -> None
        if addr % size_in_bytes != 0:
            raise ValueError(f"address not aligned: {hex(addr)} "
                             f"required alignment: {size_in_bytes}")
        for i in range(size_in_bytes):
            self.store_byte(addr + i, (value >> i * BITS_IN_BYTE) & 0xFF)

    def _memory__repr(self):
        # type: () -> str
        if len(self.memory) == 0:
            return "{}"
        keys = sorted(self.memory.keys(), reverse=True)
        CHUNK_SIZE = GPR_SIZE_IN_BYTES
        items = []  # type: list[str]
        while len(keys) != 0:
            addr = keys[-1]
            if (len(keys) >= CHUNK_SIZE
                    and addr % CHUNK_SIZE == 0
                    and keys[-CHUNK_SIZE:]
                    == list(reversed(range(addr, addr + CHUNK_SIZE)))):
                value = self.load(addr, size_in_bytes=CHUNK_SIZE)
                items.append(f"0x{addr:05x}: <0x{value:0{CHUNK_SIZE * 2}x}>")
                keys[-CHUNK_SIZE:] = ()
            else:
                items.append(f"0x{addr:05x}: 0x{self.memory[keys.pop()]:02x}")
        if len(items) == 1:
            return f"{{{items[0]}}}"
        items_str = ",\n".join(items)
        return f"{{\n{items_str}}}"

    def _ssa_vals__repr(self):
        # type: () -> str
        if len(self.ssa_vals) == 0:
            return "{}"
        items = []  # type: list[str]
        CHUNK_SIZE = 4
        for k, v in self.ssa_vals.items():
            element_strs = []  # type: list[str]
            for i, el in enumerate(v):
                if i % CHUNK_SIZE != 0:
                    element_strs.append(" " + hex(el))
                else:
                    element_strs.append("\n    " + hex(el))
            if len(element_strs) <= CHUNK_SIZE:
                element_strs[0] = element_strs[0].lstrip()
            if len(element_strs) == 1:
                element_strs.append("")
            v_str = ",".join(element_strs)
            items.append(f"{k!r}: ({v_str})")
        if len(items) == 1 and "\n" not in items[0]:
            return f"{{{items[0]}}}"
        items_str = ",\n".join(items)
        return f"{{\n{items_str},\n}}"

    def __repr__(self):
        # type: () -> str
        field_vals = []  # type: list[str]
        for name in fields(self):
            try:
                value = getattr(self, name)
            except AttributeError:
                field_vals.append(f"{name}=<not set>")
                continue
            repr_fn = getattr(self, f"_{name}__repr", None)
            if callable(repr_fn):
                field_vals.append(f"{name}={repr_fn()}")
            else:
                field_vals.append(f"{name}={value!r}")
        field_vals_str = ", ".join(field_vals)
        return f"PreRASimState({field_vals_str})"