[mlir] Add for loop specialization

[herhut]

Summary

We already has a parallel loop specialization pass that is used to
enable unrolling and consecutive vectorization by rewriting loops
whose bound is defined as a min of a constant and a dynamic value
into a loop with static bound (the constant) and the minimum as
bound, wrapped into a conditional to dispatch between the two.
This adds the same rewriting for for loops.

[herhut]

Add missing pieces...

[herhut]

formatting

[rriddle]

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

73

nit: Spell out the type here.

[ftynse]

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

16

Nit: can we take the first 'for' into single quotes for improved readability?

[herhut]

comments + rebase

[herhut]

Review comments...

mlir/include/mlir/Dialect/SCF/Passes.h

 
 #include "llvm/ADT/ArrayRef.h"
 #include <memory>
 
 namespace mlir {
 
 class Pass;
 
+/// Creates a pass that specializes for loop for unrolling and
+/// vectorization.
+std::unique_ptr<Pass> createForLoopSpecializationPass();
+
 /// Creates a loop fusion pass which fuses parallel loops.
 std::unique_ptr<Pass> createParallelLoopFusionPass();
 
 /// Creates a pass that specializes parallel loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createParallelLoopSpecializationPass();
 
 /// Creates a pass which tiles innermost parallel loops.
 std::unique_ptr<Pass>
 createParallelLoopTilingPass(llvm::ArrayRef<int64_t> tileSize = {});
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_PASSES_H_

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

-//===-- Passes.td - Loop pass definition file --------------*- tablegen -*-===//
+//===-- 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 LoopParallelLoopFusion : Pass<"parallel-loop-fusion"> {
+def SCFForLoopSpecialization
+    : FunctionPass<"for-loop-specialization"> {
+  let summary = "Specialize for loops for vectorization";
+  let constructor = "mlir::createForLoopSpecializationPass()";
+}
+
+def SCFParallelLoopFusion : Pass<"parallel-loop-fusion"> {
   let summary = "Fuse adjacent parallel loops";
   let constructor = "mlir::createParallelLoopFusionPass()";
 }
 
-def LoopParallelLoopSpecialization
+def SCFParallelLoopSpecialization
     : FunctionPass<"parallel-loop-specialization"> {
   let summary = "Specialize parallel loops for vectorization";
   let constructor = "mlir::createParallelLoopSpecializationPass()";
 }
 
-def LoopParallelLoopTiling : FunctionPass<"parallel-loop-tiling"> {
+def SCFParallelLoopTiling : FunctionPass<"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, llvm::cl::MiscFlags::CommaSeparated">
   ];
 }
 
 #endif // MLIR_DIALECT_SCF_PASSES

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

static void specializeParallelLoopForUnrolling(ParallelOp op) {
   ifOp.getElseBodyBuilder().clone(*op.getOperation());
   op.erase();
 }
 
 /// Rewrite a for loop with bounds defined by an affine.min with a constant into 2
 /// loops after checking if the bounds are equal to that constant. This is
 /// beneficial if the loop will almost always have the constant bound and that
 /// version can be fully unrolled and vectorized.
