[RTG] Add immediate type and attribute

frontends/PyRTG/test/basic.py

@@ -20,10 +20,10 @@ def entry0():
     return Set.create(Integer(0), Integer(1))
 
 
-# MLIR-LABEL: rtg.target @Tgt1 : !rtg.dict<entry0: index, entry1: !rtg.label>
+# MLIR-LABEL: rtg.target @Tgt1 : !rtg.dict<entry0: index, entry1: !rtg.isa.label>
 # MLIR-NEXT: [[C0:%.+]] = index.constant 0
 # MLIR-NEXT: [[LBL:%.+]] = rtg.label_decl "l0"
-# MLIR-NEXT: rtg.yield [[C0]], [[LBL]] : index, !rtg.label
+# MLIR-NEXT: rtg.yield [[C0]], [[LBL]] : index, !rtg.isa.label
 # MLIR-NEXT: }
 
 
@@ -40,7 +40,7 @@ def entry1():
 
 
 # MLIR-LABEL: rtg.sequence @seq0
-# MLIR-SAME: ([[SET:%.+]]: !rtg.set<!rtg.label>)
+# MLIR-SAME: ([[SET:%.+]]: !rtg.set<!rtg.isa.label>)
 # MLIR-NEXT: [[LABEL:%.+]] = rtg.set_select_random [[SET]]
 # MLIR-NEXT: rtg.label local [[LABEL]]
 # MLIR-NEXT: }
@@ -114,37 +114,37 @@ def test_args(set: Set):
 # MLIR-NEXT: rtg.label external [[L1]]
 # MLIR-NEXT: rtg.label local [[L2]]
 
-# MLIR-NEXT: [[SET0:%.+]] = rtg.set_create [[L0]], [[L1]] : !rtg.label
-# MLIR-NEXT: [[SET1:%.+]] = rtg.set_create [[L2]] : !rtg.label
-# MLIR-NEXT: [[EMPTY_SET:%.+]] = rtg.set_create  : !rtg.label
-# MLIR-NEXT: [[SET2_1:%.+]] = rtg.set_union [[SET0]], [[SET1]] : !rtg.set<!rtg.label>
-# MLIR-NEXT: [[SET2:%.+]] = rtg.set_union [[SET2_1]], [[EMPTY_SET]] : !rtg.set<!rtg.label>
-# MLIR-NEXT: [[RL0:%.+]] = rtg.set_select_random [[SET2]] : !rtg.set<!rtg.label>
+# MLIR-NEXT: [[SET0:%.+]] = rtg.set_create [[L0]], [[L1]] : !rtg.isa.label
+# MLIR-NEXT: [[SET1:%.+]] = rtg.set_create [[L2]] : !rtg.isa.label
+# MLIR-NEXT: [[EMPTY_SET:%.+]] = rtg.set_create  : !rtg.isa.label
+# MLIR-NEXT: [[SET2_1:%.+]] = rtg.set_union [[SET0]], [[SET1]] : !rtg.set<!rtg.isa.label>
+# MLIR-NEXT: [[SET2:%.+]] = rtg.set_union [[SET2_1]], [[EMPTY_SET]] : !rtg.set<!rtg.isa.label>
+# MLIR-NEXT: [[RL0:%.+]] = rtg.set_select_random [[SET2]] : !rtg.set<!rtg.isa.label>
 # MLIR-NEXT: rtg.label local [[RL0]]
