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

"""
Toom-Cook multiplication algorithm generator for SVP64
"""
import math
from abc import abstractmethod
from enum import Enum
from fractions import Fraction
from typing import Iterable, Mapping, Tuple, Union

from cached_property import cached_property
from nmutil.plain_data import plain_data

from bigint_presentation_code.compiler_ir import (GPR_SIZE_IN_BITS, BaseTy, Fn,
                                                  OpKind, SSAVal, Ty)
from bigint_presentation_code.matrix import Matrix
from bigint_presentation_code.type_util import Literal, final
from bigint_presentation_code.util import InternedMeta


@final
class PointAtInfinity(Enum):
    POINT_AT_INFINITY = "POINT_AT_INFINITY"

    def __repr__(self):
        return self.name


POINT_AT_INFINITY = PointAtInfinity.POINT_AT_INFINITY

_EvalOpPolyCoefficients = Union["Mapping[int | None, Fraction | int]",
                                "EvalOpPoly", Fraction, int, None]


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class EvalOpPoly:
    """polynomial"""
    __slots__ = "const_coeff", "var_coeffs"

    def __init__(
        self, coeffs=None,  # type: _EvalOpPolyCoefficients
        const_coeff=None,  # type: Fraction | int | None
        var_coeffs=(),  # type: Iterable[Fraction | int] | None
    ):
        if coeffs is not None:
            if const_coeff is not None or var_coeffs != ():
                raise ValueError(
                    "can't specify const_coeff or "
                    "var_coeffs along with coeffs")
            if isinstance(coeffs, EvalOpPoly):
                self.const_coeff = coeffs.const_coeff
                self.var_coeffs = coeffs.var_coeffs
                return
            if isinstance(coeffs, (int, Fraction)):
                const_coeff = Fraction(coeffs)
                final_var_coeffs = []  # type: list[Fraction]
            else:
                const_coeff = 0
                final_var_coeffs = []
                for var, coeff in coeffs.items():
                    if coeff == 0:
                        continue
                    coeff = Fraction(coeff)
                    if var is None:
                        const_coeff = coeff
                        continue
                    if var < 0:
                        raise ValueError("invalid variable index")
                    if var >= len(final_var_coeffs):
                        additional = var - len(final_var_coeffs)
                        final_var_coeffs.extend((Fraction(),) * additional)
                        final_var_coeffs.append(coeff)
                    else:
                        final_var_coeffs[var] = coeff
        else:
            if var_coeffs is None:
                final_var_coeffs = []
            else:
                final_var_coeffs = [Fraction(v) for v in var_coeffs]
                while len(final_var_coeffs) > 0 and final_var_coeffs[-1] == 0:
                    final_var_coeffs.pop()
        if const_coeff is None:
            const_coeff = 0
        self.const_coeff = Fraction(const_coeff)
        self.var_coeffs = tuple(final_var_coeffs)

    def __add__(self, rhs):
        # type: (EvalOpPoly | int | Fraction) -> EvalOpPoly
        rhs = EvalOpPoly(rhs)
        const_coeff = self.const_coeff + rhs.const_coeff
        var_coeffs = list(self.var_coeffs)
        if len(rhs.var_coeffs) > len(var_coeffs):
            var_coeffs.extend(rhs.var_coeffs[len(var_coeffs):])
        for var in range(min(len(self.var_coeffs), len(rhs.var_coeffs))):
            var_coeffs[var] += rhs.var_coeffs[var]
        return EvalOpPoly(const_coeff=const_coeff, var_coeffs=var_coeffs)

    @property
    def coefficients(self):
        # type: () -> dict[int | None, Fraction]
        retval = {}  # type: dict[int | None, Fraction]
        if self.const_coeff != 0:
            retval[None] = self.const_coeff
        for var, coeff in enumerate(self.var_coeffs):
            if coeff != 0:
                retval[var] = coeff
        return retval

    @property
    def is_const(self):
        # type: () -> bool
        return self.var_coeffs == ()

    def coeff(self, var):
        # type: (int | None) -> Fraction
        if var is None:
            return self.const_coeff
        if var < 0:
            raise ValueError("invalid variable index")
        if var < len(self.var_coeffs):
            return self.var_coeffs[var]
        return Fraction()

    __radd__ = __add__

    def __neg__(self):
        return EvalOpPoly(const_coeff=-self.const_coeff,
                          var_coeffs=(-v for v in self.var_coeffs))

    def __sub__(self, rhs):
        # type: (EvalOpPoly | int | Fraction) -> EvalOpPoly
        return self + -rhs

    def __rsub__(self, lhs):
        # type: (EvalOpPoly | int | Fraction) -> EvalOpPoly
        return lhs + -self

    def __mul__(self, rhs):
        # type: (int | Fraction | EvalOpPoly) -> EvalOpPoly
        if isinstance(rhs, EvalOpPoly):
            if self.is_const:
                self, rhs = rhs, self
            if not rhs.is_const:
                raise ValueError("can't represent exponents larger than one")
            rhs = rhs.const_coeff
        if rhs == 0:
            return EvalOpPoly()
        return EvalOpPoly(const_coeff=self.const_coeff * rhs,
                          var_coeffs=(i * rhs for i in self.var_coeffs))

    __rmul__ = __mul__

    def __truediv__(self, rhs):
        # type: (int | Fraction) -> EvalOpPoly
        if rhs == 0:
            raise ZeroDivisionError()
        return EvalOpPoly(const_coeff=self.const_coeff / rhs,
                          var_coeffs=(i / rhs for i in self.var_coeffs))

    def __repr__(self):
        return f"EvalOpPoly({self.coefficients})"


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpValueRange:
    __slots__ = ("eval_op", "inputs", "min_value", "max_value",
                 "is_signed", "output_size")

    def __init__(self, eval_op, inputs):
        # type: (EvalOp | int, tuple[EvalOpGenIrInput, ...]) -> None
        super().__init__()
        self.eval_op = eval_op
        self.inputs = inputs
        min_value = max_value = self.poly.const_coeff
        for var, coeff in enumerate(self.poly.var_coeffs):
            if coeff == 0:
                continue
            term_min = self.inputs[var].min_value * coeff
            term_max = self.inputs[var].max_value * coeff
            if term_min > term_max:
                term_min, term_max = term_max, term_min
            min_value += term_min
            max_value += term_max
        # output values are always integers, so eliminate any fractional part
        # as impossible.
        self.min_value = math.ceil(min_value)  # exclude fractional part
        self.max_value = math.floor(max_value)  # exclude fractional part
        self.is_signed = min_value < 0
        output_size = 1
        if self.is_signed:
            min_v = -1 << (GPR_SIZE_IN_BITS - 1)
            max_v = (1 << (GPR_SIZE_IN_BITS - 1)) - 1
        else:
            min_v = 0
            max_v = (1 << GPR_SIZE_IN_BITS) - 1
        while not (min_v <= self.min_value and self.max_value <= max_v):
            output_size += 1
            min_v <<= GPR_SIZE_IN_BITS
            max_v <<= GPR_SIZE_IN_BITS
        self.output_size = output_size

    @cached_property
    def poly(self):
        if isinstance(self.eval_op, int):
            return EvalOpPoly(const_coeff=self.eval_op)
        return self.eval_op.poly


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpGenIrOutput:
    __slots__ = "output", "value_range"

    def __init__(self, output, value_range):
        # type: (SSAVal, EvalOpValueRange) -> None
        super().__init__()
        if output.ty.reg_len != value_range.output_size:
            raise ValueError("wrong output size")
        self.output = output
        self.value_range = value_range

    @property
    def eval_op(self):
        # type: () -> EvalOp | int
        return self.value_range.eval_op

    @property
    def inputs(self):
        # type: () -> tuple[EvalOpGenIrInput, ...]
        return self.value_range.inputs

    @property
    def min_value(self):
        # type: () -> int
        return self.value_range.min_value

    @property
    def max_value(self):
        # type: () -> int
        return self.value_range.max_value

    @property
    def is_signed(self):
        # type: () -> bool
        return self.value_range.is_signed

    @property
    def output_size(self):
        # type: () -> int
        return self.value_range.output_size

    @property
    def current_debugging_value(self):
        # type: () -> tuple[int, ...]
        """ get the current value for debugging in pdb or similar.

        This is intended for use with
        `PreRASimState.set_current_debugging_state`.

        This is only intended for debugging, do not use in unit tests or
        production code.
        """
        return self.output.current_debugging_value


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpGenIrInput:
    __slots__ = "ssa_val", "is_signed", "min_value", "max_value"

    def __init__(self, ssa_val, is_signed, min_value=None, max_value=None):
        # type: (SSAVal, bool | None, int | None, int | None) -> None
        super().__init__()
        self.ssa_val = ssa_val
        if ssa_val.base_ty != BaseTy.I64:
            raise ValueError("input must have a base_ty of BaseTy.I64")
        if is_signed is None:
            if min_value is None or max_value is None:
                raise ValueError("must specify either is_signed or both "
                                 "min_value and max_value")
            is_signed = min_value < 0
        self.is_signed = is_signed
        if is_signed:
            if min_value is None:
                min_value = -1 << (ssa_val.ty.reg_len * GPR_SIZE_IN_BITS - 1)
            if max_value is None:
                max_value = (1 << (
                    ssa_val.ty.reg_len * GPR_SIZE_IN_BITS - 1)) - 1
        else:
            if min_value is None:
                min_value = 0
            if max_value is None:
                max_value = (1 << (ssa_val.ty.reg_len * GPR_SIZE_IN_BITS)) - 1
        self.min_value = min_value
        self.max_value = max_value
        if self.min_value > self.max_value:
            raise ValueError("invalid value range")

    @property
    def current_debugging_value(self):
        # type: () -> tuple[int, ...]
        """ get the current value for debugging in pdb or similar.

        This is intended for use with
        `PreRASimState.set_current_debugging_state`.

        This is only intended for debugging, do not use in unit tests or
        production code.
        """
        return self.ssa_val.current_debugging_value


