[MLIR] Add a foldTrait() mechanism to allow traits to define folding and test it with an Involution trait

[ahmedsabie]

Summary

This change allows folds to be done on a newly introduced involution trait rather than having to manually rewrite this optimization for every instance of an involution

[rriddle]

Can you split this commit into two? You seem to be adding core infra that isnt documented in your commit title or description, and is arguably more important than adding an involution trait.

[mehdi_amini]

This is great!

@rriddle what about renaming the revision Add a foldTrait() mechanism to allow traits to define folding and test it with an Involution trait"? I agree that the involution is completely secondary here.

[ahmedsabie]

update commit message to be more descriptive of the infrastructure change

[ahmedsabie]

I agree with Mehdi on keeping it as one PR but updating the name since making a separate PR with just the trait folding would require writing code for its test that's almost the same as what's required for involution anyways

[andyly]

mlir/test/lib/Dialect/Test/TestTraits.cpp

11–13

Should these imports be removed?

[stephenneuendorffer]

mlir/include/mlir/IR/OpBase.td

1726

I think you mean op op X == X ?

[jpienaar]

mlir/include/mlir/IR/OpBase.td

1737

These two shouldn't be needed as users don't need to write them, they are redundant with arguments and results specification and automatically inserted.

mlir/include/mlir/IR/OpDefinition.h

1438

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

mlir/test/lib/Dialect/Test/TestOps.td

801

ODS automatically inserts OneResult and OneOperand, so these need not be manually specified here. (same below)

(see OpEmitter::genTraits)

[rriddle]

Keeping them together is fine if you are using involution as a driving test case given that it is rather simple.

mlir/docs/Traits.md

69

The format of trait folding hooks should be the same as the general form for operations, unless you want to use template meta-programming to allow a trait to define either. This form here looks specific to a certain type of operation, i.e. single result operations, which does not encompass all operation types.

mlir/include/mlir/IR/OpDefinition.h

1030

Are involution operations allowed to be folded if they are side effecting?

mlir/test/lib/Dialect/Test/TestOps.td

801

Please wrap .td files at 80 characters, just like other source files.

mlir/test/lib/Dialect/Test/TestTraits.cpp

2

This should be 80 characters long, the title here is also off.

11

Please remove commented out code.

15

This file doesn't use dialect conversion or anything from FoldUtils, these shouldn't be necessary.

36

nit: Drop the newline here.

mlir/test/mlir-tblgen/trait.mlir

23

Please use a consistent format for the check lines, either all before or all after. (More preference to before than after).

[ahmedsabie]

Added no side effect requirement to involution, changed fold to two different signatures to be consistent with existing fold functionality, and fixed issues addressed by comments

[ahmedsabie]

mlir/include/mlir/IR/OpDefinition.h

1030

good catch, I added the requirement of no side effects being present

[mehdi_amini]

mlir/lib/IR/Operation.cpp

737

Nit: remove braces

mlir/test/lib/Dialect/Test/TestOps.td

806

Do we need to provide this "hasFolder"?
This seems like an unneeded limitation to me right now, can you elaborate if you're just adding it here for the fallback?
(and so could you add a test with an op without defining hasFolder = 1?)

[mehdi_amini]

The limitation on side-effect is a bit unfortunate: I suspect we'd want lmhlo involution operation on buffers. But we can leave this for future work.

[rriddle]

What do you think of adding detection for in-place folds and continuing to fold in such cases? If an operation is folded in-place, we don't necessarily have to stop the folding process prematurely.

mlir/docs/Traits.md

62

The *must* here is too strong, the second form should still be valid even if the operation isn't single result. Otherwise that limits how general a trait can be, e.g. what if we extend involution to more than just single operand/result? At that point the trait would provide a general folding implementation using the second form.

76

nit: Otherwise,

84

The OpFoldResult should be LogicalResult here.

mlir/include/mlir/IR/OpDefinition.h

450

Can you call into the general form here as a default implementation? That would allow for traits to implement the general form if they want.