-# MLIR-NEXT: [[SET2_MINUS_SET0:%.+]] = rtg.set_difference [[SET2]], [[SET0]] : !rtg.set<!rtg.label>
-# MLIR-NEXT: [[RL1:%.+]] = rtg.set_select_random [[SET2_MINUS_SET0]] : !rtg.set<!rtg.label>
+# MLIR-NEXT: [[SET2_MINUS_SET0:%.+]] = rtg.set_difference [[SET2]], [[SET0]] : !rtg.set<!rtg.isa.label>
+# MLIR-NEXT: [[RL1:%.+]] = rtg.set_select_random [[SET2_MINUS_SET0]] : !rtg.set<!rtg.isa.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: [[BAG0:%.+]] = rtg.bag_create (%idx2 x [[L0:%.+]], %idx1 x [[L1:%.+]]) : !rtg.label
-# MLIR-NEXT: [[BAG1:%.+]] = rtg.bag_create (%idx1 x [[L2:%.+]]) : !rtg.label
-# MLIR-NEXT: [[EMPTY_BAG:%.+]] = rtg.bag_create  : !rtg.label
-# MLIR-NEXT: [[BAG2_1:%.+]] = rtg.bag_union [[BAG0]], [[BAG1]] : !rtg.bag<!rtg.label>
-# MLIR-NEXT: [[BAG2:%.+]] = rtg.bag_union [[BAG2_1]], [[EMPTY_BAG]] : !rtg.bag<!rtg.label>
-# MLIR-NEXT: [[RL2:%.+]] = rtg.bag_select_random [[BAG2]] : !rtg.bag<!rtg.label>
-# MLIR-NEXT: [[SUB:%.+]] = rtg.bag_create (%idx1 x [[RL2]]) : !rtg.label
-# MLIR-NEXT: [[BAG3:%.+]] = rtg.bag_difference [[BAG2]], [[SUB]] inf : !rtg.bag<!rtg.label>
+# MLIR-NEXT: [[BAG0:%.+]] = rtg.bag_create (%idx2 x [[L0:%.+]], %idx1 x [[L1:%.+]]) : !rtg.isa.label
+# MLIR-NEXT: [[BAG1:%.+]] = rtg.bag_create (%idx1 x [[L2:%.+]]) : !rtg.isa.label
+# MLIR-NEXT: [[EMPTY_BAG:%.+]] = rtg.bag_create  : !rtg.isa.label
+# MLIR-NEXT: [[BAG2_1:%.+]] = rtg.bag_union [[BAG0]], [[BAG1]] : !rtg.bag<!rtg.isa.label>
+# MLIR-NEXT: [[BAG2:%.+]] = rtg.bag_union [[BAG2_1]], [[EMPTY_BAG]] : !rtg.bag<!rtg.isa.label>
+# MLIR-NEXT: [[RL2:%.+]] = rtg.bag_select_random [[BAG2]] : !rtg.bag<!rtg.isa.label>
+# MLIR-NEXT: [[SUB:%.+]] = rtg.bag_create (%idx1 x [[RL2]]) : !rtg.isa.label
+# MLIR-NEXT: [[BAG3:%.+]] = rtg.bag_difference [[BAG2]], [[SUB]] inf : !rtg.bag<!rtg.isa.label>
 # MLIR-NEXT: rtg.label local [[RL2]]
-# MLIR-NEXT: [[BAG4:%.+]] = rtg.bag_difference [[BAG3]], [[BAG1]] : !rtg.bag<!rtg.label>
-# MLIR-NEXT: [[RL3:%.+]] = rtg.bag_select_random [[BAG4]] : !rtg.bag<!rtg.label>
+# MLIR-NEXT: [[BAG4:%.+]] = rtg.bag_difference [[BAG3]], [[BAG1]] : !rtg.bag<!rtg.isa.label>
+# MLIR-NEXT: [[RL3:%.+]] = rtg.bag_select_random [[BAG4]] : !rtg.bag<!rtg.isa.label>
 # MLIR-NEXT: rtg.label local [[RL3]]
 
-# MLIR-NEXT: [[SEQ:%.+]] = rtg.get_sequence @seq0 : !rtg.sequence<!rtg.set<!rtg.label>>
-# MLIR-NEXT: [[SUBST:%.+]] = rtg.substitute_sequence [[SEQ]]([[SET0]]) : !rtg.sequence<!rtg.set<!rtg.label>>
+# MLIR-NEXT: [[SEQ:%.+]] = rtg.get_sequence @seq0 : !rtg.sequence<!rtg.set<!rtg.isa.label>>
+# MLIR-NEXT: [[SUBST:%.+]] = rtg.substitute_sequence [[SEQ]]([[SET0]]) : !rtg.sequence<!rtg.set<!rtg.isa.label>>
 # MLIR-NEXT: [[RAND1:%.+]] = rtg.randomize_sequence [[SUBST]]
 # MLIR-NEXT: rtg.embed_sequence [[RAND1]]
 # MLIR-NEXT: [[RAND2:%.+]] = rtg.randomize_sequence [[SUBST]]

include/circt-c/Dialect/RTG.h

@@ -79,6 +79,16 @@ MLIR_CAPI_EXPORTED MlirType rtgDictTypeGet(MlirContext ctxt,
                                            MlirAttribute const *entryNames,
                                            MlirType const *entryTypes);
 
+/// If the type is an RTG immediate.
+MLIR_CAPI_EXPORTED bool rtgTypeIsAImmediate(MlirType type);
+
+/// Creates an RTG immediate type in the context.
+MLIR_CAPI_EXPORTED MlirType rtgImmediateTypeGet(MlirContext ctx,
+                                                uint32_t width);
+
+/// Returns the width of the RTG immediate type.
+MLIR_CAPI_EXPORTED uint32_t rtgImmediateTypeGetWidth(MlirType type);
+
 //===----------------------------------------------------------------------===//
 // Attribute API.
 //===----------------------------------------------------------------------===//
