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