Add a basic tiling pass for parallel loops

[bkramer]

Summary

This exploits the fact that the iterations of parallel loops are
independent so tiling becomes just an index transformation. This pass
only tiles the innermost loop of a loop nest.

The ultimate goal is to allow vectorization of the tiled loops, but I
don't think we're there yet with the current rewriting, as the tiled
loops don't have a constant trip count.

[bkramer]

• use pass option instead of global option
• fix the recursive walker to walk recursively

[herhut]

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

25

Please use pass options.

30

nit: Doc comments need to start with ///

Use parallel loop instead of ploop.

42

tileParallelLoop

44

Maybe newSteps.reserve?

68

I do not understand this part. We want the min of either the step size or dim - upperbound_of_outer.

77

you can do something like innerLoop.region().takeBody(op.region()) which simply transfers one region to the other. As the number of block arguments does not change, no need to fix uses or anything.

95

getBlock().getOps<loop::ParallelOp>() gives an (potentially empty) iterator over all parallel loops in the block.

[ftynse]

We already several implementations of loop tiling (lib/Transforms/LoopTiling.cpp, lib/Dialect/Linalg/Transforms/Tiling.cpp, lib/Transforms/Utils/LoopUtils.cpp). Have you considered generalizing them instead of introducing yet another one? I suppose parallel loop nest as an operation removes the need for any supplementary preconditions, but the mechanics of the transformation should be very similar between parallel and sequential loops.

I agree we may want to peel off the last iterations to avoid complex conditions, but IMO it's better as an option to the transformation. @poechsel had a similar requested with loops produced by Linalg tiling, so it would make sense to try and reuse the infra here.

[ftynse]

mlir/test/Dialect/Loops/parallel-loop-tiling.mlir

20

Please don't pattern-match SSA names, and generally prefer avoiding non-essential parts of the IR from the test (e.g. repeated types in the pattern).

https://mlir.llvm.org/getting_started/TestingGuide/

[bkramer]

• One round of addressed review comments

[herhut]

In D74954#1886199, @ftynse wrote:

We already several implementations of loop tiling (lib/Transforms/LoopTiling.cpp, lib/Dialect/Linalg/Transforms/Tiling.cpp, lib/Transforms/Utils/LoopUtils.cpp). Have you considered generalizing them instead of introducing yet another one? I suppose parallel loop nest as an operation removes the need for any supplementary preconditions, but the mechanics of the transformation should be very similar between parallel and sequential loops.

The complexity of the tiling passes is in the dependency analysis and not the rewriting itself. This one is purely structural. We would need to split the actual tiling rewrite form the analysis with interfaces for generating the tiled operation. I am not sure that is worth the code we would actually get to reuse.

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

52

The replication part could be a separate pattern/pass that rewrites a loop with dynamic upper bounds specified with a constant upper bound (so an affine.min with a constant) into something like

bound = affine.min(constant, other thing)
if (bound ==constant) {
  // loop with constant upper bound
} else {
  // loop with dynamic upper bound
}

Maybe this is good enough to make LLVM recognize the potential for vectorization. That rewrite could be applied independently.

96

Can we have an ArrayRef<unsigned> for tile sizes? So we can have different ones at different levels? I would suspect that we, for vectorization, always want to tile 1x...x1xN.

[herhut]

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

68

I always get confused by this myself. I think it needs to be upper_bound - current_index_of_outer_loop to get the number of remaining iterations. And then the min with the tiling size.

[bkramer]

• Make tileSizes a list
• Fix min again

[dcaballe]

Just adding some minor nit comments!

We already several implementations of loop tiling (lib/Transforms/LoopTiling.cpp, lib/Dialect/Linalg/Transforms/Tiling.cpp, lib/Transforms/Utils/LoopUtils.cpp). Have you considered generalizing them instead of introducing yet another one? I suppose parallel loop nest as an operation removes the need for any supplementary preconditions, but the mechanics of the transformation should be very similar between parallel and sequential loops.

The complexity of the tiling passes is in the dependency analysis and not the rewriting itself. This one is purely structural. We would need to split the actual tiling rewrite form the analysis with interfaces for generating the tiled operation. I am not sure that is worth the code we would actually get to reuse.

I think it might be worth it since affine parallel for would be another candidate for tiling.

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

41

size_t i = 0; i != op.upperBound().size(); -> size_t i = 0, end = op.upperBound().size(); i != end; ?

116

please drop {} for single line for and if-else statements

mlir/test/Dialect/Loops/parallel-loop-tiling.mlir

38

Please add // ----- between tests for split-input-file to work

[herhut]

In D74954#1887027, @dcaballe wrote:

Just adding some minor nit comments!