@@ -108,6 +118,21 @@ MLIR_CAPI_EXPORTED bool rtgAttrIsADefaultContextAttr(MlirAttribute attr);
 MLIR_CAPI_EXPORTED MlirAttribute rtgDefaultContextAttrGet(MlirContext ctxt,
                                                           MlirType type);
 
+/// Checks if the attribute is an RTG immediate attribute.
+MLIR_CAPI_EXPORTED bool rtgAttrIsAImmediate(MlirAttribute attr);
+
+/// Creates an RTG immediate attribute in the context with the given width and
+/// value.
+MLIR_CAPI_EXPORTED MlirAttribute rtgImmediateAttrGet(MlirContext ctx,
+                                                     uint32_t width,
+                                                     uint64_t value);
+
+/// Returns the width of the RTG immediate attribute.
+MLIR_CAPI_EXPORTED uint32_t rtgImmediateAttrGetWidth(MlirAttribute attr);
+
+/// Returns the value of the RTG immediate attribute.
+MLIR_CAPI_EXPORTED uint64_t rtgImmediateAttrGetValue(MlirAttribute attr);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -10,6 +10,7 @@
 #define CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_H
 
 #include "circt/Dialect/RTG/IR/RTGAttrInterfaces.h"
+#include "circt/Dialect/RTG/IR/RTGTypes.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 

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

@@ -75,4 +75,25 @@ def DefaultContextAttr : RTGAttrDef<"DefaultContext", [
   let assemblyFormat = "";
 }
 
