[MLIR] Add back trait based folding and add FIXME message for adding side effect free check to Involution

[ahmedsabie]

Summary

This is the same diff as https://reviews.llvm.org/D88809/ except side effect free check is removed for involution and a FIXME is added until the dependency is resolved for shared builds. The old diff has more details on possible fixes.

[ahmedsabie]

Add the missing revert

[ahmedsabie]

I dont have push access, if you get a chance can you push for me? Thanks!

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">;

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");
+    // Involution requires the operation to be side effect free as well
+    // but currently this check is under a FIXME and is not actually done.
+    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) {
+  // FIXME: Add back check for no side effects on operation.
+  // Currently adding it would cause the shared library build
+  // to fail since there would be a dependency of IR on SideEffectInterfaces
+  // which is cyclical.
+  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();