[RTG][Elaboration] Support folders of pure operations (#8374)

include/circt/Dialect/RTGTest/IR/RTGTestOps.td

@@ -31,6 +31,14 @@ def ConstantTestOp : RTGTestOp<"constant_test", [
   let hasFolder = 1;
 }
 
+def GetHartIdOp : RTGTestOp<"get_hartid", [Pure]> {
+  let arguments = (ins CPUType:$cpu);
+  let results = (outs Index:$hartid);
+
+  let assemblyFormat = "$cpu attr-dict";
+  let hasFolder = 1;
+}
+
 //===- Instruction Formats -------------------------------------------------===//
 
 class InstFormatIOpBase<string mnemonic, int opcode7, int funct3>

lib/Dialect/RTG/Transforms/ElaborationPass.cpp

@@ -925,27 +925,38 @@ public:
     return std::get<ValueTy>(state.at(val));
   }
 
-  FailureOr<DeletionKind> visitConstantLike(Operation *op) {
-    assert(op->hasTrait<OpTrait::ConstantLike>() &&
-           "op is expected to be constant-like");
+  FailureOr<DeletionKind> visitPureOp(Operation *op) {
+    SmallVector<Attribute> operands;
+    for (auto operand : op->getOperands()) {
+      auto evalValue = state[operand];
+      if (std::holds_alternative<TypedAttr>(evalValue))
+        operands.push_back(std::get<TypedAttr>(evalValue));
+      else
+        return visitUnhandledOp(op);
+    }
 
-    SmallVector<OpFoldResult, 1> result;
-    auto foldResult = op->fold(result);
-    (void)foldResult; // Make sure there is a user when assertions are off.
-    assert(succeeded(foldResult) &&
-           "constant folder of a constant-like must always succeed");
-    auto attr = dyn_cast<TypedAttr>(result[0].dyn_cast<Attribute>());
-    if (!attr)
-      return op->emitError(
-          "only typed attributes supported for constant-like operations");
+    SmallVector<OpFoldResult> results;
+    if (failed(op->fold(operands, results)))
+      return visitUnhandledOp(op);
 
-    auto intAttr = dyn_cast<IntegerAttr>(attr);
-    if (intAttr && isa<IndexType>(attr.getType()))
-      state[op->getResult(0)] = size_t(intAttr.getInt());
-    else if (intAttr && intAttr.getType().isSignlessInteger(1))
-      state[op->getResult(0)] = bool(intAttr.getInt());
-    else
-      state[op->getResult(0)] = attr;
+    // We don't support in-place folders.
+    if (results.size() != op->getNumResults())
+      return visitUnhandledOp(op);
+
+    for (auto [res, val] : llvm::zip(results, op->getResults())) {
+      auto attr = llvm::dyn_cast_or_null<TypedAttr>(res.dyn_cast<Attribute>());
+      if (!attr)
+        return op->emitError(
+            "only typed attributes supported for constant-like operations");
+
+      auto intAttr = dyn_cast<IntegerAttr>(attr);
+      if (intAttr && isa<IndexType>(attr.getType()))
+        state[op->getResult(0)] = size_t(intAttr.getInt());
+      else if (intAttr && intAttr.getType().isSignlessInteger(1))
+        state[op->getResult(0)] = bool(intAttr.getInt());
+      else
+        state[op->getResult(0)] = attr;
+    }
 
     return DeletionKind::Delete;
   }
@@ -956,8 +967,9 @@ public:
   }
 
   FailureOr<DeletionKind> visitExternalOp(Operation *op) {
-    if (op->hasTrait<OpTrait::ConstantLike>())
-      return visitConstantLike(op);
+    auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+    if (op->hasTrait<OpTrait::ConstantLike>() || (memOp && memOp.hasNoEffect()))
+      return visitPureOp(op);
 
     // TODO: we only have this to be able to write tests for this pass without
     // having to add support for more operations for now, so it should be
@@ -968,9 +980,7 @@ public:
     return visitUnhandledOp(op);
   }
 
-  FailureOr<DeletionKind> visitOp(ConstantOp op) {
-    return visitConstantLike(op);
-  }
+  FailureOr<DeletionKind> visitOp(ConstantOp op) { return visitPureOp(op); }
 
   FailureOr<DeletionKind> visitOp(GetSequenceOp op) {
     SmallVector<ElaboratorValue> replacements;
@@ -1201,7 +1211,7 @@ public:
   }
 
   FailureOr<DeletionKind> visitOp(FixedRegisterOp op) {
-    return visitConstantLike(op);
+    return visitPureOp(op);
   }
 
   FailureOr<DeletionKind> visitOp(VirtualRegisterOp op) {

lib/Dialect/RTGTest/IR/RTGTestOps.cpp

@@ -25,6 +25,16 @@ mlir::OpFoldResult ConstantTestOp::fold(FoldAdaptor adaptor) {
   return getValueAttr();
 }
 
+//===----------------------------------------------------------------------===//
+// GetHartIdOp
+//===----------------------------------------------------------------------===//
+
+mlir::OpFoldResult GetHartIdOp::fold(FoldAdaptor adaptor) {
+  if (auto cpuAttr = dyn_cast_or_null<CPUAttr>(adaptor.getCpu()))
+    return IntegerAttr::get(IndexType::get(getContext()), cpuAttr.getId());
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen generated logic.
 //===----------------------------------------------------------------------===//

test/Dialect/RTG/Transform/elaboration.mlir

@@ -480,6 +480,11 @@ rtg.sequence @switchNestedCpuSeq(%parent: !rtgtest.cpu, %child: !rtgtest.cpu, %s
   rtg.label local %l8
 }
 
+rtg.target @singleCoreTarget : !rtg.dict<single_core: !rtgtest.cpu> {
+  %2 = rtg.constant #rtgtest.cpu<0>
+  rtg.yield %2 : !rtgtest.cpu
+}
+
 rtg.sequence @interleaveSequencesSeq0() {
   rtgtest.rv32i.ebreak
   rtgtest.rv32i.ebreak
@@ -567,6 +572,14 @@ rtg.test @tuples() {
   func.call @dummy2(%1) : (index) -> ()
 }
 
+// CHECK-LABEL: rtg.test @useFolders_singleCoreTarget
+rtg.test @useFolders(single_core = %single_core: !rtgtest.cpu) {
+  // CHECK-NEXT: index.constant 0
+  // CHECK-NEXT: call @dummy2
+  %0 = rtgtest.get_hartid %single_core
+  func.call @dummy2(%0) : (index) -> ()
+}
+
 // -----
 
 rtg.test @nestedRegionsNotSupported() {

test/Dialect/RTGTest/IR/basic.mlir

@@ -2,10 +2,12 @@
 
 // CHECK-LABEL: @cpus
 // CHECK-SAME: !rtgtest.cpu
-rtg.target @cpus : !rtg.dict<cpu: !rtgtest.cpu> {
-  // CHECK: rtg.constant #rtgtest.cpu<0> : !rtgtest.cpu
+rtg.target @cpus : !rtg.dict<cpu: !rtgtest.cpu, hartid: index> {
+  // CHECK: [[V0:%.+]] = rtg.constant #rtgtest.cpu<0> : !rtgtest.cpu
   %0 = rtg.constant #rtgtest.cpu<0>
-  rtg.yield %0 : !rtgtest.cpu
+  // CHECK: rtgtest.get_hartid [[V0]]
+  %1 = rtgtest.get_hartid %0
+  rtg.yield %0, %1 : !rtgtest.cpu, index
 }
 
 rtg.test @misc() {