[PyCDE] Clock inputs and clocking blocks (#3413)

Add a `Clock` input port type. Allow a `with` on that port to enter a
'clocking block', which the `Value.reg` function recognizes and uses as
the `clk` input. Implicitly enter a clocking block when a module has
exactly one `Clock` port.

frontends/PyCDE/src/pycde/__init__.py

@@ -3,7 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .module import (externmodule, generator, module, no_connect)
-from .module import (Input, InputChannel, Output, OutputChannel)
+from .module import (Clock, Input, InputChannel, Output, OutputChannel)
 from .system import (System)
 from .pycde_types import (dim, types)
 from .value import (Value)

frontends/PyCDE/src/pycde/module.py

@@ -4,12 +4,13 @@
 
 from __future__ import annotations
 from typing import Tuple, Union, Dict
+from pycde.pycde_types import ClockType
 
 from pycde.support import _obj_to_value
 
 from .support import (get_user_loc, _obj_to_attribute, OpOperandConnect,
                       create_type_string, create_const_zero)
-from .value import Value
+from .value import ClockValue, Value
 
 from circt import support
 from circt.dialects import esi, hw, msft
@@ -56,6 +57,13 @@ class Input(ModuleDecl):
   """Create an RTL-level input port."""
 
 
+class Clock(Input):
+  """Create a clock input"""
+
+  def __init__(self, name: str = None):
+    super().__init__(mlir.ir.IntegerType.get_signless(1), name)
+
+
 class InputChannel(Input):
   """Create an ESI input channel port."""
 
@@ -125,11 +133,23 @@ def create_output_op(args: _GeneratorPortAccess):
   with mlir.ir.InsertionPoint(
       entry_block), generator.loc, BackedgeBuilder(), bc:
     args = _GeneratorPortAccess(spec_mod)
+
+    # Enter clock block implicitly if only one clock given
+    clk = None
+    if len(spec_mod.clock_ports) == 1:
+      clk_port = list(spec_mod.clock_ports.values())[0]
+      val = entry_block.arguments[clk_port]
+      clk = ClockValue(val, ClockType())
+      clk.__enter__()
+
     outputs = generator.gen_func(args)
     if outputs is not None:
       raise ValueError("Generators must not return a value")
     create_output_op(args)
 
+    if clk is not None:
+      clk.__exit__(None, None, None)
+
 
 def create_msft_module_extern_op(sys, mod: _SpecializedModule, symbol):
   """Creation callback for creating a MSFTModuleExternOp."""
@@ -159,9 +179,9 @@ class _SpecializedModule:
   only created if said module is instantiated."""
 
   __slots__ = [
-      "name", "generators", "modcls", "loc", "input_ports", "input_port_lookup",
-      "output_ports", "output_port_lookup", "parameters", "extern_name",
-      "create_cb", "generator_cb"
+      "name", "generators", "modcls", "loc", "clock_ports", "input_ports",
+      "input_port_lookup", "output_ports", "output_port_lookup", "parameters",
+      "extern_name", "create_cb", "generator_cb"
   ]
 
   def __init__(self,
@@ -200,6 +220,7 @@ def __init__(self,
     self.output_port_lookup: Dict[str, int] = {}  # Used by 'BlockArgs' below.
     self.output_ports = []
     self.generators = {}
+    self.clock_ports: Dict[str, int] = {}
     for attr_name in dir(cls):
       if attr_name.startswith("_"):
         continue
@@ -215,6 +236,9 @@ def __init__(self,
         self.output_port_lookup[attr_name] = len(self.output_ports) - 1
       elif isinstance(attr, Generator):
         self.generators[attr_name] = attr
+
+      if isinstance(attr, Clock):
+        self.clock_ports[attr.name] = len(self.input_ports) - 1
     self.add_accessors()
 
   def add_accessors(self):
@@ -562,6 +586,8 @@ def __getattr__(self, name):
       idx = self._mod.input_port_lookup[name]
       entry_block = self._mod.circt_mod.regions[0].blocks[0]
       val = entry_block.arguments[idx]
+      if name in self._mod.clock_ports:
+        return ClockValue(val, ClockType())
       return Value(val)
     if name in self._mod.output_port_lookup:
       if name not in self._output_values:

frontends/PyCDE/src/pycde/pycde_types.py

@@ -4,8 +4,8 @@
 
 from collections import OrderedDict
 
-from .value import (BitVectorValue, ChannelValue, ListValue, StructValue,
-                    RegularValue, InOutValue, Value)
+from .value import (BitVectorValue, ChannelValue, ClockValue, ListValue,
+                    StructValue, RegularValue, InOutValue, Value)
 
 import mlir.ir
 from circt.dialects import esi, hw, sv
@@ -265,12 +265,23 @@ def _get_value_class(self):
     return BitVectorValue
 
 
+class ClockType(PyCDEType):
+  """A special single bit to represent a clock. Can't do any special operations
+  on it, except enter it as a implicit clock block."""
+
+  def __init__(self):
+    super().__init__(mlir.ir.IntegerType.get_signless(1))
+
+  def _get_value_class(self):
+    return ClockValue
+
+
 class ChannelType(PyCDEType):
   """An ESI channel type."""
 
   @property
   def inner_type(self):
-    return PyCDEType(self._type.inner)
+    return Type(self._type.inner)
 
   def _get_value_class(self):
     return ChannelValue

frontends/PyCDE/src/pycde/value.py

@@ -11,6 +11,7 @@
 
 import mlir.ir as ir
 
+from contextvars import ContextVar
 from functools import singledispatchmethod
 from typing import Union
 import re
@@ -43,13 +44,19 @@ def __new__(cls, value, type=None):
 
   _reg_name = re.compile(r"^(.*)__reg(\d+)$")
 
-  def reg(self, clk, rst=None, name=None, cycles=1, sv_attributes=None):
+  def reg(self, clk=None, rst=None, name=None, cycles=1, sv_attributes=None):
     """Register this value, returning the delayed value.
     `clk`, `rst`: the clock and reset signals.
     `name`: name this register explicitly.
     `cycles`: number of registers to add."""
+
+    if clk is None:
+      clk = ClockValue._get_current_clock_block()
+      if clk is None:
+        raise ValueError("If 'clk' not specified, must be in clock block")
     if sv_attributes is not None:
       sv_attributes = [sv.SVAttributeAttr.get(attr) for attr in sv_attributes]
+
     from .dialects import seq
     if name is None:
       basename = None
@@ -114,6 +121,27 @@ class RegularValue(Value):
   pass
 
 
+_current_clock_context = ContextVar("current_clock_context")
+
+
+class ClockValue(Value):
+  """A clock signal."""
+
+  __slots__ = ["_old_token"]
+
+  def __enter__(self):
+    self._old_token = _current_clock_context.set(self)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    if exc_value is not None:
+      return
+    _current_clock_context.reset(self._old_token)
+
+  @staticmethod
+  def _get_current_clock_block():
+    return _current_clock_context.get(None)
+
+
 class InOutValue(Value):
   # Maintain a caching of the read value.
   read_value = None

frontends/PyCDE/test/compreg.py

@@ -4,27 +4,25 @@
 # RUN: FileCheck %s --input-file %t/CompReg.tcl --check-prefix TCL
 
 import pycde
-from pycde import types, module, Input, Output
+from pycde import types, module, Clock, Input, Output
 from pycde.devicedb import LocationVector
 
-from pycde.dialects import seq
 from pycde.module import generator
 
 import sys
 
 
 @module
 class CompReg:
-  clk = Input(types.i1)
+  clk = Clock()
   input = Input(types.i8)
   output = Output(types.i8)
 
   @generator
   def build(ports):
-    compreg = ports.input.reg(clk=ports.clk,
-                              name="reg",
-                              sv_attributes=["dont_merge"])
-    ports.output = compreg
+    with ports.clk:
+      compreg = ports.input.reg(name="reg", sv_attributes=["dont_merge"])
+      ports.output = compreg
 
 
 mod = pycde.System([CompReg], name="CompReg", output_directory=sys.argv[1])

frontends/PyCDE/test/good_example.py

@@ -4,24 +4,24 @@
 # This is intended to be a simple 'tutorial' example.  Run it as a test to
 # ensure that we keep it up to date (ensure it doesn't crash).
 
-from pycde import dim, module, generator, types, Input, Output
+from pycde import dim, module, generator, types, Clock, Input, Output
 import pycde
 
 import sys
 
 
 @module
 class Mux:
-  clk = Input(types.i1)
+  clk = Clock()
   data = Input(dim(8, 14))
   sel = Input(types.i4)
 
   out = Output(types.i8)
 
   @generator
   def build(ports):
-    sel_reg = ports.sel.reg(ports.clk)
-    ports.out = ports.data.reg(ports.clk)[sel_reg].reg(ports.clk)
+    sel_reg = ports.sel.reg()
+    ports.out = ports.data.reg()[sel_reg].reg()
 
 
 t = pycde.System([Mux], name="MuxDemo", output_directory=sys.argv[1])