[nmigen-gf.git] / src / nmigen_gf / hdl / cldivrem.py

# SPDX-License-Identifier: LGPL-3-or-later
# Copyright 2022 Jacob Lifshay programmerjake@gmail.com

# Funded by NLnet Assure Programme 2021-02-052, https://nlnet.nl/assure part
# of Horizon 2020 EU Programme 957073.

""" Carry-less Division and Remainder.

https://bugs.libre-soc.org/show_bug.cgi?id=784
"""

from dataclasses import dataclass, field, fields
from nmigen.hdl.ir import Elaboratable
from nmigen.hdl.ast import Signal, Value
from nmigen.hdl.dsl import Module
from nmutil.singlepipe import ControlBase
from nmutil.clz import CLZ, clz


def equal_leading_zero_count_reference(a, b, width):
    """checks if `clz(a) == clz(b)`.
    Reference code for algorithm used in `EqualLeadingZeroCount`.
    """
    assert isinstance(width, int) and 0 <= width
    assert isinstance(a, int) and 0 <= a < (1 << width)
    assert isinstance(b, int) and 0 <= b < (1 << width)
    eq = True  # both have no leading zeros so far...
    for i in range(width):
        a_bit = (a >> i) & 1
        b_bit = (b >> i) & 1
        # `both_ones` is set if both have no leading zeros so far
        both_ones = a_bit & b_bit
        # `different` is set if there are a different number of leading
        # zeros so far
        different = a_bit != b_bit
        if both_ones:
            eq = True
        elif different:
            eq = False
        else:
            pass  # propagate from lower bits
    return eq


class EqualLeadingZeroCount(Elaboratable):
    """checks if `clz(a) == clz(b)`.

    Properties:
    width: int
        the width in bits of `a` and `b`.
    a: Signal of width `width`
        input
    b: Signal of width `width`
        input
    out: Signal of width `1`
        output, set if the number of leading zeros in `a` is the same as in
        `b`.
    """

    def __init__(self, width):
        assert isinstance(width, int)
        self.width = width
        self.a = Signal(width)
        self.b = Signal(width)
        self.out = Signal()

    def elaborate(self, platform):
        # the operation is converted into calculation of the carry-out of a
        # binary addition, allowing FPGAs to re-use their specialized
        # carry-propagation logic. This should be simplified by yosys to
        # remove the extraneous xor gates from addition when targeting
        # FPGAs/ASICs, so no efficiency is lost.
        #
        # see `equal_leading_zero_count_reference` for a Python version of
        # the algorithm, but without conversion to carry-propagation.
        # note that it's possible to do all the bits at once: a for-loop
        # (unlike in the reference-code) is not necessary

        m = Module()
        both_ones = Signal(self.width)
        different = Signal(self.width)

        # build `both_ones` and `different` such that:
        # for every bit index `i`:
        # * if `both_ones[i]` is set, then both addends bits at index `i` are
        #   set in order to set the carry bit out, since `cin + 1 + 1` always
        #   has a carry out.
        # * if `different[i]` is set, then both addends bits at index `i` are
        #   zeros in order to clear the carry bit out, since `cin + 0 + 0`
        #   never has a carry out.
        # * otherwise exactly one of the addends bits at index `i` is set and
        #   the other is clear in order to propagate the carry bit from
        #   less significant bits, since `cin + 1 + 0` has a carry out that is
        #   equal to `cin`.

        # `both_ones` is set if both have no leading zeros so far
        m.d.comb += both_ones.eq(self.a & self.b)
        # `different` is set if there are a different number of leading
        # zeros so far
        m.d.comb += different.eq(self.a ^ self.b)

        # now [ab]use add: the last bit [carry-out] is the result
        csum = Signal(self.width + 1)
        carry_in = 1  # both have no leading zeros so far, so set carry in
        m.d.comb += csum.eq(both_ones + (~different) + carry_in)
        m.d.comb += self.out.eq(csum[self.width])  # out is carry-out

        return m


