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 
20 OpFoldResult TestInvolutionTraitFailingOperationFolderOp::fold(
21     ArrayRef<Attribute> operands) {
22   // This failure should cause the trait fold to run instead.
23   return {};
24 }
25 
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 : public PassWrapper<TestTraitFolder, FunctionPass> {
35   StringRef getArgument() const final { return "test-trait-folder"; }
36   StringRef getDescription() const final { return "Run trait folding"; }
37   void runOnFunction() override {
38     (void)applyPatternsAndFoldGreedily(getFunction(),
39                                        RewritePatternSet(&getContext()));
40   }
41 };
42 } // end anonymous namespace
43 
44 namespace mlir {
45 void registerTestTraitsPass() { PassRegistration<TestTraitFolder>(); }
46 } // namespace mlir
47