@plain_data(frozen=True)
@final
class EvalOpGenIrState:
    __slots__ = "fn", "inputs", "outputs_map"

    def __init__(self, fn, inputs):
        # type: (Fn, Iterable[EvalOpGenIrInput]) -> None
        super().__init__()
        self.fn = fn
        self.inputs = tuple(inputs)
        self.outputs_map = {}  # type: dict[EvalOp | int, EvalOpGenIrOutput]

    def get_output(self, eval_op):
        # type: (EvalOp | int) -> EvalOpGenIrOutput
        retval = self.outputs_map.get(eval_op, None)
        if retval is not None:
            return retval
        value_range = EvalOpValueRange(eval_op=eval_op, inputs=self.inputs)
        if isinstance(eval_op, int):
            li = self.fn.append_new_op(OpKind.LI, immediates=[eval_op],
                                       name=f"li_{eval_op}")
            output = cast_to_size(
                fn=self.fn, ssa_val=li.outputs[0],
                dest_size=value_range.output_size,
                src_signed=value_range.is_signed, name=f"cast_{eval_op}")
            retval = EvalOpGenIrOutput(output=output, value_range=value_range)
        else:
            retval = eval_op.make_output(state=self,
                                         output_value_range=value_range)
        if retval.value_range != value_range:
            raise ValueError("wrong value_range")
        return self.outputs_map.setdefault(eval_op, retval)