1366

Is it possible for these to be made private?

mlir/test/lib/Dialect/Test/TestTraits.cpp

22

nit: Please capitalize the start of comments, and end with proper punctuation.

[ahmedsabie]

Allow traits to continue folding if an inplace fold occured

[ahmedsabie]

mlir/test/lib/Dialect/Test/TestOps.td

806

No it just emphasizes that when op folding fails the trait folding takes over. I added a separate test to clearly highlight this is not needed

[rriddle]

Thanks, LGTM after resolving the remaining comments.

mlir/include/mlir/IR/OpDefinition.h

343

Please fix the naming violations here and above.

344

Can you inline the variable definition into the if here and below?

452

nit: Drop the newline here.

455

nit: Drop else after return.

457

nit: Add a message to the assert.

1443

Please fix the naming style issues here and below.

mlir/lib/IR/Operation.cpp

685

nit: Drop the newline here.

mlir/tools/mlir-opt/mlir-opt.cpp

140

nit: Please use a registerTest* naming scheme for this method.

[ahmedsabie]

Fix last few minor comments

[ahmedsabie]

I don't have push access so if you get a chance to push my code and its good please do, thanks!

[mehdi_amini]

Pushed a revert in 16580d791f3b ; can you try to build after configuring cmake with -DBUILD_SHARED_LIBS=ON ?

[ahmedsabie]

In D88809#2320818, @mehdi_amini wrote:

Pushed a revert in 16580d791f3b ; can you try to build after configuring cmake with -DBUILD_SHARED_LIBS=ON ?

It seems there's an issue with using MemoryEffectOpInterface::getEffects in Operation.cpp where I get an undefined reference to this call. I tried adding MLIRSideEffectInterfaces under LINK_LIBS for IR/CMakeLists but I get a cyclic dependency is only allowed among static libraries error. Is there something else I can try or is there a way to check an op has no side effects with using getEffects?

[rriddle]

In D88809#2322819, @ahmedsabie wrote:

In D88809#2320818, @mehdi_amini wrote:

Pushed a revert in 16580d791f3b ; can you try to build after configuring cmake with -DBUILD_SHARED_LIBS=ON ?

It seems there's an issue with using MemoryEffectOpInterface::getEffects in Operation.cpp where I get an undefined reference to this call. I tried adding MLIRSideEffectInterfaces under LINK_LIBS for IR/CMakeLists but I get a cyclic dependency is only allowed among static libraries error. Is there something else I can try or is there a way to check an op has no side effects with using getEffects?

For now I would drop the check on non-side effecting and add a FIXME to add it back later. The actual fix, IMO, is to either:

a) Add a new Traits/ directory and start moving traits there.
b) Rename the Interfaces/ directory to something more general, and start moving traits there.

We should be funneling all of the traits into IR because it isn't scalable, and it also prevents depending on other libraries(such as in this instance). I don't have a big preference on whether we should do a) or b) in the long run.

mlir/docs/Traits.md

