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 mlir::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   void runOnFunction() override {
36     (void)applyPatternsAndFoldGreedily(getFunction(),
37                                        RewritePatternSet(&getContext()));
38   }
39 };
40 } // end anonymous namespace
41 
42 namespace mlir {
43 void registerTestTraitsPass() {
44   PassRegistration<TestTraitFolder>("test-trait-folder", "Run trait folding");
45 }
46 } // namespace mlir
47