@plain_data(frozen=True, unsafe_hash=True)
class EvalOp(metaclass=InternedMeta):
    __slots__ = "lhs", "rhs", "poly"

    @property
    def lhs_poly(self):
        # type: () -> EvalOpPoly
        if isinstance(self.lhs, int):
            return EvalOpPoly(self.lhs)
        return self.lhs.poly

    @property
    def rhs_poly(self):
        # type: () -> EvalOpPoly
        if isinstance(self.rhs, int):
            return EvalOpPoly(self.rhs)
        return self.rhs.poly

    @abstractmethod
    def _make_poly(self):
        # type: () -> EvalOpPoly
        ...

    @abstractmethod
    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        ...

    def __init__(self, lhs, rhs):
        # type: (EvalOp | int, EvalOp | int) -> None
        super().__init__()
        self.lhs = lhs
        self.rhs = rhs
        self.poly = self._make_poly()


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpAdd(EvalOp):
    __slots__ = ()

    def _make_poly(self):
        # type: () -> EvalOpPoly
        return self.lhs_poly + self.rhs_poly

    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        lhs = state.get_output(self.lhs)
        lhs_output = cast_to_size(
            fn=state.fn, ssa_val=lhs.output,
            dest_size=output_value_range.output_size, src_signed=lhs.is_signed,
            name="add_lhs_cast")
        rhs = state.get_output(self.rhs)
        rhs_output = cast_to_size(
            fn=state.fn, ssa_val=rhs.output,
            dest_size=output_value_range.output_size, src_signed=rhs.is_signed,
            name="add_rhs_cast")
        setvl = state.fn.append_new_op(
            OpKind.SetVLI, immediates=[output_value_range.output_size],
            name="setvl", maxvl=output_value_range.output_size)
        clear_ca = state.fn.append_new_op(OpKind.ClearCA, name="clear_ca")
        add = state.fn.append_new_op(
            OpKind.SvAddE, input_vals=[
                lhs_output, rhs_output, clear_ca.outputs[0], setvl.outputs[0]],
            maxvl=output_value_range.output_size, name="add")
        return EvalOpGenIrOutput(
            output=add.outputs[0], value_range=output_value_range)


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpSub(EvalOp):
    __slots__ = ()

    def _make_poly(self):
        # type: () -> EvalOpPoly
        return self.lhs_poly - self.rhs_poly

    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        lhs = state.get_output(self.lhs)
        lhs_output = cast_to_size(
            fn=state.fn, ssa_val=lhs.output,
            dest_size=output_value_range.output_size, src_signed=lhs.is_signed,
            name="add_lhs_cast")
        rhs = state.get_output(self.rhs)
        rhs_output = cast_to_size(
            fn=state.fn, ssa_val=rhs.output,
            dest_size=output_value_range.output_size, src_signed=rhs.is_signed,
            name="add_rhs_cast")
        setvl = state.fn.append_new_op(
            OpKind.SetVLI, immediates=[output_value_range.output_size],
            name="setvl", maxvl=output_value_range.output_size)
        set_ca = state.fn.append_new_op(OpKind.SetCA, name="set_ca")
        sub = state.fn.append_new_op(
            OpKind.SvSubFE, input_vals=[
                rhs_output, lhs_output, set_ca.outputs[0], setvl.outputs[0]],
            maxvl=output_value_range.output_size, name="sub")
        return EvalOpGenIrOutput(
            output=sub.outputs[0], value_range=output_value_range)


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpMul(EvalOp):
    __slots__ = ()
    rhs: int

    def _make_poly(self):
        # type: () -> EvalOpPoly
        if not isinstance(self.rhs, int):  # type: ignore
            raise TypeError("invalid rhs type")
        return self.lhs_poly * self.rhs

    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        raise NotImplementedError  # FIXME: finish


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpExactDiv(EvalOp):
    __slots__ = ()
    rhs: int

    def _make_poly(self):
        # type: () -> EvalOpPoly
        if not isinstance(self.rhs, int):  # type: ignore
            raise TypeError("invalid rhs type")
        return self.lhs_poly / self.rhs

    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        raise NotImplementedError  # FIXME: finish


@plain_data(frozen=True, unsafe_hash=True)
@final
class EvalOpInput(EvalOp):
    __slots__ = ()
    lhs: int
    rhs: Literal[0]

    def __init__(self, lhs, rhs=0):
        # type: (int, int) -> None
        if lhs < 0:
            raise ValueError("Input part_index (lhs) must be >= 0")
        if rhs != 0:
            raise ValueError("Input rhs must be 0")
        super().__init__(lhs, rhs)

    @property
    def part_index(self):
        return self.lhs

    def _make_poly(self):
        # type: () -> EvalOpPoly
        return EvalOpPoly({self.part_index: 1})

    def make_output(self, state, output_value_range):
        # type: (EvalOpGenIrState, EvalOpValueRange) -> EvalOpGenIrOutput
        inp = state.inputs[self.part_index]
        output = cast_to_size(
            fn=state.fn, ssa_val=inp.ssa_val, src_signed=inp.is_signed,
            dest_size=output_value_range.output_size,
            name=f"input_{self.part_index}_cast")
        return EvalOpGenIrOutput(output=output, value_range=output_value_range)