def cldivrem_shifting(n, d, width):
    """ Carry-less Division and Remainder based on shifting at start and end
        allowing us to get away with checking a single bit each iteration
        rather than checking for equal degrees every iteration.
        `n` and `d` are integers, `width` is the number of bits needed to hold
        each input/output.
        Returns a tuple `q, r` of the quotient and remainder.
    """
    assert isinstance(width, int) and width >= 1
    assert isinstance(n, int) and 0 <= n < 1 << width
    assert isinstance(d, int) and 0 <= d < 1 << width
    assert d != 0, "TODO: decide what happens on division by zero"

    shape = CLDivRemShape(width)

    # `clz(d, width)`, but maxes out at `width - 1` instead of `width` in
    # order to both fit in `shape.shift_width` bits and to not shift by more
    # than needed.
    shift = clz(d >> 1, width - 1)
    assert 0 <= shift < 1 << shape.shift_width, "shift overflow"
    d <<= shift
    assert 0 <= d < 1 << shape.d_width, "d overflow"
    r = n << shift
    assert 0 <= r < 1 << shape.r_width, "r overflow"
    q = 0
    for step in range(width):
        q <<= 1
        r <<= 1
        if r >> (width * 2 - 1) != 0:
            r ^= d << width
            q |= 1
        assert 0 <= q < 1 << shape.q_width, "q overflow"
        assert 0 <= r < 1 << shape.r_width, "r overflow"
    r >>= width
    r >>= shift
    return q, r


@dataclass(frozen=True, unsafe_hash=True)
class CLDivRemShape:
    width: int

    def __post_init__(self):
        assert isinstance(self.width, int) and self.width >= 1, "invalid width"

    @property
    def done_step(self):
        """the step number when iteration is finished
        -- the largest `CLDivRemState.step` will get
        """
        return self.width

    @property
    def step_range(self):
        """the range that `CLDivRemState.step` will fall in.

        returns: range
        """
        return range(self.done_step + 1)

    @property
    def d_width(self):
        """bit-width of the internal signal `CLDivRemState.d`"""
        return self.width

    @property
    def r_width(self):
        """bit-width of the internal signal `CLDivRemState.r`"""
        return self.width * 2

    @property
    def q_width(self):
        """bit-width of the internal signal `CLDivRemState.q`"""
        return self.width

    @property
    def shift_width(self):
        """bit-width of the internal signal `CLDivRemState.shift`"""
        return (self.width - 1).bit_length()


@dataclass(frozen=True, eq=False)
class CLDivRemState:
    shape: CLDivRemShape
    name: str
    step: Signal = field(init=False)
    d: Signal = field(init=False)
    r: Signal = field(init=False)
    q: Signal = field(init=False)
    shift: Signal = field(init=False)

    def __init__(self, shape, *, name=None, src_loc_at=0):
        assert isinstance(shape, CLDivRemShape)
        if name is None:
            name = Signal(src_loc_at=1 + src_loc_at).name
        assert isinstance(name, str)
        step = Signal(shape.step_range, name=f"{name}_step")
        d = Signal(shape.d_width, name=f"{name}_d")
        r = Signal(shape.r_width, name=f"{name}_r")
        q = Signal(shape.q_width, name=f"{name}_q")
        shift = Signal(shape.shift_width, name=f"{name}_shift")
        object.__setattr__(self, "shape", shape)
        object.__setattr__(self, "name", name)
        object.__setattr__(self, "step", step)
        object.__setattr__(self, "d", d)
        object.__setattr__(self, "r", r)
        object.__setattr__(self, "q", q)
        object.__setattr__(self, "shift", shift)

    def eq(self, rhs):
        assert isinstance(rhs, CLDivRemState)
        for f in fields(CLDivRemState):
            if f.name in ("shape", "name"):
                continue
            l = getattr(self, f.name)
            r = getattr(rhs, f.name)
            yield l.eq(r)

    @staticmethod
    def like(other, *, name=None, src_loc_at=0):
        assert isinstance(other, CLDivRemState)
        return CLDivRemState(other.shape, name=name, src_loc_at=1 + src_loc_at)

    @property
    def done(self):
        return self.will_be_done_after(steps=0)

    def will_be_done_after(self, steps):
        """ Returns True if this state will be done after
            another `steps` passes through `set_to_next`."""
        assert isinstance(steps, int) and steps >= 0
        return self.step >= max(0, self.shape.done_step - steps)

    def get_output(self):
        return self.q, (self.r >> self.shape.width) >> self.shift

    def set_to_initial(self, m, n, d):
        assert isinstance(m, Module)
        n = Value.cast(n)  # convert to Value
        d = Value.cast(d)  # convert to Value
        clz_mod = CLZ(self.shape.width - 1)
        # can't name submodule since it would conflict if this function is
        # called multiple times in a Module
        m.submodules += clz_mod
        assert clz_mod.lz.width == self.shape.shift_width, \
            "internal inconsistency -- mismatched shift signal width"
        m.d.comb += [
            clz_mod.sig_in.eq(d >> 1),
            self.shift.eq(clz_mod.lz),
            self.d.eq(d << self.shift),
            self.r.eq(n << self.shift),
            self.q.eq(0),
            self.step.eq(0),
        ]

    def set_to_next(self, m, state_in):
        assert isinstance(m, Module)
        assert isinstance(state_in, CLDivRemState)
        assert state_in.shape == self.shape
        assert self is not state_in, "a.set_to_next(m, a) is not allowed"
        width = self.shape.width

        with m.If(state_in.done):
            m.d.comb += self.eq(state_in)
        with m.Else():
            m.d.comb += [
                self.step.eq(state_in.step + 1),
                self.d.eq(state_in.d),
                self.shift.eq(state_in.shift),
            ]
            q = state_in.q << 1
            r = state_in.r << 1
            with m.If(r[width * 2 - 1]):
                m.d.comb += [
                    self.q.eq(q | 1),
                    self.r.eq(r ^ (state_in.d << width)),
                ]
            with m.Else():
                m.d.comb += [
                    self.q.eq(q),
                    self.r.eq(r),
                ]