We already several implementations of loop tiling (lib/Transforms/LoopTiling.cpp, lib/Dialect/Linalg/Transforms/Tiling.cpp, lib/Transforms/Utils/LoopUtils.cpp). Have you considered generalizing them instead of introducing yet another one? I suppose parallel loop nest as an operation removes the need for any supplementary preconditions, but the mechanics of the transformation should be very similar between parallel and sequential loops.

The complexity of the tiling passes is in the dependency analysis and not the rewriting itself. This one is purely structural. We would need to split the actual tiling rewrite form the analysis with interfaces for generating the tiled operation. I am not sure that is worth the code we would actually get to reuse.

I think it might be worth it since affine parallel for would be another candidate for tiling.

I guess we would need to compute the new step sizes as affine expressions, as well, in such case. My point was how much code reuse we would get.

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

34

This comment does not state the actual rewriting, does it?

[herhut]

Thanks!

LGTM with comments addressed.

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

40

mega-nit: tileSizeConstants.reserve

[bkramer]

• Address moar comments

[bkramer]

• WTF phab?

[rriddle]

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

113

nit: Please drop the trivial braces. Here and below.

mlir/include/mlir/Dialect/LoopOps/Passes.h

 
 namespace mlir {
 
 class Pass;
 
 /// Creates a loop fusion pass which fuses parallel loops.
 std::unique_ptr<Pass> createParallelLoopFusionPass();
 
+/// Creates a pass which tiles innermost parallel loops.
+std::unique_ptr<Pass> createParallelLoopTilingPass(int64_t tileSize = 4);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LOOPOPS_PASSES_H_

mlir/include/mlir/InitAllPasses.h

#endif
   createLinalgPromotionPass(0);
   createConvertLinalgToLoopsPass();
   createConvertLinalgToParallelLoopsPass();
   createConvertLinalgToAffineLoopsPass();
   createConvertLinalgToLLVMPass();
 
   // LoopOps
   createParallelLoopFusionPass();
+  createParallelLoopTilingPass();
 
   // QuantOps
   quant::createConvertSimulatedQuantPass();
   quant::createConvertConstPass();
   quantizer::createAddDefaultStatsPass();
   quantizer::createRemoveInstrumentationPass();
   quantizer::registerInferQuantizedTypesPass();
 

mlir/lib/Dialect/LoopOps/Transforms/CMakeLists.txt

 add_llvm_library(MLIRLoopOpsTransforms
   ParallelLoopFusion.cpp
+  ParallelLoopTiling.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LoopOps
   )
 
 target_link_libraries(MLIRLoopOpsTransforms
   MLIRPass
   MLIRLoopOps
   )

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

+//===- ParallelLoopTiling.cpp - Tiles loop.parallel ---------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements loop tiling on parallel loops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/LoopOps/LoopOps.h"
+#include "mlir/Dialect/LoopOps/Passes.h"
+#include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace mlir;
+using loop::ParallelOp;
+
+/// Tile a parallel loop of the form
+///   loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
+///                                             step (%arg4, %arg5)
+///
+/// into
+///   loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
+///                                             step (%arg4*tileSize,
+///                                                   %arg5*tileSize)
+///     loop.parallel (%i0, %i1) = (0, 0) to (min(%arg2, %arg0),
+///                                           min(%arg3, %arg1))
+///                                        step (%arg4, %arg5)
+/// The old loop is replaced with the new one.
+static void tileParallelLoop(ParallelOp op, int64_t tileSize) {
+  OpBuilder b(op);
+  auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0);
+  auto tileSizeConstant = b.create<ConstantIndexOp>(op.getLoc(), tileSize);
+
+  // Create the outer loop with adjusted steps.
+  SmallVector<Value, 2> newSteps;
+  newSteps.reserve(op.step().size());
+  for (Value step : op.step()) {
+    newSteps.push_back(b.create<MulIOp>(op.getLoc(), step, tileSizeConstant));
+  }
+  auto outerLoop = b.create<ParallelOp>(op.getLoc(), op.lowerBound(),
+                                        op.upperBound(), newSteps);
+  b.setInsertionPointToStart(outerLoop.getBody());
+
+  // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
+  // FIXME: Instead of using min, we want to replicate the tail. This would give
+  // the inner loop constant bounds for easy vectorization.
+  auto minMap = AffineMap::get(
+      /*dimCount=*/3, /*symbolCount=*/0,
+      {getAffineDimExpr(/*position=*/0, b.getContext()),
+       getAffineDimExpr(/*position=*/1, b.getContext()) -
+           getAffineDimExpr(/*position=*/2, b.getContext())});
+
+  // Create the inner loop with adjusted bounds.
+  SmallVector<Value, 2> newBounds;
+  newBounds.reserve(op.upperBound().size());
+  for (auto bounds : llvm::zip(op.upperBound(), op.lowerBound())) {
+    newBounds.push_back(
+        b.create<AffineMinOp>(op.getLoc(), b.getIndexType(), minMap,
+                              ValueRange{tileSizeConstant, std::get<0>(bounds),
+                                         std::get<1>(bounds)}));
+  }
+  auto innerLoop = b.create<ParallelOp>(
+      op.getLoc(), SmallVector<Value, 2>(newBounds.size(), zero), newBounds,
+      op.step());
+
+  // Steal the body of the old parallel loop and erase it.
+  innerLoop.region().takeBody(op.region());
+  op.erase();
+}
+
+/// Get a list of most nested parallel loops. Assumes that ParallelOps are only
+/// directly nested.
+static bool getInnermostNestedLoops(Block *block,
+                                    SmallVectorImpl<ParallelOp> &loops) {
+  bool hasInnerLoop = false;
+  for (auto parallelOp : block->getOps<ParallelOp>()) {
+    hasInnerLoop = true;
+    if (!getInnermostNestedLoops(parallelOp.getBody(), loops)) {
+      loops.push_back(parallelOp);
+    }
+  }
+  return hasInnerLoop;
+}
+
+namespace {
+struct ParallelLoopTiling : public FunctionPass<ParallelLoopTiling> {
+  ParallelLoopTiling() = default;
+  ParallelLoopTiling(const ParallelLoopTiling &) {} // tileSize is non-copyable.
+  explicit ParallelLoopTiling(int64_t tileSize) { this->tileSize = tileSize; }
+
+  void runOnFunction() override {
+    SmallVector<ParallelOp, 2> mostNestedParallelOps;
+    for (Block &block : getFunction()) {
+      getInnermostNestedLoops(&block, mostNestedParallelOps);
+    }
+    for (ParallelOp pLoop : mostNestedParallelOps) {
+      tileParallelLoop(pLoop, tileSize);
+    }
+  }
+
+  Option<int64_t> tileSize{
+      *this, "parallel-loop-tile-size",
+      llvm::cl::desc("factor to tile innermost parallel loops by"),
+      llvm::cl::init(4)};
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createParallelLoopTilingPass(int64_t tileSize) {
+  return std::make_unique<ParallelLoopTiling>(tileSize);
+}
+
+static PassRegistration<ParallelLoopTiling> pass("parallel-loop-tiling",
+                                                 "Tile parallel loops.");

mlir/test/Dialect/Loops/parallel-loop-tiling.mlir

+// RUN: mlir-opt %s -pass-pipeline='func(parallel-loop-tiling)' -split-input-file | FileCheck %s --dump-input-on-failure
+
+func @parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
+                    %arg3 : index, %arg4 : index, %arg5 : index,
+        %A: memref<?x?xf32>, %B: memref<?x?xf32>,
+                    %C: memref<?x?xf32>, %result: memref<?x?xf32>) {
+  loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
+    %B_elem = load %B[%i0, %i1] : memref<?x?xf32>
+    %C_elem = load %C[%i0, %i1] : memref<?x?xf32>
+    %sum_elem = addf %B_elem, %C_elem : f32
+    store %sum_elem, %result[%i0, %i1] : memref<?x?xf32>
+  }
+  return
+}
+
+// CHECK:       #map0 = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
+
+// CHECK-LABEL:   func @parallel_loop(
+// CHECK-SAME:                        [[VAL_0:%.*]]: index, [[VAL_1:%.*]]: index, [[VAL_2:%.*]]: index, [[VAL_3:%.*]]: index, [[VAL_4:%.*]]: index, [[VAL_5:%.*]]: index, [[VAL_6:%.*]]: memref<?x?xf32>, [[VAL_7:%.*]]: memref<?x?xf32>, [[VAL_8:%.*]]: memref<?x?xf32>, [[VAL_9:%.*]]: memref<?x?xf32>) {
+// CHECK:           [[VAL_10:%.*]] = constant 0 : index
+// CHECK:           [[VAL_11:%.*]] = constant 4 : index
+// CHECK:           [[VAL_12:%.*]] = muli [[VAL_4]], [[VAL_11]] : index
+// CHECK:           [[VAL_13:%.*]] = muli [[VAL_5]], [[VAL_11]] : index
+// CHECK:           loop.parallel ([[VAL_14:%.*]], [[VAL_15:%.*]]) = ([[VAL_0]], [[VAL_1]]) to ([[VAL_2]], [[VAL_3]]) step ([[VAL_12]], [[VAL_13]]) {
+// CHECK:             [[VAL_16:%.*]] = affine.min #map0([[VAL_11]], [[VAL_2]], [[VAL_0]])
+// CHECK:             [[VAL_17:%.*]] = affine.min #map0([[VAL_11]], [[VAL_3]], [[VAL_1]])
+// CHECK:             loop.parallel ([[VAL_18:%.*]], [[VAL_19:%.*]]) = ([[VAL_10]], [[VAL_10]]) to ([[VAL_16]], [[VAL_17]]) step ([[VAL_4]], [[VAL_5]]) {
+// CHECK:               [[VAL_20:%.*]] = load [[VAL_7]]{{\[}}[[VAL_18]], [[VAL_19]]] : memref<?x?xf32>
+// CHECK:               [[VAL_21:%.*]] = load [[VAL_8]]{{\[}}[[VAL_18]], [[VAL_19]]] : memref<?x?xf32>
+// CHECK:               [[VAL_22:%.*]] = addf [[VAL_20]], [[VAL_21]] : f32
+// CHECK:               store [[VAL_22]], [[VAL_9]]{{\[}}[[VAL_18]], [[VAL_19]]] : memref<?x?xf32>
+// CHECK:               "loop.terminator"() : () -> ()
+// CHECK:             }
+// CHECK:             "loop.terminator"() : () -> ()
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+
+
+func @tile_nested_innermost() {
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  loop.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+    loop.parallel (%k, %l) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+    }
+  }
+  loop.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+  }
+  return
+}
+
+// CHECK-LABEL:   func @tile_nested_innermost() {
+// CHECK:           [[VAL_23:%.*]] = constant 2 : index
+// CHECK:           [[VAL_24:%.*]] = constant 0 : index
+// CHECK:           [[VAL_25:%.*]] = constant 1 : index
+// CHECK:           loop.parallel ([[VAL_26:%.*]], [[VAL_27:%.*]]) = ([[VAL_24]], [[VAL_24]]) to ([[VAL_23]], [[VAL_23]]) step ([[VAL_25]], [[VAL_25]]) {
+// CHECK:             [[VAL_28:%.*]] = constant 0 : index
+// CHECK:             [[VAL_29:%.*]] = constant 4 : index
+// CHECK:             [[VAL_30:%.*]] = muli [[VAL_25]], [[VAL_29]] : index
+// CHECK:             [[VAL_31:%.*]] = muli [[VAL_25]], [[VAL_29]] : index
+// CHECK:             loop.parallel ([[VAL_32:%.*]], [[VAL_33:%.*]]) = ([[VAL_24]], [[VAL_24]]) to ([[VAL_23]], [[VAL_23]]) step ([[VAL_30]], [[VAL_31]]) {
+// CHECK:               [[VAL_34:%.*]] = affine.min #map0([[VAL_29]], [[VAL_23]], [[VAL_24]])
+// CHECK:               [[VAL_35:%.*]] = affine.min #map0([[VAL_29]], [[VAL_23]], [[VAL_24]])
+// CHECK:               loop.parallel ([[VAL_36:%.*]], [[VAL_37:%.*]]) = ([[VAL_28]], [[VAL_28]]) to ([[VAL_34]], [[VAL_35]]) step ([[VAL_25]], [[VAL_25]]) {
+// CHECK:                 "loop.terminator"() : () -> ()
+// CHECK:               }
+// CHECK:               "loop.terminator"() : () -> ()
+// CHECK:             }
+// CHECK:             "loop.terminator"() : () -> ()
+// CHECK:           }
+// CHECK:           [[VAL_38:%.*]] = constant 0 : index
+// CHECK:           [[VAL_39:%.*]] = constant 4 : index
+// CHECK:           [[VAL_40:%.*]] = muli [[VAL_25]], [[VAL_39]] : index
+// CHECK:           [[VAL_41:%.*]] = muli [[VAL_25]], [[VAL_39]] : index
+// CHECK:           loop.parallel ([[VAL_42:%.*]], [[VAL_43:%.*]]) = ([[VAL_24]], [[VAL_24]]) to ([[VAL_23]], [[VAL_23]]) step ([[VAL_40]], [[VAL_41]]) {
+// CHECK:             [[VAL_44:%.*]] = affine.min #map0([[VAL_39]], [[VAL_23]], [[VAL_24]])
+// CHECK:             [[VAL_45:%.*]] = affine.min #map0([[VAL_39]], [[VAL_23]], [[VAL_24]])
+// CHECK:             loop.parallel ([[VAL_46:%.*]], [[VAL_47:%.*]]) = ([[VAL_38]], [[VAL_38]]) to ([[VAL_44]], [[VAL_45]]) step ([[VAL_25]], [[VAL_25]]) {
+// CHECK:               "loop.terminator"() : () -> ()
+// CHECK:             }
+// CHECK:             "loop.terminator"() : () -> ()
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }