[RTG] Add on_context and context_switch operations (#8150)

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

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)

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

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

@@ -16,6 +16,7 @@
 include "circt/Dialect/RTG/IR/RTGInterfaces.td"
 include "circt/Dialect/RTG/IR/RTGDialect.td"
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/Interfaces.td"
 
 class RTGAttrDef<string name, list<Trait> traits = []> 
   : AttrDef<RTGDialect, name, traits>;
@@ -56,4 +57,22 @@ class ImmediateAttrBase<int width> : RTGAttrDef<"Imm" # width, [
   let genVerifyDecl = 1;
 }
 
+
+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

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();
   }];
 }

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">,
 ]>;
 

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]> {

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

@@ -35,6 +35,8 @@ public:
             // Bags
             BagCreateOp, BagSelectRandomOp, BagDifferenceOp, BagUnionOp,
             BagUniqueSizeOp,
+            // Contexts
+            OnContextOp, ContextSwitchOp,
             // Labels
             LabelDeclOp, LabelUniqueDeclOp, LabelOp,
             // Registers
@@ -51,10 +53,6 @@ public:
             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 @@ public:
   /// 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 @@ public:
   HANDLE(RandomizeSequenceOp, Unhandled);
   HANDLE(EmbedSequenceOp, Unhandled);
   HANDLE(RandomNumberInRangeOp, Unhandled);
+  HANDLE(OnContextOp, Unhandled);
+  HANDLE(ContextSwitchOp, Unhandled);
   HANDLE(SetCreateOp, Unhandled);
   HANDLE(SetSelectRandomOp, Unhandled);
   HANDLE(SetDifferenceOp, Unhandled);

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);

integration_test/Bindings/Python/dialects/rtg.py

@@ -208,3 +208,11 @@ with Context() as ctx, Location.unknown():
   # 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)

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);
 }

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)));
+}

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

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"

lib/Dialect/RTG/IR/RTGDialect.cpp

@@ -25,6 +25,7 @@ using namespace rtg;
 //===----------------------------------------------------------------------===//
 
 void RTGDialect::initialize() {
+  registerAttributes();
   registerTypes();
   // Register operations.
   addOperations<

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
 //===----------------------------------------------------------------------===//

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(); }
 
 //===----------------------------------------------------------------------===//

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

test/CAPI/rtg.c

@@ -11,6 +11,7 @@
 #include <stdio.h>
 
 #include "circt-c/Dialect/RTG.h"
+#include "circt-c/Dialect/RTGTest.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 
@@ -120,9 +121,23 @@ static void testLabelVisibilityAttr(MlirContext ctx) {
   }
 }
 