@plain_data(frozen=True, unsafe_hash=True)
@final
class ToomCookInstance:
    __slots__ = ("lhs_part_count", "rhs_part_count", "eval_points",
                 "lhs_eval_ops", "rhs_eval_ops", "prod_eval_ops")

    @property
    def prod_part_count(self):
        return self.lhs_part_count + self.rhs_part_count - 1

    @staticmethod
    def make_eval_matrix(width, eval_points):
        # type: (int, tuple[PointAtInfinity | int, ...]) -> Matrix[Fraction]
        retval = Matrix(height=len(eval_points), width=width)
        for row, col in retval.indexes():
            eval_point = eval_points[row]
            if eval_point is POINT_AT_INFINITY:
                retval[row, col] = int(col == width - 1)
            else:
                retval[row, col] = eval_point ** col
        return retval

    def get_lhs_eval_matrix(self):
        # type: () -> Matrix[Fraction]
        return self.make_eval_matrix(self.lhs_part_count, self.eval_points)

    @staticmethod
    def make_input_poly_vector(height):
        # type: (int) -> Matrix[EvalOpPoly]
        return Matrix(height=height, width=1, element_type=EvalOpPoly,
                      data=(EvalOpPoly({i: 1}) for i in range(height)))

    def get_lhs_eval_polys(self):
        # type: () -> list[EvalOpPoly]
        return list(self.get_lhs_eval_matrix().cast(EvalOpPoly)
                    @ self.make_input_poly_vector(self.lhs_part_count))

    def get_rhs_eval_matrix(self):
        # type: () -> Matrix[Fraction]
        return self.make_eval_matrix(self.rhs_part_count, self.eval_points)

    def get_rhs_eval_polys(self):
        # type: () -> list[EvalOpPoly]
        return list(self.get_rhs_eval_matrix().cast(EvalOpPoly)
                    @ self.make_input_poly_vector(self.rhs_part_count))

    def get_prod_inverse_eval_matrix(self):
        # type: () -> Matrix[Fraction]
        return self.make_eval_matrix(self.prod_part_count, self.eval_points)

    def get_prod_eval_matrix(self):
        # type: () -> Matrix[Fraction]
        return self.get_prod_inverse_eval_matrix().inverse()

    def get_prod_eval_polys(self):
        # type: () -> list[EvalOpPoly]
        return list(self.get_prod_eval_matrix().cast(EvalOpPoly)
                    @ self.make_input_poly_vector(self.prod_part_count))

    def __init__(
        self, lhs_part_count,  # type: int
        rhs_part_count,  # type: int
        eval_points,  # type: Iterable[PointAtInfinity | int]
        lhs_eval_ops,  # type: Iterable[EvalOp]
        rhs_eval_ops,  # type: Iterable[EvalOp]
        prod_eval_ops,  # type: Iterable[EvalOp]
    ):
        # type: (...) -> None
        self.lhs_part_count = lhs_part_count
        if self.lhs_part_count < 2:
            raise ValueError("lhs_part_count must be at least 2")
        self.rhs_part_count = rhs_part_count
        if self.rhs_part_count < 2:
            raise ValueError("rhs_part_count must be at least 2")
        eval_points = list(eval_points)
        self.eval_points = tuple(eval_points)
        if len(self.eval_points) != len(set(self.eval_points)):
            raise ValueError("duplicate eval points")
        self.lhs_eval_ops = tuple(lhs_eval_ops)
        if len(self.lhs_eval_ops) != self.prod_part_count:
            raise ValueError("wrong number of lhs_eval_ops")
        self.rhs_eval_ops = tuple(rhs_eval_ops)
        if len(self.rhs_eval_ops) != self.prod_part_count:
            raise ValueError("wrong number of rhs_eval_ops")
        if len(self.eval_points) != self.prod_part_count:
            raise ValueError("wrong number of eval_points")
        self.prod_eval_ops = tuple(prod_eval_ops)
        if len(self.prod_eval_ops) != self.prod_part_count:
            raise ValueError("wrong number of prod_eval_ops")

        lhs_eval_polys = self.get_lhs_eval_polys()
        for i, eval_op in enumerate(self.lhs_eval_ops):
            if lhs_eval_polys[i] != eval_op.poly:
                raise ValueError(
                    f"lhs_eval_ops[{i}] is incorrect: expected polynomial: "
                    f"{lhs_eval_polys[i]} found polynomial: {eval_op.poly}")

        rhs_eval_polys = self.get_rhs_eval_polys()
        for i, eval_op in enumerate(self.rhs_eval_ops):
            if rhs_eval_polys[i] != eval_op.poly:
                raise ValueError(
                    f"rhs_eval_ops[{i}] is incorrect: expected polynomial: "
                    f"{rhs_eval_polys[i]} found polynomial: {eval_op.poly}")

        prod_eval_polys = self.get_prod_eval_polys()  # also checks matrix
        for i, eval_op in enumerate(self.prod_eval_ops):
            if prod_eval_polys[i] != eval_op.poly:
                raise ValueError(
                    f"prod_eval_ops[{i}] is incorrect: expected polynomial: "
                    f"{prod_eval_polys[i]} found polynomial: {eval_op.poly}")

    def reversed(self):
        # type: () -> ToomCookInstance
        """return a ToomCookInstance where lhs/rhs are reversed"""
        return ToomCookInstance(
            lhs_part_count=self.rhs_part_count,
            rhs_part_count=self.lhs_part_count,
            eval_points=self.eval_points,
            lhs_eval_ops=self.rhs_eval_ops,
            rhs_eval_ops=self.lhs_eval_ops,
            prod_eval_ops=self.prod_eval_ops)

    @staticmethod
    def make_toom_2():
        # type: () -> ToomCookInstance
        """make an instance of Toom-2 aka Karatsuba multiplication"""
        return ToomCookInstance(
            lhs_part_count=2,
            rhs_part_count=2,
            eval_points=[0, 1, POINT_AT_INFINITY],
            lhs_eval_ops=[
                EvalOpInput(0),
                EvalOpAdd(EvalOpInput(0), EvalOpInput(1)),
                EvalOpInput(1),
            ],
            rhs_eval_ops=[
                EvalOpInput(0),
                EvalOpAdd(EvalOpInput(0), EvalOpInput(1)),
                EvalOpInput(1),
            ],
            prod_eval_ops=[
                EvalOpInput(0),
                EvalOpSub(EvalOpSub(EvalOpInput(1), EvalOpInput(0)),
                          EvalOpInput(2)),
                EvalOpInput(2),
            ],
        )

    @staticmethod
    def make_toom_2_5():
        # type: () -> ToomCookInstance
        """makes an instance of Toom-2.5"""
        inp_0_plus_inp_2 = EvalOpAdd(EvalOpInput(0), EvalOpInput(2))
        inp_1_minus_inp_2 = EvalOpSub(EvalOpInput(1), EvalOpInput(2))
        inp_1_plus_inp_2 = EvalOpAdd(EvalOpInput(1), EvalOpInput(2))
        inp_1_minus_inp_2_all_div_2 = EvalOpExactDiv(inp_1_minus_inp_2, 2)
        inp_1_plus_inp_2_all_div_2 = EvalOpExactDiv(inp_1_plus_inp_2, 2)
        return ToomCookInstance(
            lhs_part_count=3,
            rhs_part_count=2,
            eval_points=[0, 1, -1, POINT_AT_INFINITY],
            lhs_eval_ops=[
                EvalOpInput(0),
                EvalOpAdd(inp_0_plus_inp_2, EvalOpInput(1)),
                EvalOpSub(inp_0_plus_inp_2, EvalOpInput(1)),
                EvalOpInput(2),
            ],
            rhs_eval_ops=[
                EvalOpInput(0),
                EvalOpAdd(EvalOpInput(0), EvalOpInput(1)),
                EvalOpSub(EvalOpInput(0), EvalOpInput(1)),
                EvalOpInput(1),
            ],
            prod_eval_ops=[
                EvalOpInput(0),
                EvalOpSub(inp_1_minus_inp_2_all_div_2, EvalOpInput(3)),
                EvalOpSub(inp_1_plus_inp_2_all_div_2, EvalOpInput(0)),
                EvalOpInput(3),
            ],
        )

    # TODO: add make_toom_3


