In [1]:
from functools import reduce
import magma as m
import mantle

def one_hot_mux(conds, inputs):
    outputs = []
    for cond, inp in zip(conds, inputs):
        outputs.append(inp & m.bits([cond for _ in range(len(inp))]))
    return reduce(lambda x, y: x | y, outputs)


class SimpleALU(m.Circuit):
    name = "SimpleALU"
    IO = ["a", m.In(m.UInt[4]), "b", m.In(m.UInt[4]), 
          "opcode", m.In(m.UInt[2]), "out", m.Out(m.UInt[4])]
    
    @classmethod
    def definition(io):
        is_op0 = io.opcode == m.uint(0, n=2)
        is_op1 = io.opcode == m.uint(1, n=2)
        is_op2 = io.opcode == m.uint(2, n=2)
        is_op3 = io.opcode == m.uint(3, n=2)
        op0_out = io.a + io.b
        op1_out = io.a - io.b
        op2_out = io.a
        op3_out = io.b
        m.wire(io.out, one_hot_mux([is_op0, is_op1, is_op2, is_op3], [op0_out, op1_out, op2_out, op3_out]))

In [2]:
m.compile("build/SimpleALU", SimpleALU, output="coreir")
%cat build/SimpleALU.json


{"top":"global.SimpleALU",
"namespaces":{
  "global":{
    "modules":{
      "Add4_cin":{
        "type":["Record",[
          ["I0",["Array",4,"BitIn"]],
          ["I1",["Array",4,"BitIn"]],
          ["O",["Array",4,"Bit"]],
          ["CIN","BitIn"]
        ]],
        "instances":{
          "bit_const_0_None":{
            "modref":"corebit.const",
            "modargs":{"value":["Bool",false]}
          },
          "coreir_add4_inst0":{
            "genref":"coreir.add",
            "genargs":{"width":["Int",4]}
          },
          "coreir_add4_inst1":{
            "genref":"coreir.add",
            "genargs":{"width":["Int",4]}
          }
        },
        "connections":[
          ["coreir_add4_inst1.in0.1","bit_const_0_None.out"],
          ["coreir_add4_inst1.in0.2","bit_const_0_None.out"],
          ["coreir_add4_inst1.in0.3","bit_const_0_None.out"],
          ["coreir_add4_inst1.out","coreir_add4_inst0.in0"],
          ["self.I1","coreir_add4_inst0.in1"],
          ["self.O","coreir_add4_inst0.out"],
          ["self.CIN","coreir_add4_inst1.in0.0"],
          ["self.I0","coreir_add4_inst1.in1"]
        ]
      },
      "SimpleALU":{
        "type":["Record",[
          ["a",["Array",4,"BitIn"]],
          ["b",["Array",4,"BitIn"]],
          ["opcode",["Array",2,"BitIn"]],
          ["out",["Array",4,"Bit"]]
        ]],
        "instances":{
          "Sub4_inst0":{
            "modref":"global.Sub4"
          },
          "and4_inst0":{
            "genref":"coreir.and",
            "genargs":{"width":["Int",4]}
          },
          "and4_inst1":{
            "genref":"coreir.and",
            "genargs":{"width":["Int",4]}
          },
          "and4_inst2":{
            "genref":"coreir.and",
            "genargs":{"width":["Int",4]}
          },
          "and4_inst3":{
            "genref":"coreir.and",
            "genargs":{"width":["Int",4]}
          },
          "const_0_2":{
            "genref":"coreir.const",
            "genargs":{"width":["Int",2]},
            "modargs":{"value":[["BitVector",2],"2'h0"]}
          },
          "const_1_2":{
            "genref":"coreir.const",
            "genargs":{"width":["Int",2]},
            "modargs":{"value":[["BitVector",2],"2'h1"]}
          },
          "const_2_2":{
            "genref":"coreir.const",
            "genargs":{"width":["Int",2]},
            "modargs":{"value":[["BitVector",2],"2'h2"]}
          },
          "const_3_2":{
            "genref":"coreir.const",
            "genargs":{"width":["Int",2]},
            "modargs":{"value":[["BitVector",2],"2'h3"]}
          },
          "coreir_add4_inst0":{
            "genref":"coreir.add",
            "genargs":{"width":["Int",4]}
          },
          "coreir_eq_2_inst0":{
            "genref":"coreir.eq",
            "genargs":{"width":["Int",2]}
          },
          "coreir_eq_2_inst1":{
            "genref":"coreir.eq",
            "genargs":{"width":["Int",2]}
          },
          "coreir_eq_2_inst2":{
            "genref":"coreir.eq",
            "genargs":{"width":["Int",2]}
          },
          "coreir_eq_2_inst3":{
            "genref":"coreir.eq",
            "genargs":{"width":["Int",2]}
          },
          "or4_inst0":{
            "genref":"coreir.or",
            "genargs":{"width":["Int",4]}
          },
          "or4_inst1":{
            "genref":"coreir.or",
            "genargs":{"width":["Int",4]}
          },
          "or4_inst2":{
            "genref":"coreir.or",
            "genargs":{"width":["Int",4]}
          }
        },
        "connections":[
          ["self.a","Sub4_inst0.I0"],
          ["self.b","Sub4_inst0.I1"],
          ["and4_inst1.in0","Sub4_inst0.O"],
          ["coreir_add4_inst0.out","and4_inst0.in0"],
          ["coreir_eq_2_inst0.out","and4_inst0.in1.0"],
          ["coreir_eq_2_inst0.out","and4_inst0.in1.1"],
          ["coreir_eq_2_inst0.out","and4_inst0.in1.2"],
          ["coreir_eq_2_inst0.out","and4_inst0.in1.3"],
          ["or4_inst0.in0","and4_inst0.out"],
          ["coreir_eq_2_inst1.out","and4_inst1.in1.0"],
          ["coreir_eq_2_inst1.out","and4_inst1.in1.1"],
          ["coreir_eq_2_inst1.out","and4_inst1.in1.2"],
          ["coreir_eq_2_inst1.out","and4_inst1.in1.3"],
          ["or4_inst0.in1","and4_inst1.out"],
          ["self.a","and4_inst2.in0"],
          ["coreir_eq_2_inst2.out","and4_inst2.in1.0"],
          ["coreir_eq_2_inst2.out","and4_inst2.in1.1"],
          ["coreir_eq_2_inst2.out","and4_inst2.in1.2"],
          ["coreir_eq_2_inst2.out","and4_inst2.in1.3"],
          ["or4_inst1.in1","and4_inst2.out"],
          ["self.b","and4_inst3.in0"],
          ["coreir_eq_2_inst3.out","and4_inst3.in1.0"],
          ["coreir_eq_2_inst3.out","and4_inst3.in1.1"],
          ["coreir_eq_2_inst3.out","and4_inst3.in1.2"],
          ["coreir_eq_2_inst3.out","and4_inst3.in1.3"],
          ["or4_inst2.in1","and4_inst3.out"],
          ["coreir_eq_2_inst0.in1","const_0_2.out"],
          ["coreir_eq_2_inst1.in1","const_1_2.out"],
          ["coreir_eq_2_inst2.in1","const_2_2.out"],
          ["coreir_eq_2_inst3.in1","const_3_2.out"],
          ["self.a","coreir_add4_inst0.in0"],
          ["self.b","coreir_add4_inst0.in1"],
          ["self.opcode","coreir_eq_2_inst0.in0"],
          ["self.opcode","coreir_eq_2_inst1.in0"],
          ["self.opcode","coreir_eq_2_inst2.in0"],
          ["self.opcode","coreir_eq_2_inst3.in0"],
          ["or4_inst1.in0","or4_inst0.out"],
          ["or4_inst2.in0","or4_inst1.out"],
          ["self.out","or4_inst2.out"]
        ]
      },
      "Sub4":{
        "type":["Record",[
          ["I0",["Array",4,"BitIn"]],
          ["I1",["Array",4,"BitIn"]],
          ["O",["Array",4,"Bit"]]
        ]],
        "instances":{
          "Add4_cin_inst0":{
            "modref":"global.Add4_cin"
          },
          "Invert4_inst0":{
            "genref":"coreir.not",
            "genargs":{"width":["Int",4]}
          },
          "bit_const_1_None":{
            "modref":"corebit.const",
            "modargs":{"value":["Bool",true]}
          }
        },
        "connections":[
          ["bit_const_1_None.out","Add4_cin_inst0.CIN"],
          ["self.I0","Add4_cin_inst0.I0"],
          ["Invert4_inst0.out","Add4_cin_inst0.I1"],
          ["self.O","Add4_cin_inst0.O"],
          ["self.I1","Invert4_inst0.in"]
        ]
      }
    }
  }
}
}

In [3]:
from magma.simulator import PythonSimulator

simulator = PythonSimulator(SimpleALU)
simulator.set_value(SimpleALU.a, 3)
simulator.set_value(SimpleALU.b, 2)

simulator.set_value(SimpleALU.opcode, 0)
simulator.evaluate()
assert simulator.get_value(SimpleALU.out) == 3 + 2, simulator.get_value(SimpleALU.out)

simulator.set_value(SimpleALU.opcode, 1)
simulator.evaluate()
assert simulator.get_value(SimpleALU.out) == 3 - 2

simulator.set_value(SimpleALU.opcode, 2)
simulator.evaluate()
assert simulator.get_value(SimpleALU.out) == 3

simulator.set_value(SimpleALU.opcode, 3)
simulator.evaluate()
assert simulator.get_value(SimpleALU.out) == 2
print("Success!")


Success!