Transpose Reshape Elimination pass.

This is a new pass that is designed to replace the existing Transpose
and Reshape Elemination passes.  This pass adds a lot of new rewrite
rules which enable pushing the transposes and reshapes around so that
they can be combined and then eliminated.

Signed-off-by: Matthew Francis-Landau <[email protected]>

fix transpose lit test

Signed-off-by: Matthew Francis-Landau <[email protected]>

cp

cp

cp

Signed-off-by: Matthew Francis-Landau <[email protected]>

cp

cp

some improvements for matching matmul with reshapes, and update lit tests

Signed-off-by: Matthew Francis-Landau <[email protected]>

Author: Matthew Francis-Landau

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/TensorRT/IR/TensorRTOps.td

@@ -3790,6 +3790,7 @@ def TensorRT_ReshapeOp : TensorRT_Op<"reshape",
   let extraClassDeclaration = [{
     /// Returns true if created op is valid for TensorRT major version.
     bool isValidForTensorRTVersion(int64_t trtMajorVersion);
+    static void getCanonicalizationPatternsSameOp(RewritePatternSet &results, MLIRContext *context);
   }] # baseClassDeclaration;
 }
 

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/TensorRT/Transforms/Passes.td

@@ -256,4 +256,18 @@ def ReshapeEliminationPass : Pass<"tensorrt-reshape-elimination"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// TransposeReshapeEliminationPass
+//===----------------------------------------------------------------------===//
+def TransposeReshapeEliminationPass : Pass<"tensorrt-transpose-reshape-elimination"> {
+  let summary = "try to eliminate tensorrt.transpose, tensorrt.reshape, and tensorrt.shuffle operations";
+
+  let description = [{
+    Push tensorrt.transpose and tensorrt.reshape operations around to attempt to eleminate them
+    and merge them with other ops such as matrix multiply.  The intention is to improve
+    pattern matching and fusion inside of TensorRT.
+  }];
+}
+
+
 #endif // MLIR_TENSORRT_DIALECT_TENSORRT_TRANSFORMS_PASSES

mlir-tensorrt/tensorrt/lib/TensorRT/IR/TensorRT.cpp

@@ -1850,6 +1850,13 @@ void tensorrt::ReshapeOp::getCanonicalizationPatterns(
                  SimplifyReshapeToRankExpandCollapse>(context);
 }
 
+void tensorrt::ReshapeOp::getCanonicalizationPatternsSameOp(
+      RewritePatternSet &results, MLIRContext *context) {
+  results.insert<ConstFoldReshapePattern<ReshapeOp>, SimplifyReshapeReshape
+                 // NOT INCLUDED: SimplifyReshapeToRankExpandCollapse
+                 >(context);
+}
+
 void tensorrt::ReshapeOp::build(OpBuilder &builder, OperationState &state,
                                 Type result, Value input) {
   ReshapeOp::build(builder, state, result, input, Value());

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/CMakeLists.txt

@@ -21,6 +21,7 @@ add_mtrtd_library(MLIRTensorRTTransforms
   RaiseNormalizations.cpp
   ReshapeElimination.cpp
   TransposeElimination.cpp
+  TransposeReshapeElimination.cpp
 
   DEPENDS
   MLIRTensorRTTransformsActivationsPdllGen

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/Passes.cpp

@@ -86,9 +86,7 @@ void tensorrt::buildTensorRTModuleSimplificationPipeline(OpPassManager &pm) {
   // Try to eliminate as many `tensorrt.broadcast` ops as possible.
   pm.addPass(tensorrt::createBroadcastEliminationPass());
   addCleanupPasses(pm);
-  pm.addPass(tensorrt::createTransposeEliminationPass());
-  addCleanupPasses(pm);
-  pm.addPass(tensorrt::createReshapeEliminationPass());
+  pm.addPass(tensorrt::createTransposeReshapeEliminationPass());
   addCleanupPasses(pm);
   pm.addPass(tensorrt::createRaiseNormalizationsPass());
   addCleanupPasses(pm);