@plain_data(frozen=True, unsafe_hash=True)
@final
class PartialProduct:
    __slots__ = "ssa_val_spread", "shift_in_words", "is_signed", "subtract"

    def __init__(self, ssa_val_spread, shift_in_words, is_signed, subtract):
        # type: (Iterable[SSAVal], int, bool, bool) -> None
        if shift_in_words < 0:
            raise ValueError("invalid shift_in_words")
        self.ssa_val_spread = tuple(ssa_val_spread)
        for ssa_val in ssa_val_spread:
            if ssa_val.ty != Ty(base_ty=BaseTy.I64, reg_len=1):
                raise ValueError("invalid ssa_val.ty")
        self.shift_in_words = shift_in_words
        self.is_signed = is_signed
        self.subtract = subtract


def sum_partial_products(fn, partial_products, retval_size, name):
    # type: (Fn, Iterable[PartialProduct], int, str) -> SSAVal
    retval_spread = []  # type: list[SSAVal]
    retval_signed = False
    zero = fn.append_new_op(OpKind.LI, immediates=[0],
                            name=f"{name}_zero").outputs[0]
    has_carry_word = False
    for idx, partial_product in enumerate(partial_products):
        shift_in_words = partial_product.shift_in_words
        spread = list(partial_product.ssa_val_spread)
        if (not retval_signed and shift_in_words >= len(retval_spread)
                and not partial_product.subtract):
            retval_spread.extend(
                [zero] * (shift_in_words - len(retval_spread)))
            retval_spread.extend(spread)
            retval_signed = partial_product.is_signed
            has_carry_word = False
            continue
        assert len(retval_spread) != 0, "logic error"
        retval_hi_len = len(retval_spread) - shift_in_words
        if retval_hi_len <= len(spread):
            maxvl = len(spread) + 1
            has_carry_word = True
        elif has_carry_word:
            maxvl = retval_hi_len
        else:
            maxvl = retval_hi_len + 1
            has_carry_word = True
        if not has_carry_word:
            maxvl += 1
            has_carry_word = True
        if maxvl > retval_size - shift_in_words:
            maxvl = retval_size - shift_in_words
            has_carry_word = False
        retval_spread = cast_to_size_spread(
            fn=fn, ssa_vals=retval_spread, src_signed=retval_signed,
            dest_size=maxvl + shift_in_words, name=f"{name}_{idx}_cast_retval")
        spread = cast_to_size_spread(
            fn=fn, ssa_vals=spread, src_signed=partial_product.is_signed,
            dest_size=maxvl, name=f"{name}_{idx}_cast_pp")
        setvl = fn.append_new_op(
            OpKind.SetVLI, immediates=[maxvl],
            maxvl=maxvl, name=f"{name}_{idx}_setvl")
        retval_concat = fn.append_new_op(
            kind=OpKind.Concat,
            input_vals=[*retval_spread[shift_in_words:], setvl.outputs[0]],
            name=f"{name}_{idx}_retval_concat", maxvl=maxvl)
        pp_concat = fn.append_new_op(
            kind=OpKind.Concat,
            input_vals=[*spread, setvl.outputs[0]],
            name=f"{name}_{idx}_pp_concat", maxvl=maxvl)
        if partial_product.subtract:
            set_ca = fn.append_new_op(kind=OpKind.SetCA,
                                      name=f"{name}_{idx}_set_ca")
            add_sub = fn.append_new_op(
                kind=OpKind.SvSubFE, input_vals=[
                    pp_concat.outputs[0], retval_concat.outputs[0],
                    set_ca.outputs[0], setvl.outputs[0]],
                maxvl=maxvl, name=f"{name}_{idx}_sub")
        else:
            clear_ca = fn.append_new_op(kind=OpKind.ClearCA,
                                        name=f"{name}_{idx}_clear_ca")
            add_sub = fn.append_new_op(
                kind=OpKind.SvAddE, input_vals=[
                    retval_concat.outputs[0], pp_concat.outputs[0],
                    clear_ca.outputs[0], setvl.outputs[0]],
                maxvl=maxvl, name=f"{name}_{idx}_add")
        retval_spread[shift_in_words:] = fn.append_new_op(
            kind=OpKind.Spread,
            input_vals=[add_sub.outputs[0], setvl.outputs[0]],
            name=f"{name}_{idx}_sum_spread", maxvl=maxvl).outputs
    retval_spread = cast_to_size_spread(
        fn=fn, ssa_vals=retval_spread, src_signed=retval_signed,
        dest_size=retval_size, name=f"{name}_retval_cast")
    retval_setvl = fn.append_new_op(
        OpKind.SetVLI, immediates=[retval_size],
        maxvl=retval_size, name=f"{name}_setvl")
    retval_concat = fn.append_new_op(
        kind=OpKind.Concat,
        input_vals=[*retval_spread, retval_setvl.outputs[0]],
        name=f"{name}_concat", maxvl=retval_size)
    return retval_concat.outputs[0]


