diff --git a/lib/Dialect/RTG/Transforms/ElaborationPass.cpp b/lib/Dialect/RTG/Transforms/ElaborationPass.cpp index fdc06a82a042..4000095245d4 100644 --- a/lib/Dialect/RTG/Transforms/ElaborationPass.cpp +++ b/lib/Dialect/RTG/Transforms/ElaborationPass.cpp @@ -25,6 +25,7 @@ #include "mlir/IR/IRMapping.h" #include "mlir/IR/PatternMatch.h" #include "llvm/ADT/DenseMapInfoVariant.h" +#include "llvm/ADT/ScopeExit.h" #include "llvm/Support/Debug.h" #include @@ -108,6 +109,7 @@ struct MemoryBlockStorage; struct SymbolicComputationWithIdentityStorage; struct SymbolicComputationWithIdentityValue; struct SymbolicComputationStorage; +struct OpaqueExternalStorage; /// The abstract base class for elaborated values. using ElaboratorValue = @@ -117,7 +119,7 @@ using ElaboratorValue = ArrayStorage *, TupleStorage *, MemoryStorage *, MemoryBlockStorage *, SymbolicComputationWithIdentityStorage *, SymbolicComputationWithIdentityValue *, - SymbolicComputationStorage *>; + SymbolicComputationStorage *, OpaqueExternalStorage *>; // NOLINTNEXTLINE(readability-identifier-naming) llvm::hash_code hash_value(const ElaboratorValue &val) { @@ -512,6 +514,15 @@ struct SymbolicComputationWithIdentityValue : IdentityValue { const unsigned idx; }; +/// Storage for SSA values produced by external ops with regions — block +/// arguments and results. Carries identity only; the corresponding new-IR +/// `Value` is registered via `Materializer::map()` immediately upon creation so +/// `Materializer::materialize()` always hits its cache and never dispatches +// through `visit()`. +struct OpaqueExternalStorage : IdentityValue { + OpaqueExternalStorage(Type type, Location loc) : IdentityValue(type, loc) {} +}; + /// An 'Internalizer' object internalizes storages and takes ownership of them. /// When the initializer object is destroyed, all owned storages are also /// deallocated and thus must not be accessed anymore. @@ -717,6 +728,10 @@ static void print(const SymbolicComputationStorage *val, os << ">"; } +static void print(const OpaqueExternalStorage *val, llvm::raw_ostream &os) { + os << "type << ">"; +} + static llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const ElaboratorValue &value) { std::visit([&](auto val) { print(val, os); }, value); @@ -874,6 +889,7 @@ class ElaboratorValueToAttributeConverter { VISIT_UNSUPPORTED(SymbolicComputationWithIdentityStorage) VISIT_UNSUPPORTED(SymbolicComputationWithIdentityValue) VISIT_UNSUPPORTED(SymbolicComputationStorage) + VISIT_UNSUPPORTED(OpaqueExternalStorage) #undef VISIT_UNSUPPORTED @@ -922,8 +938,9 @@ class Materializer { Materializer(OpBuilder builder, TestState &testState, SharedState &sharedState, SmallVector &blockArgs) - : builder(builder), testState(testState), sharedState(sharedState), - blockArgs(blockArgs), attrConverter(builder.getContext()) {} + : builder(builder), rootBlock(builder.getBlock()), testState(testState), + sharedState(sharedState), blockArgs(blockArgs), + attrConverter(builder.getContext()) {} /// Materialize IR representing the provided `ElaboratorValue` and return the /// `Value` or a null value on failure. @@ -984,6 +1001,11 @@ class Materializer { /// before the insertion point. LogicalResult materialize(Operation *op, DenseMap &state) { + // Region-bearing ops must be elaborated away before reaching here (either + // by a dedicated visitor, or by visitExternalOp which performs its own + // structural rewrite). Routing them through this generic operand-only + // materialize path would silently keep un-elaborated regions in the + // output. if (op->getNumRegions() > 0) return op->emitOpError("ops with nested regions must be elaborated away"); @@ -1056,6 +1078,8 @@ class Materializer { void map(ElaboratorValue eval, Value val) { materializedValues[eval] = val; } + OpBuilder &getBuilder() { return builder; } + template OpTy create(Location location, Args &&...args) { return OpTy::create(builder, location, std::forward(args)...); @@ -1064,6 +1088,13 @@ class Materializer { private: Value tryMaterializeAsConstant(ElaboratorValue val, Location loc) { if (auto attr = attrConverter.convert(val)) { + // Hoist constants to the materializer's root block when the builder's + // current insertion point is in a nested region. This ensures the + // constant's defining region dominates every potential use, including + // sibling regions that share this materializer's value cache. + OpBuilder::InsertionGuard guard(builder); + if (builder.getBlock() != rootBlock) + builder.setInsertionPointToStart(rootBlock); Value res = ConstantOp::create(builder, loc, attr); materializedValues[val] = res; return res; @@ -1336,6 +1367,16 @@ class Materializer { return op->getResult(0); } + Value visit(OpaqueExternalStorage *val, Location loc, + function_ref emitError) { + // Opaque external values (block arguments and results of external + // region-bearing ops) are always pre-mapped to a concrete SSA value by + // visitExternalOp, so materialize() short-circuits via the cache before + // reaching here. + emitError() << "cannot materialize opaque external value"; + return {}; + } + private: /// Cache values we have already materialized to reuse them later. We start /// with an insertion point at the start of the block and cache the (updated) @@ -1348,6 +1389,10 @@ class Materializer { /// for the reason stated above. OpBuilder builder; + /// The block this materializer was constructed at. Used as the hoist target + /// for constants that would otherwise be sunk into nested sibling regions. + Block *rootBlock; + SmallVector toDelete; TestState &testState; @@ -1401,7 +1446,95 @@ class Elaborator : public RTGOpVisitor> { } FailureOr 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. + // + // The new op is a structural shell: + // - Its operands are re-materialized from the elaborator state so that + // references to RTG-elaborated values become the corresponding new + // SSA values (and don't dangle after finalize() erases the old ops). + // - Each region gets a fresh block whose arguments mirror the old + // block's argument types/locations. Old block args are registered as + // opaque IdentityValues mapped to the new block args, so RTG ops + // inside the region can resolve them via state lookup. + // - Each result is registered as an opaque IdentityValue mapped to the + // corresponding new result, so downstream RTG ops consuming them + // find them in state. + auto *newOp = op->cloneWithoutRegions(); + materializer.getBuilder().insert(newOp); + + // Re-map operands: substitute each operand with the SSA value + // corresponding to its current ElaboratorValue. Materialization of + // operands must happen *before* `newOp` so the defs dominate the use; + // anchor the builder at `newOp` for the duration of operand re-map, then + // advance past it. + materializer.getBuilder().setInsertionPoint(newOp); + for (auto &operand : newOp->getOpOperands()) { + auto emitError = [&]() { + auto diag = newOp->emitError(); + diag.attachNote(newOp->getLoc()) + << "while materializing operand#" << operand.getOperandNumber() + << " of external region op"; + return diag; + }; + auto elabVal = state.at(operand.get()); + Value val = materializer.materialize(elabVal, newOp->getLoc(), emitError); + if (!val) + return failure(); + operand.set(val); + } + materializer.getBuilder().setInsertionPointAfter(newOp); + + // Register an old value as an opaque external value mapped to the + // corresponding new SSA value. RTG ops that consume `oldVal` will find the + // opaque value in `state` (it is treated as symbolic), and materializing it + // yields `newVal` directly via the materializer's cache. + auto mapOpaque = [&](Value oldVal, Value newVal) { + auto *storage = sharedState.internalizer.create( + oldVal.getType(), oldVal.getLoc()); + state[oldVal] = storage; + materializer.map(storage, newVal); + }; + + for (auto [oldRegion, newRegion] : + llvm::zip(op->getRegions(), newOp->getRegions())) { + if (oldRegion.empty()) + continue; + + // Give the new region a block whose arguments mirror the old block's + // argument types/locations, and register the old args as opaque values + // mapped to the new args so region-internal RTG ops can resolve them. + Block &oldBlock = oldRegion.front(); + Block &newBlock = newRegion.emplaceBlock(); + for (auto oldArg : oldBlock.getArguments()) { + Value newArg = newBlock.addArgument(oldArg.getType(), oldArg.getLoc()); + mapOpaque(oldArg, newArg); + } + + { + OpBuilder::InsertionGuard guard(materializer.getBuilder()); + materializer.getBuilder().setInsertionPoint(&newBlock, + newBlock.begin()); + SmallVector 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(oldRegion, {}, + /*keepTerminator=*/true, unused))) + return failure(); + } + } + + // Register each result as an opaque value mapped to the new op's result so + // downstream RTG ops consuming them find them in `state`. + for (auto [oldRes, newRes] : + llvm::zip(op->getResults(), newOp->getResults())) + mapOpaque(oldRes, newRes); + + return DeletionKind::Delete; } FailureOr visitOp(GetSequenceOp op) { @@ -2180,7 +2313,8 @@ class Elaborator : public RTGOpVisitor> { val) || std::holds_alternative( val) || - std::holds_alternative(val); + std::holds_alternative(val) || + std::holds_alternative(val); } bool isSymbolic(Operation *op) { @@ -2312,6 +2446,31 @@ class Elaborator : public RTGOpVisitor> { return region.getParentOp()->emitOpError( "regions with more than one block are not supported"); + // Save any prior bindings for this region's block arguments so that a + // recursive re-entry of the same region (e.g. nested invocation of the + // same handler region during a multi-shot resume whose continuation body + // re-performs the handled effect) does not clobber the outer frame's + // bindings of those args. Without this, the outer handler's continuation + // SSA value is rebound to the inner frame's continuation, and the outer + // handler's subsequent `rtg.resume %k, ...` resumes the inner frame. + SmallVector>> savedArgs; + savedArgs.reserve(region.getNumArguments()); + for (auto arg : region.getArguments()) { + auto it = state.find(arg); + if (it != state.end()) + savedArgs.emplace_back(arg, it->second); + else + savedArgs.emplace_back(arg, std::nullopt); + } + llvm::scope_exit restoreArgs([&] { + for (auto &[arg, prev] : savedArgs) { + if (prev.has_value()) + state[arg] = *prev; + else + state.erase(arg); + } + }); + for (auto [arg, elabArg] : llvm::zip(region.getArguments(), regionArguments)) state[arg] = elabArg; diff --git a/test/Dialect/RTG/Transform/elaboration.mlir b/test/Dialect/RTG/Transform/elaboration.mlir index 154f97da589d..1222f0a7ece1 100644 --- a/test/Dialect/RTG/Transform/elaboration.mlir +++ b/test/Dialect/RTG/Transform/elaboration.mlir @@ -1125,14 +1125,88 @@ 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 { + %0 = index.constant 0 + rtg.yield %0 : index +} + +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 + 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 + +// ----- + +// External region op produces a result that is consumed by an RTG op. + rtg.target @singletonTarget : !rtg.dict { %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_b(%arg0: index) -> () { return } + +// CHECK-LABEL: rtg.test @externalOpResultUsedByRTG +rtg.test @externalOpResultUsedByRTG(singleton = %none: index) { + // CHECK: [[C5:%.+]] = rtg.constant 5 + // CHECK: [[R:%.+]] = scf.execute_region -> index { + // CHECK-NEXT: scf.yield [[C5]] + %r = scf.execute_region -> index { + %c = index.constant 5 + scf.yield %c : index + } + // The result [[R]] of the external op is an opaque value that downstream ops + // resolve to the new op's SSA result. + // CHECK: func.call @consume_index_b([[R]]) + func.call @consume_index_b(%r) : (index) -> () +} + +// ----- + +// External region op with block arguments. + +rtg.target @singletonTarget : !rtg.dict { + %0 = index.constant 0 + rtg.yield %0 : index +} + +func.func @consume_index_c(%arg0: index) -> () { return } + +// CHECK-LABEL: rtg.test @externalOpRegionWithBlockArgs +rtg.test @externalOpRegionWithBlockArgs(singleton = %none: index) { + %c0 = index.constant 0 + %c10 = index.constant 10 + %c1 = index.constant 1 + // CHECK: scf.while ({{.+}}) : (index) -> index + %r = scf.while (%i = %c0) : (index) -> index { + %cond = index.cmp ult(%i, %c10) + scf.condition(%cond) %i : index + } do { + ^bb0(%j: index): + %next = index.add %j, %c1 + scf.yield %next : index + } + func.call @consume_index_c(%r) : (index) -> () } // -----