public:
   }
 };
 ```
 
 Note: It is generally good practice to define the implementation of the
 `verifyTrait` hook out-of-line as a free function when possible to avoid
 instantiating the implementation for every concrete operation type.
 
+Operation traits may also provide a `foldTrait` hook that is called when
+folding the concrete operation. The trait folders will only be invoked if
+the concrete operation fold is either not implemented, fails, or performs
+an in-place fold.
+
+The following signature of fold will be called if it is implemented
+and the op has a single result.
+
+```c++
+template <typename ConcreteType>
+class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
+public:
+  /// Override the 'foldTrait' hook to support trait based folding on the
+  /// concrete operation.
+  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) { {
+    // ...
+  }
+};
+```
+
+Otherwise, if the operation has a single result and the above signature is
+not implemented, or the operation has multiple results, then the following signature
+will be used (if implemented):
+
+```c++
+template <typename ConcreteType>
+class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
+public:
+  /// Override the 'foldTrait' hook to support trait based folding on the
+  /// concrete operation.
+  static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
+                                 SmallVectorImpl<OpFoldResult> &results) { {
+    // ...
+  }
+};
+```
+
+Note: It is generally good practice to define the implementation of the
+`foldTrait` hook out-of-line as a free function when possible to avoid
+instantiating the implementation for every concrete operation type.
+
 ### Parametric Traits
 
 The above demonstrates the definition of a simple self-contained trait. It is
 also often useful to provide some static parameters to the trait to control its
 behavior. Given that the definition of the trait class is rigid, i.e. we must
 have a single template argument for the concrete object, the templates for the
 parameters will need to be split out. An example is shown below:
 

mlir/include/mlir/IR/OpBase.td

 def AffineScope : NativeOpTrait<"AffineScope">;
 // Op defines an automatic allocation scope.
 def AutomaticAllocationScope : NativeOpTrait<"AutomaticAllocationScope">;
 // Op supports operand broadcast behavior.
 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
+// op op X == X
+def Involution  : NativeOpTrait<"IsInvolution">;
 // Op behaves like a constant.
 def ConstantLike : NativeOpTrait<"ConstantLike">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
 // Op has the same operand type.
 def SameTypeOperands : NativeOpTrait<"SameTypeOperands">;
 // Op has same shape for all operands.
 def SameOperandsShape : NativeOpTrait<"SameOperandsShape">;
 // Op has same operand and result shape.
 def SameOperandsAndResultShape : NativeOpTrait<"SameOperandsAndResultShape">;
 // Op has the same operand and result type.
 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 // Op has the same element type (or type itself, if scalar) for all operands.
 def SameOperandsElementType : NativeOpTrait<"SameOperandsElementType">;

mlir/include/mlir/IR/OpDefinition.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_OPDEFINITION_H
 #define MLIR_IR_OPDEFINITION_H
 
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
+
 #include <type_traits>
 
 namespace mlir {
 class Builder;
 class OpBuilder;
 
 namespace OpTrait {
 template <typename ConcreteType> class OneResult;

 /// custom ops which can return multiple results.
 template <typename ConcreteType, bool isSingleResult, typename = void>
 class FoldingHook {
 public:
   /// This is an implementation detail of the constant folder hook for
   /// AbstractOperation.
   static LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
                                 SmallVectorImpl<OpFoldResult> &results) {
-    return cast<ConcreteType>(op).fold(operands, results);
+    auto operationFoldResult = cast<ConcreteType>(op).fold(operands, results);
+    // Failure to fold or in place fold both mean we can continue folding.
+    if (failed(operationFoldResult) || results.empty()) {
+      auto traitFoldResult = ConcreteType::foldTraits(op, operands, results);
+      // Only return the trait fold result if it is a success since
+      // operationFoldResult might have been a success originally.
+      if (succeeded(traitFoldResult))
+        return traitFoldResult;
+    }
+    return operationFoldResult;
   }
 
   /// This hook implements a generalized folder for this operation.  Operations
   /// can implement this to provide simplifications rules that are applied by
   /// the Builder::createOrFold API and the canonicalization pass.
   ///
   /// This is an intentionally limited interface - implementations of this hook
   /// can only perform the following changes to the operation:
operator Value() {
     return static_cast<ConcreteType *>(this)->getOperation()->getResult(0);
   }
 
   /// This is an implementation detail of the constant folder hook for
   /// AbstractOperation.
   static LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
                                 SmallVectorImpl<OpFoldResult> &results) {
     auto result = cast<ConcreteType>(op).fold(operands);
+    // Failure to fold or in place fold both mean we can continue folding.
+    if (!result || result.template dyn_cast<Value>() == op->getResult(0)) {
+      // Only consider the trait fold result if it is a success since
+      // the operation fold might have been a success originally.
+      if (auto traitFoldResult = ConcreteType::foldTraits(op, operands))
+        result = traitFoldResult;
+    }
+
     if (!result)
       return failure();
 
     // Check if the operation was folded in place. In this case, the operation
     // returns itself.
     if (result.template dyn_cast<Value>() != op->getResult(0))
       results.push_back(result);
     return success();

 //===----------------------------------------------------------------------===//
 
 namespace OpTrait {
 
 // These functions are out-of-line implementations of the methods in the
 // corresponding trait classes.  This avoids them being template
 // instantiated/duplicated.
 namespace impl {
+OpFoldResult foldInvolution(Operation *op);
 LogicalResult verifyZeroOperands(Operation *op);
 LogicalResult verifyOneOperand(Operation *op);
 LogicalResult verifyNOperands(Operation *op, unsigned numOperands);
+LogicalResult verifyIsInvolution(Operation *op);
 LogicalResult verifyAtLeastNOperands(Operation *op, unsigned numOperands);
 LogicalResult verifyOperandsAreFloatLike(Operation *op);
 LogicalResult verifyOperandsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifySameTypeOperands(Operation *op);
 LogicalResult verifyZeroRegion(Operation *op);
 LogicalResult verifyOneRegion(Operation *op);
 LogicalResult verifyNRegions(Operation *op, unsigned numRegions);
 LogicalResult verifyAtLeastNRegions(Operation *op, unsigned numRegions);
protected:
   }
 
   /// Provide default implementations of trait hooks.  This allows traits to
   /// provide exactly the overrides they care about.
   static LogicalResult verifyTrait(Operation *op) { return success(); }
   static AbstractOperation::OperationProperties getTraitProperties() {
     return 0;
   }
+
+  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+    SmallVector<OpFoldResult, 1> results;
+    if (failed(foldTrait(op, operands, results)))
+      return {};
+    if (results.empty())
+      return op->getResult(0);
+    assert(results.size() == 1 &&
+           "Single result op cannot return multiple fold results");
+
+    return results[0];
+  }
+
+  static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
+                                 SmallVectorImpl<OpFoldResult> &results) {
+    return failure();
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // Operand Traits
 
 namespace detail {
 /// Utility trait base that provides accessors for derived traits that have
 /// multiple operands.

 class IsCommutative : public TraitBase<ConcreteType, IsCommutative> {
 public:
   static AbstractOperation::OperationProperties getTraitProperties() {
     return static_cast<AbstractOperation::OperationProperties>(
         OperationProperty::Commutative);
   }
 };
 
+/// This class adds property that the operation is an involution.
+/// This means a unary to unary operation "f" that satisfies f(f(x)) = f(x)
+template <typename ConcreteType>
+class IsInvolution : public TraitBase<ConcreteType, IsInvolution> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<OneResult>(),
+                  "expected operation to produce one result");
+    static_assert(ConcreteType::template hasTrait<OneOperand>(),
+                  "expected operation to take one operand");
+    static_assert(ConcreteType::template hasTrait<SameOperandsAndResultType>(),
+                  "expected operation to preserve type");
+    return impl::verifyIsInvolution(op);
+  }
+
+  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+    return impl::foldInvolution(op);
+  }
+};
+
 /// This class verifies that all operands of the specified op have a float type,
 /// a vector thereof, or a tensor thereof.
 template <typename ConcreteType>
 class OperandsAreFloatLike
     : public TraitBase<ConcreteType, OperandsAreFloatLike> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyOperandsAreFloatLike(op);
#endif
   /// On success this returns false; on failure it emits an error to the
   /// diagnostic subsystem and returns true.
   static LogicalResult verifyInvariants(Operation *op) {
     return failure(
         failed(BaseVerifier<Traits<ConcreteType>...>::verifyTrait(op)) ||
         failed(cast<ConcreteType>(op).verify()));
   }
 
+  /// This is the hook that tries to fold the given operation according to its
+  /// traits. It delegates to the Traits for their policy implementations, and
+  /// allows the user to specify their own fold() method.
+  static OpFoldResult foldTraits(Operation *op, ArrayRef<Attribute> operands) {
+    return BaseFolder<Traits<ConcreteType>...>::foldTraits(op, operands);
+  }
+
+  static LogicalResult foldTraits(Operation *op, ArrayRef<Attribute> operands,
+                                  SmallVectorImpl<OpFoldResult> &results) {
+    return BaseFolder<Traits<ConcreteType>...>::foldTraits(op, operands,
+                                                           results);
+  }
+
   // Returns the properties of an operation by combining the properties of the
   // traits of the op.
   static AbstractOperation::OperationProperties getOperationProperties() {
     return BaseProperties<Traits<ConcreteType>...>::getTraitProperties();
   }
 
   /// Expose the type we are instantiated on to template machinery that may want
   /// to introspect traits on this operation.
private:
   template <typename First, typename... Rest>
   struct BaseProperties<First, Rest...> {
     static AbstractOperation::OperationProperties getTraitProperties() {
       return First::getTraitProperties() |
              BaseProperties<Rest...>::getTraitProperties();
     }
   };
 
+  template <typename... Types>
+  struct BaseFolder;
+
+  template <typename First, typename... Rest>
+  struct BaseFolder<First, Rest...> {
+    static OpFoldResult foldTraits(Operation *op,
+                                   ArrayRef<Attribute> operands) {
+      auto result = First::foldTrait(op, operands);
+      // Failure to fold or in place fold both mean we can continue folding.
+      if (!result || result.template dyn_cast<Value>() == op->getResult(0)) {
+        // Only consider the trait fold result if it is a success since
+        // the operation fold might have been a success originally.
+        auto resultRemaining = BaseFolder<Rest...>::foldTraits(op, operands);
+        if (resultRemaining)
+          result = resultRemaining;
+      }
+
+      return result;
+    }
+
+    static LogicalResult foldTraits(Operation *op, ArrayRef<Attribute> operands,
+                                    SmallVectorImpl<OpFoldResult> &results) {
+      auto result = First::foldTrait(op, operands, results);
+      // Failure to fold or in place fold both mean we can continue folding.
+      if (failed(result) || results.empty()) {
+        auto resultRemaining =
+            BaseFolder<Rest...>::foldTraits(op, operands, results);
+        if (succeeded(resultRemaining))
+          result = resultRemaining;
+      }
+
+      return result;
+    }
+  };
+
+  template <typename...>
+  struct BaseFolder {
+    static OpFoldResult foldTraits(Operation *op,
+                                   ArrayRef<Attribute> operands) {
+      return {};
+    }
+    static LogicalResult foldTraits(Operation *op, ArrayRef<Attribute> operands,
+                                    SmallVectorImpl<OpFoldResult> &results) {
+      return failure();
+    }
+  };
+
   template <typename...> struct BaseProperties {
     static AbstractOperation::OperationProperties getTraitProperties() {
       return 0;
     }
   };
 
   /// Returns true if this operation contains the trait for the given typeID.
   static bool hasTrait(TypeID traitID) {

mlir/lib/IR/Operation.cpp

 //===- Operation.cpp - Operation support code -----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include <numeric>
 
 using namespace mlir;
 
 OpAsmParser::~OpAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // OperationName

 InFlightDiagnostic OpState::emitRemark(const Twine &message) {
   return getOperation()->emitRemark(message);
 }
 
 //===----------------------------------------------------------------------===//
 // Op Trait implementations
 //===----------------------------------------------------------------------===//
 
+OpFoldResult OpTrait::impl::foldInvolution(Operation *op) {
+  auto *argumentOp = op->getOperand(0).getDefiningOp();
+  if (argumentOp && op->getName() == argumentOp->getName()) {
+    // Replace the outer involutions output with inner's input.
+    return argumentOp->getOperand(0);
+  }
+
+  return {};
+}
+
 LogicalResult OpTrait::impl::verifyZeroOperands(Operation *op) {
   if (op->getNumOperands() != 0)
     return op->emitOpError() << "requires zero operands";
   return success();
 }
 
 LogicalResult OpTrait::impl::verifyOneOperand(Operation *op) {
   if (op->getNumOperands() != 1)
if (auto vec = type.dyn_cast<VectorType>())
     return vec.getElementType();
 
   // Look through tensor<vector<...>> to find the underlying element type.
   if (auto tensor = type.dyn_cast<TensorType>())
     return getTensorOrVectorElementType(tensor.getElementType());
   return type;
 }
 
+LogicalResult OpTrait::impl::verifyIsInvolution(Operation *op) {
+  if (!MemoryEffectOpInterface::hasNoEffect(op))
+    return op->emitOpError() << "requires operation to have no side effects";
+  return success();
+}
+
 LogicalResult
 OpTrait::impl::verifyOperandsAreSignlessIntegerLike(Operation *op) {
   for (auto opType : op->getOperandTypes()) {
     auto type = getTensorOrVectorElementType(opType);
     if (!type.isSignlessIntOrIndex())
       return op->emitOpError() << "requires an integer or index type";
   }
   return success();

mlir/test/lib/Dialect/Test/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   TestDialect.cpp
   TestPatterns.cpp
+  TestTraits.cpp
 )
 
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
 mlir_tablegen(TestTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TestTypeInterfaces.cpp.inc -gen-type-interface-defs)
 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestOps.td)
 mlir_tablegen(TestOps.h.inc -gen-op-decls)
 mlir_tablegen(TestOps.cpp.inc -gen-op-defs)
 mlir_tablegen(TestOpsDialect.h.inc -gen-dialect-decls)
 mlir_tablegen(TestOpEnums.h.inc -gen-enum-decls)
 mlir_tablegen(TestOpEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(TestOpStructs.h.inc -gen-struct-attr-decls)
 mlir_tablegen(TestOpStructs.cpp.inc -gen-struct-attr-defs)
 mlir_tablegen(TestPatterns.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestOpsIncGen)
 
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestDialect
   TestDialect.cpp
   TestPatterns.cpp
+  TestTraits.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   DEPENDS
   MLIRTestInterfaceIncGen
   MLIRTestOpsIncGen
 
   LINK_LIBS PUBLIC

mlir/test/lib/Dialect/Test/TestOps.td

def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
   let hasFolder = 1;
 }
 
 def TestCommutativeOp : TEST_Op<"op_commutative", [Commutative]> {
   let arguments = (ins I32:$op1, I32:$op2, I32:$op3, I32:$op4);
   let results = (outs I32);
 }
 
+def TestInvolutionTraitNoOperationFolderOp
+ : TEST_Op<"op_involution_trait_no_operation_fold",
+           [SameOperandsAndResultType, NoSideEffect, Involution]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+}
+
+def TestInvolutionTraitFailingOperationFolderOp
+ : TEST_Op<"op_involution_trait_failing_operation_fold",
+           [SameOperandsAndResultType, NoSideEffect, Involution]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+  let hasFolder = 1;
+}
+
+def TestInvolutionTraitSuccesfulOperationFolderOp
+ : TEST_Op<"op_involution_trait_succesful_operation_fold",
+           [SameOperandsAndResultType, NoSideEffect, Involution]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+  let hasFolder = 1;
+}
+
 def TestOpInPlaceFoldAnchor : TEST_Op<"op_in_place_fold_anchor"> {
   let arguments = (ins I32);
   let results = (outs I32);
 }
 
 def TestOpInPlaceFold : TEST_Op<"op_in_place_fold"> {
   let arguments = (ins I32:$op, I32Attr:$attr);
   let results = (outs I32);

mlir/test/lib/Dialect/Test/TestTraits.cpp

+//===- TestTraits.cpp - Test trait folding --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/FoldUtils.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Trait Folder.
+//===----------------------------------------------------------------------===//
+
+OpFoldResult TestInvolutionTraitFailingOperationFolderOp::fold(
+    ArrayRef<Attribute> operands) {
+  // This failure should cause the trait fold to run instead.
+  return {};
+}
+
+OpFoldResult TestInvolutionTraitSuccesfulOperationFolderOp::fold(
+    ArrayRef<Attribute> operands) {
+  auto argument_op = getOperand();
+  // The success case should cause the trait fold to be supressed.
+  return argument_op.getDefiningOp() ? argument_op : OpFoldResult{};
+}
+
+namespace {
+struct TestTraitFolder : public PassWrapper<TestTraitFolder, FunctionPass> {
+  void runOnFunction() override {
+    applyPatternsAndFoldGreedily(getFunction(), {});
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+void registerTestTraitsPass() {
+  PassRegistration<TestTraitFolder>("test-trait-folder", "Run trait folding");
+}
+} // namespace mlir

mlir/test/mlir-tblgen/trait.mlir

+// RUN: mlir-opt -test-trait-folder %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Test that involutions fold correctly
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @testSingleInvolution
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testSingleInvolution(%arg0 : i32) -> i32 {
+  // CHECK: [[INVOLUTION:%.+]] = "test.op_involution_trait_no_operation_fold"([[ARG0]])
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  // CHECK: return [[INVOLUTION]]
+  return %0: i32
+}
+
+// CHECK-LABEL: func @testDoubleInvolution
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testDoubleInvolution(%arg0: i32) -> i32 {
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  %1 = "test.op_involution_trait_no_operation_fold"(%0) : (i32) -> i32
+  // CHECK: return [[ARG0]]
+  return %1: i32
+}
+
+// CHECK-LABEL: func @testTripleInvolution
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testTripleInvolution(%arg0: i32) -> i32 {
+  // CHECK: [[INVOLUTION:%.+]] = "test.op_involution_trait_no_operation_fold"([[ARG0]])
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  %1 = "test.op_involution_trait_no_operation_fold"(%0) : (i32) -> i32
+  %2 = "test.op_involution_trait_no_operation_fold"(%1) : (i32) -> i32
+  // CHECK: return [[INVOLUTION]]
+  return %2: i32
+}
+
+//===----------------------------------------------------------------------===//
+// Test that involutions fold occurs if operation fold fails
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @testFailingOperationFolder
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testFailingOperationFolder(%arg0: i32) -> i32 {
+  %0 = "test.op_involution_trait_failing_operation_fold"(%arg0) : (i32) -> i32
+  %1 = "test.op_involution_trait_failing_operation_fold"(%0) : (i32) -> i32
+  // CHECK: return [[ARG0]]
+  return %1: i32
+}
+
+//===----------------------------------------------------------------------===//
+// Test that involution fold does not occur if operation fold succeeds
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @testInhibitInvolution
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testInhibitInvolution(%arg0: i32) -> i32 {
+  // CHECK: [[OP:%.+]] = "test.op_involution_trait_succesful_operation_fold"([[ARG0]])
+  %0 = "test.op_involution_trait_succesful_operation_fold"(%arg0) : (i32) -> i32
+  %1 = "test.op_involution_trait_succesful_operation_fold"(%0) : (i32) -> i32
+  // CHECK: return [[OP]]
+  return %1: i32
+}

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestOpaqueLoc();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestPrintDefUsePass();
 void registerTestPrintNestingPass();
 void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSCFUtilsPass();
+void registerTestTraitsPass();
 void registerTestVectorConversions();
 void registerVectorizerTestPass();
 } // namespace mlir
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerConvertCallOpPass();
   registerConvertToTargetEnvPass();
#endif
   registerTestPreparationPassWithAllowedMemrefResults();
   registerTestPrintDefUsePass();
   registerTestPrintNestingPass();
   registerTestRecursiveTypesPass();
   registerTestReducer();
   registerTestGpuParallelLoopMappingPass();
   registerTestSpirvEntryPointABIPass();
   registerTestSCFUtilsPass();
+  registerTestTraitsPass();
   registerTestVectorConversions();
   registerVectorizerTestPass();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllDialects();
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();