X86 scf

@irdl_op_definition
class YieldOp(AbstractYieldOperation[X86RegisterType]):
    name = "x86_scf.yield"

    traits = lazy_traits_def(
        lambda: (
            IsTerminator(),
            HasParent(ForRofOperation),
            NoMemoryEffect(),
        )
    )

class ForRofOperation(X86HasRegisterConstraints, ABC):
    lb = operand_def(GeneralRegisterType)
    ub_val = opt_operand_def(GeneralRegisterType)
    ub_attr = opt_prop_def(IntegerAttr[SI32])
    step_val = opt_operand_def(GeneralRegisterType)
    step_attr = opt_prop_def(IntegerAttr[SI32])

    iter_args = var_operand_def(X86RegisterType)

    lb_end = result_def(GeneralRegisterType)
    """Final value of the lower-bound / induction-variable register (inout with `lb`)."""

    res = var_result_def(X86RegisterType)

    body = region_def("single_block")

    traits = traits_def(SingleBlockImplicitTerminator(YieldOp), RecursiveMemoryEffect())
    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    @property
    def ub(self) -> IntegerAttr[SI32] | SSAValue:
        """Static upper bound (typed integer) or dynamic register SSA value."""
        match (ub_attr := self.ub_attr, ub_val := self.ub_val):
            case (None, None):
                raise ValueError("Exactly one of ub_attr or ub_val must be set")
            case (None, _):
                return ub_val
            case (_, None):
                return ub_attr
            case (_, _):
                raise ValueError("Exactly one of ub_attr or ub_val must be set")

    @property
    def step(self) -> IntegerAttr[SI32] | SSAValue:
        """Static step (typed integer) or dynamic register SSA value."""
        match (step_attr := self.step_attr, step_val := self.step_val):
            case (None, None):
                raise ValueError("Exactly one of step_attr or step_val must be set")
            case (None, _):
                return step_val
            case (_, None):
                return step_attr
            case (_, _):
                raise ValueError("Exactly one of step_attr or step_val must be set")

    def __init__(
        self,
        lb: SSAValue | Operation,
        ub: SSAValue | Operation | IntegerAttr,
        step: SSAValue | Operation | IntegerAttr,
        iter_args: Sequence[SSAValue],
        body: Region | Sequence[Operation] | Sequence[Block] | Block | None = None,
    ):
        lb = SSAValue.get(lb)
        if body is None:
            body = Region(
                Block(arg_types=(lb.type, *(iter_arg.type for iter_arg in iter_args)))
            )

        if isinstance(body, Block):
            body = [body]

        if isinstance(ub, IntegerAttr):
            ub_attr = ub
            ub_val = None
        else:
            ub_attr = None
            ub_val = ub

        if isinstance(step, IntegerAttr):
            step_attr = step
            step_val = None
        else:
            step_attr = None
            step_val = step

        super().__init__(
            operands=[lb, ub_val, step_val, iter_args],
            properties={"ub_attr": ub_attr, "step_attr": step_attr},
            result_types=[lb.type, [SSAValue.get(a).type for a in iter_args]],
            regions=[body],
        )

    def verify_(self):
        try:
            _ = self.ub
            _ = self.step
        except ValueError as exc:
            raise VerifyException(exc.args[0])

        if (len(self.iter_args) + 1) != len(self.body.block.args):
            raise VerifyException(
                f"Wrong number of block arguments, expected {len(self.iter_args) + 1}, got "
                f"{len(self.body.block.args)}. The body must have the induction "
                f"variable and loop-carried variables as arguments."
            )
        if self.body.block.args and (iter_var := self.body.block.args[0]):
            if not isinstance(iter_var.type, GeneralRegisterType):
                raise VerifyException(
                    f"The first block argument of the body is of type {iter_var.type}"
                    " instead of x86 GeneralRegisterType"
                )
            if iter_var.type != self.lb.type:
                raise VerifyException(
                    f"Expected induction var to be same type as lb, "
                    f"got {iter_var.type} and {self.lb.type}"
                )
            if iter_var.type != self.lb_end.type:
                raise VerifyException(
                    f"Expected induction var to be same type as lb_end result, "
                    f"got {iter_var.type} and {self.lb_end.type}"
                )
        for idx, (arg, block_arg) in enumerate(
            zip(self.iter_args, self.body.block.args[1:])
        ):
            if block_arg.type != arg.type:
                raise VerifyException(
                    f"Block argument {idx + 1} has wrong type, expected {arg.type}, "
                    f"got {block_arg.type}. Arguments after the "
                    f"induction variable must match the carried variables."
                )
        if len(self.body.ops) > 0 and isinstance(
            yieldop := self.body.block.last_op, YieldOp
        ):
            if len(yieldop.arguments) != len(self.iter_args):
                raise VerifyException(
                    f"Expected {len(self.iter_args)} args, got {len(yieldop.arguments)}. "
                    f"The {self.name} must yield its carried variables."
                )
            for iter_arg, yield_arg in zip(self.iter_args, yieldop.arguments):
                if iter_arg.type != yield_arg.type:
                    raise VerifyException(
                        f"Expected {iter_arg.type}, got {yield_arg.type}. The "
                        f"{self.name}'s {YieldOp.name} must match carried "
                        f"variables types."
                    )

    def allocate_registers(self, allocator: BlockAllocator) -> None:
        """Allocate loop-carried and IV registers, then the body under those reservations."""
        # Allocate values used inside the body but defined outside.
        # Their scope lasts for the whole body execution scope
        live_ins = allocator.live_ins_per_block[self.body.block]
        for live_in in live_ins:
            allocator.allocate_value(live_in)

        yield_op = self.body.block.last_op
        assert yield_op is not None, (
            f"last op of {self.name} is guaranteed to be {YieldOp.name}"
        )
        block_args = self.body.block.args

        # The loop-carried variables are trickier
        # The for op operand, block arg, and yield operand must have the same type
        for block_arg, operand, yield_operand, op_result in zip(
            block_args[1:], self.iter_args, yield_op.operands, self.res, strict=True
        ):
            allocator.allocate_values_same_reg(
                (block_arg, operand, yield_operand, op_result)
            )

        allocator.allocate_values_same_reg((block_args[0], self.lb, self.lb_end))

        # ub and step are used throughout loop when dynamic
        if self.ub_val is not None:
            allocator.allocate_value(self.ub_val)
        if self.step_val is not None:
            allocator.allocate_value(self.step_val)

        # Reserve the loop carried variables for allocation within the body
        regs = self.iter_args.types
        assert all(isinstance(reg, X86RegisterType) for reg in regs)
        regs = cast(tuple[X86RegisterType, ...], regs)
        with allocator.available_registers.reserve_registers(regs):
            allocator.allocate_block(self.body.block)

    def get_register_constraints(self) -> RegisterConstraints:
        """`lb` and each iter_arg are inout; dynamic `ub`/`step` are in-only."""
        ins: list[SSAValue] = []
        if self.ub_val is not None:
            ins.append(self.ub_val)
        if self.step_val is not None:
            ins.append(self.step_val)
        inouts = ((self.lb, self.lb_end), *zip(self.iter_args, self.res, strict=True))
        return RegisterConstraints(ins, (), inouts)

    def _body_live_outs(self, live_after: AbstractSet[SSAValue]) -> set[SSAValue]:
        """
        Values live at the end of the loop body.
        """
        block = self.body.block
        res = set(live_after)
        # The body runs repeatedly, so every value defined outside it and used inside
        # must survive a whole iteration. This covers the dynamic bounds: the loop
        # re-reads them on the back edge, and a clobber in the body is itself a use.
        res.update(live_ins_per_block(block)[block])
        # The induction variable is a block argument rather than a live-in, but the
        # loop reads it on the back edge to compute the next value.
        res.add(block.args[0])
        return res

    def update_liveness(self, ctx: LivenessContext) -> None:
        # Create a new context to use inside the loop
        body_ctx = ctx.copy(self._body_live_outs(ctx.alive))
        body_ctx.process_block(self.body.block)
        # Update the outer context with all the values that are alive coming into the
        # body
        ctx.alive.update(body_ctx.alive)
        # HasRegisterConstraints default implementation
        super().update_liveness(ctx)