def simple_mul(fn, lhs, lhs_signed, rhs, rhs_signed, name, retval_size=None):
    # type: (Fn, SSAVal, bool, SSAVal, bool, str, int | None) -> SSAVal
    """ simple O(n^2) big-int multiply """
    if retval_size is None:
        retval_size = lhs.ty.reg_len + rhs.ty.reg_len
    if lhs.ty.reg_len < rhs.ty.reg_len:
        lhs, rhs = rhs, lhs
        lhs_signed, rhs_signed = rhs_signed, lhs_signed
    # split rhs into elements
    rhs_setvl = fn.append_new_op(
        kind=OpKind.SetVLI, immediates=[rhs.ty.reg_len],
        name=f"{name}_rhs_setvl")
    rhs_spread = fn.append_new_op(
        kind=OpKind.Spread, input_vals=[rhs, rhs_setvl.outputs[0]],
        maxvl=rhs.ty.reg_len, name=f"{name}_rhs_spread")
    rhs_words = rhs_spread.outputs
    zero = fn.append_new_op(
        kind=OpKind.LI, immediates=[0], name=f"{name}_zero").outputs[0]
    maxvl = lhs.ty.reg_len
    lhs_setvl = fn.append_new_op(
        kind=OpKind.SetVLI, immediates=[maxvl], name=f"{name}_lhs_setvl",
        maxvl=maxvl)
    vl = lhs_setvl.outputs[0]

    def partial_products():
        # type: () -> Iterable[PartialProduct]
        for shift_in_words, rhs_word in enumerate(rhs_words):
            mul = fn.append_new_op(
                kind=OpKind.SvMAddEDU, input_vals=[lhs, rhs_word, zero, vl],
                maxvl=maxvl, name=f"{name}_{shift_in_words}_mul")
            mul_rt_spread = fn.append_new_op(
                kind=OpKind.Spread, input_vals=[mul.outputs[0], vl],
                name=f"{name}_{shift_in_words}_mul_rt_spread", maxvl=maxvl)
            yield PartialProduct(
                ssa_val_spread=[*mul_rt_spread.outputs, mul.outputs[1]],
                shift_in_words=shift_in_words,
                is_signed=False, subtract=False)
        if lhs_signed:
            lhs_spread = fn.append_new_op(
                kind=OpKind.Spread, input_vals=[lhs, lhs_setvl.outputs[0]],
                maxvl=lhs.ty.reg_len, name=f"{name}_lhs_spread")
            rhs_mask = fn.append_new_op(
                kind=OpKind.SRADI, input_vals=[lhs_spread.outputs[-1]],
                immediates=[GPR_SIZE_IN_BITS - 1], name=f"{name}_rhs_mask")
            lhs_and = fn.append_new_op(
                kind=OpKind.SvAndVS,
                input_vals=[rhs, rhs_mask.outputs[0], rhs_setvl.outputs[0]],
                maxvl=rhs.ty.reg_len, name=f"{name}_rhs_and")
            rhs_and_spread = fn.append_new_op(
                kind=OpKind.Spread,
                input_vals=[lhs_and.outputs[0], rhs_setvl.outputs[0]],
                name=f"{name}_rhs_and_spread", maxvl=rhs.ty.reg_len)
            yield PartialProduct(
                ssa_val_spread=rhs_and_spread.outputs,
                shift_in_words=lhs.ty.reg_len, is_signed=False, subtract=True)
        if rhs_signed:
            rhs_spread = fn.append_new_op(
                kind=OpKind.Spread, input_vals=[rhs, rhs_setvl.outputs[0]],
                maxvl=rhs.ty.reg_len, name=f"{name}_rhs_spread")
            lhs_mask = fn.append_new_op(
                kind=OpKind.SRADI, input_vals=[rhs_spread.outputs[-1]],
                immediates=[GPR_SIZE_IN_BITS - 1], name=f"{name}_lhs_mask")
            rhs_and = fn.append_new_op(
                kind=OpKind.SvAndVS,
                input_vals=[lhs, lhs_mask.outputs[0], lhs_setvl.outputs[0]],
                maxvl=lhs.ty.reg_len, name=f"{name}_lhs_and")
            lhs_and_spread = fn.append_new_op(
                kind=OpKind.Spread,
                input_vals=[rhs_and.outputs[0], lhs_setvl.outputs[0]],
                name=f"{name}_lhs_and_spread", maxvl=lhs.ty.reg_len)
            yield PartialProduct(
                ssa_val_spread=lhs_and_spread.outputs,
                shift_in_words=rhs.ty.reg_len, is_signed=False, subtract=True)
    return sum_partial_products(
        fn=fn, partial_products=partial_products(),
        retval_size=retval_size, name=name)


