473 lines · cpp
1//===- Rewrite.cpp - C API for Rewrite Patterns ---------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir-c/Rewrite.h"10 11#include "mlir-c/Transforms.h"12#include "mlir/CAPI/IR.h"13#include "mlir/CAPI/Rewrite.h"14#include "mlir/CAPI/Support.h"15#include "mlir/CAPI/Wrap.h"16#include "mlir/IR/Attributes.h"17#include "mlir/IR/PDLPatternMatch.h.inc"18#include "mlir/IR/PatternMatch.h"19#include "mlir/Rewrite/FrozenRewritePatternSet.h"20#include "mlir/Transforms/GreedyPatternRewriteDriver.h"21 22using namespace mlir;23 24//===----------------------------------------------------------------------===//25/// RewriterBase API inherited from OpBuilder26//===----------------------------------------------------------------------===//27 28MlirContext mlirRewriterBaseGetContext(MlirRewriterBase rewriter) {29 return wrap(unwrap(rewriter)->getContext());30}31 32//===----------------------------------------------------------------------===//33/// Insertion points methods34//===----------------------------------------------------------------------===//35 36void mlirRewriterBaseClearInsertionPoint(MlirRewriterBase rewriter) {37 unwrap(rewriter)->clearInsertionPoint();38}39 40void mlirRewriterBaseSetInsertionPointBefore(MlirRewriterBase rewriter,41 MlirOperation op) {42 unwrap(rewriter)->setInsertionPoint(unwrap(op));43}44 45void mlirRewriterBaseSetInsertionPointAfter(MlirRewriterBase rewriter,46 MlirOperation op) {47 unwrap(rewriter)->setInsertionPointAfter(unwrap(op));48}49 50void mlirRewriterBaseSetInsertionPointAfterValue(MlirRewriterBase rewriter,51 MlirValue value) {52 unwrap(rewriter)->setInsertionPointAfterValue(unwrap(value));53}54 55void mlirRewriterBaseSetInsertionPointToStart(MlirRewriterBase rewriter,56 MlirBlock block) {57 unwrap(rewriter)->setInsertionPointToStart(unwrap(block));58}59 60void mlirRewriterBaseSetInsertionPointToEnd(MlirRewriterBase rewriter,61 MlirBlock block) {62 unwrap(rewriter)->setInsertionPointToEnd(unwrap(block));63}64 65MlirBlock mlirRewriterBaseGetInsertionBlock(MlirRewriterBase rewriter) {66 return wrap(unwrap(rewriter)->getInsertionBlock());67}68 69MlirBlock mlirRewriterBaseGetBlock(MlirRewriterBase rewriter) {70 return wrap(unwrap(rewriter)->getBlock());71}72 73MlirOperation74mlirRewriterBaseGetOperationAfterInsertion(MlirRewriterBase rewriter) {75 mlir::RewriterBase *base = unwrap(rewriter);76 mlir::Block *block = base->getInsertionBlock();77 mlir::Block::iterator it = base->getInsertionPoint();78 if (it == block->end())79 return {nullptr};80 81 return wrap(std::addressof(*it));82}83 84//===----------------------------------------------------------------------===//85/// Block and operation creation/insertion/cloning86//===----------------------------------------------------------------------===//87 88MlirBlock mlirRewriterBaseCreateBlockBefore(MlirRewriterBase rewriter,89 MlirBlock insertBefore,90 intptr_t nArgTypes,91 MlirType const *argTypes,92 MlirLocation const *locations) {93 SmallVector<Type, 4> args;94 ArrayRef<Type> unwrappedArgs = unwrapList(nArgTypes, argTypes, args);95 SmallVector<Location, 4> locs;96 ArrayRef<Location> unwrappedLocs = unwrapList(nArgTypes, locations, locs);97 return wrap(unwrap(rewriter)->createBlock(unwrap(insertBefore), unwrappedArgs,98 unwrappedLocs));99}100 101MlirOperation mlirRewriterBaseInsert(MlirRewriterBase rewriter,102 MlirOperation op) {103 return wrap(unwrap(rewriter)->insert(unwrap(op)));104}105 106// Other methods of OpBuilder107 108MlirOperation mlirRewriterBaseClone(MlirRewriterBase rewriter,109 MlirOperation op) {110 return wrap(unwrap(rewriter)->clone(*unwrap(op)));111}112 113MlirOperation mlirRewriterBaseCloneWithoutRegions(MlirRewriterBase rewriter,114 MlirOperation op) {115 return wrap(unwrap(rewriter)->cloneWithoutRegions(*unwrap(op)));116}117 118void mlirRewriterBaseCloneRegionBefore(MlirRewriterBase rewriter,119 MlirRegion region, MlirBlock before) {120 121 unwrap(rewriter)->cloneRegionBefore(*unwrap(region), unwrap(before));122}123 124//===----------------------------------------------------------------------===//125/// RewriterBase API126//===----------------------------------------------------------------------===//127 128void mlirRewriterBaseInlineRegionBefore(MlirRewriterBase rewriter,129 MlirRegion region, MlirBlock before) {130 unwrap(rewriter)->inlineRegionBefore(*unwrap(region), unwrap(before));131}132 133void mlirRewriterBaseReplaceOpWithValues(MlirRewriterBase rewriter,134 MlirOperation op, intptr_t nValues,135 MlirValue const *values) {136 SmallVector<Value, 4> vals;137 ArrayRef<Value> unwrappedVals = unwrapList(nValues, values, vals);138 unwrap(rewriter)->replaceOp(unwrap(op), unwrappedVals);139}140 141void mlirRewriterBaseReplaceOpWithOperation(MlirRewriterBase rewriter,142 MlirOperation op,143 MlirOperation newOp) {144 unwrap(rewriter)->replaceOp(unwrap(op), unwrap(newOp));145}146 147void mlirRewriterBaseEraseOp(MlirRewriterBase rewriter, MlirOperation op) {148 unwrap(rewriter)->eraseOp(unwrap(op));149}150 151void mlirRewriterBaseEraseBlock(MlirRewriterBase rewriter, MlirBlock block) {152 unwrap(rewriter)->eraseBlock(unwrap(block));153}154 155void mlirRewriterBaseInlineBlockBefore(MlirRewriterBase rewriter,156 MlirBlock source, MlirOperation op,157 intptr_t nArgValues,158 MlirValue const *argValues) {159 SmallVector<Value, 4> vals;160 ArrayRef<Value> unwrappedVals = unwrapList(nArgValues, argValues, vals);161 162 unwrap(rewriter)->inlineBlockBefore(unwrap(source), unwrap(op),163 unwrappedVals);164}165 166void mlirRewriterBaseMergeBlocks(MlirRewriterBase rewriter, MlirBlock source,167 MlirBlock dest, intptr_t nArgValues,168 MlirValue const *argValues) {169 SmallVector<Value, 4> args;170 ArrayRef<Value> unwrappedArgs = unwrapList(nArgValues, argValues, args);171 unwrap(rewriter)->mergeBlocks(unwrap(source), unwrap(dest), unwrappedArgs);172}173 174void mlirRewriterBaseMoveOpBefore(MlirRewriterBase rewriter, MlirOperation op,175 MlirOperation existingOp) {176 unwrap(rewriter)->moveOpBefore(unwrap(op), unwrap(existingOp));177}178 179void mlirRewriterBaseMoveOpAfter(MlirRewriterBase rewriter, MlirOperation op,180 MlirOperation existingOp) {181 unwrap(rewriter)->moveOpAfter(unwrap(op), unwrap(existingOp));182}183 184void mlirRewriterBaseMoveBlockBefore(MlirRewriterBase rewriter, MlirBlock block,185 MlirBlock existingBlock) {186 unwrap(rewriter)->moveBlockBefore(unwrap(block), unwrap(existingBlock));187}188 189void mlirRewriterBaseStartOpModification(MlirRewriterBase rewriter,190 MlirOperation op) {191 unwrap(rewriter)->startOpModification(unwrap(op));192}193 194void mlirRewriterBaseFinalizeOpModification(MlirRewriterBase rewriter,195 MlirOperation op) {196 unwrap(rewriter)->finalizeOpModification(unwrap(op));197}198 199void mlirRewriterBaseCancelOpModification(MlirRewriterBase rewriter,200 MlirOperation op) {201 unwrap(rewriter)->cancelOpModification(unwrap(op));202}203 204void mlirRewriterBaseReplaceAllUsesWith(MlirRewriterBase rewriter,205 MlirValue from, MlirValue to) {206 unwrap(rewriter)->replaceAllUsesWith(unwrap(from), unwrap(to));207}208 209void mlirRewriterBaseReplaceAllValueRangeUsesWith(MlirRewriterBase rewriter,210 intptr_t nValues,211 MlirValue const *from,212 MlirValue const *to) {213 SmallVector<Value, 4> fromVals;214 ArrayRef<Value> unwrappedFromVals = unwrapList(nValues, from, fromVals);215 SmallVector<Value, 4> toVals;216 ArrayRef<Value> unwrappedToVals = unwrapList(nValues, to, toVals);217 unwrap(rewriter)->replaceAllUsesWith(unwrappedFromVals, unwrappedToVals);218}219 220void mlirRewriterBaseReplaceAllOpUsesWithValueRange(MlirRewriterBase rewriter,221 MlirOperation from,222 intptr_t nTo,223 MlirValue const *to) {224 SmallVector<Value, 4> toVals;225 ArrayRef<Value> unwrappedToVals = unwrapList(nTo, to, toVals);226 unwrap(rewriter)->replaceAllOpUsesWith(unwrap(from), unwrappedToVals);227}228 229void mlirRewriterBaseReplaceAllOpUsesWithOperation(MlirRewriterBase rewriter,230 MlirOperation from,231 MlirOperation to) {232 unwrap(rewriter)->replaceAllOpUsesWith(unwrap(from), unwrap(to));233}234 235void mlirRewriterBaseReplaceOpUsesWithinBlock(MlirRewriterBase rewriter,236 MlirOperation op,237 intptr_t nNewValues,238 MlirValue const *newValues,239 MlirBlock block) {240 SmallVector<Value, 4> vals;241 ArrayRef<Value> unwrappedVals = unwrapList(nNewValues, newValues, vals);242 unwrap(rewriter)->replaceOpUsesWithinBlock(unwrap(op), unwrappedVals,243 unwrap(block));244}245 246void mlirRewriterBaseReplaceAllUsesExcept(MlirRewriterBase rewriter,247 MlirValue from, MlirValue to,248 MlirOperation exceptedUser) {249 unwrap(rewriter)->replaceAllUsesExcept(unwrap(from), unwrap(to),250 unwrap(exceptedUser));251}252 253//===----------------------------------------------------------------------===//254/// IRRewriter API255//===----------------------------------------------------------------------===//256 257MlirRewriterBase mlirIRRewriterCreate(MlirContext context) {258 return wrap(new IRRewriter(unwrap(context)));259}260 261MlirRewriterBase mlirIRRewriterCreateFromOp(MlirOperation op) {262 return wrap(new IRRewriter(unwrap(op)));263}264 265void mlirIRRewriterDestroy(MlirRewriterBase rewriter) {266 delete static_cast<IRRewriter *>(unwrap(rewriter));267}268 269//===----------------------------------------------------------------------===//270/// RewritePatternSet and FrozenRewritePatternSet API271//===----------------------------------------------------------------------===//272 273MlirFrozenRewritePatternSet274mlirFreezeRewritePattern(MlirRewritePatternSet set) {275 auto *m = new mlir::FrozenRewritePatternSet(std::move(*unwrap(set)));276 set.ptr = nullptr;277 return wrap(m);278}279 280void mlirFrozenRewritePatternSetDestroy(MlirFrozenRewritePatternSet set) {281 delete unwrap(set);282 set.ptr = nullptr;283}284 285MlirLogicalResult286mlirApplyPatternsAndFoldGreedily(MlirModule op,287 MlirFrozenRewritePatternSet patterns,288 MlirGreedyRewriteDriverConfig) {289 return wrap(mlir::applyPatternsGreedily(unwrap(op), *unwrap(patterns)));290}291 292MlirLogicalResult293mlirApplyPatternsAndFoldGreedilyWithOp(MlirOperation op,294 MlirFrozenRewritePatternSet patterns,295 MlirGreedyRewriteDriverConfig) {296 return wrap(mlir::applyPatternsGreedily(unwrap(op), *unwrap(patterns)));297}298 299//===----------------------------------------------------------------------===//300/// PatternRewriter API301//===----------------------------------------------------------------------===//302 303MlirRewriterBase mlirPatternRewriterAsBase(MlirPatternRewriter rewriter) {304 return wrap(static_cast<mlir::RewriterBase *>(unwrap(rewriter)));305}306 307//===----------------------------------------------------------------------===//308/// RewritePattern API309//===----------------------------------------------------------------------===//310 311namespace mlir {312 313class ExternalRewritePattern : public mlir::RewritePattern {314public:315 ExternalRewritePattern(MlirRewritePatternCallbacks callbacks, void *userData,316 StringRef rootName, PatternBenefit benefit,317 MLIRContext *context,318 ArrayRef<StringRef> generatedNames)319 : RewritePattern(rootName, benefit, context, generatedNames),320 callbacks(callbacks), userData(userData) {321 if (callbacks.construct)322 callbacks.construct(userData);323 }324 325 ~ExternalRewritePattern() {326 if (callbacks.destruct)327 callbacks.destruct(userData);328 }329 330 LogicalResult matchAndRewrite(Operation *op,331 PatternRewriter &rewriter) const override {332 return unwrap(callbacks.matchAndRewrite(333 wrap(static_cast<const mlir::RewritePattern *>(this)), wrap(op),334 wrap(&rewriter), userData));335 }336 337private:338 MlirRewritePatternCallbacks callbacks;339 void *userData;340};341 342} // namespace mlir343 344MlirRewritePattern mlirOpRewritePatternCreate(345 MlirStringRef rootName, unsigned benefit, MlirContext context,346 MlirRewritePatternCallbacks callbacks, void *userData,347 size_t nGeneratedNames, MlirStringRef *generatedNames) {348 std::vector<mlir::StringRef> generatedNamesVec;349 generatedNamesVec.reserve(nGeneratedNames);350 for (size_t i = 0; i < nGeneratedNames; ++i) {351 generatedNamesVec.push_back(unwrap(generatedNames[i]));352 }353 return wrap(new mlir::ExternalRewritePattern(354 callbacks, userData, unwrap(rootName), PatternBenefit(benefit),355 unwrap(context), generatedNamesVec));356}357 358//===----------------------------------------------------------------------===//359/// RewritePatternSet API360//===----------------------------------------------------------------------===//361 362MlirRewritePatternSet mlirRewritePatternSetCreate(MlirContext context) {363 return wrap(new mlir::RewritePatternSet(unwrap(context)));364}365 366void mlirRewritePatternSetDestroy(MlirRewritePatternSet set) {367 delete unwrap(set);368}369 370void mlirRewritePatternSetAdd(MlirRewritePatternSet set,371 MlirRewritePattern pattern) {372 std::unique_ptr<mlir::RewritePattern> patternPtr(373 const_cast<mlir::RewritePattern *>(unwrap(pattern)));374 pattern.ptr = nullptr;375 unwrap(set)->add(std::move(patternPtr));376}377 378//===----------------------------------------------------------------------===//379/// PDLPatternModule API380//===----------------------------------------------------------------------===//381 382#if MLIR_ENABLE_PDL_IN_PATTERNMATCH383MlirPDLPatternModule mlirPDLPatternModuleFromModule(MlirModule op) {384 return wrap(new mlir::PDLPatternModule(385 mlir::OwningOpRef<mlir::ModuleOp>(unwrap(op))));386}387 388void mlirPDLPatternModuleDestroy(MlirPDLPatternModule op) {389 delete unwrap(op);390 op.ptr = nullptr;391}392 393MlirRewritePatternSet394mlirRewritePatternSetFromPDLPatternModule(MlirPDLPatternModule op) {395 auto *m = new mlir::RewritePatternSet(std::move(*unwrap(op)));396 op.ptr = nullptr;397 return wrap(m);398}399 400MlirValue mlirPDLValueAsValue(MlirPDLValue value) {401 return wrap(unwrap(value)->dyn_cast<mlir::Value>());402}403 404MlirType mlirPDLValueAsType(MlirPDLValue value) {405 return wrap(unwrap(value)->dyn_cast<mlir::Type>());406}407 408MlirOperation mlirPDLValueAsOperation(MlirPDLValue value) {409 return wrap(unwrap(value)->dyn_cast<mlir::Operation *>());410}411 412MlirAttribute mlirPDLValueAsAttribute(MlirPDLValue value) {413 return wrap(unwrap(value)->dyn_cast<mlir::Attribute>());414}415 416void mlirPDLResultListPushBackValue(MlirPDLResultList results,417 MlirValue value) {418 unwrap(results)->push_back(unwrap(value));419}420 421void mlirPDLResultListPushBackType(MlirPDLResultList results, MlirType value) {422 unwrap(results)->push_back(unwrap(value));423}424 425void mlirPDLResultListPushBackOperation(MlirPDLResultList results,426 MlirOperation value) {427 unwrap(results)->push_back(unwrap(value));428}429 430void mlirPDLResultListPushBackAttribute(MlirPDLResultList results,431 MlirAttribute value) {432 unwrap(results)->push_back(unwrap(value));433}434 435inline std::vector<MlirPDLValue> wrap(ArrayRef<PDLValue> values) {436 std::vector<MlirPDLValue> mlirValues;437 mlirValues.reserve(values.size());438 for (auto &value : values) {439 mlirValues.push_back(wrap(&value));440 }441 return mlirValues;442}443 444void mlirPDLPatternModuleRegisterRewriteFunction(445 MlirPDLPatternModule pdlModule, MlirStringRef name,446 MlirPDLRewriteFunction rewriteFn, void *userData) {447 unwrap(pdlModule)->registerRewriteFunction(448 unwrap(name),449 [userData, rewriteFn](PatternRewriter &rewriter, PDLResultList &results,450 ArrayRef<PDLValue> values) -> LogicalResult {451 std::vector<MlirPDLValue> mlirValues = wrap(values);452 return unwrap(rewriteFn(wrap(&rewriter), wrap(&results),453 mlirValues.size(), mlirValues.data(),454 userData));455 });456}457 458void mlirPDLPatternModuleRegisterConstraintFunction(459 MlirPDLPatternModule pdlModule, MlirStringRef name,460 MlirPDLConstraintFunction constraintFn, void *userData) {461 unwrap(pdlModule)->registerConstraintFunction(462 unwrap(name),463 [userData, constraintFn](PatternRewriter &rewriter,464 PDLResultList &results,465 ArrayRef<PDLValue> values) -> LogicalResult {466 std::vector<MlirPDLValue> mlirValues = wrap(values);467 return unwrap(constraintFn(wrap(&rewriter), wrap(&results),468 mlirValues.size(), mlirValues.data(),469 userData));470 });471}472#endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH473