70 lines · cpp
1//===- BubbleDownMemorySpaceCasts.cpp - Bubble down casts transform -------===//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/Transforms/BubbleDownMemorySpaceCasts.h"10#include "mlir/IR/PatternMatch.h"11#include "mlir/Interfaces/MemOpInterfaces.h"12#include "mlir/Pass/Pass.h"13#include "mlir/Transforms/GreedyPatternRewriteDriver.h"14#include "mlir/Transforms/Passes.h"15#include "llvm/Support/Debug.h"16 17using namespace mlir;18 19namespace mlir {20#define GEN_PASS_DEF_BUBBLEDOWNMEMORYSPACECASTS21#include "mlir/Transforms/Passes.h.inc"22} // namespace mlir23 24namespace {25//===----------------------------------------------------------------------===//26// BubbleDownCastsPattern pattern27//===----------------------------------------------------------------------===//28/// Pattern to bubble down casts into consumer operations.29struct BubbleDownCastsPattern30 : public OpInterfaceRewritePattern<MemorySpaceCastConsumerOpInterface> {31 using OpInterfaceRewritePattern::OpInterfaceRewritePattern;32 33 LogicalResult matchAndRewrite(MemorySpaceCastConsumerOpInterface op,34 PatternRewriter &rewriter) const override {35 FailureOr<std::optional<SmallVector<Value>>> results =36 op.bubbleDownCasts(rewriter);37 if (failed(results))38 return failure();39 if (!results->has_value()) {40 rewriter.modifyOpInPlace(op, []() {});41 return success();42 }43 rewriter.replaceOp(op, **results);44 return success();45 }46};47 48//===----------------------------------------------------------------------===//49// BubbleDownMemorySpaceCasts pass50//===----------------------------------------------------------------------===//51 52struct BubbleDownMemorySpaceCasts53 : public impl::BubbleDownMemorySpaceCastsBase<BubbleDownMemorySpaceCasts> {54 using impl::BubbleDownMemorySpaceCastsBase<55 BubbleDownMemorySpaceCasts>::BubbleDownMemorySpaceCastsBase;56 57 void runOnOperation() override {58 RewritePatternSet patterns(&getContext());59 populateBubbleDownMemorySpaceCastPatterns(patterns, PatternBenefit(1));60 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))61 signalPassFailure();62 }63};64} // namespace65 66void mlir::populateBubbleDownMemorySpaceCastPatterns(67 RewritePatternSet &patterns, PatternBenefit benefit) {68 patterns.add<BubbleDownCastsPattern>(patterns.getContext(), benefit);69}70