def cast_to_size(fn, ssa_val, src_signed, dest_size, name):
    # type: (Fn, SSAVal, bool, int, str) -> SSAVal
    if dest_size <= 0:
        raise ValueError("invalid dest_size -- must be a positive integer")
    if ssa_val.ty.reg_len == dest_size:
        return ssa_val
    in_setvl = fn.append_new_op(
        OpKind.SetVLI, immediates=[ssa_val.ty.reg_len],
        maxvl=ssa_val.ty.reg_len, name=f"{name}_in_setvl")
    spread = fn.append_new_op(
        OpKind.Spread, input_vals=[ssa_val, in_setvl.outputs[0]],
        name=f"{name}_spread", maxvl=ssa_val.ty.reg_len)
    spread_values = cast_to_size_spread(
        fn=fn, ssa_vals=spread.outputs, src_signed=src_signed,
        dest_size=dest_size, name=name)
    out_setvl = fn.append_new_op(
        OpKind.SetVLI, immediates=[dest_size], maxvl=dest_size,
        name=f"{name}_out_setvl")
    concat = fn.append_new_op(
        OpKind.Concat, input_vals=[*spread_values, out_setvl.outputs[0]],
        name=f"{name}_concat", maxvl=dest_size)
    return concat.outputs[0]


def cast_to_size_spread(fn, ssa_vals, src_signed, dest_size, name):
    # type: (Fn, Iterable[SSAVal], bool, int, str) -> list[SSAVal]
    if dest_size <= 0:
        raise ValueError("invalid dest_size -- must be a positive integer")
    spread_values = list(ssa_vals)
    for ssa_val in ssa_vals:
        if ssa_val.ty != Ty(base_ty=BaseTy.I64, reg_len=1):
            raise ValueError("invalid ssa_val.ty")
    if len(spread_values) == dest_size:
        return spread_values
    if len(spread_values) > dest_size:
        spread_values[dest_size:] = []
    elif src_signed:
        sign = fn.append_new_op(
            OpKind.SRADI, input_vals=[spread_values[-1]],
            immediates=[GPR_SIZE_IN_BITS - 1], name=f"{name}_sign")
        spread_values += [sign.outputs[0]] * (dest_size - len(spread_values))
    else:
        zero = fn.append_new_op(
            OpKind.LI, immediates=[0], name=f"{name}_zero")
        spread_values += [zero.outputs[0]] * (dest_size - len(spread_values))
    return spread_values


def split_into_exact_sized_parts(fn, ssa_val, part_count, part_size, name):
    # type: (Fn, SSAVal, int, int, str) -> tuple[SSAVal, ...]
    """split ssa_val into part_count parts, where all but the last part have
    `part.ty.reg_len == part_size`.
    """
    if part_size <= 0:
        raise ValueError("invalid part size, must be positive")
    if part_count <= 0:
        raise ValueError("invalid part count, must be positive")
    if part_count == 1:
        return (ssa_val,)
    too_short_reg_len = (part_count - 1) * part_size
    if ssa_val.ty.reg_len <= too_short_reg_len:
        raise ValueError(f"ssa_val is too short to split, must have "
                         f"reg_len > {too_short_reg_len}: {ssa_val}")
    maxvl = ssa_val.ty.reg_len
    setvl = fn.append_new_op(OpKind.SetVLI, immediates=[maxvl],
                             maxvl=maxvl, name=f"{name}_setvl")
    spread = fn.append_new_op(
        OpKind.Spread, input_vals=[ssa_val, setvl.outputs[0]],
        name=f"{name}_spread", maxvl=maxvl)
    retval = []  # type: list[SSAVal]
    for part in range(part_count):
        start = part * part_size
        stop = min(maxvl, start + part_size)
        part_maxvl = stop - start
        part_setvl = fn.append_new_op(
            OpKind.SetVLI, immediates=[part_size], maxvl=part_size,
            name=f"{name}_{part}_setvl")
        concat = fn.append_new_op(
            OpKind.Concat,
            input_vals=[*spread.outputs[start:stop], part_setvl.outputs[0]],
            name=f"{name}_{part}_concat", maxvl=part_maxvl)
        retval.append(concat.outputs[0])
    return tuple(retval)


_TCIs = Tuple[ToomCookInstance, ...]


