Circuit sublasses with arguments or verilog modules with parameters

[rkshthrmsh]

Hi @leonardt, this is a longer question. For the bootcamp, I am trying to show how higher order programming can be used to make the code concise. The log_shifter.py seems like a good candidate (https://github.com/leonardt/magma_examples/blob/master/magma_examples/log_shifter.py).

The idea is to create a new Circuit subclass to define a Register-like circuit that takes a shift_amount argument. The shift_amount determines the left shift which is hard-coded in the original example to 8, 4, 2, and 1. Ex.: s0.I @= io.I << 8

My attempt:

import magma as m

class ModifiedReg(m.Circuit):
    # How to make a Circuit subclass with arguments that can be passed? 
    # Similar to magma Primitives: m.Register(..., has_enable=True)()
    # The first comment in LogShifter class definition shows a use case
    def __init__(self, shift_amount: int=1):
        self.shift_amount = shift_amount
    # This could be compiled to verilog module parameters? 
    # Ex.: module ModifiedReg #(parameter shift_amount=1) (//ports);

    io = m.IO(
            # Is it possible to get m.Register IOs here?
            # Something like get_ports(m.Register(m.UInt[16])) from magma Primitives 
            # or other Circuit subclasses
            # Else add them manually
            I=m.In(m.UInt[16]),
            O=m.Out(m.UInt[16]),
            # and then add a new port
            shift_en=m.In(m.Bit),
    )
    reg = m.Register(m.UInt[16])()

    with m.when(io.shift_en):
        # use argument to determine shift amount
        reg.I @= reg.O << self.shift_amount # NameError: name 'self' is not defined

    io.O @= reg.O

class LogShifter(m.Circuit):
    io = m.IO(
        I=m.In(m.UInt[16]),
        shift_amount=m.In(m.UInt[4]),
        O=m.Out(m.UInt[16])
    ) + m.ClockIO()

    regs = [ModifiedReg(shift_amount=2**i)() for i in range(3,0,-1)]  # since shift amount is 'hard-coded'
    shift_en_ = io.shift_amount[1:]
    O = m.braid(regs, foldargs={"I":"O"}, joinargs=["shift_en"])(I=io.I, shift_en=shift_en_)

    with m.when(io.shift_amount[0]):
        io.O @= O << 1
    with m.otherwise:
        io.O @= O

I suppose this should have been two posts:

1. How to make circuit subclasses with arguments (or verilog modules with parameters)? I am aware of Generator2, but in this situation, it will generate definitions for multiple verilog modules instead of defining a single module with parameters. Here is the same example re-written with generator:

class ModifiedReg(m.Generator2):
    def __init__(self, shift_amount: int=1):
        self.shift_amount = shift_amount
        self.io = io = m.IO(
            I=m.In(m.UInt[16]),
            O=m.Out(m.UInt[16]),
            shift_en=m.In(m.Bit),
        ) + m.ClockIO()
        reg = m.Register(m.UInt[16])()

        with m.when(io.shift_en):
            reg.I @= reg.O << self.shift_amount

        io.O @= reg.O

Which generates:

module ModifiedReg(
// skipping...
  always_comb begin
    _GEN = Register_inst0;
    if (shift_en)
      _GEN = {Register_inst0[7:0], 8'h0}; // << 8
  end // always_comb
// skipping...
endmodule

module ModifiedReg_unq1(
// skipping...
  always_comb begin
    _GEN = Register_inst0;
    if (shift_en)
      _GEN = {Register_inst0[11:0], 4'h0}; // << 4
  end // always_comb
// skipping...
endmodule

module ModifiedReg_unq2(
// skipping...
  always_comb begin
    _GEN = Register_inst0;
    if (shift_en)
      _GEN = {Register_inst0[13:0], 2'h0}; // << 2
  end // always_comb
// skipping...
endmodule

A parameterized module could have been:

module ModifiedReg # (
  parameter shift_amount=1)(
//skipping...
  always_comb begin
    _GEN = Register_inst0;
    if (shift_en)
      _GEN = {Register_inst0[15-shift_amount:0], {shift_amount{1'h0}}}; // << shift_amount
  end // always_comb

Do you see this as a worthwhile distinction?

2. Is there a way to derive ports from other circuit's or primitive's ports? Which would be very useful for verilog wrapper modules.

[leonardt]

1. We currently don't support parametrized modules without using Generator. There's some practical reasons, for example inside the m.Circuit class definition there's no way to reference a parameter without wrapping it in either a factory function or something like Generator2. We'd need to introduce a notion of "symbolic" values that are used in the circuit definition and this seems like a lot of work to replicate the exact same frontend feature as Generator2. Is there a specific reason you need the modules to be a single module in the generated verilog? It's not difficult to add logic to Generator2 to emit a list of module names corresponding to a Generator2, which can be used for pattern matching in tools downstream.

2. Can you explain more on what you mean by derive ports? You could subclass a circuit and introspect the inherited io object. In general in magma the interface ports can be introspected, so you can easily write code that loops over the contents of the ports dictionary, for example.

[rkshthrmsh]

1. I suspect having multiple module definitions and instances may have an impact on the performance of simulation and synthesis tools using the verilog output, especially in larger designs. Will try to run a small experiment and see if this holds. Secondly, the verilog module parameterization gives concise module definitions and I was wondering what sort of effort would be needed to enable this in magma or whether this was tricky because of CIRCT/MLIR.
2. Maybe I can explain better with a use case: Consider a 'core' verilog module that is being wrapped to add context specific logic like IO synchronizers. The wrapper module shares the same IOs as the core. My idea was that since the core is already defined (i.e., IOs are fixed), having native support to populate m.IO for the wrapper from the core's m.IO would be quite enabling. A very simple, contrived example:

class Xor(m.Circuit):
    io = m.IO(
        A=m.In(m.Bit),
        B=m.In(m.Bit),
        C=m.Out(m.Bit)
    )
    io.C @= io.A ^ io.B

class ShiftRegisterXor(m.Circuit):
    io = m.IO(
        # This circuit "wraps" some Registers and a Xor instance
        # Has one input wired to the first m.Register(m.Bit) and 
        # another wired to Xor. Output is the output of Xor.
        get_io(m.Register(m.Bit), i), # only get inputs of m.Register(m.Bit):
        #I=m.In(m.Bit),
        get_io(Xor, io, in_prefix='A'), # get IOs but only inputs with prefix:
        #A=m.In(m.Bit),
        #C=m.Out(m.Bit)
    ) + m.ClockIO()
    O = m.fold([m.Register(m.Bit)() for _ in range(4)], foldargs={"I":"O"})(io.I)
    io.C @= Xor()(io.A, O)