[mlir] Make test pass that runs on Module use a ModulePass

[mravishankar]

Summary

The TestGreedyFusionPass runs a pass pipeline using the module, but is
declared as a FunctionPass. Instead, just define the pass as a
ModulePass.

[mehdi_amini]

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

189

This looks likes a different behavior than before in how it'll repeat the nested pipeline multiple times per function even if only a single function gets a fusion. Is this really needed?

194

Actually, why are we even doing this? Can we just remove all these passes and keep this pass dedicated to fuseLinalgOpsGreedily alone?

197

Please use the dynamic pass pipeline: we should not run nested pipeline this way: https://mlir.llvm.org/docs/PassManagement/#dynamic-pass-pipelines

201

I don't quite get why you can't run this on the current function instead?
Keeping this a FunctionPass would be better in general.

[mravishankar]

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

189

Ah your right. Not really needed I think, but at the least need to move the while loop here to encompass the pipeline

201

I think loop invariant code motion is a ModulePass? Is it valid to take the parent module of a pass labeled as a FunctionPass? Just using the funcOp here threw an error that it was running a module pass on a function pass.

[mravishankar]

(push)

[mehdi_amini]

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

201

LICM is even more generic: it is an Operation pass. It can be scheduled on a FuncOp or on a spv.func, or a gpu.module!

If you can model after it here that'd be ideal!

[mehdi_amini]

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

194

Ping on this comment: can we remove this please?

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

static LogicalResult fuseLinalgOpsGreedily(FuncOp f) {
   for (auto *e : eraseSet)
     e->erase();
 
   return changed ? success() : failure();
 }
 
 namespace {
 struct TestLinalgGreedyFusion
-    : public PassWrapper<TestLinalgGreedyFusion, FunctionPass> {
-  void runOnFunction() override {
+    : public PassWrapper<TestLinalgGreedyFusion, OperationPass<ModuleOp>> {
+  void runOnOperation() override {
     MLIRContext *context = &getContext();
     OwningRewritePatternList patterns =
         linalg::getLinalgTilingCanonicalizationPatterns(context);
     patterns.insert<AffineMinSCFCanonicalizationPattern>(context);
     FrozenRewritePatternList frozenPatterns(std::move(patterns));
-    while (succeeded(fuseLinalgOpsGreedily(getFunction()))) {
-      (void)applyPatternsAndFoldGreedily(getFunction(), frozenPatterns);
+    bool repeat = true;
+    while (repeat) {
+      repeat = false;
+      for (FuncOp funcOp : getOperation().getOps<FuncOp>()) {
+        if (succeeded(fuseLinalgOpsGreedily(funcOp))) {
+          repeat = true;
+          (void)applyPatternsAndFoldGreedily(funcOp, frozenPatterns);
+        }
+      }
       PassManager pm(context);
       pm.addPass(createLoopInvariantCodeMotionPass());
       pm.addPass(createCanonicalizerPass());
       pm.addPass(createCSEPass());
-      LogicalResult res = pm.run(getFunction()->getParentOfType<ModuleOp>());
+      LogicalResult res = pm.run(getOperation());
       if (failed(res))
         this->signalPassFailure();
     }
   }
 };
 
 /// Pass to test tile and fuse of sequence of operations. Intended only for
 /// testing.