[mlir][linalg] Run CSE after every CodegenStrategy transformation.

[gysit]

Summary

Add CSE after every transformation. Transformations such as tiling introduce redundant computation, for example, one AffineMinOp for every operand dimension pair. Follow up transformations such as Padding and Hoisting benefit from CSE since comparing slice sizes simplifies to comparing SSA values instead of analyzing affine expressions.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp

49 ↗

(On Diff #389741)

Can this be folded inside the createLinalgStrategyEnablePass ?
This is the "one size fits all" for all enabling stuff (hoisting etc).
This way you'd also have more control if you need to run some enabling stuff multiple times.

[gysit]

mlir/lib/Dialect/Linalg/Transforms/CodegenStrategy.cpp

49 ↗

(On Diff #389741)

I think this would require another dynamic pass manager nested inside the LinalgStrategyEnablePass? My understanding is that we should minimize the use of dynamic pass managers?

@mehdi_amini do you have a recommendation on how to address this problem? Our goal is to "call one thing" that runs the LinalgStrategyEnablePass and performs CSE. The LinalgStrategyEnablePass applies canonicalization patterns while CSE is a pass.

[gysit]

Run CSE from within the enabeling pass using a dynamic pass pipeline.

mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp

 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/VectorTransforms.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 using namespace linalg;
 
 namespace {
 
void runOnFunction() override {
     }
 
     promoteSingleIterationLoops(funcOp);
     if (options.hoistRedundantVectorTransfers)
       hoistRedundantVectorTransfers(funcOp);
 
     if (options.hoistRedundantVectorTransfersOnTensor)
       hoistRedundantVectorTransfersOnTensor(funcOp);
+
+    // Run CSE to cleanup after canonicalization.
+    OpPassManager dynamicPM("builtin.func");
+    dynamicPM.addPass(createCSEPass());
+    if (failed(runPipeline(dynamicPM, funcOp)))
+      return signalPassFailure();
   }
 
   LinalgEnablingOptions options;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to lower vector operations.
 struct LinalgStrategyLowerVectorsPass

mlir/test/Dialect/Linalg/codegen-strategy.mlir

func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
   //      CHECK-INTERCHANGE:      linalg.generic
   // CHECK-INTERCHANGE-SAME:      iterator_types = ["parallel", "reduction", "parallel"]
   %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
   return %0 : tensor<72x72xf32>
 }
 
 // -----
 
+//     CHECK-PAD-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (16, -d0 + 72)>
+
 //         CHECK-PAD: func @matmul(
 func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
 
   // Check the padding of the input operands has been hoisted out of the tile loop nest.
   //      CHECK-PAD-COUNT=2: linalg.pad_tensor %{{.*}} nofold
-  //      CHECK-PAD-COUNT=3: scf.for
+  //              CHECK-PAD: scf.for
+  // Check CSE eliminates the duplicate min operations introduced by tiling.
+  //              CHECK-PAD: affine.min #[[MAP0]]
+  //          CHECK-PAD-NOT: affine.min #[[MAP0]]
+  //      CHECK-PAD-COUNT=2: scf.for
   //              CHECK-PAD: linalg.matmul
   %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
   return %0 : tensor<72x72xf32>
 }
-