[RTG][Elaborator] Elaborate external ops with regions

lib/Dialect/RTG/Transforms/ElaborationPass.cpp

@@ -984,9 +984,6 @@ class Materializer {
   /// before the insertion point.
   LogicalResult materialize(Operation *op,
                             DenseMap<Value, ElaboratorValue> &state) {
-    if (op->getNumRegions() > 0)
-      return op->emitOpError("ops with nested regions must be elaborated away");
-
     // We don't support opaque values. If there is an SSA value that has a
     // use-site it needs an equivalent ElaborationValue representation.
     // NOTE: We could support cases where there is initially a use-site but that
@@ -1056,6 +1053,8 @@ class Materializer {
 
   void map(ElaboratorValue eval, Value val) { materializedValues[eval] = val; }
 
+  OpBuilder &getBuilder() { return builder; }
+
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&...args) {
     return OpTy::create(builder, location, std::forward<Args>(args)...);
@@ -1401,7 +1400,31 @@ class Elaborator : public RTGOpVisitor<Elaborator, FailureOr<DeletionKind>> {
   }
 
   FailureOr<DeletionKind> visitExternalOp(Operation *op) {
-    return visitOpGeneric(op);
+    if (op->getNumRegions() == 0)
+      return visitOpGeneric(op);
+
+    // Elaborate all regions of unknown external ops. Any op appearing inside
+    // an RTG test body with regions is expected to contain RTG constructs.
+    auto *newOp = op->cloneWithoutRegions();
+    materializer.getBuilder().insert(newOp);
+
+    for (unsigned i = 0; i < op->getNumRegions(); ++i) {
+      Block &newBlock = newOp->getRegion(i).emplaceBlock();
+      {
+        OpBuilder::InsertionGuard guard(materializer.getBuilder());
+        materializer.getBuilder().setInsertionPoint(&newBlock,
+                                                    newBlock.begin());
+        SmallVector<ElaboratorValue> unused;
+        // keepTerminator=true: the original terminator (rtg.yield, scf.yield,
+        // etc.) is cloned into the new block, preserving the op's expected
+        // terminator type.
+        if (failed(elaborate(op->getRegion(i), {},
+                             /*keepTerminator=*/true, unused)))
+          return failure();
+      }
+    }
+
+    return DeletionKind::Delete;
   }
 
   FailureOr<DeletionKind> visitOp(GetSequenceOp op) {

test/Dialect/RTG/Transform/elaboration.mlir

@@ -1125,15 +1125,35 @@ rtg.test @registerIndexConversions(singleton = %none: index) {
 
 // -----
 
+// Test: external ops (from non-RTG dialects) that have regions have their
+// regions elaborated in-place by visitExternalOp.  RTG constructs inside the
+// region are resolved; the op itself and its (dialect-appropriate) terminator
+// survive unchanged.
+
 rtg.target @singletonTarget : !rtg.dict<singleton: index> {
   %0 = index.constant 0
   rtg.yield %0 : index
 }
 
-rtg.test @nestedRegionsNotSupported(singleton = %none: index) {
-  // expected-error @below {{ops with nested regions must be elaborated away}}
-  scf.execute_region { scf.yield }
+func.func @consume_index(%arg0: index) -> () { return }
+
+// CHECK-LABEL: rtg.test @externalOpRegionElaborated
+rtg.test @externalOpRegionElaborated(singleton = %none: index) {
+  %c1 = index.constant 1
+  %c2 = index.constant 2
+  %set = rtg.set_create %c1, %c2 : index
+  // scf.execute_region is an external (non-RTG) op with a region.
+  scf.execute_region {
+    %sel = rtg.set_select_random %set : !rtg.set<index>
+    func.call @consume_index(%sel) : (index) -> ()
+    scf.yield
+  }
 }
+// The external op survives; RTG constructs inside its region are resolved.
+// CHECK: scf.execute_region
+// CHECK-NOT: rtg.set_select_random
+// CHECK: func.call @consume_index
+// CHECK-NOT: scf.execute_region
 
 // -----