@irdl_op_definition
class ForOp(ForRofOperation):
    """
    A for loop, counting up from lb to ub by step each iteration.
    """

    name = "x86_scf.for"

    def print(self, printer: Printer):
        print_for_op_like(
            printer,
            self.lb,
            self.ub,
            self.step,
            self.iter_args,
            self.body,
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        lb, ub, step, iter_arg_operands, body = parse_for_op_like(
            parser, allow_static_stop=True, allow_static_step=True
        )
        _, *iter_args = body.block.args

        for_op = cls(lb, ub, step, iter_arg_operands, body)

        if not iter_args:
            for trait in for_op.get_traits_of_type(SingleBlockImplicitTerminator):
                ensure_terminator(for_op, trait)

        return for_op

@irdl_op_definition
class RofOp(ForRofOperation):
    """
    Reverse Order For loop.

    MLIR's for loops have the constraint of always executing from lb to ub,
    so in order to express loops that count down from ub to lb, the rof op
    is needed.

    Rof has the semantics of going from ub to lb, decrementing by step each time.
    The implicit constraints are that lb < ub, and step > 0.

    In order to convert a for to a rof, one needs to switch lb and ub.
    (for the normalized case that (ub - lb) % step == 0)
    """

    name = "x86_scf.rof"

    def print(self, printer: Printer):
        print_for_op_like(
            printer,
            self.ub,
            self.lb,
            self.step,
            self.iter_args,
            self.body,
            bound_words=["down", "to"],
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        ub, lb, step, iter_arg_operands, body = parse_for_op_like(
            parser,
            bound_words=["down", "to"],
            allow_static_stop=True,
            allow_static_step=True,
        )
        _, *iter_args = body.block.args

        if isinstance(lb, IntegerAttr):
            parser.raise_error("Expected an operand.")

        rof_op = cls(lb, ub, step, iter_arg_operands, body)

        if not iter_args:
            for trait in rof_op.get_traits_of_type(SingleBlockImplicitTerminator):
                ensure_terminator(rof_op, trait)

        return rof_op