[RTG] Add operation and types to represent labels (#7964)

include/circt-c/Dialect/RTG.h

@@ -31,6 +31,12 @@ MLIR_CAPI_EXPORTED bool rtgTypeIsASequence(MlirType type);
 /// Creates an RTG sequence type in the context.
 MLIR_CAPI_EXPORTED MlirType rtgSequenceTypeGet(MlirContext ctxt);
 
+/// If the type is an RTG label.
+MLIR_CAPI_EXPORTED bool rtgTypeIsALabel(MlirType type);
+
+/// Creates an RTG mode type in the context.
+MLIR_CAPI_EXPORTED MlirType rtgLabelTypeGet(MlirContext ctxt);
+
 /// If the type is an RTG set.
 MLIR_CAPI_EXPORTED bool rtgTypeIsASet(MlirType type);
 

include/circt/Dialect/RTG/IR/CMakeLists.txt

@@ -2,6 +2,11 @@ add_circt_dialect(RTG rtg)
 add_circt_doc(RTG Dialects/RTGOps -gen-op-doc)
 add_circt_doc(RTG Dialects/RTGTypes -gen-typedef-doc -dialect rtg)
 
+set(LLVM_TARGET_DEFINITIONS RTG.td)
+mlir_tablegen(RTGEnums.h.inc -gen-enum-decls)
+mlir_tablegen(RTGEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(CIRCTRTGEnumsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS RTGInterfaces.td)
 mlir_tablegen(RTGOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(RTGOpInterfaces.cpp.inc -gen-op-interface-defs)

include/circt/Dialect/RTG/IR/RTGDialect.h

@@ -16,6 +16,9 @@
 #include "circt/Support/LLVM.h"
 #include "mlir/IR/Dialect.h"
 
+// Pull in all enum type definitions and utility function declarations.
+#include "circt/Dialect/RTG/IR/RTGEnums.h.inc"
+
 // Pull in the Dialect definition.
 #include "circt/Dialect/RTG/IR/RTGDialect.h.inc"
 

include/circt/Dialect/RTG/IR/RTGOps.td

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/CommonTypeConstraints.td"
 include "mlir/IR/CommonAttrConstraints.td"
 include "mlir/IR/Properties.td"
@@ -97,6 +98,62 @@ def InvokeSequenceOp : RTGOp<"invoke_sequence", []> {
   let assemblyFormat = "$sequence attr-dict";
 }
 
+//===- Label Operations ---------------------------------------------------===//
+
+class LabelDeclBase<string mnemonic,
+                    list<Trait> traits> : RTGOp<mnemonic, traits> {
+  let description = [{
+    Declares a label that can then be placed by an `rtg.label` operation in an
+    instruction sequence, passed on to sequences via their arguments, and used
+    by instructions (e.g., as jump targets) by allowing ISA dialects to use them
+    directly as an operand of an instruction or by casting it to a value
+    representing an immediate.
+
+    The format string may contain placeholders of the form `{i}` where `i`
+    refers to the i-th element in `args`.
+    The declared label is uniqued by the compiler to no collide with any other
+    label declarations.
+  }];
+
+  // TODO: 'args' can be generalized to more types
+  let arguments = (ins StrAttr:$formatString, Variadic<Index>:$args);
+
+  let assemblyFormat = [{
+    $formatString (`,` $args^)? attr-dict
+  }];
+}
+
+def LabelDeclOp : LabelDeclBase<"label_decl", [Pure]> {
+  let summary = "declares a label for an instruction sequence";
+  let results = (outs LabelType:$label);
+}
+
+def LabelUniqueDeclOp : LabelDeclBase<"label_unique_decl", []> {
+  let summary = "declares a unique label for an instruction sequence";
+  let results = (outs Res<LabelType, "", [MemAlloc]>:$label);
+}
+
+def LabelVisibilityAttr : I32EnumAttr<"LabelVisibility",
+                                      "visibility specifiers for labels", [
+  I32EnumAttrCase<"local",  0>,
+  I32EnumAttrCase<"global", 1>,
+  I32EnumAttrCase<"external", 2>,
+]> {
+  let cppNamespace = "::circt::rtg";
+}
+
+def LabelOp : RTGOp<"label", []> {
+  let summary = "places a label in an instruction sequence";
+  let description = [{
+    Any declared label must only be placed at most once in any fully elaborated
+    instruction sequence.
+  }];
+
+  let arguments = (ins LabelVisibilityAttr:$visibility, LabelType:$label);
+
+  let assemblyFormat = "$visibility $label attr-dict";
+}
+
 //===- Set Operations ------------------------------------------------------===//
 
 def SetCreateOp : RTGOp<"set_create", [Pure, SameTypeOperands]> {

include/circt/Dialect/RTG/IR/RTGTypes.td

@@ -29,6 +29,18 @@ def SequenceType : RTGTypeDef<"Sequence"> {
   let assemblyFormat = "";
 }
 
+def LabelType : RTGTypeDef<"Label"> {
+  let summary = "a reference to a label";
+  let description = [{
+    This type represents a label. Payload dialects can add operations to cast
+    from this type to an immediate type they can use as an operand to an
+    instruction.
+  }];
+
+  let mnemonic = "label";
+  let assemblyFormat = "";
+}
+
 def SetType : RTGTypeDef<"Set"> {
   let summary = "a set of values";
   let description = [{

include/circt/Dialect/RTG/IR/RTGVisitors.h

@@ -35,10 +35,10 @@ class RTGOpVisitor {
                        SetSelectRandomOp, SetDifferenceOp, SetUnionOp,
                        SetSizeOp, TestOp, InvokeSequenceOp, BagCreateOp,
                        BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
-                       BagUniqueSizeOp, TargetOp, YieldOp>(
-            [&](auto expr) -> ResultType {
-              return thisCast->visitOp(expr, args...);
-            })
+                       BagUniqueSizeOp, LabelDeclOp, LabelUniqueDeclOp, LabelOp,
+                       TargetOp, YieldOp>([&](auto expr) -> ResultType {
+          return thisCast->visitOp(expr, args...);
+        })
         .template Case<RegisterOpInterface>([&](auto expr) -> ResultType {
           return thisCast->visitRegisterOp(expr, args...);
         })
@@ -97,6 +97,9 @@ class RTGOpVisitor {
   HANDLE(BagDifferenceOp, Unhandled);
   HANDLE(BagUnionOp, Unhandled);
   HANDLE(BagUniqueSizeOp, Unhandled);
+  HANDLE(LabelDeclOp, Unhandled);
+  HANDLE(LabelUniqueDeclOp, Unhandled);
+  HANDLE(LabelOp, Unhandled);
   HANDLE(TestOp, Unhandled);
   HANDLE(TargetOp, Unhandled);
   HANDLE(YieldOp, Unhandled);
@@ -111,8 +114,8 @@ class RTGTypeVisitor {
   ResultType dispatchTypeVisitor(Type type, ExtraArgs... args) {
     auto *thisCast = static_cast<ConcreteType *>(this);
     return TypeSwitch<Type, ResultType>(type)
-        .template Case<SequenceType, SetType, BagType, DictType, IndexType,
-                       IntegerType>([&](auto expr) -> ResultType {
+        .template Case<SequenceType, SetType, BagType, DictType, LabelType,
+                       IndexType, IntegerType>([&](auto expr) -> ResultType {
           return thisCast->visitType(expr, args...);
         })
         .template Case<ContextResourceTypeInterface>(
@@ -160,6 +163,7 @@ class RTGTypeVisitor {
   HANDLE(DictType, Unhandled);
   HANDLE(IndexType, Unhandled);
   HANDLE(IntegerType, Unhandled);
+  HANDLE(LabelType, Unhandled);
 #undef HANDLE
 };
 

integration_test/Bindings/Python/dialects/rtg.py

@@ -84,11 +84,12 @@
   with InsertionPoint(m.body):
     indexTy = IndexType.get()
     sequenceTy = rtg.SequenceType.get()
+    labelTy = rtg.LabelType.get()
     setTy = rtg.SetType.get(indexTy)
     bagTy = rtg.BagType.get(indexTy)
     seq = rtg.SequenceOp('seq')
-    Block.create_at_start(seq.bodyRegion, [sequenceTy, setTy, bagTy])
+    Block.create_at_start(seq.bodyRegion, [sequenceTy, labelTy, setTy, bagTy])
 
   # CHECK: rtg.sequence @seq
-  # CHECK: (%{{.*}}: !rtg.sequence, %{{.*}}: !rtg.set<index>, %{{.*}}: !rtg.bag<index>):
+  # CHECK: (%{{.*}}: !rtg.sequence, %{{.*}}: !rtg.label, %{{.*}}: !rtg.set<index>, %{{.*}}: !rtg.bag<index>):
   print(m)

lib/Bindings/Python/RTGModule.cpp

@@ -32,6 +32,14 @@ void circt::python::populateDialectRTGSubmodule(py::module &m) {
           },
           py::arg("self"), py::arg("ctxt") = nullptr);
 
+  mlir_type_subclass(m, "LabelType", rtgTypeIsALabel)
+      .def_classmethod(
+          "get",
+          [](py::object cls, MlirContext ctxt) {
+            return cls(rtgLabelTypeGet(ctxt));
+          },
+          py::arg("self"), py::arg("ctxt") = nullptr);
+
   mlir_type_subclass(m, "SetType", rtgTypeIsASet)
       .def_classmethod(
           "get",

lib/CAPI/Dialect/RTG.cpp

@@ -36,6 +36,15 @@ MlirType rtgSequenceTypeGet(MlirContext ctxt) {
   return wrap(SequenceType::get(unwrap(ctxt)));
 }
 
+// LabelType
+//===----------------------------------------------------------------------===//
+
+bool rtgTypeIsALabel(MlirType type) { return isa<LabelType>(unwrap(type)); }
+
+MlirType rtgLabelTypeGet(MlirContext ctxt) {
+  return wrap(LabelType::get(unwrap(ctxt)));
+}
+
 // SetType
 //===----------------------------------------------------------------------===//
 

lib/Dialect/RTG/IR/CMakeLists.txt

@@ -11,6 +11,7 @@ add_circt_dialect_library(CIRCTRTGDialect
 
   DEPENDS
   MLIRRTGIncGen
+  CIRCTRTGEnumsIncGen
   CIRCTRTGOpInterfacesIncGen
   CIRCTRTGISAAssemblyOpInterfacesIncGen
   CIRCTRTGTypeInterfacesIncGen

lib/Dialect/RTG/IR/RTGDialect.cpp

@@ -33,4 +33,6 @@ void RTGDialect::initialize() {
       >();
 }
 
+#include "circt/Dialect/RTG/IR/RTGEnums.cpp.inc"
+
 #include "circt/Dialect/RTG/IR/RTGDialect.cpp.inc"

test/CAPI/rtg.c

@@ -24,6 +24,15 @@ static void testSequenceType(MlirContext ctx) {
   mlirTypeDump(sequenceTy);
 }
 
+static void testLabelType(MlirContext ctx) {
+  MlirType labelTy = rtgLabelTypeGet(ctx);
+
+  // CHECK: is_label
+  fprintf(stderr, rtgTypeIsALabel(labelTy) ? "is_label\n" : "isnot_label\n");
+  // CHECK: !rtg.label
+  mlirTypeDump(labelTy);
+}
+
 static void testSetType(MlirContext ctx) {
   MlirType elTy = mlirIntegerTypeGet(ctx, 32);
   MlirType setTy = rtgSetTypeGet(elTy);
@@ -71,6 +80,7 @@ int main(int argc, char **argv) {
   mlirDialectHandleLoadDialect(mlirGetDialectHandle__rtg__(), ctx);
 
   testSequenceType(ctx);
+  testLabelType(ctx);
   testSetType(ctx);
   testBagType(ctx);
   testDictType(ctx);

test/Dialect/RTG/IR/basic.mlir

@@ -3,6 +3,17 @@
 // CHECK-LABEL: rtg.sequence @seq
 // CHECK-SAME: attributes {rtg.some_attr} {
 rtg.sequence @seq0 attributes {rtg.some_attr} {
+  %arg = arith.constant 1 : index
+  // CHECK: [[LBL0:%.*]] = rtg.label_decl "label_string_{0}_{1}", %{{.*}}, %{{.*}}
+  %0 = rtg.label_decl "label_string_{0}_{1}", %arg, %arg
+  // CHECK: [[LBL1:%.+]] = rtg.label_unique_decl "label_string"
+  %1 = rtg.label_unique_decl "label_string"
+  // CHECK: rtg.label local [[LBL0]]
+  rtg.label local %0
+  // CHECK: rtg.label global [[LBL1]]
+  rtg.label global %1
+  // CHECK: rtg.label external [[LBL0]]
+  rtg.label external %0
 }
 
 // CHECK-LABEL: rtg.sequence @seq1

test/Dialect/RTG/IR/cse.mlir

@@ -0,0 +1,23 @@
+// RUN: circt-opt --cse %s | FileCheck %s
+
+// CHECK-LABEL: rtg.sequence @seq
+// CHECK-SAME: attributes {rtg.some_attr} {
+rtg.sequence @seq0 attributes {rtg.some_attr} {
+  // CHECK-NEXT: arith.constant
+  %arg = arith.constant 1 : index
+  // CHECK-NEXT: rtg.label_decl "label_string_{0}_{1}", %{{.*}}, %{{.*}}
+  // They are CSE'd and DCE'd 
+  %0 = rtg.label_decl "label_string_{0}_{1}", %arg, %arg
+  %1 = rtg.label_decl "label_string_{0}_{1}", %arg, %arg
+  // CHECK-NEXT: rtg.label_unique_decl "label_string"
+  // CHECK-NEXT: rtg.label_unique_decl "label_string"
+  // They are DCE'd but not CSE'd 
+  %2 = rtg.label_unique_decl "label_string"
+  %3 = rtg.label_unique_decl "label_string"
+  %4 = rtg.label_unique_decl "label_string"
+  // CHECK-NEXT: rtg.label global
+  rtg.label global %0
+  rtg.label global %1
+  rtg.label global %2
+  rtg.label global %3
+}