TheRock/compiler/amd-llvm/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsTiling.cpp

//===- 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/Arith/Utils/Utils.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.
// Returns a value with the same type as targetType.
static mlir::Value resolveAndCastTileSize(mlir::Value tileSize,
                                          int32_t defaultTileSize,
                                          mlir::Type targetType,
                                          mlir::RewriterBase &rewriter,
                                          mlir::Location loc) {
  auto constVal = mlir::getConstantIntValue(tileSize);
  if (constVal && *constVal < 0) {
    // Create constant with the target type directly
    return mlir::arith::ConstantOp::create(
        rewriter, loc, targetType,
        rewriter.getIntegerAttr(targetType, defaultTileSize));
  }
  return mlir::getValueOrCreateCastToIndexLike(rewriter, loc, targetType,
                                               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) {
  // nOps includes the terminator; move all ops except the terminator:
  // [begin, begin + nOps - 1)
  mlir::Block::iterator begin = sourceLoop.getBody().begin();
  mlir::Block::iterator end = std::next(begin, nOps - 1);

  // Notify the rewriter about all ops being moved (and their nested ops).
  // Directly moved ops have their parent block changed (rewriter fingerprint
  // tracking invalidated). Nested ops may have operands replaced by
  // replaceAllUsesInRegionWith below.
  llvm::SmallVector<mlir::Operation *> movedOps;
  for (mlir::Block::iterator it = begin; it != end; ++it)
    it->walk([&](mlir::Operation *op) {
      movedOps.push_back(op);
      rewriter.startOpModification(op);
    });

  targetLoop.getBody().getOperations().splice(
      targetLoop.getBody().getOperations().begin(),
      sourceLoop.getBody().getOperations(), begin, end);

  // Replace uses of origIV with newIV
  for (auto [i, newIV] : llvm::enumerate(newIVs))
    mlir::replaceAllUsesInRegionWith(origIVs[i], newIV, targetLoop.getRegion());

  for (mlir::Operation *op : movedOps)
    rewriter.finalizeOpModification(op);
}

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();

  mlir::acc::LoopOp innerLoop = tileLoops[tileLoops.size() - 1];
  llvm::SmallVector<mlir::Value, 3> origIVs;
  llvm::SmallVector<mlir::Value, 3> origSteps;
  llvm::SmallVector<mlir::Value, 3> origUBs;
  llvm::SmallVector<mlir::Value, 3> newSteps;
  llvm::SmallVector<mlir::Value, 3> newUBs;
  llvm::SmallVector<mlir::Value, 3> newIVs;
  size_t nOps = innerLoop.getBody().getOperations().size();

  // Extract original inclusiveUBs
  llvm::SmallVector<bool> inclusiveUBs;
  for (auto tileLoop : tileLoops) {
    for (auto [j, step] : llvm::enumerate(tileLoop.getStep())) {
      // inclusiveUBs are present on the IR from Fortran frontend for DO loops
      // but might not be present from other frontends (python)
      // So check if it exists
      if (tileLoop.getInclusiveUpperboundAttr())
        inclusiveUBs.push_back(
            tileLoop.getInclusiveUpperboundAttr().asArrayRef()[j]);
      else
        inclusiveUBs.push_back(false);
    }
  }

  // Extract original ivs, UBs, steps, and calculate new steps
  rewriter.setInsertionPoint(outerLoop);
  for (auto [i, tileLoop] : llvm::enumerate(tileLoops)) {
    for (auto arg : tileLoop.getBody().getArguments())
      origIVs.push_back(arg);
    for (auto ub : tileLoop.getUpperbound())
      origUBs.push_back(ub);

    llvm::SmallVector<mlir::Value, 3> currentLoopSteps;
    for (auto [j, step] : llvm::enumerate(tileLoop.getStep())) {
      origSteps.push_back(step);
      if (i + j >= tileSizes.size()) {
        currentLoopSteps.push_back(step);
      } else {
        mlir::Value tileSize = resolveAndCastTileSize(
            tileSizes[i + j], defaultTileSize, step.getType(), rewriter, loc);
        auto newLoopStep =
            mlir::arith::MulIOp::create(rewriter, loc, step, tileSize);
        currentLoopSteps.push_back(newLoopStep);
        newSteps.push_back(newLoopStep);
      }
    }

    rewriter.startOpModification(tileLoop);
    tileLoop.getStepMutable().clear();
    tileLoop.getStepMutable().append(currentLoopSteps);
    rewriter.finalizeOpModification(tileLoop);
  }

  // Calculate new upper bounds for element loops
  for (size_t i = 0; i < newSteps.size(); i++) {
    rewriter.setInsertionPoint(innerLoop.getBody().getTerminator());
    // UpperBound: min(origUB, origIV+(originalStep*tile_size))
    auto stepped =
        mlir::arith::AddIOp::create(rewriter, loc, origIVs[i], newSteps[i]);
    mlir::Value newUB = stepped;
    if (inclusiveUBs[i]) {
      // Handle InclusiveUB
      // UpperBound: min(origUB, origIV+(originalStep*tile_size - 1))
      auto c1 = mlir::arith::ConstantOp::create(
          rewriter, loc, newSteps[i].getType(),
          rewriter.getIntegerAttr(newSteps[i].getType(), 1));
      newUB = mlir::arith::SubIOp::create(rewriter, loc, stepped, c1);
    }
    newUBs.push_back(
        mlir::arith::MinSIOp::create(rewriter, loc, origUBs[i], newUB));
  }

  // Create and insert nested elementLoopOps before terminator of outer loopOp
  mlir::acc::LoopOp currentLoop = innerLoop;
  for (size_t i = 0; i < tileSizes.size(); i++) {
    rewriter.setInsertionPoint(currentLoop.getBody().getTerminator());
    mlir::DenseBoolArrayAttr inclusiveUBAttr = mlir::DenseBoolArrayAttr{};
    if (inclusiveUBs[i])
      inclusiveUBAttr = rewriter.getDenseBoolArrayAttr({true});

    mlir::acc::LoopOp elementLoop =
        createInnerLoop(innerLoop, rewriter, mlir::ValueRange{origIVs[i]},
                        mlir::ValueRange{newUBs[i]},
                        mlir::ValueRange{origSteps[i]}, inclusiveUBAttr, loc);

    // Remove vector/worker attributes from inner element loops except
    // outermost element loop
    if (i > 0) {
      rewriter.startOpModification(elementLoop);
      removeWorkerVectorFromLoop(elementLoop);
      rewriter.finalizeOpModification(elementLoop);
    }
    newIVs.push_back(elementLoop.getBody().getArgument(0));
    currentLoop = elementLoop;
  }

  // Remove vector/worker attributes from outer tile loops
  for (auto tileLoop : tileLoops) {
    rewriter.startOpModification(tileLoop);
    removeWorkerVectorFromLoop(tileLoop);
    rewriter.finalizeOpModification(tileLoop);
  }

  // Move ops from inner tile loop to inner element loop and replace IV uses
  moveOpsAndReplaceIVs(innerLoop, currentLoop, newIVs, origIVs, nOps, rewriter);

  return outerLoop;
}

