Files
clang-p2996/mlir/unittests/Rewrite/PatternBenefit.cpp
River Riddle 76f3c2f3f3 [mlir][Pattern] Add better support for using interfaces/traits to match root operations in rewrite patterns
To match an interface or trait, users currently have to use the `MatchAny` tag. This tag can be quite problematic for compile time for things like the canonicalizer, as the `MatchAny` patterns may get applied to  *every* operation. This revision adds better support by bucketing interface/trait patterns based on which registered operations have them registered. This means that moving forward we will only attempt to match these patterns to operations that have this interface registered. Two simplify defining patterns that match traits and interfaces, two new utility classes have been added: OpTraitRewritePattern and OpInterfaceRewritePattern.

Differential Revision: https://reviews.llvm.org/D98986
2021-03-23 14:05:33 -07:00

80 lines
2.2 KiB
C++

//===- PatternBenefit.cpp - RewritePattern benefit unit tests -------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "mlir/IR/PatternMatch.h"
#include "mlir/Rewrite/PatternApplicator.h"
#include "gtest/gtest.h"
using namespace mlir;
namespace {
TEST(PatternBenefitTest, BenefitOrder) {
// There was a bug which caused low-benefit op-specific patterns to never be
// called in presence of high-benefit op-agnostic pattern
MLIRContext context;
OpBuilder builder(&context);
auto module = ModuleOp::create(builder.getUnknownLoc());
struct Pattern1 : public OpRewritePattern<ModuleOp> {
Pattern1(mlir::MLIRContext *context, bool *called)
: OpRewritePattern<ModuleOp>(context, /*benefit*/ 1), called(called) {}
mlir::LogicalResult
matchAndRewrite(ModuleOp /*op*/,
mlir::PatternRewriter & /*rewriter*/) const override {
*called = true;
return failure();
}
private:
bool *called;
};
struct Pattern2 : public RewritePattern {
Pattern2(MLIRContext *context, bool *called)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/2, context),
called(called) {}
mlir::LogicalResult
matchAndRewrite(Operation * /*op*/,
mlir::PatternRewriter & /*rewriter*/) const override {
*called = true;
return failure();
}
private:
bool *called;
};
RewritePatternSet patterns(&context);
bool called1 = false;
bool called2 = false;
patterns.add<Pattern1>(&context, &called1);
patterns.add<Pattern2>(&context, &called2);
FrozenRewritePatternSet frozenPatterns(std::move(patterns));
PatternApplicator pa(frozenPatterns);
pa.applyDefaultCostModel();
class MyPatternRewriter : public PatternRewriter {
public:
MyPatternRewriter(MLIRContext *ctx) : PatternRewriter(ctx) {}
};
MyPatternRewriter rewriter(&context);
(void)pa.matchAndRewrite(module, rewriter);
EXPECT_TRUE(called1);
EXPECT_TRUE(called2);
}
} // namespace