[RTG] Add any_context attribute (#8373)

frontends/PyRTG/src/pyrtg/contexts.py

@@ -9,6 +9,7 @@ from .rtgtest import rtgtest
 from .core import Value
 from .circt import ir
 from .support import _FromCirctValue
+from .sequences import Sequence
 
 import ctypes
 import sys
@@ -34,12 +35,39 @@ class CPUCore(Value):
   when generating randomized tests.
   """
 
-  def __init__(self, hartid: Union[int, ir.Value]) -> CPUCore:
+  def __init__(self, hartid: Union[int, ir.Attribute, ir.Value]) -> CPUCore:
     """
     Creates a CPUCore instance for a specific hardware thread ID.
     """
 
-    self._value = hartid
+    if isinstance(hartid, ir.Value):
+      self._attr = None
+      self._value = hartid
+      return
+
+    self._attr = hartid if isinstance(
+        hartid, ir.Attribute) else rtgtest.CPUAttr.get(hartid)
+    self._value = self._attr
+
+  @staticmethod
+  def any() -> CPUCore:
+    """
+    Creates a CPUCore instance referring to any core. This is useful when
+    specifying a sequence to use for switching from one core to another.
+    """
+
+    return CPUCore(rtg.AnyContextAttr.get(rtgtest.CPUType.get()))
+
+  @staticmethod
+  def register_switch(from_core: CPUCore, to_core: CPUCore,
+                      seq: Sequence) -> None:
+    """
+    Register a sequence with arguments of types [CPUCore, CPUCore, Sequence] to
+    be usable for switching from context 'from_core' to context 'to_core' and back.
+    """
+
+    assert from_core._attr is not None and to_core._attr is not None, "must have attribute available"
+    rtg.ContextSwitchOp(from_core._attr, to_core._attr, seq)
 
   def __enter__(self):
     # TODO: just adding all variables in the context is not particularly nice.
@@ -115,8 +143,8 @@ class CPUCore(Value):
     ctypes.pythonapi.PyFrame_LocalsToFast(ctypes.py_object(s), ctypes.c_int(1))
 
   def _get_ssa_value(self) -> ir.Value:
-    if isinstance(self._value, int):
-      self = rtg.ConstantOp(rtgtest.CPUAttr.get(self._value))
+    if isinstance(self._value, ir.Attribute):
+      self = rtg.ConstantOp(self._value)
     return self._value
 
   def get_type(self) -> ir.Type:

frontends/PyRTG/test/test_contexts.py

@@ -1,6 +1,6 @@
 # RUN: %rtgtool% %s --seed=0 --output-format=mlir | FileCheck %s
 
-from pyrtg import test, sequence, Integer, CPUCore
+from pyrtg import test, sequence, Integer, CPUCore, Sequence, target, entry
 
 
 @sequence(Integer.type())
@@ -8,6 +8,27 @@ def consumer(arg):
   pass
 
 
+@sequence(CPUCore.type(), CPUCore.type(), Sequence.type())
+def switch(from_ctxt, to_ctxt, seq):
+  pass
+
+
+# MLIR-LABEL: rtg.target @Tgt0 : !rtg.dict<cpu: !rtgtest.cpu>
+# MLIR-NEXT: [[V0:%.+]] = rtg.constant #rtgtest.cpu<0> : !rtgtest.cpu
+# MLIR-NEXT: [[V1:%.+]] = rtg.get_sequence @switch : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+# MLIR-NEXT: rtg.context_switch #rtg.any_context : !rtgtest.cpu -> #rtgtest.cpu<0> : !rtgtest.cpu, [[V1]] : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+# MLIR-NEXT: rtg.yield [[V0]] : !rtgtest.cpu
+
+
+@target
+class Tgt0:
+
+  @entry
+  def cpu():
+    CPUCore.register_switch(CPUCore.any(), CPUCore(0), switch)
+    return CPUCore(0)
+
+
 # CHECK-LABEL: rtg.test @test0_context_args
 # CHECK-NEXT:   [[IDX4:%.+]] = index.constant 4
 # CHECK-NEXT:   [[SEQ0:%.+]] = rtg.get_sequence @_context_seq_0 : !rtg.sequence<index, index, !rtgtest.cpu, index>

include/circt-c/Dialect/RTG.h

@@ -162,6 +162,13 @@ MLIR_CAPI_EXPORTED uint32_t rtgImmediateAttrGetWidth(MlirAttribute attr);
 /// Returns the value of the RTG immediate attribute.
 MLIR_CAPI_EXPORTED uint64_t rtgImmediateAttrGetValue(MlirAttribute attr);
 
+/// If the attribute is an RTG any context attribute.
+MLIR_CAPI_EXPORTED bool rtgAttrIsAAnyContextAttr(MlirAttribute attr);
+
+/// Creates an RTG any context attribute in the context.
+MLIR_CAPI_EXPORTED MlirAttribute rtgAnyContextAttrGet(MlirContext ctxt,
+                                                      MlirType type);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -38,6 +38,19 @@ def DefaultContextAttr : RTGAttrDef<"DefaultContext", [
   let assemblyFormat = "";
 }
 
+def AnyContextAttr : RTGAttrDef<"AnyContext", [
+  DeclareAttrInterfaceMethods<ContextResourceAttrInterface>,
+]> {
+  let summary = "any single context of its type";
+  let description = [{
+    This attribute can be used to refer to any context of its type.
+  }];
+
+  let mnemonic = "any_context";
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type);
+  let assemblyFormat = "";
+}
+
 //===----------------------------------------------------------------------===//
 // Attributes for ISA targets
 //===----------------------------------------------------------------------===//

integration_test/Bindings/Python/dialects/rtg.py

@@ -205,6 +205,12 @@ with Context() as ctx, Location.unknown():
   # CHECK: #rtg.default : !rtgtest.cpu
   print(attr)
 
+  attr = rtg.AnyContextAttr.get(rtgtest.CPUType.get())
+  # CHECK: !rtgtest.cpu
+  print(attr.type)
+  # CHECK: #rtg.any_context : !rtgtest.cpu
+  print(attr)
+
   immediate_type = rtg.ImmediateType.get(32)
   # CHECK: width=32
   print(f"width={immediate_type.width}")

lib/Bindings/Python/RTGModule.cpp

@@ -167,6 +167,14 @@ void circt::python::populateDialectRTGSubmodule(nb::module_ &m) {
           },
           nb::arg("self"), nb::arg("type"), nb::arg("ctxt") = nullptr);
 
+  mlir_attribute_subclass(m, "AnyContextAttr", rtgAttrIsAAnyContextAttr)
+      .def_classmethod(
+          "get",
+          [](nb::object cls, MlirType type, MlirContext ctxt) {
+            return cls(rtgAnyContextAttrGet(ctxt, type));
+          },
+          nb::arg("self"), nb::arg("type"), nb::arg("ctxt") = nullptr);
+
   // Attributes for ISA targets
   //===--------------------------------------------------------------------===//
 

lib/CAPI/Dialect/RTG.cpp

@@ -242,3 +242,14 @@ uint32_t rtgImmediateAttrGetWidth(MlirAttribute attr) {
 uint64_t rtgImmediateAttrGetValue(MlirAttribute attr) {
   return cast<ImmediateAttr>(unwrap(attr)).getValue().getZExtValue();
 }
+
+// AnyContexts
+//===----------------------------------------------------------------------===//
+
+bool rtgAttrIsAAnyContextAttr(MlirAttribute attr) {
+  return isa<AnyContextAttr>(unwrap(attr));
+}
+
+MlirAttribute rtgAnyContextAttrGet(MlirContext ctxt, MlirType type) {
+  return wrap(AnyContextAttr::get(unwrap(ctxt), unwrap(type)));
+}

lib/Dialect/RTG/Transforms/ElaborationPass.cpp

@@ -1433,7 +1433,28 @@ public:
     }
 
     // Switch to the desired context.
+    // First, check if a context switch is registered that has the concrete
+    // context as source and target.
     auto *iter = testState.contextSwitches.find({from, to});
+
+    // Try with 'any' context as target and the concrete context as source.
+    if (iter == testState.contextSwitches.end())
+      iter = testState.contextSwitches.find(
+          {from, AnyContextAttr::get(op->getContext(), to.getType())});
+
+    // Try with 'any' context as source and the concrete context as target.
+    if (iter == testState.contextSwitches.end())
+      iter = testState.contextSwitches.find(
+          {AnyContextAttr::get(op->getContext(), from.getType()), to});
+
+    // Try with 'any' context for both the source and the target.
+    if (iter == testState.contextSwitches.end())
+      iter = testState.contextSwitches.find(
+          {AnyContextAttr::get(op->getContext(), from.getType()),
+           AnyContextAttr::get(op->getContext(), to.getType())});
+
+    // Otherwise, fail with an error because we couldn't find a user
+    // specification on how to switch between the requested contexts.
     // NOTE: we could think about supporting context switching via intermediate
     // context, i.e., treat it as a transitive relation.
     if (iter == testState.contextSwitches.end())

test/CAPI/rtg.c

@@ -138,17 +138,26 @@ static void testLabelVisibilityAttr(MlirContext ctx) {
   }
 }
 
-static void testDefaultContextAttr(MlirContext ctx) {
+static void testContextAttrs(MlirContext ctx) {
   MlirType cpuTy = rtgtestCPUTypeGet(ctx);
-  MlirAttribute defaultCtxtAttr = rtgDefaultContextAttrGet(ctx, cpuTy);
 
+  // DefaultContext
+  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));
+
+  // AnyContext
+  MlirAttribute anyCtxtAttr = rtgAnyContextAttrGet(ctx, cpuTy);
+  // CHECK: is_any_context
+  fprintf(stderr, rtgAttrIsAAnyContextAttr(anyCtxtAttr)
+                      ? "is_any_context\n"
+                      : "isnot_any_context\n");
+  // CHECK: !rtgtest.cpu
+  mlirTypeDump(mlirAttributeGetType(anyCtxtAttr));
 }
 
 static void testImmediate(MlirContext ctx) {
@@ -212,7 +221,7 @@ int main(int argc, char **argv) {
   testArrayType(ctx);
 
   testLabelVisibilityAttr(ctx);
-  testDefaultContextAttr(ctx);
+  testContextAttrs(ctx);
   testImmediate(ctx);
   testMemories(ctx);
 

test/Dialect/RTG/IR/basic.mlir

@@ -139,6 +139,9 @@ rtg.test @contexts(ctxt0 = %ctxt0: !rtgtest.cpu) {
   // CHECK: rtg.on_context {{%.+}}, {{%.+}} : !rtgtest.cpu
   %seq = rtg.get_sequence @seq0 : !rtg.sequence
   rtg.on_context %ctxt0, %seq : !rtgtest.cpu
+
+  // CHECK: rtg.constant #rtg.any_context : !rtgtest.cpu
+  rtg.constant #rtg.any_context : !rtgtest.cpu
 }
 
 // CHECK-LABEL: rtg.test @test0

test/Dialect/RTG/Transform/elaboration.mlir

@@ -524,7 +524,7 @@ rtg.sequence @switchCpuSeq(%parent: !rtgtest.cpu, %child: !rtgtest.cpu, %seq: !r
 // CHECK:  rtg.sequence @switchNestedCpuSeq_0() {
 // CHECK-NEXT:    [[L12:%.+]] = rtg.label_decl "label7"
 // CHECK-NEXT:    rtg.label local [[L12]]
-// CHECK-NEXT:    [[SEQ8:%.+]] = rtg.get_sequence @nestedCpuSeq_0 : !rtg.sequence
+// CHECK-NEXT:    [[SEQ8:%.+]] = rtg.get_sequence @nestedCpuSeq{{.*}} : !rtg.sequence
 // CHECK-NEXT:    [[SEQ9:%.+]] = rtg.randomize_sequence [[SEQ8]]
 // CHECK-NEXT:    rtg.embed_sequence [[SEQ9]]
 // CHECK-NEXT:    [[L13:%.+]] = rtg.label_decl "label8"
@@ -540,10 +540,24 @@ rtg.sequence @switchNestedCpuSeq(%parent: !rtgtest.cpu, %child: !rtgtest.cpu, %s
 }
 
 rtg.target @singleCoreTarget : !rtg.dict<single_core: !rtgtest.cpu> {
+  %0 = rtg.get_sequence @switchCpuSeq : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+  %1 = rtg.get_sequence @switchNestedCpuSeq : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+  rtg.context_switch #rtg.any_context : !rtgtest.cpu -> #rtg.any_context : !rtgtest.cpu, %1 : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
+  rtg.context_switch #rtg.default : !rtgtest.cpu -> #rtg.any_context : !rtgtest.cpu, %0 : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtg.sequence>
   %2 = rtg.constant #rtgtest.cpu<0>
   rtg.yield %2 : !rtgtest.cpu
 }
 
+// CHECK-LABEL: rtg.test @anyContextSwitch_singleCoreTarget
+rtg.test @anyContextSwitch(single_core = %single_core: !rtgtest.cpu) {
+  // CHECK-NEXT: [[V0:%.+]] = rtg.get_sequence @switchCpuSeq{{.*}} : !rtg.sequence
+  // CHECK-NEXT: [[V1:%.+]] = rtg.randomize_sequence [[V0]]
+  // CHECK-NEXT: rtg.embed_sequence [[V1]]
+  // CHECK-NEXT: }
+  %0 = rtg.get_sequence @nestedCpuSeq : !rtg.sequence
+  rtg.on_context %single_core, %0 : !rtgtest.cpu
+}
+
 rtg.sequence @interleaveSequencesSeq0() {
   rtgtest.rv32i.ebreak
   rtgtest.rv32i.ebreak