[Mlir-commits] [mlir] [mlir][acc] Add loop tiling utilities for OpenACC (PR #171490)

[Razvan Lupusoru]

================
@@ -0,0 +1,349 @@
+//===- OpenACCUtilsTilingTest.cpp - Unit tests for loop tiling utilities --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/OpenACC/OpenACCUtilsTiling.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::acc;
+
+//===----------------------------------------------------------------------===//
+// Test Fixture
+//===----------------------------------------------------------------------===//
+
+class OpenACCUtilsTilingTest : public ::testing::Test {
+protected:
+  OpenACCUtilsTilingTest() : b(&context), loc(UnknownLoc::get(&context)) {
+    context.loadDialect<acc::OpenACCDialect, arith::ArithDialect,
+                        memref::MemRefDialect, func::FuncDialect>();
+  }
+
+  // Create a simple LoopOp with specified bounds using the simple builder
+  acc::LoopOp createLoopOp(OpBuilder &builder, ValueRange lbs, ValueRange ubs,
+                           ValueRange steps) {
+    auto loopOp = acc::LoopOp::create(builder, loc, lbs, ubs, steps,
+                                      acc::LoopParMode::loop_independent);
+
+    // Add body block with IV arguments and yield
+    Region &region = loopOp.getRegion();
+    Block *block = builder.createBlock(&region, region.begin());
+    for (Value lb : lbs)
+      block->addArgument(lb.getType(), loc);
+    builder.setInsertionPointToEnd(block);
+    acc::YieldOp::create(builder, loc);
+
+    return loopOp;
+  }
+
+  // Helper to count nested acc.loop ops within a loop
+  unsigned countNestedLoops(acc::LoopOp loop) {
+    unsigned count = 0;
+    loop.getBody().walk([&](acc::LoopOp) { ++count; });
+    return count;
+  }
+
+  // Helper to collect all nested acc.loop ops in order
+  SmallVector<acc::LoopOp> collectNestedLoops(acc::LoopOp loop) {
+    SmallVector<acc::LoopOp> loops;
+    loop.getBody().walk([&](acc::LoopOp nestedLoop) {
+      loops.push_back(nestedLoop);
+    });
+    return loops;
+  }
+
+  MLIRContext context;
+  OpBuilder b;
+  Location loc;
+};
+
+//===----------------------------------------------------------------------===//
+// tileACCLoops Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsTilingTest, tileACCLoopsSingleLoop) {
+  // Create a module to hold the function
+  OwningOpRef<ModuleOp> module = ModuleOp::create(loc);
+  Block *moduleBlock = module->getBody();
+
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPointToStart(moduleBlock);
+
+  // Create a function
+  auto funcType = b.getFunctionType({}, {});
+  OwningOpRef<func::FuncOp> funcOp =
+      func::FuncOp::create(b, loc, "test_func", funcType);
+  Block *funcBlock = funcOp->addEntryBlock();
+
+  b.setInsertionPointToStart(funcBlock);
+
+  // Create loop bounds
+  Value lb =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
+  Value ub =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(100));
+  Value step =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
+  Value tileSize =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(4));
+
+  // Create the loop
+  acc::LoopOp loopOp = createLoopOp(b, {lb}, {ub}, {step});
+
+  // Tile the loop using IRRewriter
+  IRRewriter rewriter(&context);
+  rewriter.setInsertionPoint(loopOp);
+
+  SmallVector<acc::LoopOp> loopsToTile = {loopOp};
+  SmallVector<Value> tileSizes = {tileSize};
+
+  acc::LoopOp tiledLoop = tileACCLoops(loopsToTile, tileSizes, /*defaultTileSize=*/128, rewriter);
+
+  // Verify the tiled loop was created
+  EXPECT_TRUE(tiledLoop != nullptr);
+  EXPECT_FALSE(tiledLoop.getBody().empty());
+
+  // After tiling a single loop with tile(4), we should have:
+  // - 1 tile loop (the outer loop)
+  // - 1 element loop nested inside
+  // Total: 1 nested loop inside the tile loop
+  EXPECT_EQ(countNestedLoops(tiledLoop), 1u);
+
+  // The tile loop (outer) should have 1 IV
+  EXPECT_EQ(tiledLoop.getBody().getNumArguments(), 1u);
+
+  // Collect nested loops and verify
+  auto nestedLoops = collectNestedLoops(tiledLoop);
+  EXPECT_EQ(nestedLoops.size(), 1u);
+  if (!nestedLoops.empty()) {
+    // The element loop should have 1 IV
+    EXPECT_EQ(nestedLoops[0].getBody().getNumArguments(), 1u);
+  }
+}
+
+TEST_F(OpenACCUtilsTilingTest, tileACCLoopsNestedLoops) {
+  // Create a module to hold the function
+  OwningOpRef<ModuleOp> module = ModuleOp::create(loc);
+  Block *moduleBlock = module->getBody();
+
+  OpBuilder::InsertionGuard guard(b);
+  b.setInsertionPointToStart(moduleBlock);
+
+  // Create a function
+  auto funcType = b.getFunctionType({}, {});
+  OwningOpRef<func::FuncOp> funcOp =
+      func::FuncOp::create(b, loc, "test_func", funcType);
+  Block *funcBlock = funcOp->addEntryBlock();
+
+  b.setInsertionPointToStart(funcBlock);
+
+  // Create loop bounds for outer loop
+  Value lb1 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
+  Value ub1 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(100));
+  Value step1 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
+
+  // Create loop bounds for inner loop
+  Value lb2 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
+  Value ub2 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(50));
+  Value step2 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
+
+  // Tile sizes
+  Value tileSize1 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(4));
+  Value tileSize2 =
+      arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(8));
+
+  // Create outer loop
+  acc::LoopOp outerLoop = createLoopOp(b, {lb1}, {ub1}, {step1});
+
+  // Create inner loop inside outer loop
+  b.setInsertionPoint(outerLoop.getBody().getTerminator());
+  acc::LoopOp innerLoop = createLoopOp(b, {lb2}, {ub2}, {step2});
+
+  // Tile the loops
+  IRRewriter rewriter(&context);
+  rewriter.setInsertionPoint(outerLoop);
+
+  SmallVector<acc::LoopOp> loopsToTile = {outerLoop, innerLoop};
+  SmallVector<Value> tileSizes = {tileSize1, tileSize2};
+
+  acc::LoopOp tiledLoop = tileACCLoops(loopsToTile, tileSizes, /*defaultTileSize=*/128, rewriter);
+
+  // Verify the tiled loop nest was created
+  EXPECT_TRUE(tiledLoop != nullptr);
+  EXPECT_FALSE(tiledLoop.getBody().empty());
+
+  // After tiling a 2-level nested loop with tile(4,8), we should have:
+  // tile_loop_1 -> tile_loop_2 -> element_loop_1 -> element_loop_2
+  // Total: 3 nested loops inside the outermost tile loop
+  unsigned nestedCount = countNestedLoops(tiledLoop);
+  EXPECT_EQ(nestedCount, 3u);
+
+  // The outermost tile loop should have 1 IV
+  EXPECT_EQ(tiledLoop.getBody().getNumArguments(), 1u);
+
+  // Collect all nested loops and verify each has 1 IV
+  auto nestedLoops = collectNestedLoops(tiledLoop);
+  EXPECT_EQ(nestedLoops.size(), 3u);
+  for (auto loop : nestedLoops) {
+    EXPECT_EQ(loop.getBody().getNumArguments(), 1u);
----------------
razvanlupusoru wrote:

Done

https://github.com/llvm/llvm-project/pull/171490