1 //===- TestTraits.cpp - Test trait folding --------------------------------===//
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "TestDialect.h"
10 #include "mlir/Pass/Pass.h"
11 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
12 
13 using namespace mlir;
14 using namespace test;
15 
16 //===----------------------------------------------------------------------===//
17 // Trait Folder.
18 //===----------------------------------------------------------------------===//
19 
fold(ArrayRef<Attribute> operands)20 OpFoldResult TestInvolutionTraitFailingOperationFolderOp::fold(
21     ArrayRef<Attribute> operands) {
22   // This failure should cause the trait fold to run instead.
23   return {};
24 }
25 
fold(ArrayRef<Attribute> operands)26 OpFoldResult TestInvolutionTraitSuccesfulOperationFolderOp::fold(
27     ArrayRef<Attribute> operands) {
28   auto argumentOp = getOperand();
29   // The success case should cause the trait fold to be supressed.
30   return argumentOp.getDefiningOp() ? argumentOp : OpFoldResult{};
31 }
32 
33 namespace {
34 struct TestTraitFolder
35     : public PassWrapper<TestTraitFolder, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon864ddab30111::TestTraitFolder36   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTraitFolder)
37 
38   StringRef getArgument() const final { return "test-trait-folder"; }
getDescription__anon864ddab30111::TestTraitFolder39   StringRef getDescription() const final { return "Run trait folding"; }
runOnOperation__anon864ddab30111::TestTraitFolder40   void runOnOperation() override {
41     (void)applyPatternsAndFoldGreedily(getOperation(),
42                                        RewritePatternSet(&getContext()));
43   }
44 };
45 } // namespace
46 
47 namespace mlir {
registerTestTraitsPass()48 void registerTestTraitsPass() { PassRegistration<TestTraitFolder>(); }
49 } // namespace mlir
50