nmigen-alu.py

from nmigen import Cat, Module, Repl, Signal
from nmigen.cli import main


class Adder:
    def __init__(self, width):
        self.a = Signal(width)
        self.b = Signal(width)
        self.subtract = Signal()

        self.o = Signal(width)

    def elaborate(self, platform):
        m = Module()

        with m.If(self.subtract):
            m.d.comb += self.o.eq(self.a - self.b)
        with m.Else():
            m.d.comb += self.o.eq(self.a + self.b)

        return m


class ShiftLeft:
    def __init__(self, width):
        assert width == 32  # For now, anyway.

        self.a = Signal(width)
        self.b = Signal(5)  # Needs to be ceil(log2(width)) for 64-bit

        self.o = Signal(width)

    def elaborate(self, platform):
        m = Module()
        m.d.comb += self.o.eq(self.a << self.b)
        return m

class SignedComparator:
    def __init__(self, width):
        self.a = Signal(width, True)
        self.b = Signal(width, True)

        self.o = Signal()

    def elaborate(self, platform):
        m = Module()
        m.d.comb += self.o.eq(self.a < self.b)
        return m


class ALU:
    def __init__(self, width):
        self.a = Signal(width)
        self.b = Signal(width)
        self.op = Signal(3)
        self.flag = Signal()

        self.o = Signal(width)

        self.add = Adder(width)
        self.shl = ShiftLeft(width)
        self.scmp = SignedComparator(width)

        self.width = width

    def elaborate(self, platform):
        m = Module()

        m.submodules.add = self.add
        m.submodules.shl = self.shl
        m.submodules.scmp = self.scmp

        m.d.comb += [
            self.add.a.eq(self.a),
            self.add.b.eq(self.b),
            self.add.subtract.eq(self.flag),
            self.shl.a.eq(self.a),
            self.shl.b.eq(self.b[:4]),
            self.scmp.a.eq(self.a),
            self.scmp.b.eq(self.b)
        ]

        with m.Switch(self.op):
            with m.Case(0b000):  # Add
                m.d.comb += self.o.eq(self.add.o)
            with m.Case(0b001):
                m.d.comb += self.o.eq(self.shl.o)
            with m.Case(0b010):
                m.d.comb += self.o.eq(Cat(Repl(0, self.width-1), self.scmp.o))

        return m


if __name__ == "__main__":
    alu = ALU(width=32)
    main(alu)