From 6c82ca141afa81e9ead9aaf3e7259f28f491179f Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Fri, 31 Jul 2026 18:00:45 -0700 Subject: [PATCH] [mlir-c] Add RewriterBase insertion point save/restore (#206531) Exposes `OpBuilder::saveInsertionPoint` / `restoreInsertionPoint` through the MLIR C API for `MlirRewriterBase`, continuing the buildout of the rewrite/conversion C bindings. Assisted by: Claude --- mlir/include/mlir-c/Rewrite.h | 23 ++++++++- mlir/lib/CAPI/Transforms/Rewrite.cpp | 26 ++++++++++ mlir/test/CAPI/rewrite.c | 77 ++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir-c/Rewrite.h b/mlir/include/mlir-c/Rewrite.h index a77214ebb986..0e5562477d4f 100644 --- a/mlir/include/mlir-c/Rewrite.h +++ b/mlir/include/mlir-c/Rewrite.h @@ -78,8 +78,8 @@ mlirRewriterBaseGetContext(MlirRewriterBase rewriter); //===----------------------------------------------------------------------===// // These do not include functions using Block::iterator or Region::iterator, as -// they are not exposed by the C API yet. Similarly for methods using -// `InsertPoint` directly. +// they are not exposed by the C API yet. `InsertPoint` is exposed as +// `MlirRewriterBaseInsertPoint` below. /// Reset the insertion point to no location. Creating an operation without a /// set insertion point is an error, but this can still be useful when the @@ -133,6 +133,25 @@ mlirRewriterBaseGetBlock(MlirRewriterBase rewriter); MLIR_CAPI_EXPORTED MlirOperation mlirRewriterBaseGetOperationAfterInsertion(MlirRewriterBase rewriter); +/// A saved insertion point: a (block, operationAfter) pair. `operationAfter` is +/// the operation that subsequent insertions go before. If `operationAfter` is +/// null, the insertion point is at the end of `block`. If `block` is null, the +/// insertion point is not set (cleared). +typedef struct MlirRewriterBaseInsertPoint { + MlirBlock block; + MlirOperation operationAfter; +} MlirRewriterBaseInsertPoint; + +/// Returns the current insertion point of the rewriter so that it can be +/// restored later with mlirRewriterBaseRestoreInsertionPoint. +MLIR_CAPI_EXPORTED MlirRewriterBaseInsertPoint +mlirRewriterBaseSaveInsertionPoint(MlirRewriterBase rewriter); + +/// Restores a previously saved insertion point. +MLIR_CAPI_EXPORTED void +mlirRewriterBaseRestoreInsertionPoint(MlirRewriterBase rewriter, + MlirRewriterBaseInsertPoint insertPoint); + //===----------------------------------------------------------------------===// /// Block and operation creation/insertion/cloning //===----------------------------------------------------------------------===// diff --git a/mlir/lib/CAPI/Transforms/Rewrite.cpp b/mlir/lib/CAPI/Transforms/Rewrite.cpp index cd040dea124c..e54c7dcd7271 100644 --- a/mlir/lib/CAPI/Transforms/Rewrite.cpp +++ b/mlir/lib/CAPI/Transforms/Rewrite.cpp @@ -87,6 +87,32 @@ mlirRewriterBaseGetOperationAfterInsertion(MlirRewriterBase rewriter) { return wrap(std::addressof(*it)); } +MlirRewriterBaseInsertPoint +mlirRewriterBaseSaveInsertionPoint(MlirRewriterBase rewriter) { + OpBuilder::InsertPoint ip = unwrap(rewriter)->saveInsertionPoint(); + if (!ip.isSet()) + return {{nullptr}, {nullptr}}; + Block *block = ip.getBlock(); + MlirOperation operationAfter = ip.getPoint() == block->end() + ? MlirOperation{nullptr} + : wrap(&*ip.getPoint()); + return {wrap(block), operationAfter}; +} + +void mlirRewriterBaseRestoreInsertionPoint( + MlirRewriterBase rewriter, MlirRewriterBaseInsertPoint insertPoint) { + if (mlirBlockIsNull(insertPoint.block)) { + unwrap(rewriter)->clearInsertionPoint(); + return; + } + Block *block = unwrap(insertPoint.block); + if (mlirOperationIsNull(insertPoint.operationAfter)) + unwrap(rewriter)->setInsertionPointToEnd(block); + else + unwrap(rewriter)->setInsertionPoint( + block, Block::iterator(unwrap(insertPoint.operationAfter))); +} + //===----------------------------------------------------------------------===// /// Block and operation creation/insertion/cloning //===----------------------------------------------------------------------===// diff --git a/mlir/test/CAPI/rewrite.c b/mlir/test/CAPI/rewrite.c index 321ab6c1f039..60f5dca31b9e 100644 --- a/mlir/test/CAPI/rewrite.c +++ b/mlir/test/CAPI/rewrite.c @@ -623,6 +623,82 @@ void testCloneWithMapping(MlirContext ctx) { fprintf(stderr, "testCloneWithMapping: PASSED\n"); } +void testInsertionPointSaveRestore(MlirContext ctx) { + // CHECK-LABEL: @testInsertionPointSaveRestore + fprintf(stderr, "@testInsertionPointSaveRestore\n"); + + const char *moduleString = "\"dialect.op1\"() : () -> ()\n" + "\"dialect.op2\"() : () -> ()\n"; + MlirModule module = + mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleString)); + MlirOperation op = mlirModuleGetOperation(module); + MlirBlock body = mlirModuleGetBody(module); + MlirOperation op1 = mlirBlockGetFirstOperation(body); + MlirOperation op2 = mlirOperationGetNextInBlock(op1); + + MlirRewriterBase rewriter = mlirIRRewriterCreate(ctx); + + // Save an insertion point that points right before op2. + mlirRewriterBaseSetInsertionPointBefore(rewriter, op2); + MlirRewriterBaseInsertPoint saved = + mlirRewriterBaseSaveInsertionPoint(rewriter); + assert(!mlirBlockIsNull(saved.block)); + assert(mlirOperationEqual(saved.operationAfter, op2)); + + // Move the insertion point to the end of the block. An end-of-block insertion + // point round-trips with a null `operationAfter`. + mlirRewriterBaseSetInsertionPointToEnd(rewriter, body); + MlirRewriterBaseInsertPoint endIp = + mlirRewriterBaseSaveInsertionPoint(rewriter); + assert(!mlirBlockIsNull(endIp.block)); + assert(mlirOperationIsNull(endIp.operationAfter)); + + // Restoring the first saved insertion point makes subsequent insertions land + // before op2 again, not at the end where we just were. + mlirRewriterBaseRestoreInsertionPoint(rewriter, saved); + MlirOperation opRestored = + createOperationWithName(ctx, "dialect.op_restored"); + mlirRewriterBaseInsert(rewriter, opRestored); + + // Restoring the null-`operationAfter` point re-establishes end-of-block, even + // though the insertion point currently sits in the middle of the block. + mlirRewriterBaseRestoreInsertionPoint(rewriter, endIp); + assert(!mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter))); + assert(mlirOperationIsNull( + mlirRewriterBaseGetOperationAfterInsertion(rewriter))); + MlirOperation opEnd = createOperationWithName(ctx, "dialect.op_end"); + mlirRewriterBaseInsert(rewriter, opEnd); + + // A cleared insertion point round-trips as a null block. + mlirRewriterBaseClearInsertionPoint(rewriter); + MlirRewriterBaseInsertPoint clearedIp = + mlirRewriterBaseSaveInsertionPoint(rewriter); + assert(mlirBlockIsNull(clearedIp.block)); + assert(mlirOperationIsNull(clearedIp.operationAfter)); + + // Restoring a cleared insertion point clears the current one. + mlirRewriterBaseSetInsertionPointToStart(rewriter, body); + assert(!mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter))); + mlirRewriterBaseRestoreInsertionPoint(rewriter, clearedIp); + assert(mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter))); + + mlirOperationDump(op); + // clang-format off + // CHECK: module { + // CHECK-NEXT: "dialect.op1"() : () -> () + // CHECK-NEXT: %{{.*}} = "dialect.op_restored"() : () -> index + // CHECK-NEXT: "dialect.op2"() : () -> () + // CHECK-NEXT: %{{.*}} = "dialect.op_end"() : () -> index + // CHECK-NEXT: } + // clang-format on + + mlirIRRewriterDestroy(rewriter); + mlirModuleDestroy(module); + + // CHECK: testInsertionPointSaveRestore: PASSED + fprintf(stderr, "testInsertionPointSaveRestore: PASSED\n"); +} + static MlirConversionTargetLegality dynamicLegalityAlwaysLegal(MlirOperation op, void *userData) { (void)op; @@ -1899,6 +1975,7 @@ int main(void) { testReplaceUses(ctx); testGreedyRewriteDriverConfig(ctx); testCloneWithMapping(ctx); + testInsertionPointSaveRestore(ctx); testConversionTargetDynamicLegality(ctx); testTypeConverterSourceMaterialization(ctx); testTypeConverterTargetMaterialization(ctx);