[RTG] Add on_context and context_switch operations

maerhart · maerhart · commit b2a9f5265f53 · 2025-02-05T11:43:33.000Z
diff --git a/include/circt-c/Dialect/RTG.h b/include/circt-c/Dialect/RTG.h
@@ -95,6 +95,13 @@ rtgLabelVisibilityAttrGetValue(MlirAttribute attr);
 MLIR_CAPI_EXPORTED MlirAttribute
 rtgLabelVisibilityAttrGet(MlirContext ctxt, RTGLabelVisibility visibility);
 
+/// If the attribute is an RTG default context.
+MLIR_CAPI_EXPORTED bool rtgAttrIsADefaultContextAttr(MlirAttribute attr);
+
+/// Creates an RTG default context attribute in the context.
+MLIR_CAPI_EXPORTED MlirAttribute rtgDefaultContextAttrGet(MlirContext ctxt,
+                                                          MlirType type);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/include/circt/Dialect/RTG/IR/CMakeLists.txt b/include/circt/Dialect/RTG/IR/CMakeLists.txt
@@ -38,3 +38,8 @@ mlir_tablegen(RTGISAAssemblyOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(RTGISAAssemblyOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(CIRCTRTGISAAssemblyOpInterfacesIncGen)
 add_dependencies(circt-headers CIRCTRTGISAAssemblyOpInterfacesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS RTGAttributes.td)
+mlir_tablegen(RTGAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(RTGAttributes.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(CIRCTRTGAttributeIncGen)
diff --git a/include/circt/Dialect/RTG/IR/RTGAttributes.h b/include/circt/Dialect/RTG/IR/RTGAttributes.h
@@ -0,0 +1,19 @@
+//===- RTGAttributes.h - RTG dialect attributes -----------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_H
+#define CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_H
+
+#include "circt/Dialect/RTG/IR/RTGAttrInterfaces.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "circt/Dialect/RTG/IR/RTGAttributes.h.inc"
+
+#endif // CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_H
diff --git a/include/circt/Dialect/RTG/IR/RTGAttributes.td b/include/circt/Dialect/RTG/IR/RTGAttributes.td
@@ -0,0 +1,42 @@
+//===- RTGAttributes.td - RTG attributes -------------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This describes the RTG attributes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_TD
+#define CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_TD
+
+include "circt/Dialect/RTG/IR/RTGDialect.td"
+include "circt/Dialect/RTG/IR/RTGInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/Interfaces.td"
+
+class RTGAttrDef<string name, list<Trait> traits = []> 
+  : AttrDef<RTGDialect, name, traits>;
+
+
+def DefaultContextAttr : RTGAttrDef<"DefaultContext", [
+  DeclareAttrInterfaceMethods<ContextResourceAttrInterface>,
+]> {
+  let summary = "the default context of its type";
+  let description = [{
+    A target must specify a default value for each context resource type it
+    uses. The tests matched against that target then start in that context
+    initially and 'on_context' operations can be used to switch to another
+    context within a test. Essentially, this attribute is used within tests to
+    refer to the default context set in the target.
+  }];
+
+  let mnemonic = "default";
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type);
+  let assemblyFormat = "";
+}
+
+#endif // CIRCT_DIALECT_RTG_IR_RTGATTRIBUTES_TD
diff --git a/include/circt/Dialect/RTG/IR/RTGDialect.td b/include/circt/Dialect/RTG/IR/RTGDialect.td
@@ -25,10 +25,12 @@ def RTGDialect : Dialect {
     test can, for example, be a sequence of ISA instructions to test a CPU core.
   }];
 
+  let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
   let cppNamespace = "::circt::rtg";
 
   let extraClassDeclaration = [{
+    void registerAttributes();
     void registerTypes();
   }];
 }
diff --git a/include/circt/Dialect/RTG/IR/RTGInterfaces.td b/include/circt/Dialect/RTG/IR/RTGInterfaces.td
@@ -13,28 +13,8 @@ include "mlir/IR/Interfaces.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 
-def ContextResourceOpInterface : OpInterface<"ContextResourceOpInterface"> {
-  let description = [{
-    This interface should be implemented by operations that define
-    context resources. The operation should define at least one SSA value of a type
-    implementing the `ContextResourceTypeInterface`.
-  }];
-  let cppNamespace = "::circt::rtg";
-
-  let methods = [
-    InterfaceMethod<[{
-        Provides a unique identifier for the defined context resource at `idx`
-        (not counting op results that are not of a type implementing the
-        `ContextResourceTypeInterface`).
-        For example, if the context resource are CPUs it could be the core ID.
-      }],
-      "size_t", "getIdentifier", (ins "size_t":$idx)>,
-  ];
-}
-
 /// Context resources can only be defined inside the `rtg.target` operation.
 def ContextResourceDefining : TraitList<[
-  DeclareOpInterfaceMethods<ContextResourceOpInterface>,
   HasParent<"::circt::rtg::TargetOp">,
 ]>;
 
diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td
@@ -197,6 +197,57 @@ def LabelOp : RTGOp<"label", []> {
   let assemblyFormat = "$visibility $label attr-dict";
 }
 
+//===- Context Operations -------------------------------------------------===//
+
+def OnContextOp : RTGOp<"on_context", []> {
+  let summary = "places a sequence on a context";
+  let description = [{
+    This operation takes a context and a fully substituted, but not yet
+    randomized sequence and inserts the necessary instructions to switch from
+    the current context to the provided context, randomizes and embeds the given
+    sequence under the given context, and inserts instructions to switch back to
+    the original context.
+
+    These instructions are provided by the 'rtg.context_switch' operation. If no
+    'rtg.context_switch' for this transition is provided, the compiler will
+    error out. If multiple such context switches apply, the most recently
+    registered one takes precedence.
+  }];
+
+  let arguments = (ins ContextResourceTypeInterface:$context,
+                       FullySubstitutedSequenceType:$sequence);
+
+  let assemblyFormat = [{
+    $context `,` $sequence `:` qualified(type($context)) attr-dict
+  }];
+}
+
+def ContextSwitchOp : RTGOp<"context_switch", [
+  HasParent<"rtg::TargetOp">,
+]> {
+  let summary = "a specification of how to switch contexts";
+  let description = [{
+    This operation allows the user to specify a sequence of instructions to
+    switch from context 'from' to context 'to', randomize and embed a provided
+    sequence, and switch back from context 'to' to context 'from'.  This
+    sequence of instructions should be provided as the 'sequence' operand which
+    is a sequence of the type '!rtg.sequence<context-type-interface,
+    context-type-interface, !rtg.sequence>'.  The first parameter is the 'from'
+    context, the second one the 'to' context, and the third is the sequence to
+    randomize and embed under the 'to' context.
+  }];
+
+  let arguments = (ins ContextResourceAttrInterface:$from,
+                       ContextResourceAttrInterface:$to,
+                       SequenceType:$sequence);
+
+  let assemblyFormat = [{
+    $from `->` $to `,` $sequence `:` qualified(type($sequence)) attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
+
 //===- Set Operations ------------------------------------------------------===//
 
 def SetCreateOp : RTGOp<"set_create", [Pure, SameTypeOperands]> {
diff --git a/include/circt/Dialect/RTG/IR/RTGVisitors.h b/include/circt/Dialect/RTG/IR/RTGVisitors.h
@@ -35,6 +35,8 @@ class RTGOpVisitor {
             // Bags
             BagCreateOp, BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
             BagUniqueSizeOp,
+            // Contexts
+            OnContextOp, ContextSwitchOp,
             // Labels
             LabelDeclOp, LabelUniqueDeclOp, LabelOp,
             // Registers
@@ -51,10 +53,6 @@ class RTGOpVisitor {
             SetSizeOp>([&](auto expr) -> ResultType {
           return thisCast->visitOp(expr, args...);
         })
-        .template Case<ContextResourceOpInterface>(
-            [&](auto expr) -> ResultType {
-              return thisCast->visitContextResourceOp(expr, args...);
-            })
         .Default([&](auto expr) -> ResultType {
           if (op->getDialect() ==
               op->getContext()->getLoadedDialect<RTGDialect>())
@@ -75,11 +73,6 @@ class RTGOpVisitor {
   /// handled by the concrete visitor.
   ResultType visitUnhandledOp(Operation *op, ExtraArgs... args);
 
-  ResultType visitContextResourceOp(ContextResourceOpInterface op,
-                                    ExtraArgs... args) {
-    return static_cast<ConcreteType *>(this)->visitUnhandledOp(op, args...);
-  }
-
   ResultType visitExternalOp(Operation *op, ExtraArgs... args) {
     return ResultType();
   }
@@ -95,6 +88,8 @@ class RTGOpVisitor {
   HANDLE(RandomizeSequenceOp, Unhandled);
   HANDLE(EmbedSequenceOp, Unhandled);
   HANDLE(RandomNumberInRangeOp, Unhandled);
+  HANDLE(OnContextOp, Unhandled);
+  HANDLE(ContextSwitchOp, Unhandled);
   HANDLE(SetCreateOp, Unhandled);
   HANDLE(SetSelectRandomOp, Unhandled);
   HANDLE(SetDifferenceOp, Unhandled);
diff --git a/include/circt/Dialect/RTGTest/IR/RTGTestOps.td b/include/circt/Dialect/RTGTest/IR/RTGTestOps.td
@@ -29,8 +29,8 @@ def CPUDeclOp : RTGTestOp<"cpu_decl", [
 ]> {
   let summary = "declare a CPU";
   let description = [{
-    This operation is used to test the `ContextResourceOpInterface` and passes
-    taking advantage of it.
+    This operation is used to test the `ContextResourceAttrInterface` and
+    `ContextResourceTypeInterface` passes taking advantage of it.
   }];
 
   let arguments = (ins CPUAttr:$id);
diff --git a/integration_test/Bindings/Python/dialects/rtg.py b/integration_test/Bindings/Python/dialects/rtg.py
@@ -208,3 +208,11 @@
   # CHECK: rtg.label_decl "label"
   # CHECK: rtg.label global {{%.+}}
   print(m)
+
+with Context() as ctx, Location.unknown():
+  circt.register_dialects(ctx)
+  attr = rtg.DefaultContextAttr.get(rtgtest.CPUType.get())
+  # CHECK: !rtgtest.cpu
+  print(attr.type)
+  # CHECK: #rtg.default : !rtgtest.cpu
+  print(attr)
diff --git a/lib/Bindings/Python/RTGModule.cpp b/lib/Bindings/Python/RTGModule.cpp
@@ -107,4 +107,12 @@ void circt::python::populateDialectRTGSubmodule(nb::module_ &m) {
       .def_property_readonly("value", [](MlirAttribute self) {
         return rtgLabelVisibilityAttrGetValue(self);
       });
+
+  mlir_attribute_subclass(m, "DefaultContextAttr", rtgAttrIsADefaultContextAttr)
+      .def_classmethod(
+          "get",
+          [](nb::object cls, MlirType type, MlirContext ctxt) {
+            return cls(rtgDefaultContextAttrGet(ctxt, type));
+          },
+          nb::arg("self"), nb::arg("type"), nb::arg("ctxt") = nullptr);
 }
diff --git a/lib/CAPI/Dialect/RTG.cpp b/lib/CAPI/Dialect/RTG.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "circt-c/Dialect/RTG.h"
+#include "circt/Dialect/RTG/IR/RTGAttributes.h"
 #include "circt/Dialect/RTG/IR/RTGDialect.h"
 #include "circt/Dialect/RTG/IR/RTGTypes.h"
 
@@ -142,3 +143,11 @@ MlirAttribute rtgLabelVisibilityAttrGet(MlirContext ctxt,
   };
   return wrap(LabelVisibilityAttr::get(unwrap(ctxt), convert(visibility)));
 }
+
+bool rtgAttrIsADefaultContextAttr(MlirAttribute attr) {
+  return isa<DefaultContextAttr>(unwrap(attr));
+}
+
+MlirAttribute rtgDefaultContextAttrGet(MlirContext ctxt, MlirType type) {
+  return wrap(DefaultContextAttr::get(unwrap(ctxt), unwrap(type)));
+}
diff --git a/lib/Dialect/RTG/IR/CMakeLists.txt b/lib/Dialect/RTG/IR/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_circt_dialect_library(CIRCTRTGDialect
+  RTGAttributes.cpp
   RTGAttrInterfaces.cpp
   RTGDialect.cpp
   RTGISAAssemblyAttrInterfaces.cpp
@@ -13,6 +14,7 @@ add_circt_dialect_library(CIRCTRTGDialect
   ${CIRCT_MAIN_INCLUDE_DIR}/circt/Dialect/RTG/IR
 
   DEPENDS
+  CIRCTRTGAttributeIncGen
   CIRCTRTGAttrInterfacesIncGen
   CIRCTRTGEnumsIncGen
   CIRCTRTGISAAssemblyAttrInterfacesIncGen
diff --git a/lib/Dialect/RTG/IR/RTGAttributes.cpp b/lib/Dialect/RTG/IR/RTGAttributes.cpp
@@ -0,0 +1,30 @@
+//===- RTGAttributes.cpp --------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/RTG/IR/RTGAttributes.h"
+#include "circt/Dialect/RTG/IR/RTGDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace circt;
+using namespace rtg;
+
+//===----------------------------------------------------------------------===//
+// TableGen generated logic.
+//===----------------------------------------------------------------------===//
+
+void RTGDialect::registerAttributes() {
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "circt/Dialect/RTG/IR/RTGAttributes.cpp.inc"
+      >();
+}
+
+#define GET_ATTRDEF_CLASSES
+#include "circt/Dialect/RTG/IR/RTGAttributes.cpp.inc"
diff --git a/lib/Dialect/RTG/IR/RTGDialect.cpp b/lib/Dialect/RTG/IR/RTGDialect.cpp
@@ -25,6 +25,7 @@ using namespace rtg;
 //===----------------------------------------------------------------------===//
 
 void RTGDialect::initialize() {
+  registerAttributes();
   registerTypes();
   // Register operations.
   addOperations<
diff --git a/lib/Dialect/RTG/IR/RTGOps.cpp b/lib/Dialect/RTG/IR/RTGOps.cpp
@@ -363,6 +363,31 @@ LogicalResult VirtualRegisterOp::inferReturnTypes(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ContextSwitchOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ContextSwitchOp::verify() {
+  auto elementTypes = getSequence().getType().getElementTypes();
+  if (elementTypes.size() != 3)
+    return emitOpError("sequence type must have exactly 3 element types");
+
+  if (getFrom().getType() != elementTypes[0])
+    return emitOpError(
+        "first sequence element type must match 'from' attribute type");
+
+  if (getTo().getType() != elementTypes[1])
+    return emitOpError(
+        "second sequence element type must match 'to' attribute type");
+
+  auto seqTy = dyn_cast<SequenceType>(elementTypes[2]);
+  if (!seqTy || !seqTy.getElementTypes().empty())
+    return emitOpError(
+        "third sequence element type must be a fully substituted sequence");
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TestOp
 //===----------------------------------------------------------------------===//
diff --git a/lib/Dialect/RTGTest/IR/RTGTestOps.cpp b/lib/Dialect/RTGTest/IR/RTGTestOps.cpp
@@ -21,8 +21,6 @@ using namespace rtgtest;
 // CPUDeclOp
 //===----------------------------------------------------------------------===//
 
-size_t CPUDeclOp::getIdentifier(size_t idx) { return getId().getId(); }
-
 mlir::OpFoldResult CPUDeclOp::fold(FoldAdaptor adaptor) { return getId(); }
 
 //===----------------------------------------------------------------------===//
diff --git a/test/CAPI/CMakeLists.txt b/test/CAPI/CMakeLists.txt
@@ -95,6 +95,7 @@ target_link_libraries(circt-capi-rtg-test
 
   MLIRCAPIIR
   CIRCTCAPIRTG
+  CIRCTCAPIRTGTest
 )
 
 add_llvm_executable(circt-capi-rtgtest-test
diff --git a/test/CAPI/rtg.c b/test/CAPI/rtg.c
diff --git a/test/Dialect/RTG/IR/basic.mlir b/test/Dialect/RTG/IR/basic.mlir
diff --git a/test/Dialect/RTG/IR/errors.mlir b/test/Dialect/RTG/IR/errors.mlir
