[MLIR][Loops] Expose transformations on loop.parallel in header (NFC).

[herhut]

Summary

This change does not add any functionality but merely exposes existing
static functions to make the associated transformations available
outside of their testing passes.

[herhut]

PTAL.

We probably have other transforms that should ultimately move to the header. This is a first step as the added transforms are not available otherwise.

[herhut]

fix linter errors.

[bkramer]

looks good, thanks

[ftynse]

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

171

Ultra-nit: drop trivial braces

[rriddle]

mlir/include/mlir/Dialect/LoopOps/Transforms.h

49

Could you please fix this.

[herhut]

Comments...

[herhut]

Thanks!

mlir/include/mlir/Dialect/LoopOps/Transforms.h

49

Of course. Still not used to all the C++ cruft :(

mlir/include/mlir/Dialect/LoopOps/Transforms.h

+//===- Transforms.h - Pass Entrypoints --------------------------*- C++ -*-===//
+//
+// 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 header file defines transformations on loop operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LOOPOPS_TRANSFORMS_H_
+#define MLIR_DIALECT_LOOPOPS_TRANSFORMS_H_
+
+#include "llvm/ADT/ArrayRef.h"
+
+namespace mlir {
+
+class Region;
+
+namespace loop {
+
+class ParallelOp;
+
+/// Fuses all adjacent loop.parallel operations with identical bounds and step
+/// into one loop.parallel operations. Uses a naive aliasing and dependency
+/// analysis.
+void naivelyFuseParallelOps(Region &region);
+
+/// 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[0],
+///                                                   %arg5*tileSize[1])
+///     loop.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%j0)
+///                                           min(tileSize[1], %arg3-%j1))
+///                                        step (%arg4, %arg5)
+/// The old loop is replaced with the new one.
+void tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes);
+
+} // namespace loop
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LOOPOPS_TRANSFORMS_H_

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

 //===- ParallelLoopFusion.cpp - Code to perform loop fusion ---------------===//
 //
 // 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 fusion on parallel loops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/LoopOps/Passes.h"
+#include "mlir/Dialect/LoopOps/Transforms.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
-using loop::ParallelOp;
+using namespace mlir::loop;
 
 /// Verify there are no nested ParallelOps.
 static bool hasNestedParallelOp(ParallelOp ploop) {
   auto walkResult =
       ploop.getBody()->walk([](ParallelOp) { return WalkResult::interrupt(); });
   return walkResult.wasInterrupted();
 }
 
if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices))
     return;
 
   b.setInsertionPointToStart(secondPloop.getBody());
   for (auto &op : firstPloop.getBody()->without_terminator())
     b.clone(op, firstToSecondPloopIndices);
   firstPloop.erase();
 }
 
-static void naivelyFuseParallelOps(Operation *op) {
-  OpBuilder b(op);
+void mlir::loop::naivelyFuseParallelOps(Region &region) {
+  OpBuilder b(region);
   // Consider every single block and attempt to fuse adjacent loops.
-  for (auto &region : op->getRegions()) {
-    for (auto &block : region.getBlocks()) {
-      SmallVector<SmallVector<ParallelOp, 8>, 1> ploop_chains{{}};
-      // Not using `walk()` to traverse only top-level parallel loops and also
-      // make sure that there are no side-effecting ops between the parallel
-      // loops.
-      bool noSideEffects = true;
-      for (auto &op : block.getOperations()) {
-        if (auto ploop = dyn_cast<ParallelOp>(op)) {
-          if (noSideEffects) {
-            ploop_chains.back().push_back(ploop);
-          } else {
-            ploop_chains.push_back({ploop});
-            noSideEffects = true;
-          }
-          continue;
-        }
-        noSideEffects &= op.hasNoSideEffect();
-      }
-      for (ArrayRef<ParallelOp> ploops : ploop_chains) {
-        for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
-          fuseIfLegal(ploops[i], ploops[i + 1], b);
-      }
-    }
-  }
-}
+    SmallVector<SmallVector<ParallelOp, 8>, 1> ploopChains{{}};
+          ploopChains.back().push_back(ploop);
+          ploopChains.push_back({ploop});
+    for (ArrayRef<ParallelOp> ploops : ploopChains) {
 
 namespace {
 
 struct ParallelLoopFusion : public OperationPass<ParallelLoopFusion> {
-  void runOnOperation() override { naivelyFuseParallelOps(getOperation()); }
+  void runOnOperation() override {
+    for (Region &region : getOperation()->getRegions())
+      naivelyFuseParallelOps(region);
+  }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopFusionPass() {
   return std::make_unique<ParallelLoopFusion>();
 }
 
 static PassRegistration<ParallelLoopFusion>
     pass("parallel-loop-fusion", "Fuse adjacent parallel loops.");

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/LoopOps/Transforms.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;
-using loop::ParallelOp;
+using namespace mlir::loop;
 
 /// 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[0],
 ///                                                   %arg5*tileSize[1])
 ///     loop.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%j0)
 ///                                           min(tileSize[1], %arg3-%j1))
 ///                                        step (%arg4, %arg5)
 /// The old loop is replaced with the new one.
-static void tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
+void mlir::loop::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
   OpBuilder b(op);
   auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0);
   SmallVector<Value, 2> tileSizeConstants;
   tileSizeConstants.reserve(op.upperBound().size());
   for (size_t i = 0, end = op.upperBound().size(); i != end; ++i) {
     if (i < tileSizes.size())
       tileSizeConstants.push_back(
           b.create<ConstantIndexOp>(op.getLoc(), tileSizes[i]));