[mlir] Generalize SCF passes to not have to run on FuncOp.

[stellaraccident]

Summary

Seems to have been an accident of history and none of these had any reason to be restricted to FuncOp.

[jpienaar]

Thanks, yes, seems like oversight/nothing func.func specific here.

[Mogball]

Hah! Beat you

[stellaraccident]

Rebase

mlir/include/mlir/Dialect/SCF/Transforms/Passes.td

 //===-- Passes.td - SCF pass definition file ---------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_PASSES
 #define MLIR_DIALECT_SCF_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def SCFBufferize : Pass<"scf-bufferize", "func::FuncOp"> {
+def SCFBufferize : Pass<"scf-bufferize"> {
   let summary = "Bufferize the scf dialect.";
   let constructor = "mlir::createSCFBufferizePass()";
   let dependentDialects = ["bufferization::BufferizationDialect",
                            "memref::MemRefDialect"];
 }
 
 // Note: Making these canonicalization patterns would require a dependency
 // of the SCF dialect on the Affine/Tensor/MemRef dialects or vice versa.
 def SCFForLoopCanonicalization
-    : Pass<"scf-for-loop-canonicalization", "func::FuncOp"> {
+    : Pass<"scf-for-loop-canonicalization"> {
   let summary = "Canonicalize operations within scf.for loop bodies";
   let constructor = "mlir::createSCFForLoopCanonicalizationPass()";
   let dependentDialects = ["AffineDialect", "tensor::TensorDialect",
                            "memref::MemRefDialect"];
 }
 
-def SCFForLoopPeeling : Pass<"scf-for-loop-peeling", "func::FuncOp"> {
+def SCFForLoopPeeling : Pass<"scf-for-loop-peeling"> {
   let summary = "Peel `for` loops at their upper bounds.";
   let constructor = "mlir::createForLoopPeelingPass()";
   let options = [
     Option<"skipPartial", "skip-partial", "bool",
            /*default=*/"true",
            "Do not peel loops inside of the last, partial iteration of another "
            "already peeled loop.">
   ];
   let dependentDialects = ["AffineDialect"];
 }
 
-def SCFForLoopSpecialization : Pass<"scf-for-loop-specialization", "func::FuncOp"> {
+def SCFForLoopSpecialization : Pass<"scf-for-loop-specialization"> {
   let summary = "Specialize `for` loops for vectorization";
   let constructor = "mlir::createForLoopSpecializationPass()";
 }
 
 def SCFParallelLoopFusion : Pass<"scf-parallel-loop-fusion"> {
   let summary = "Fuse adjacent parallel loops";
   let constructor = "mlir::createParallelLoopFusionPass()";
 }
 
 def SCFParallelLoopCollapsing : Pass<"scf-parallel-loop-collapsing"> {
   let summary = "Collapse parallel loops to use less induction variables";
   let constructor = "mlir::createParallelLoopCollapsingPass()";
   let options = [
     ListOption<"clCollapsedIndices0", "collapsed-indices-0", "unsigned",
                "Which loop indices to combine 0th loop index">,
     ListOption<"clCollapsedIndices1", "collapsed-indices-1", "unsigned",
                "Which loop indices to combine into the position 1 loop index">,
     ListOption<"clCollapsedIndices2", "collapsed-indices-2", "unsigned",
                "Which loop indices to combine into the position 2 loop index">,
   ];
 }
 
 def SCFParallelLoopSpecialization
-    : Pass<"scf-parallel-loop-specialization", "func::FuncOp"> {
+    : Pass<"scf-parallel-loop-specialization"> {
   let summary = "Specialize parallel loops for vectorization";
   let constructor = "mlir::createParallelLoopSpecializationPass()";
 }
 
-def SCFParallelLoopTiling : Pass<"scf-parallel-loop-tiling", "func::FuncOp"> {
+def SCFParallelLoopTiling : Pass<"scf-parallel-loop-tiling"> {
   let summary = "Tile parallel loops";
   let constructor = "mlir::createParallelLoopTilingPass()";
   let options = [
     ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t",
                "Factors to tile parallel loops by", "llvm::cl::ZeroOrMore">,
     Option<"noMinMaxBounds", "no-min-max-bounds", "bool",
            /*default=*/"false",
            "Perform tiling with fixed upper bound with inbound check "
            "inside the internal loops">
   ];
   let dependentDialects = ["AffineDialect"];
 }
 
 def SCFForLoopRangeFolding : Pass<"scf-for-loop-range-folding"> {
   let summary = "Fold add/mul ops into loop range";
   let constructor = "mlir::createForLoopRangeFoldingPass()";
 }
 
-def SCFForToWhileLoop : Pass<"scf-for-to-while", "func::FuncOp"> {
+def SCFForToWhileLoop : Pass<"scf-for-to-while"> {
   let summary = "Convert SCF for loops to SCF while loops";
   let constructor = "mlir::createForToWhileLoopPass()";
   let description = [{
     This pass transforms SCF.ForOp operations to SCF.WhileOp. The For loop
     condition is placed in the 'before' region of the while operation, and the
     induction variable incrementation and loop body in the 'after' region.
     The loop carried values of the while op are the induction variable (IV) of
     the for-loop + any iter_args specified for the for-loop.

mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp

LogicalResult matchAndRewrite(ForOp forOp,
 
     rewriter.eraseOp(forOp);
     return success();
   }
 };
 
 struct ForToWhileLoop : public SCFForToWhileLoopBase<ForToWhileLoop> {
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
-    MLIRContext *ctx = funcOp.getContext();
+    auto *parentOp = getOperation();
+    MLIRContext *ctx = parentOp->getContext();
     RewritePatternSet patterns(ctx);
     patterns.add<ForLoopLoweringPattern>(ctx);
-    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+    (void)applyPatternsAndFoldGreedily(parentOp, std::move(patterns));
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createForToWhileLoopPass() {
   return std::make_unique<ForToWhileLoop>();
 }

mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp

LogicalResult matchAndRewrite(OpTy op,
     return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),
                                            op.operands(), IsMin, loopMatcher);
   }
 };
 
 struct SCFForLoopCanonicalization
     : public SCFForLoopCanonicalizationBase<SCFForLoopCanonicalization> {
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
-    MLIRContext *ctx = funcOp.getContext();
+    auto *parentOp = getOperation();
+    MLIRContext *ctx = parentOp->getContext();
     RewritePatternSet patterns(ctx);
     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
-    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
+    if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 void mlir::scf::populateSCFForLoopCanonicalizationPatterns(
     RewritePatternSet &patterns) {
   MLIRContext *ctx = patterns.getContext();

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
   bool skipPartial;
 };
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnOperation() override {
-    getOperation().walk(
+    getOperation()->walk(
         [](ParallelOp op) { specializeParallelLoopForUnrolling(op); });
   }
 };
 
 struct ForLoopSpecialization
     : public SCFForLoopSpecializationBase<ForLoopSpecialization> {
   void runOnOperation() override {
-    getOperation().walk([](ForOp op) { specializeForLoopForUnrolling(op); });
+    getOperation()->walk([](ForOp op) { specializeForLoopForUnrolling(op); });
   }
 };
 
 struct ForLoopPeeling : public SCFForLoopPeelingBase<ForLoopPeeling> {
   void runOnOperation() override {
-    func::FuncOp funcOp = getOperation();
-    MLIRContext *ctx = funcOp.getContext();
+    auto *parentOp = getOperation();
+    MLIRContext *ctx = parentOp->getContext();
     RewritePatternSet patterns(ctx);
     patterns.add<ForLoopPeelingPattern>(ctx, skipPartial);
-    (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+    (void)applyPatternsAndFoldGreedily(parentOp, std::move(patterns));
 
     // Drop the markers.
-    funcOp.walk([](Operation *op) {
+    parentOp->walk([](Operation *op) {
       op->removeAttr(kPeeledLoopLabel);
       op->removeAttr(kPartialIterationLabel);
     });
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {

mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp

struct ParallelLoopTiling
   ParallelLoopTiling() = default;
   explicit ParallelLoopTiling(ArrayRef<int64_t> tileSizes,
                               bool noMinMaxBounds = false) {
     this->tileSizes = tileSizes;
     this->noMinMaxBounds = noMinMaxBounds;
   }
 
   void runOnOperation() override {
+    auto *parentOp = getOperation();
     SmallVector<ParallelOp, 2> innermostPloops;
-    getInnermostParallelLoops(getOperation().getOperation(), innermostPloops);
+    getInnermostParallelLoops(parentOp, innermostPloops);
     for (ParallelOp ploop : innermostPloops) {
       // FIXME: Add reduction support.
       if (ploop.getNumReductions() == 0)
         tileParallelLoop(ploop, tileSizes, noMinMaxBounds);
     }
   }
 };
 } // namespace
 
 std::unique_ptr<Pass>
 mlir::createParallelLoopTilingPass(ArrayRef<int64_t> tileSizes,
                                    bool noMinMaxBounds) {
   return std::make_unique<ParallelLoopTiling>(tileSizes, noMinMaxBounds);
 }