-static void specializeParallelLoopForUnrolling(ForOp op) {
-  auto bound = op.upper();
-    auto minOp = bound.getDefiningOp<AffineMinOp>();
-    if (!minOp)
-      return;
-    int64_t minConstant = std::numeric_limits<int64_t>::max();
-    for (auto expr : minOp.map().getResults()) {
-      if (auto constantIndex = expr.dyn_cast<AffineConstantExpr>())
-        minConstant = std::min(minConstant, constantIndex.getValue());
-    }
-    if (minConstant == std::numeric_limits<int64_t>::max())
-      return;
-    constantIndices.push_back(minConstant);
+static void specializeForLoopForUnrolling(ForOp op) {
+  auto bound = op.upperBound();
 
   OpBuilder b(op);
   BlockAndValueMapping map;
   Value constant = b.create<ConstantIndexOp>(op.getLoc(), minConstant);
   Value cond = b.create<CmpIOp>(op.getLoc(), CmpIPredicate::eq,
                                  bound, constant);
   map.map(bound, constant);
   auto ifOp = b.create<scf::IfOp>(op.getLoc(), cond, /*withElseRegion=*/true);
   ifOp.getThenBodyBuilder().clone(*op.getOperation(), map);
   ifOp.getElseBodyBuilder().clone(*op.getOperation());
   op.erase();
 }
 
 namespace {
 struct ParallelLoopSpecialization
-    : public LoopParallelLoopSpecializationBase<ParallelLoopSpecialization> {
+    : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnFunction() override {
     getFunction().walk([](ParallelOp op) { specializeParallelLoopForUnrolling(op); });
   }
 };
+
+struct ForLoopSpecialization
+    : public SCFForLoopSpecializationBase<ForLoopSpecialization> {
+  void runOnFunction() override {
+    getFunction().walk([](ForOp op) { specializeForLoopForUnrolling(op); });
+  }
+};
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {
   return std::make_unique<ParallelLoopSpecialization>();
 }
+
+std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {
+  return std::make_unique<ForLoopSpecialization>();
+}

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

for (ArrayRef<ParallelOp> ploops : ploopChains) {
       for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
         fuseIfLegal(ploops[i], ploops[i + 1], b);
     }
   }
 }
 
 namespace {
 struct ParallelLoopFusion
-    : public LoopParallelLoopFusionBase<ParallelLoopFusion> {
+    : public SCFParallelLoopFusionBase<ParallelLoopFusion> {
   void runOnOperation() override {
     getOperation()->walk([&](Operation *child) {
       for (Region &region : child->getRegions())
         naivelyFuseParallelOps(region);
     });
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopFusionPass() {
   return std::make_unique<ParallelLoopFusion>();
 }

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

for (auto parallelOp : block->getOps<ParallelOp>()) {
     if (!getInnermostNestedLoops(parallelOp.getBody(), loops))
       loops.push_back(parallelOp);
   }
   return hasInnerLoop;
 }
 
 namespace {
 struct ParallelLoopTiling
-    : public LoopParallelLoopTilingBase<ParallelLoopTiling> {
+    : public SCFParallelLoopTilingBase<ParallelLoopTiling> {
   ParallelLoopTiling() = default;
   explicit ParallelLoopTiling(ArrayRef<int64_t> tileSizes) {
     this->tileSizes = tileSizes;
   }
 
   void runOnFunction() override {
     SmallVector<ParallelOp, 2> mostNestedParallelOps;
     for (Block &block : getFunction()) {

mlir/test/Dialect/SCF/for-loop-specialization.mlir

+// RUN: mlir-opt %s -for-loop-specialization -split-input-file | FileCheck %s
+
+#map0 = affine_map<()[s0, s1] -> (1024, s0 - s1)>
+#map1 = affine_map<()[s0, s1] -> (64, s0 - s1)>
+
+func @for(%outer: index, %A: memref<?xf32>, %B: memref<?xf32>,
+          %C: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %d0 = dim %A, %c0 : memref<?xf32>
+  %b0 = affine.min #map0()[%d0, %outer]
+  scf.for %i0 = %c0 to %b0 step %c1 {
+    %B_elem = load %B[%i0] : memref<?xf32>
+    %C_elem = load %C[%i0] : memref<?xf32>
+    %sum_elem = addf %B_elem, %C_elem : f32
+    store %sum_elem, %result[%i0] : memref<?xf32>
+  }
+  return
+}
+
+// CHECK-LABEL:   func @for(
+// CHECK-SAME:                        [[ARG0:%.*]]: index, [[VAL_2:%.*]]: memref<?xf32>, [[VAL_3:%.*]]: memref<?xf32>, [[VAL_4:%.*]]: memref<?xf32>, [[VAL_5:%.*]]: memref<?xf32>) {
+// CHECK:           [[CST_0:%.*]] = constant 0 : index
+// CHECK:           [[CST_1:%.*]] = constant 1 : index
+// CHECK:           [[DIM_0:%.*]] = dim [[VAL_2]], [[CST_0]] : memref<?xf32>
+// CHECK:           [[MIN:%.*]] = affine.min #map0(){{\[}}[[DIM_0]], [[ARG0]]]
+// CHECK:           [[CST_1024:%.*]] = constant 1024 : index
+// CHECK:           [[PRED:%.*]] = cmpi "eq", [[MIN]], [[CST_1024]] : index
+// CHECK:           scf.if [[PRED]] {
+// CHECK:             scf.for [[IDX0:%.*]] = [[CST_0]] to [[CST_1024]] step [[CST_1]] {
+// CHECK:               store
+// CHECK:             }
+// CHECK:           } else {
+// CHECK:             scf.for [[IDX0:%.*]] = [[CST_0]] to [[MIN]] step [[CST_1]] {
+// CHECK:               store
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }