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/IR/PatternMatch.h"
11 #include "mlir/Pass/Pass.h"
12 #include "mlir/Transforms/FoldUtils.h"
13 
14 using namespace mlir;
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 argument_op = getOperand();
29   // The success case should cause the trait fold to be supressed.
30   return argument_op.getDefiningOp() ? argument_op : OpFoldResult{};
31 }
32 
33 namespace {
34 struct TestTraitFolder : public PassWrapper<TestTraitFolder, FunctionPass> {
35   void runOnFunction() override {
36     applyPatternsAndFoldGreedily(getFunction(), {});
37   }
38 };
39 } // end anonymous namespace
40 
41 namespace mlir {
42 void registerTestTraitsPass() {
43   PassRegistration<TestTraitFolder>("test-trait-folder", "Run trait folding");
44 }
45 } // namespace mlir
46