[MLIR][Affine] Add parametric tile size support for affine.for tiling

Add support to tile affine.for ops with parametric sizes (i.e., SSA
values). Currently supports hyper-rectangular loop nests with constant
lower bounds only. Move methods

  - moveLoopBody(*)
  - getTileableBands(*)
  - checkTilingLegality(*)
  - tilePerfectlyNested(*)
  - constructTiledIndexSetHyperRect(*)

to allow reuse with constant tile size API. Add a test pass -test-affine
-parametric-tile to test parametric tiling.

Differential Revision: https://reviews.llvm.org/D87353

2 files changed, 91 insertions, 0 deletions

diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt
index 3ac1e7c..99424f1 100644
--- a/mlir/test/lib/Transforms/CMakeLists.txt
+++ b/mlir/test/lib/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
+  TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
diff --git a/mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp b/mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp
new file mode 100644
index 0000000..5d369e6
--- /dev/null
+++ b/mlir/test/lib/Transforms/TestAffineLoopParametricTiling.cpp
@@ -0,0 +1,90 @@
+//= TestAffineLoopParametricTiling.cpp -- Parametric Affine loop tiling pass =//
+//
+// 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 a test pass to test parametric tiling of perfectly
+// nested affine for loops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Transforms/LoopUtils.h"
+
+using namespace mlir;
+
+#define DEBUG_TYPE "test-affine-parametric-tile"
+
+namespace {
+
+struct TestAffineLoopParametricTiling
+    : public PassWrapper<TestAffineLoopParametricTiling, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+/// Checks if the function enclosing the loop nest has any arguments passed to
+/// it, which can be used as tiling parameters. Assumes that atleast 'n'
+/// arguments are passed, where 'n' is the number of loops in the loop nest.
+static void checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {
+  assert(!band.empty() && "no loops in input band");
+  AffineForOp topLoop = band[0];
+
+  if (FuncOp funcOp = dyn_cast<FuncOp>(topLoop.getParentOp()))
+    assert(funcOp.getNumArguments() >= band.size() && "Too few tile sizes");
+}
+
+/// Captures tiling parameters, which are expected to be passed as arguments
+/// to the function enclosing the loop nest. Also checks if the required
+/// parameters are of index type. This approach is temporary for testing
+/// purposes.
+static void getTilingParameters(ArrayRef<AffineForOp> band,
+                                SmallVectorImpl<Value> &tilingParameters) {
+  AffineForOp topLoop = band[0];
+  Region *funcOpRegion = topLoop.getParentRegion();
+  unsigned nestDepth = band.size();
+
+  for (BlockArgument blockArgument :
+       funcOpRegion->getArguments().take_front(nestDepth)) {
+    if (blockArgument.getArgNumber() < nestDepth) {
+      assert(blockArgument.getType().isIndex() &&
+             "expected tiling parameters to be of index type.");
+      tilingParameters.push_back(blockArgument);
+    }
+  }
+}
+
+void TestAffineLoopParametricTiling::runOnFunction() {
+  // Bands of loops to tile.
+  std::vector<SmallVector<AffineForOp, 6>> bands;
+  getTileableBands(getFunction(), &bands);
+
+  // Tile each band.
+  for (SmallVectorImpl<AffineForOp> &band : bands) {
+    // Capture the tiling parameters from the arguments to the function
+    // enclosing this loop nest.
+    SmallVector<AffineForOp, 6> tiledNest;
+    SmallVector<Value, 6> tilingParameters;
+    // Check if tiling parameters are present.
+    checkIfTilingParametersExist(band);
+
+    // Get function arguments as tiling parameters.
+    getTilingParameters(band, tilingParameters);
+
+    if (failed(
+            tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest)))
+      return signalPassFailure();
+  }
+}
+
+namespace mlir {
+void registerTestAffineLoopParametricTilingPass() {
+  PassRegistration<TestAffineLoopParametricTiling>(
+      "test-affine-parametric-tile",
+      "Tile affine loops using SSA values as tile sizes");
+}
+} // namespace mlir