[PyRTG] Support integers

Author: maerhart

frontends/PyRTG/src/CMakeLists.txt

@@ -15,6 +15,8 @@ declare_mlir_python_sources(PyRTGSources
   SOURCES
   pyrtg/__init__.py
   pyrtg/core.py
+  pyrtg/index.py
+  pyrtg/integers.py
   pyrtg/labels.py
   pyrtg/rtg.py
   pyrtg/sets.py

frontends/PyRTG/src/pyrtg/__init__.py

@@ -7,4 +7,6 @@
 from .tests import test
 from .labels import Label
 from .rtg import rtg
+from .index import index
 from .sets import Set
+from .integers import Integer

frontends/PyRTG/src/pyrtg/index.py

@@ -0,0 +1,8 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .support import wrap_opviews_with_values
+from .circt.dialects import index
+
+wrap_opviews_with_values(index, index.__name__)

frontends/PyRTG/src/pyrtg/integers.py

@@ -0,0 +1,51 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from __future__ import annotations
+
+from .circt import ir
+from .core import Value
+from .index import index
+
+import typing
+
+
+class Integer(Value):
+  """
+  This represents an integer with the same number of bits as a 'size_t' in C.
+  It is used to provide parameter values to meta-level constructs such as the
+  multiple of an element in a Bag. These integers will be fully constant folded
+  away during randomization.
+  """
+
+  def __init__(self, value: typing.Union[ir.Value, int]) -> Integer:
+    """
+    Use this constructor to create an Integer from a builtin Python int.
+    """
+
+    self._value = value
+
+  def __add__(self, other: Integer) -> Integer:
+    return index.AddOp(self._get_ssa_value(), other._get_ssa_value())
+
+  def __sub__(self, other: Integer) -> Integer:
+    return index.SubOp(self._get_ssa_value(), other._get_ssa_value())
+
+  def __and__(self, other: Integer) -> Integer:
+    return index.AndOp(self._get_ssa_value(), other._get_ssa_value())
+
+  def __or__(self, other: Integer) -> Integer:
+    return index.OrOp(self._get_ssa_value(), other._get_ssa_value())
+
+  def __xor__(self, other: Integer) -> Integer:
+    return index.XOrOp(self._get_ssa_value(), other._get_ssa_value())
+
+  def get_type(self) -> ir.Type:
+    return ir.IndexType.get()
+
+  def _get_ssa_value(self) -> ir.Value:
+    if isinstance(self._value, int):
+      self = index.ConstantOp(self._value)
+
+    return self._value

frontends/PyRTG/src/pyrtg/labels.py

@@ -7,6 +7,9 @@
 from .circt import ir
 from .core import Value
 from .rtg import rtg
+from .integers import Integer
+
+from typing import Union
 
 
 class Label(Value):
@@ -22,23 +25,27 @@ class Label(Value):
   def __init__(self, value: ir.Value):
     self._value = value
 
-  def declare(string: str) -> Label:
+  def declare(string: str, *args: Union[Integer, int]) -> Label:
     """
     Declares a label with a fixed name. Labels returned by different calls to
     this function but with the same arguments refer to the same label.
     """
 
-    return rtg.LabelDeclOp(string, [])
+    return rtg.LabelDeclOp(
+        string,
+        [(arg if isinstance(arg, Integer) else Integer(arg)) for arg in args])
 
-  def declare_unique(string: str) -> Label:
+  def declare_unique(string: str, *args: Union[Integer, int]) -> Label:
     """
     Declares a unique label. This means, all usages of the value returned by this
     function will refer to the same label, but no other label declarations can
     conflict with this label, including labels returned by other calls to this
     function or fixed labels declared with 'declare_label'.
     """
 
-    return rtg.LabelUniqueDeclOp(string, [])
+    return rtg.LabelUniqueDeclOp(
+        string,
+        [(arg if isinstance(arg, Integer) else Integer(arg)) for arg in args])
 
   def place(
       self,

frontends/PyRTG/src/pyrtg/support.py

@@ -15,6 +15,9 @@ def _FromCirctValue(value: ir.Value) -> Value:
   if isinstance(type, rtg.SetType):
     from .sets import Set
     return Set(value)
+  if isinstance(type, ir.IndexType):
+    from .integers import Integer
+    return Integer(value)
   assert False, "Unsupported value"
 
 

frontends/PyRTG/test/basic.py

@@ -2,7 +2,7 @@
 # RUN: %rtgtool% %s --seed=0 --output-format=elaborated | FileCheck %s --check-prefix=ELABORATED
 # RUN: %rtgtool% %s --seed=0 -o %t --output-format=asm && FileCheck %s --input-file=%t --check-prefix=ASM
 
-from pyrtg import test, rtg, Label, Set
+from pyrtg import test, rtg, Label, Set, Integer
 
 # MLIR-LABEL: rtg.test @test0
 # MLIR-NEXT: }
@@ -20,6 +20,8 @@ def test0():
 
 
 # MLIR-LABEL: rtg.test @test_labels
+# MLIR-NEXT: index.constant 5
+# MLIR-NEXT: index.constant 3
 # MLIR-NEXT: [[L0:%.+]] = rtg.label_decl "l0"
 # MLIR-NEXT: [[L1:%.+]] = rtg.label_unique_decl "l1"
 # MLIR-NEXT: [[L2:%.+]] = rtg.label_unique_decl "l1"
@@ -38,6 +40,11 @@ def test0():
 # MLIR-NEXT: [[RL1:%.+]] = rtg.set_select_random [[SET2_MINUS_SET0]] : !rtg.set<!rtg.label>
 # MLIR-NEXT: rtg.label local [[RL1]]
 
+# MLIR-NEXT: rtg.label_decl "L_{{[{][{]0[}][}]}}", %idx5
+# MLIR-NEXT: rtg.label local
+# MLIR-NEXT: rtg.label_decl "L_{{[{][{]0[}][}]}}", %idx3
+# MLIR-NEXT: rtg.label local
+
 # MLIR-NEXT: }
 
 # ELABORATED-LABEL: rtg.test @test_labels
@@ -51,6 +58,11 @@ def test0():
 # ELABORATED-NEXT: rtg.label local [[L0]]
 # ELABORATED-NEXT: rtg.label local [[L2]]
 
+# ELABORATED-NEXT: rtg.label_decl "L_5"
+# ELABORATED-NEXT: rtg.label local
+# ELABORATED-NEXT: rtg.label_decl "L_3"
+# ELABORATED-NEXT: rtg.label local
+
 # ELABORATED-NEXT: }
 
 # ASM-LABEL: Begin of test_labels
@@ -63,6 +75,9 @@ def test0():
 # ASM-NEXT: l0:
 # ASM-NEXT: l1_1:
 
+# ASM-NEXT: L_5:
+# ASM-NEXT: L_3:
+
 # ASM-EMPTY:
 # ASM: End of test_labels
 
@@ -86,3 +101,11 @@ def test_labels():
   set2 -= set0
   rl1 = set2.get_random_and_exclude()
   rl1.place()
+
+  sub = Integer(1) - Integer(2)
+  add = (sub & Integer(4) | Integer(3) ^ Integer(5))
+  add += sub
+  l3 = Label.declare(r"L_{{0}}", add)
+  l3.place()
+  l4 = Label.declare(r"L_{{0}}", 3)
+  l4.place()

lib/Bindings/Python/CIRCTModule.cpp

@@ -33,6 +33,7 @@
 #include "circt-c/Dialect/Verif.h"
 #include "circt-c/ExportVerilog.h"
 #include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir-c/Dialect/Index.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Transforms.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
@@ -101,6 +102,10 @@ NB_MODULE(_circt, m) {
         mlirDialectHandleRegisterDialect(hwarith, context);
         mlirDialectHandleLoadDialect(hwarith, context);
 
+        MlirDialectHandle index = mlirGetDialectHandle__index__();
+        mlirDialectHandleRegisterDialect(index, context);
+        mlirDialectHandleLoadDialect(index, context);
+
         MlirDialectHandle om = mlirGetDialectHandle__om__();
         mlirDialectHandleRegisterDialect(om, context);
         mlirDialectHandleLoadDialect(om, context);

lib/Bindings/Python/CMakeLists.txt

@@ -52,6 +52,7 @@ set(PYTHON_BINDINGS_LINK_LIBS
   CIRCTCAPISV
   CIRCTCAPIVerif
   CIRCTCAPITransforms
+  MLIRCAPIIndex
   # needed for mlirFrozenRewritePatternSetDestroy
   # but not the actual passes
   MLIRCAPITransforms
@@ -139,6 +140,15 @@ declare_mlir_dialect_python_bindings(
     dialects/hw.py
   DIALECT_NAME hw)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT CIRCTBindingsPythonSources.Dialects
+  ROOT_DIR "${MLIR_MAIN_SRC_DIR}/python/mlir"
+  TD_FILE dialects/IndexOps.td
+  SOURCES
+    dialects/index.py
+  DIALECT_NAME index
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT CIRCTBindingsPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"