+//===----------------------------------------------------------------------===//
+// Attributes for ISA targets
+//===----------------------------------------------------------------------===//
+
+class RTGISAAttrDef<string name, list<Trait> traits = []> 
+  : RTGAttrDef<name, traits> { let mnemonic = "isa." # !tolower(name); }
+
+def ImmediateAttr : RTGISAAttrDef<"Immediate", [
+  DeclareAttrInterfaceMethods<TypedAttrInterface>,
+]> {
+  let summary = "an ISA immediate value";
+  let description = [{
+    This represents an ISA immediate of arbitrary but fixed bit-width. The type
+    of this attribute must always be an `ImmediateType` of matching bit-width.
+  }];
+
+  let parameters = (ins "llvm::APInt":$value);
+
+  let hasCustomAssemblyFormat = true;
+}
+
 #endif // CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_TD

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

@@ -29,6 +29,8 @@ def RTGDialect : Dialect {
   let useDefaultTypePrinterParser = 1;
   let cppNamespace = "::circt::rtg";
 
+  let hasConstantMaterializer = 1;
+
   let extraClassDeclaration = [{
     void registerAttributes();
     void registerTypes();

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

@@ -25,6 +25,21 @@ include "circt/Dialect/RTG/IR/RTGISAAssemblyInterfaces.td"
 class RTGOp<string mnemonic, list<Trait> traits = []> :
     Op<RTGDialect, mnemonic, traits>;
 
+
+def ConstantOp : RTGOp<"constant", [
+  Pure,
+  ConstantLike,
+  DeclareOpInterfaceMethods<InferTypeOpInterface>
+]> {
+  let summary = "create an SSA value from an attribute";
+
+  let arguments = (ins TypedAttrInterface:$value);
+  let results = (outs AnyType:$result);
+
+  let assemblyFormat = "$value attr-dict";
+  let hasFolder = 1;
+}
+
 //===- Sequence Handling Operations ---------------------------------------===//
 
 def SequenceOp : RTGOp<"sequence", [

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

@@ -16,7 +16,12 @@
 include "circt/Dialect/RTG/IR/RTGDialect.td"
 include "mlir/IR/AttrTypeBase.td"
 
-class RTGTypeDef<string name> : TypeDef<RTGDialect, name>;
+class RTGTypeDef<string name, list<Trait> traits = []>
+  : TypeDef<RTGDialect, name, traits>;
+
+//===----------------------------------------------------------------------===//
+// Sequence Types
+//===----------------------------------------------------------------------===//
 
 def SequenceType : RTGTypeDef<"Sequence"> {
   let summary = "handle to a sequence or sequence family";
@@ -51,17 +56,9 @@ def FullySubstitutedSequenceType : DialectType<RTGDialect,
     "::circt::rtg::SequenceType::get($_builder.getContext(), " #
                                     "llvm::ArrayRef<::mlir::Type>{})">;
 
-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 = "";
-}
+//===----------------------------------------------------------------------===//
+// Types for common datastructures
+//===----------------------------------------------------------------------===//
 
 def SetType : RTGTypeDef<"Set"> {
   let summary = "a set of values";
@@ -123,6 +120,43 @@ def DictType : RTGTypeDef<"Dict"> {
   let genVerifyDecl = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Types for ISA targets
+//===----------------------------------------------------------------------===//
+
+class RTGISATypeDef<string name, list<Trait> traits = []>
+  : RTGTypeDef<name, traits> { let mnemonic = "isa." # !tolower(name); }
+
+
+def LabelType : RTGISATypeDef<"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 or allow an operand of this type directly.
+  }];
+
+  let assemblyFormat = "";
+}
+
+def ImmediateType : RTGISATypeDef<"Immediate"> {
+  let summary = "an ISA immediate";
+  let description = [{
+    This type represents immediates of arbitrary but fixed bit-width.
+    The RTG dialect provides this type to avoid duplication in ISA payload
+    dialects.
+  }];
+
+  let parameters = (ins "uint32_t":$width);
+  let assemblyFormat = "`<` $width `>`";
+}
+
+class ImmediateOfWidth<int width> : Type<
+  And<[CPred<"llvm::isa<::circt::rtg::ImmediateType>($_self)">,
+       CPred<"llvm::cast<::circt::rtg::ImmediateType>($_self).getWidth() == " # width>]>,
+  "a " # width # "-bit immediate">,
+  BuildableType<"::circt::rtg::ImmediateType::get($_builder.getContext(), " # width # ")">;
+
 class ImmTypeBase<int width> : TypeDef<RTGDialect, "Imm" # width, []> {
   let summary = "represents a " # width # "-bit immediate";
   let mnemonic = "imm" # width;

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

@@ -32,6 +32,8 @@ class RTGOpVisitor {
     auto *thisCast = static_cast<ConcreteType *>(this);
     return TypeSwitch<Operation *, ResultType>(op)
         .template Case<
+            // Constants
+            ConstantOp,
             // Bags
             BagCreateOp, BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
             BagUniqueSizeOp,
@@ -82,6 +84,7 @@ class RTGOpVisitor {
     return static_cast<ConcreteType *>(this)->visit##OPKIND##Op(op, args...);  \
   }
 
+  HANDLE(ConstantOp, Unhandled);
   HANDLE(SequenceOp, Unhandled);
   HANDLE(GetSequenceOp, Unhandled);
   HANDLE(SubstituteSequenceOp, Unhandled);
@@ -120,10 +123,11 @@ 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, LabelType,
-                       IndexType, IntegerType>([&](auto expr) -> ResultType {
-          return thisCast->visitType(expr, args...);
-        })
+        .template Case<ImmediateType, SequenceType, SetType, BagType, DictType,
+                       LabelType, IndexType, IntegerType>(
+            [&](auto expr) -> ResultType {
+              return thisCast->visitType(expr, args...);
+            })
         .template Case<ContextResourceTypeInterface>(
             [&](auto expr) -> ResultType {
               return thisCast->visitContextResourceType(expr, args...);
@@ -163,6 +167,7 @@ class RTGTypeVisitor {
                                                                     args...);  \
   }
 
+  HANDLE(ImmediateType, Unhandled);
   HANDLE(SequenceType, Unhandled);
   HANDLE(SetType, Unhandled);
   HANDLE(BagType, Unhandled);

integration_test/Bindings/Python/dialects/rtg.py

@@ -106,7 +106,7 @@
     # CHECK: index{{$}}
     print(bagTy.element_type)
 
-  # CHECK: rtg.sequence @seq(%{{.*}}: !rtg.sequence, %{{.*}}: !rtg.label, %{{.*}}: !rtg.set<index>, %{{.*}}: !rtg.bag<index>, %{{.*}}: !rtgtest.ireg, %{{.*}}: !rtg.randomized_sequence)
+  # CHECK: rtg.sequence @seq(%{{.*}}: !rtg.sequence, %{{.*}}: !rtg.isa.label, %{{.*}}: !rtg.set<index>, %{{.*}}: !rtg.bag<index>, %{{.*}}: !rtgtest.ireg, %{{.*}}: !rtg.randomized_sequence)
   print(m)
 
 with Context() as ctx, Location.unknown():
@@ -221,3 +221,17 @@
   print(attr.type)
   # CHECK: #rtg.default : !rtgtest.cpu
   print(attr)
+
+  immediate_type = rtg.ImmediateType.get(32)
+  # CHECK: width=32
+  print(f"width={immediate_type.width}")
+  # CHECK: !rtg.isa.immediate<32>
+  print(immediate_type)
+
+  immediate_attr = rtg.ImmediateAttr.get(32, 42)
+  # CHECK: width=32
+  print(f"width={immediate_attr.width}")
+  # CHECK: value=42
+  print(f"value={immediate_attr.value}")
+  # CHECK: #rtg.isa.immediate<32, 42>
+  print(immediate_attr)