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

[Razvan Lupusoru]

================
@@ -0,0 +1,313 @@
+//===- OpenACCUtilsTiling.cpp - OpenACC 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains utility functions for tiling OpenACC loops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/OpenACC/OpenACCUtilsTiling.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Transforms/RegionUtils.h"
+
+// Resolve unknown tile sizes (represented as -1 for tile(*)) to the default.
+static mlir::Value resolveUnknownTileSize(mlir::Value tileSize,
+                                          int32_t defaultTileSize,
+                                          mlir::RewriterBase &rewriter,
+                                          mlir::Location loc) {
+  auto constVal = mlir::getConstantIntValue(tileSize);
+  if (constVal && *constVal < 0) {
+    return mlir::arith::ConstantOp::create(
+        rewriter, loc, rewriter.getI32Type(),
+        rewriter.getI32IntegerAttr(defaultTileSize));
+  }
+  return tileSize;
+}
+
+// Remove vector/worker attributes from loop
+static void removeWorkerVectorFromLoop(mlir::acc::LoopOp loop) {
+  if (loop.hasVector() || loop.getVectorValue()) {
+    loop.removeVectorAttr();
+    loop.removeVectorOperandsDeviceTypeAttr();
+  } else if (loop.hasWorker() || loop.getWorkerValue()) {
+    loop.removeWorkerAttr();
+    loop.removeWorkerNumOperandsDeviceTypeAttr();
+  }
+}
+
+// Create a new ACC loop with new steps, lb, ub from original loop
+static mlir::acc::LoopOp
+createACCLoopFromOriginal(mlir::acc::LoopOp origLoop,
+                          mlir::RewriterBase &rewriter, mlir::ValueRange lb,
+                          mlir::ValueRange ub, mlir::ValueRange step,
+                          mlir::DenseBoolArrayAttr inclusiveUBAttr,
+                          mlir::acc::CombinedConstructsTypeAttr combinedAttr,
+                          mlir::Location loc, bool preserveCollapse) {
+  mlir::ArrayAttr collapseAttr = mlir::ArrayAttr{};
+  mlir::ArrayAttr collapseDeviceTypeAttr = mlir::ArrayAttr{};
+  if (preserveCollapse) {
+    collapseAttr = origLoop.getCollapseAttr();
+    collapseDeviceTypeAttr = origLoop.getCollapseDeviceTypeAttr();
+  }
+  auto newLoop = mlir::acc::LoopOp::create(
+      rewriter, loc, origLoop->getResultTypes(), lb, ub, step, inclusiveUBAttr,
+      collapseAttr, collapseDeviceTypeAttr, origLoop.getGangOperands(),
+      origLoop.getGangOperandsArgTypeAttr(),
+      origLoop.getGangOperandsSegmentsAttr(),
+      origLoop.getGangOperandsDeviceTypeAttr(), origLoop.getWorkerNumOperands(),
+      origLoop.getWorkerNumOperandsDeviceTypeAttr(),
+      origLoop.getVectorOperands(), origLoop.getVectorOperandsDeviceTypeAttr(),
+      origLoop.getSeqAttr(), origLoop.getIndependentAttr(),
+      origLoop.getAuto_Attr(), origLoop.getGangAttr(), origLoop.getWorkerAttr(),
+      origLoop.getVectorAttr(), mlir::ValueRange{}, mlir::DenseI32ArrayAttr{},
+      mlir::ArrayAttr{}, origLoop.getCacheOperands(),
+      origLoop.getPrivateOperands(), origLoop.getFirstprivateOperands(),
+      origLoop.getReductionOperands(), combinedAttr);
+  return newLoop;
+}
+
+// Create inner loop inside input loop
+static mlir::acc::LoopOp
+createInnerLoop(mlir::acc::LoopOp inputLoop, mlir::RewriterBase &rewriter,
+                mlir::ValueRange lb, mlir::ValueRange ub, mlir::ValueRange step,
+                mlir::DenseBoolArrayAttr inclusiveUBAttr, mlir::Location loc) {
+  mlir::acc::LoopOp elementLoop = createACCLoopFromOriginal(
+      inputLoop, rewriter, lb, ub, step, inclusiveUBAttr,
+      mlir::acc::CombinedConstructsTypeAttr{}, loc, /*preserveCollapse*/ false);
+
+  // Remove gang/worker attributes from inner loops
+  rewriter.startOpModification(elementLoop);
+  if (inputLoop.hasGang() ||
+      inputLoop.getGangValue(mlir::acc::GangArgType::Num) ||
+      inputLoop.getGangValue(mlir::acc::GangArgType::Dim) ||
+      inputLoop.getGangValue(mlir::acc::GangArgType::Static)) {
+    elementLoop.removeGangAttr();
+    elementLoop.removeGangOperandsArgTypeAttr();
+    elementLoop.removeGangOperandsSegmentsAttr();
+    elementLoop.removeGangOperandsDeviceTypeAttr();
+  }
+  if (inputLoop.hasVector() || inputLoop.getVectorValue()) {
+    elementLoop.removeWorkerAttr();
+    elementLoop.removeWorkerNumOperandsDeviceTypeAttr();
+  }
+  rewriter.finalizeOpModification(elementLoop);
+
+  // Create empty block in elementLoop and add IV argument
+  mlir::Block *blk = rewriter.createBlock(&elementLoop.getRegion(),
+                                          elementLoop.getRegion().begin());
+  rewriter.setInsertionPointToEnd(blk);
+  mlir::acc::YieldOp::create(rewriter, loc);
+  elementLoop.getBody().addArgument(
+      inputLoop.getBody().getArgument(0).getType(), loc);
+
+  return elementLoop;
+}
+
+// Move ops from source to target Loop and replace uses of IVs
+static void moveOpsAndReplaceIVs(mlir::acc::LoopOp sourceLoop,
+                                 mlir::acc::LoopOp targetLoop,
+                                 llvm::ArrayRef<mlir::Value> newIVs,
+                                 llvm::ArrayRef<mlir::Value> origIVs,
+                                 size_t nOps, mlir::RewriterBase &rewriter) {
+  // Move ops from source to target loop [begin, begin + nOps - 1)
+  mlir::Block::iterator begin = sourceLoop.getBody().begin();
+  targetLoop.getBody().getOperations().splice(
+      targetLoop.getBody().getOperations().begin(),
+      sourceLoop.getBody().getOperations(), begin, std::next(begin, nOps - 1));
+
+  // Replace uses of origIV with newIV
+  for (auto [i, newIV] : llvm::enumerate(newIVs))
+    mlir::replaceAllUsesInRegionWith(origIVs[i], newIV, targetLoop.getRegion());
+}
+
+mlir::acc::LoopOp
+mlir::acc::tileACCLoops(llvm::SmallVector<mlir::acc::LoopOp> &tileLoops,
+                        const llvm::SmallVector<mlir::Value> &tileSizes,
+                        int32_t defaultTileSize, mlir::RewriterBase &rewriter) {
+  // Tile collapsed and/or nested loops
+  mlir::acc::LoopOp outerLoop = tileLoops[0];
+  const mlir::Location loc = outerLoop.getLoc();
+
+  // Resolve unknown tile sizes (tile(*) represented as -1)
+  llvm::SmallVector<mlir::Value> resolvedTileSizes;
+  rewriter.setInsertionPoint(outerLoop);
+  for (mlir::Value tileSize : tileSizes) {
+    resolvedTileSizes.push_back(
+        resolveUnknownTileSize(tileSize, defaultTileSize, rewriter, loc));
+  }
----------------
razvanlupusoru wrote:

Done

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