class CLDivRemInputData:
    def __init__(self, shape):
        assert isinstance(shape, CLDivRemShape)
        self.shape = shape
        self.n = Signal(shape.width)
        self.d = Signal(shape.width)

    def __iter__(self):
        """ Get member signals. """
        yield self.n
        yield self.d

    def eq(self, rhs):
        """ Assign member signals. """
        return [
            self.n.eq(rhs.n),
            self.d.eq(rhs.d),
        ]


class CLDivRemOutputData:
    def __init__(self, shape):
        assert isinstance(shape, CLDivRemShape)
        self.shape = shape
        self.q = Signal(shape.width)
        self.r = Signal(shape.width)

    def __iter__(self):
        """ Get member signals. """
        yield self.q
        yield self.r

    def eq(self, rhs):
        """ Assign member signals. """
        return [
            self.q.eq(rhs.q),
            self.r.eq(rhs.r),
        ]


class CLDivRemFSMStage(ControlBase):
    """carry-less div/rem

    Attributes:
    shape: CLDivRemShape
        the shape
    steps_per_clock: int
        number of steps that should be taken per clock cycle
    pspec:
        pipe-spec
    empty: Signal()
        true if nothing is stored in `self.saved_state`
    saved_state: CLDivRemState()
        the saved state that is currently being worked on.
    """

    def __init__(self, pspec, shape, *, steps_per_clock=8):
        assert isinstance(shape, CLDivRemShape)
        assert isinstance(steps_per_clock, int) and steps_per_clock >= 1
        self.shape = shape
        self.steps_per_clock = steps_per_clock
        self.pspec = pspec  # store now: used in ispec and ospec
        super().__init__(stage=self)
        self.empty = Signal(reset=1)
        self.saved_state = CLDivRemState(shape)

    def ispec(self):
        return CLDivRemInputData(self.shape)

    def ospec(self):
        return CLDivRemOutputData(self.shape)

    def setup(self, m, i):
        pass

    def elaborate(self, platform):
        m = super().elaborate(platform)
        i_data: CLDivRemInputData = self.p.i_data
        o_data: CLDivRemOutputData = self.n.o_data

        # TODO: handle cancellation

        m.d.comb += self.n.o_valid.eq(~self.empty & self.saved_state.done)
        m.d.comb += self.p.o_ready.eq(self.empty)

        def make_nc(i):
            return CLDivRemState(self.shape, name=f"next_chain_{i}")
        next_chain = [make_nc(i) for i in range(self.steps_per_clock + 1)]
        for i in range(self.steps_per_clock):
            next_chain[i + 1].set_to_next(m, next_chain[i])
        m.d.comb += next_chain[0].eq(self.saved_state)
        out_q, out_r = self.saved_state.get_output()
        m.d.comb += o_data.q.eq(out_q)
        m.d.comb += o_data.r.eq(out_r)
        initial_state = CLDivRemState(self.shape)
        initial_state.set_to_initial(m, n=i_data.n, d=i_data.d)

        with m.If(self.empty):
            m.d.sync += self.saved_state.eq(initial_state)
            with m.If(self.p.i_valid):
                m.d.sync += self.empty.eq(0)
        with m.Else():
            m.d.sync += self.saved_state.eq(next_chain[-1])
            with m.If(self.n.i_ready & self.n.o_valid):
                m.d.sync += self.empty.eq(1)
        return m

    def __iter__(self):
        yield from self.p
        yield from self.n

    def ports(self):
        return list(self)