llvm::SmallVector<mlir::acc::LoopOp>
mlir::acc::uncollapseLoops(mlir::acc::LoopOp origLoop, unsigned tileCount,
                           unsigned collapseCount,
                           mlir::RewriterBase &rewriter) {
  llvm::SmallVector<mlir::acc::LoopOp> newLoops;
  llvm::SmallVector<mlir::Value, 3> newIVs;
  mlir::Location loc = origLoop.getLoc();
  llvm::SmallVector<bool> newInclusiveUBs;
  llvm::SmallVector<mlir::Value, 3> lbs, ubs, steps;
  for (unsigned i = 0; i < collapseCount; i++) {
    // inclusiveUpperbound attribute might not be set, default to false
    bool inclusiveUB = false;
    if (origLoop.getInclusiveUpperboundAttr())
      inclusiveUB = origLoop.getInclusiveUpperboundAttr().asArrayRef()[i];
    newInclusiveUBs.push_back(inclusiveUB);
    lbs.push_back(origLoop.getLowerbound()[i]);
    ubs.push_back(origLoop.getUpperbound()[i]);
    steps.push_back(origLoop.getStep()[i]);
  }
  mlir::acc::LoopOp outerLoop = createACCLoopFromOriginal(
      origLoop, rewriter, lbs, ubs, steps,
      rewriter.getDenseBoolArrayAttr(newInclusiveUBs),
      origLoop.getCombinedAttr(), loc, /*preserveCollapse*/ true);
  mlir::Block *blk = rewriter.createBlock(&outerLoop.getRegion(),
                                          outerLoop.getRegion().begin());
  rewriter.setInsertionPointToEnd(blk);
  mlir::acc::YieldOp::create(rewriter, loc);
  for (unsigned i = 0; i < collapseCount; i++) {
    outerLoop.getBody().addArgument(origLoop.getBody().getArgument(i).getType(),
                                    loc);
    newIVs.push_back(outerLoop.getBody().getArgument(i));
  }
  newLoops.push_back(outerLoop);

  mlir::acc::LoopOp currentLoopOp = outerLoop;
  for (unsigned i = collapseCount; i < tileCount; i++) {
    rewriter.setInsertionPoint(currentLoopOp.getBody().getTerminator());
    bool inclusiveUB = false;
    if (origLoop.getInclusiveUpperboundAttr())
      inclusiveUB = origLoop.getInclusiveUpperboundAttr().asArrayRef()[i];
    mlir::DenseBoolArrayAttr inclusiveUBAttr =
        rewriter.getDenseBoolArrayAttr({inclusiveUB});
    mlir::acc::LoopOp innerLoop = createInnerLoop(
        origLoop, rewriter, mlir::ValueRange{origLoop.getLowerbound()[i]},
        mlir::ValueRange{origLoop.getUpperbound()[i]},
        mlir::ValueRange{origLoop.getStep()[i]}, inclusiveUBAttr, loc);
    newIVs.push_back(innerLoop.getBody().getArgument(0));
    newLoops.push_back(innerLoop);
    currentLoopOp = innerLoop;
  }
  // Move ops from origLoop to innermost loop and replace uses of IVs
  size_t nOps = origLoop.getBody().getOperations().size();
  llvm::SmallVector<mlir::Value, 3> origIVs;
  for (auto arg : origLoop.getBody().getArguments())
    origIVs.push_back(arg);
  moveOpsAndReplaceIVs(origLoop, currentLoopOp, newIVs, origIVs, nOps,
                       rewriter);

  return newLoops;
}