[PyCDE] Add 'Mux' construct

Multi-bit select signal mux. Detects error cases.

frontends/PyCDE/src/pycde/constructs.py

@@ -5,11 +5,25 @@
 from __future__ import annotations
 
 from .pycde_types import dim
-from .value import Value
+from .value import BitVectorValue, ListValue, Value
 from circt.support import get_value
 from circt.dialects import msft, hw
 import mlir.ir as ir
 
+import typing
+
+
+def Mux(sel: BitVectorValue, *data_inputs: typing.List[Value]):
+  """Create a single mux from a list of values."""
+  num_inputs = len(data_inputs)
+  if num_inputs == 0:
+    raise ValueError("'Mux' must have 1 or more data input")
+  if num_inputs == 1:
+    return data_inputs[0]
+  if sel.type.width != (num_inputs - 1).bit_length():
+    raise TypeError("'Sel' bit width must be clog2 of number of inputs")
+  return ListValue(data_inputs)[sel]
+
 
 def SystolicArray(row_inputs, col_inputs, pe_builder):
   """Build a systolic array."""

frontends/PyCDE/test/muxing.py

@@ -1,8 +1,8 @@
 # RUN: %PYTHON% py-split-input-file.py %s | FileCheck %s
 
-from pycde import System, generator, dim, Input, Output, Value, types
+from pycde import generator, dim, Clock, Input, Output, Value, types
+from pycde.constructs import Mux
 from pycde.testing import unittestmodule
-import sys
 
 
 def array_from_tuple(*input):
@@ -36,7 +36,7 @@ def array_from_tuple(*input):
 @unittestmodule()
 class ComplexMux:
 
-  Clk = Input(dim(1))
+  Clk = Clock()
   In = Input(dim(3, 4, 5))
   Sel = Input(dim(1))
   Out = Output(dim(3, 4))
@@ -46,9 +46,7 @@ class ComplexMux:
 
   @generator
   def create(ports):
-    clk = ports.Clk
-    select_from = Value([ports.In[3].reg(clk).reg(clk, cycles=2), ports.In[1]])
-    ports.Out = select_from[ports.Sel]
+    ports.Out = Mux(ports.Sel, ports.In[3].reg().reg(cycles=2), ports.In[1])
 
     ports.OutArr = array_from_tuple(ports.In[0], ports.In[1])
     ports.OutSlice = ports.In[0:3]