@plain_data(frozen=True)
@final
class ToomCookMul:
    __slots__ = (
        "fn", "lhs", "lhs_signed", "rhs", "rhs_signed", "instances",
        "retval_size", "start_instance_index", "instance", "part_size",
        "lhs_parts", "lhs_inputs", "lhs_eval_state", "lhs_outputs",
        "rhs_parts", "rhs_inputs", "rhs_eval_state", "rhs_outputs",
        "prod_inputs", "prod_eval_state", "prod_parts",
        "partial_products", "retval",
    )

    def __init__(self, fn, lhs, lhs_signed, rhs, rhs_signed, instances,
                 retval_size=None, start_instance_index=0):
        # type: (Fn, SSAVal, bool, SSAVal, bool, _TCIs, None | int, int) -> None
        self.fn = fn
        self.lhs = lhs
        self.lhs_signed = lhs_signed
        self.rhs = rhs
        self.rhs_signed = rhs_signed
        self.instances = instances
        if retval_size is None:
            retval_size = lhs.ty.reg_len + rhs.ty.reg_len
        self.retval_size = retval_size
        if start_instance_index < 0:
            raise ValueError("start_instance_index must be non-negative")
        self.start_instance_index = start_instance_index
        self.instance = None
        self.part_size = 0  # type: int
        while start_instance_index < len(instances):
            self.instance = instances[start_instance_index]
            self.part_size = max(
                lhs.ty.reg_len // self.instance.lhs_part_count,
                rhs.ty.reg_len // self.instance.rhs_part_count)
            if self.part_size <= 0:
                self.instance = None
                start_instance_index += 1
            else:
                break
        if self.instance is None:
            self.retval = simple_mul(fn=fn,
                                     lhs=lhs, lhs_signed=lhs_signed,
                                     rhs=rhs, rhs_signed=rhs_signed,
                                     name="toom_cook_base_case")
            return
        self.lhs_parts = split_into_exact_sized_parts(
            fn=fn, ssa_val=lhs, part_count=self.instance.lhs_part_count,
            part_size=self.part_size, name="lhs")
        self.lhs_inputs = []  # type: list[EvalOpGenIrInput]
        for part, ssa_val in enumerate(self.lhs_parts):
            self.lhs_inputs.append(EvalOpGenIrInput(
                ssa_val=ssa_val,
                is_signed=lhs_signed and part == len(self.lhs_parts) - 1))
        self.lhs_eval_state = EvalOpGenIrState(fn=fn, inputs=self.lhs_inputs)
        lhs_eval_ops = self.instance.lhs_eval_ops
        self.lhs_outputs = [
            self.lhs_eval_state.get_output(i) for i in lhs_eval_ops]
        self.rhs_parts = split_into_exact_sized_parts(
            fn=fn, ssa_val=rhs, part_count=self.instance.rhs_part_count,
            part_size=self.part_size, name="rhs")
        self.rhs_inputs = []  # type: list[EvalOpGenIrInput]
        for part, ssa_val in enumerate(self.rhs_parts):
            self.rhs_inputs.append(EvalOpGenIrInput(
                ssa_val=ssa_val,
                is_signed=rhs_signed and part == len(self.rhs_parts) - 1))
        self.rhs_eval_state = EvalOpGenIrState(fn=fn, inputs=self.rhs_inputs)
        rhs_eval_ops = self.instance.rhs_eval_ops
        self.rhs_outputs = [
            self.rhs_eval_state.get_output(i) for i in rhs_eval_ops]
        self.prod_inputs = []  # type: list[EvalOpGenIrInput]
        for lhs_output, rhs_output in zip(self.lhs_outputs, self.rhs_outputs):
            ssa_val = toom_cook_mul(
                fn=fn,
                lhs=lhs_output.output, lhs_signed=lhs_output.is_signed,
                rhs=rhs_output.output, rhs_signed=rhs_output.is_signed,
                instances=instances,
                start_instance_index=start_instance_index + 1)
            products = (lhs_output.min_value * rhs_output.min_value,
                        lhs_output.min_value * rhs_output.max_value,
                        lhs_output.max_value * rhs_output.min_value,
                        lhs_output.max_value * rhs_output.max_value)
            self.prod_inputs.append(EvalOpGenIrInput(
                ssa_val=ssa_val,
                is_signed=None,
                min_value=min(products),
                max_value=max(products)))
        self.prod_eval_state = EvalOpGenIrState(fn=fn, inputs=self.prod_inputs)
        prod_eval_ops = self.instance.prod_eval_ops
        self.prod_parts = [
            self.prod_eval_state.get_output(i) for i in prod_eval_ops]

        def partial_products():
            # type: () -> Iterable[PartialProduct]
            for part, prod_part in enumerate(self.prod_parts):
                part_maxvl = prod_part.output.ty.reg_len
                part_setvl = fn.append_new_op(
                    OpKind.SetVLI, immediates=[part_maxvl],
                    name=f"prod_{part}_setvl", maxvl=part_maxvl)
                spread_part = fn.append_new_op(
                    OpKind.Spread,
                    input_vals=[prod_part.output, part_setvl.outputs[0]],
                    name=f"prod_{part}_spread", maxvl=part_maxvl)
                yield PartialProduct(
                    spread_part.outputs, shift_in_words=part * self.part_size,
                    is_signed=prod_part.is_signed, subtract=False)
        self.partial_products = tuple(partial_products())
        self.retval = sum_partial_products(
            fn=fn, partial_products=self.partial_products,
            retval_size=retval_size, name="prod")


def toom_cook_mul(fn, lhs, lhs_signed, rhs, rhs_signed, instances,
                  retval_size=None, start_instance_index=0):
    # type: (Fn, SSAVal, bool, SSAVal, bool, _TCIs, None | int, int) -> SSAVal
    return ToomCookMul(
        fn=fn, lhs=lhs, lhs_signed=lhs_signed, rhs=rhs, rhs_signed=rhs_signed,
        instances=instances, retval_size=retval_size,
        start_instance_index=start_instance_index).retval