[mlir][linalg] Remove template parameter from loop lowering.

[gysit]

Summary

Replace the templated linalgLowerOpToLoops method by three specialized methods linalgOpToLoops, LinalgOpToParallelLoops, and linalgOpToAffineLoops.

[gysit]

Note that the replaceIndexOpsByInductionVariables moves and is otherwise unchanged.

[ThomasRaoux]

LGTM

[gysit]

rebase

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 Optional<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                    LinalgPromotionOptions options,
                                    OperationFolder *folder = nullptr);
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
 LogicalResult vectorizeLinalgOp(OpBuilder &builder, Operation *op,
                                 SmallVectorImpl<Value> &newResults);
 
-/// Emits a loop nest of `LoopTy` with the proper body for `linalgOp`.
-template <typename LoopTy>
-Optional<LinalgLoops> linalgLowerOpToLoops(PatternRewriter &rewriter,
-                                           LinalgOp linalgOp);
-
 /// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
-LogicalResult linalgOpToLoops(PatternRewriter &rewriter, LinalgOp linalgOp);
+Optional<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
+                                      LinalgOp linalgOp);
 
 /// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
-LogicalResult linalgOpToParallelLoops(PatternRewriter &rewriter,
+Optional<LinalgLoops> linalgOpToParallelLoops(PatternRewriter &rewriter,
                                               LinalgOp linalgOp);
 
 /// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
-LogicalResult linalgOpToAffineLoops(PatternRewriter &rewriter,
+Optional<LinalgLoops> linalgOpToAffineLoops(PatternRewriter &rewriter,
                                             LinalgOp linalgOp);
 
 //===----------------------------------------------------------------------===//
 // Preconditions that ensure the corresponding transformation succeeds and can
 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
 /// Emits a `generic` or `indexed_generic` operation with the `indexing_maps`
 /// and `iterator_types` permutated according to `permutation`.
 LogicalResult
LogicalResult matchAndRewrite(Operation *op,
     if (failed(filter.checkAndNotify(rewriter, linalgOp)))
       return failure();
 
     switch (loweringType) {
     case LinalgLoweringType::LibraryCall:
       // TODO: Move lowering to library calls here.
       return failure();
     case LinalgLoweringType::Loops:
-      if (failed(linalgOpToLoops(rewriter, op)))
+      if (!linalgOpToLoops(rewriter, op))
         return failure();
       break;
     case LinalgLoweringType::AffineLoops:
-      if (failed(linalgOpToAffineLoops(rewriter, op)))
+      if (!linalgOpToAffineLoops(rewriter, op))
         return failure();
       break;
     case LinalgLoweringType::ParallelLoops:
-      if (failed(linalgOpToParallelLoops(rewriter, op)))
+      if (!linalgOpToParallelLoops(rewriter, op))
         return failure();
       break;
     }
 
     rewriter.eraseOp(op);
     return success();
   }
 

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

static void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingSumOp op) {
   auto indices = getInputAndOutputIndices(allIvs, op);
   IndexedValueType output(op.output());
 
   // Emit scalar form.
   output(indices.outputs) +=
       getPoolingInput<IndexedValueType>(op, indices.inputs);
 }
 
+/// Replace the index operations in the body of the loop nest by the matching
+/// induction variables.
+static void replaceIndexOpsByInductionVariables(LinalgOp linalgOp,
+                                                PatternRewriter &rewriter,
+                                                ArrayRef<Operation *> loopOps) {
+  // Extract the induction variables of the loop nest from outer to inner.
+  SmallVector<Value> allIvs;
+  for (Operation *loopOp : loopOps) {
+    llvm::TypeSwitch<Operation *>(loopOp)
+        .Case([&](scf::ParallelOp parallelOp) {
+          allIvs.append(parallelOp.getInductionVars().begin(),
+                        parallelOp.getInductionVars().end());
+        })
+        .Case([&](scf::ForOp forOp) {
+          allIvs.push_back(forOp.getInductionVar());
+        })
+        .Case([&](AffineForOp affineForOp) {
+          allIvs.push_back(affineForOp.getInductionVar());
+        })
+        .Default([&](Operation *op) { assert(false && "unexpected op"); });
+  }
+  assert(linalgOp.getNumLoops() == allIvs.size() &&
+         "expected the number of loops and induction variables to match");
+  // Replace the index operations in the body of the innermost loop op.
+  if (!loopOps.empty()) {
+    LoopLikeOpInterface loopOp = loopOps.back();
+    for (IndexOp indexOp :
+         llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
+      rewriter.replaceOp(indexOp, allIvs[indexOp.dim()]);
+  }
+}
+
 template <typename LoopTy>
-static Optional<LinalgLoops> linalgOpToLoopsImpl(LinalgOp linalgOp,
-                                                 OpBuilder &builder) {
+static Optional<LinalgLoops> linalgOpToLoopsImpl(PatternRewriter &rewriter,
+                                                 LinalgOp linalgOp) {
   using IndexedValueTy = typename GenerateLoopNest<LoopTy>::IndexedValueTy;
-  ScopedContext scope(builder, linalgOp.getLoc());
+  ScopedContext scope(rewriter, linalgOp.getLoc());
+
+  // Canonicalize indexed_generic operations before lowering them to loops.
+  if (isa<IndexedGenericOp>(linalgOp))
+    return llvm::None;
 
   // The flattened loopToOperandRangesMaps is expected to be an invertible
   // permutation map (which is asserted in the inverse calculation).
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
 
-  auto loopRanges = linalgOp.createLoopRanges(builder, linalgOp.getLoc());
+  auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc());
   auto iteratorTypes = llvm::to_vector<4>(linalgOp.iterator_types().getValue());
 
   SmallVector<Value, 4> allIvs;
   GenerateLoopNest<LoopTy>::doit(
       loopRanges, linalgOp, iteratorTypes,
       [&](ValueRange ivs, ValueRange iterArgs) -> scf::ValueVector {
         assert(iterArgs.empty() && "unexpected iterArgs");
         allIvs.append(ivs.begin(), ivs.end());
for (Value iv : allIvs) {
     // The induction variable is a block argument of the entry block of the
     // loop operation.
     BlockArgument ivVal = iv.dyn_cast<BlockArgument>();
     if (!ivVal)
       return {};
     loopSet.insert(ivVal.getOwner()->getParentOp());
   }
   LinalgLoops loops(loopSet.begin(), loopSet.end());
+  // Replace all index operations in the loop body.
+  replaceIndexOpsByInductionVariables(linalgOp, rewriter, loops);
   return loops;
 }
 
-/// Replace the index operations in the body of the loop nest by the matching
-/// induction variables.
-static void replaceIndexOpsByInductionVariables(LinalgOp linalgOp,
-                                                PatternRewriter &rewriter,
-                                                ArrayRef<Operation *> loopOps) {
-  // Extract the induction variables of the loop nest from outer to inner.
-  SmallVector<Value> allIvs;
-  for (Operation *loopOp : loopOps) {
-    llvm::TypeSwitch<Operation *>(loopOp)
-        .Case([&](scf::ParallelOp parallelOp) {
-          allIvs.append(parallelOp.getInductionVars().begin(),
-                        parallelOp.getInductionVars().end());
-        })
-        .Case([&](scf::ForOp forOp) {
-          allIvs.push_back(forOp.getInductionVar());
-        })
-        .Case([&](AffineForOp affineForOp) {
-          allIvs.push_back(affineForOp.getInductionVar());
-        })
-        .Default([&](Operation *op) { assert(false && "unexpected op"); });
-  }
-  assert(linalgOp.getNumLoops() == allIvs.size() &&
-         "expected the number of loops and induction variables to match");
-  // Replace the index operations in the body of the innermost loop op.
-  if (!loopOps.empty()) {
-    LoopLikeOpInterface loopOp = loopOps.back();
-    for (IndexOp indexOp :
-         llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
-      rewriter.replaceOp(indexOp, allIvs[indexOp.dim()]);
-  }
-}
-
 namespace {
 template <typename LoopType>
 class LinalgRewritePattern : public RewritePattern {
 public:
   LinalgRewritePattern(MLIRContext *context)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     auto linalgOp = dyn_cast<LinalgOp>(op);
     if (!isa<LinalgOp>(op))
       return failure();
-    if (!linalgLowerOpToLoops<LoopType>(rewriter, linalgOp))
+    if (!linalgOpToLoopsImpl<LoopType>(rewriter, linalgOp))
       return failure();
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 struct TiledLoopToSCFPattern : public OpRewritePattern<TiledLoopOp> {
   using OpRewritePattern<TiledLoopOp>::OpRewritePattern;
mlir::createConvertLinalgToParallelLoopsPass() {
   return std::make_unique<LowerToParallelLoops>();
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createConvertLinalgToAffineLoopsPass() {
   return std::make_unique<LowerToAffineLoops>();
 }
 
-/// Emits a loop nest with the proper body for `linalgOp`.
-template <typename LoopTy>
-Optional<LinalgLoops>
-mlir::linalg::linalgLowerOpToLoops(PatternRewriter &rewriter,
-                                   LinalgOp linalgOp) {
-  // Convert indexed_generic ops to generic ops before lowering them to loops.
-  if (isa<IndexedGenericOp>(linalgOp))
-    return llvm::None;
-
-  Optional<LinalgLoops> loopOps =
-      linalgOpToLoopsImpl<LoopTy>(linalgOp.getOperation(), rewriter);
-  if (loopOps.hasValue())
-    replaceIndexOpsByInductionVariables(linalgOp, rewriter, loopOps.getValue());
-  return loopOps;
-}
-
-template Optional<LinalgLoops>
-mlir::linalg::linalgLowerOpToLoops<AffineForOp>(PatternRewriter &rewriter,
-                                                LinalgOp linalgOp);
-template Optional<LinalgLoops>
-mlir::linalg::linalgLowerOpToLoops<scf::ForOp>(PatternRewriter &rewriter,
-                                               LinalgOp linalgOp);
-template Optional<LinalgLoops>
-mlir::linalg::linalgLowerOpToLoops<scf::ParallelOp>(PatternRewriter &rewriter,
-                                                    LinalgOp linalgOp);
-
 /// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
-LogicalResult mlir::linalg::linalgOpToAffineLoops(PatternRewriter &rewriter,
-                                                  LinalgOp linalgOp) {
-  Optional<LinalgLoops> loops =
-      linalgLowerOpToLoops<AffineForOp>(rewriter, linalgOp);
-  return loops ? success() : failure();
+Optional<LinalgLoops>
+mlir::linalg::linalgOpToAffineLoops(PatternRewriter &rewriter,
+  return linalgOpToLoopsImpl<AffineForOp>(rewriter, linalgOp);
 }
 
 /// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
-LogicalResult mlir::linalg::linalgOpToLoops(PatternRewriter &rewriter,
+Optional<LinalgLoops> mlir::linalg::linalgOpToLoops(PatternRewriter &rewriter,
                                                     LinalgOp linalgOp) {
-  Optional<LinalgLoops> loops =
-      linalgLowerOpToLoops<scf::ForOp>(rewriter, linalgOp);
-  return loops ? success() : failure();
+  return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp);
 }
 
 /// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
-LogicalResult mlir::linalg::linalgOpToParallelLoops(PatternRewriter &rewriter,
-                                                    LinalgOp linalgOp) {
-  Optional<LinalgLoops> loops =
-      linalgLowerOpToLoops<scf::ParallelOp>(rewriter, linalgOp);
-  return loops ? success() : failure();
+Optional<LinalgLoops>
+mlir::linalg::linalgOpToParallelLoops(PatternRewriter &rewriter,
+  return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp);
 }