+static void testDefaultContextAttr(MlirContext ctx) {
+  MlirType cpuTy = rtgtestCPUTypeGet(ctx);
+  MlirAttribute defaultCtxtAttr = rtgDefaultContextAttrGet(ctx, cpuTy);
+
+  // CHECK: is_default_context
+  fprintf(stderr, rtgAttrIsADefaultContextAttr(defaultCtxtAttr)
+                      ? "is_default_context\n"
+                      : "isnot_default_context\n");
+
+  // CHECK: !rtgtest.cpu
+  mlirTypeDump(mlirAttributeGetType(defaultCtxtAttr));
+}
+
 int main(int argc, char **argv) {
   MlirContext ctx = mlirContextCreate();
   mlirDialectHandleLoadDialect(mlirGetDialectHandle__rtg__(), ctx);
+  mlirDialectHandleLoadDialect(mlirGetDialectHandle__rtgtest__(), ctx);
 
   testSequenceType(ctx);
   testRandomizedSequenceType(ctx);
@@ -132,6 +147,7 @@ int main(int argc, char **argv) {
   testDictType(ctx);
 
   testLabelVisibilityAttr(ctx);
+  testDefaultContextAttr(ctx);
 
   mlirContextDestroy(ctx);
 

test/Dialect/RTG/IR/basic.mlir

@@ -100,6 +100,28 @@ rtg.target @target : !rtg.dict<num_cpus: i32, num_modes: i32> {
   rtg.yield %1, %1 : i32, i32
 }
 
+// CHECK-LABEL: rtg.sequence @switch_seq
+// CHECK-SAME: (%{{.*}}: !rtgtest.cpu, %{{.*}}: !rtgtest.cpu, %{{.*}}: !rtg.sequence)
+rtg.sequence @switch_seq(%from: !rtgtest.cpu, %to: !rtgtest.cpu, %seq: !rtg.sequence) { }
+
+// CHECK-LABEL: rtg.target @context_switch
+rtg.target @context_switch : !rtg.dict<> {
+  // CHECK: [[V0:%.+]] = rtg.get_sequence
+  // CHECK: rtg.context_switch #rtg.default : !rtgtest.cpu -> #rtgtest.cpu<1> : !rtgtest.cpu, [[V0]] : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+  %0 = rtg.get_sequence @switch_seq : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+  rtg.context_switch #rtg.default : !rtgtest.cpu -> #rtgtest.cpu<1>, %0 : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+
+  rtg.yield
+}
+
+// CHECK-LABEL: @contexts
+rtg.test @contexts : !rtg.dict<ctxt0: !rtgtest.cpu> {
+^bb0(%arg0: !rtgtest.cpu):
+  // CHECK: rtg.on_context {{%.+}}, {{%.+}} : !rtgtest.cpu
+  %seq = rtg.get_sequence @seq0 : !rtg.sequence
+  rtg.on_context %arg0, %seq : !rtgtest.cpu
+}
+
 // CHECK-LABEL: rtg.test @test : !rtg.dict<num_cpus: i32, num_modes: i32> {
 // CHECK: ^bb0(%arg0: i32, %arg1: i32):
 // CHECK: }

test/Dialect/RTG/IR/errors.mlir

@@ -134,3 +134,43 @@ rtg.sequence @seq() {
   // expected-error @below {{expected 1 or more operands, but found 0}}
   rtg.bag_union : !rtg.bag<i32>
 }
+
+// -----
+
+rtg.sequence @seq() {}
+
+rtg.target @target : !rtg.dict<> {
+  %0 = rtg.get_sequence @seq : !rtg.sequence
+  // expected-error @below {{sequence type must have exactly 3 element types}}
+  rtg.context_switch #rtgtest.cpu<0> -> #rtgtest.cpu<1>, %0 : !rtg.sequence
+}
+
+// -----
+
+rtg.sequence @seq(%arg0: !rtg.sequence, %arg1: !rtg.sequence, %arg2: !rtg.sequence) {}
+
+rtg.target @target : !rtg.dict<> {
+  %0 = rtg.get_sequence @seq : !rtg.sequence<!rtg.sequence, !rtg.sequence, !rtg.sequence>
+  // expected-error @below {{first sequence element type must match 'from' attribute type}}
+  rtg.context_switch #rtgtest.cpu<0> -> #rtgtest.cpu<1>, %0 : !rtg.sequence<!rtg.sequence, !rtg.sequence, !rtg.sequence>
+}
+
+// -----
+
+rtg.sequence @seq(%arg0: !rtgtest.cpu, %arg1: !rtg.sequence, %arg2: !rtg.sequence) {}
+
+rtg.target @target : !rtg.dict<> {
+  %0 = rtg.get_sequence @seq : !rtg.sequence<!rtgtest.cpu, !rtg.sequence, !rtg.sequence>
+  // expected-error @below {{second sequence element type must match 'to' attribute type}}
+  rtg.context_switch #rtgtest.cpu<0> -> #rtgtest.cpu<1>, %0 : !rtg.sequence<!rtgtest.cpu, !rtg.sequence, !rtg.sequence>
+}
+
+// -----
+
+rtg.sequence @seq(%arg0: !rtgtest.cpu, %arg1: !rtgtest.cpu, %arg2: !rtgtest.cpu) {}
+
+rtg.target @target : !rtg.dict<> {
+  %0 = rtg.get_sequence @seq : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtgtest.cpu>
+  // expected-error @below {{third sequence element type must be a fully substituted sequence}}
+  rtg.context_switch #rtgtest.cpu<0> -> #rtgtest.cpu<1>, %0 : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtgtest.cpu>
+}