OpenACCUtilsTiling.cpp

Go to the documentation of this file.

//===- 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"
#include "llvm/ADT/STLExtras.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) {
  mlir::ArrayAttr collapseAttr = mlir::ArrayAttr{};
  mlir::ArrayAttr collapseDeviceTypeAttr = mlir::ArrayAttr{};
  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;
}
 
// 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);
}
 
// Create a single "element group" loop nested in `tileLoop`, carrying
// `ivTypes.size()` induction variables so the whole element space is one
// multi-IV loop. The element group carries vector or worker but not gang.
static mlir::acc::LoopOp
createElementGroupLoop(mlir::acc::LoopOp tileLoop, mlir::RewriterBase &rewriter,
                       mlir::ValueRange lbs, mlir::ValueRange ubs,
                       mlir::ValueRange steps,
                       mlir::DenseBoolArrayAttr inclusiveUBAttr,
                       llvm::ArrayRef<mlir::Type> ivTypes, mlir::Location loc) {
  mlir::acc::LoopOp elementLoop = createACCLoopFromOriginal(
      tileLoop, rewriter, lbs, ubs, steps, inclusiveUBAttr,
      mlir::acc::CombinedConstructsTypeAttr{}, loc);
 
  // Drop gang from the element group, keeping vector/worker. The operand
  // values must be cleared too, not just the attributes.
  rewriter.startOpModification(elementLoop);
  if (tileLoop.hasGang() ||
      tileLoop.getGangValue(mlir::acc::GangArgType::Num) ||
      tileLoop.getGangValue(mlir::acc::GangArgType::Dim) ||
      tileLoop.getGangValue(mlir::acc::GangArgType::Static)) {
    elementLoop.removeGangAttr();
    elementLoop.removeGangOperandsArgTypeAttr();
    elementLoop.removeGangOperandsSegmentsAttr();
    elementLoop.removeGangOperandsDeviceTypeAttr();
    elementLoop.getGangOperandsMutable().clear();
  }
  if (tileLoop.hasVector() || tileLoop.getVectorValue()) {
    elementLoop.removeWorkerAttr();
    elementLoop.removeWorkerNumOperandsDeviceTypeAttr();
    elementLoop.getWorkerNumOperandsMutable().clear();
  }
  rewriter.finalizeOpModification(elementLoop);
 
  // Create the element loop body: one block argument per IV plus a terminator.
  mlir::Block *blk = rewriter.createBlock(&elementLoop.getRegion(),
                                          elementLoop.getRegion().begin());
  rewriter.setInsertionPointToEnd(blk);
  mlir::acc::YieldOp::create(rewriter, loc);
  for (mlir::Type ivType : ivTypes)
    elementLoop.getBody().addArgument(ivType, loc);
 
  return elementLoop;
}
 
mlir::acc::LoopOp
mlir::acc::tileACCLoops(mlir::acc::LoopOp tileLoop,
                        const llvm::SmallVector<mlir::Value> &tileSizes,
                        int32_t defaultTileSize, mlir::RewriterBase &rewriter) {
  // Tile a single fused acc.loop that carries all associated induction
  // variables. This keeps the tile iterations as one multi-IV "tile group"
  // loop and the in-tile iterations as one multi-IV "element group" loop, each
  // spanning all of its induction variables.
  const mlir::Location loc = tileLoop.getLoc();
  const unsigned tileCount = tileSizes.size();
 
  llvm::SmallVector<mlir::Value, 3> origIVs(tileLoop.getBody().getArguments());
  llvm::SmallVector<mlir::Value, 3> origUBs(tileLoop.getUpperbound());
  llvm::SmallVector<mlir::Value, 3> origSteps(tileLoop.getStep());
  const unsigned numIVs = origIVs.size();
  const size_t nOps = tileLoop.getBody().getOperations().size();
 
  // Original inclusive-UB flags (default false when the attribute is absent).
  llvm::SmallVector<bool> inclusiveUBs;
  for (unsigned i = 0; i < numIVs; ++i) {
    if (tileLoop.getInclusiveUpperboundAttr())
      inclusiveUBs.push_back(
          tileLoop.getInclusiveUpperboundAttr().asArrayRef()[i]);
    else
      inclusiveUBs.push_back(false);
  }
 
  // Scale each tiled dimension's step by its tile size to form the tile group
  // loop steps.
  rewriter.setInsertionPoint(tileLoop);
  llvm::SmallVector<mlir::Value, 3> scaledSteps;
  llvm::SmallVector<mlir::Value, 3> tileLoopSteps;
  for (unsigned i = 0; i < numIVs; ++i) {
    if (i < tileCount) {
      mlir::Value tileSize = resolveAndCastTileSize(
          tileSizes[i], defaultTileSize, origSteps[i].getType(), rewriter, loc);
      mlir::Value scaled =
          mlir::arith::MulIOp::create(rewriter, loc, origSteps[i], tileSize);
      scaledSteps.push_back(scaled);
      tileLoopSteps.push_back(scaled);
    } else {
      tileLoopSteps.push_back(origSteps[i]);
    }
  }
 
  // Compute each element-loop upper bound by clamping the tile's far edge,
  // origIV + scaledStep, to origUB. An ascending loop takes
  // min(origUB, edge - 1), a descending one max(origUB, edge + 1). The -1 and
  // +1 apply in the presence of inclusive bounds and are omitted otherwise.
  rewriter.setInsertionPoint(tileLoop.getBody().getTerminator());
  llvm::SmallVector<mlir::Value, 3> elemLBs, elemUBs, elemSteps;
  llvm::SmallVector<mlir::Type, 3> elemIVTypes;
  llvm::SmallVector<bool> elemInclusiveUBs;
  for (unsigned i = 0; i < tileCount; ++i) {
    mlir::Value stepped =
        mlir::arith::AddIOp::create(rewriter, loc, origIVs[i], scaledSteps[i]);
    mlir::Type stepTy = scaledSteps[i].getType();
    std::optional<int64_t> constStep = mlir::getConstantIntValue(origSteps[i]);
 
    // The inclusive edge is one iteration short of the next tile's start, in
    // whichever direction the loop runs.
    auto inclusiveEdge = [&](bool descending) -> mlir::Value {
      if (!inclusiveUBs[i])
        return stepped;
      mlir::Value c1 = mlir::arith::ConstantOp::create(
          rewriter, loc, stepTy, rewriter.getIntegerAttr(stepTy, 1));
      if (descending)
        return mlir::arith::AddIOp::create(rewriter, loc, stepped, c1);
      return mlir::arith::SubIOp::create(rewriter, loc, stepped, c1);
    };
 
    mlir::Value elemUB;
    if (constStep && *constStep < 0) {
      elemUB = mlir::arith::MaxSIOp::create(rewriter, loc, origUBs[i],
                                            inclusiveEdge(/*descending=*/true));
    } else if (constStep) {
      elemUB = mlir::arith::MinSIOp::create(
          rewriter, loc, origUBs[i], inclusiveEdge(/*descending=*/false));
    } else {
      // The sign is only known at run time, so compute both and select.
      mlir::Type origStepTy = origSteps[i].getType();
      mlir::Value zero = mlir::arith::ConstantOp::create(
          rewriter, loc, origStepTy, rewriter.getIntegerAttr(origStepTy, 0));
      mlir::Value isDescending = mlir::arith::CmpIOp::create(
          rewriter, loc, mlir::arith::CmpIPredicate::slt, origSteps[i], zero);
      mlir::Value descUB = mlir::arith::MaxSIOp::create(
          rewriter, loc, origUBs[i], inclusiveEdge(/*descending=*/true));
      mlir::Value ascUB = mlir::arith::MinSIOp::create(
          rewriter, loc, origUBs[i], inclusiveEdge(/*descending=*/false));
      elemUB = mlir::arith::SelectOp::create(rewriter, loc, isDescending,
                                             descUB, ascUB);
    }
    elemUBs.push_back(elemUB);
    elemLBs.push_back(origIVs[i]);
    elemSteps.push_back(origSteps[i]);
    elemIVTypes.push_back(origIVs[i].getType());
    elemInclusiveUBs.push_back(inclusiveUBs[i]);
  }
 
  // Only attach an inclusiveUpperbound attribute if at least one element
  // dimension is inclusive.
  mlir::DenseBoolArrayAttr elemInclAttr = mlir::DenseBoolArrayAttr{};
  if (llvm::is_contained(elemInclusiveUBs, true))
    elemInclAttr = rewriter.getDenseBoolArrayAttr(elemInclusiveUBs);
 
  // Create the element group loop from the unmodified tile loop.
  mlir::acc::LoopOp elementLoop =
      createElementGroupLoop(tileLoop, rewriter, elemLBs, elemUBs, elemSteps,
                             elemInclAttr, elemIVTypes, loc);
 
  // Move the original body into the element loop and remap the tiled IVs to the
  // element IVs.
  llvm::SmallVector<mlir::Value, 3> newIVs(
      elementLoop.getBody().getArguments());
  llvm::SmallVector<mlir::Value, 3> tiledOrigIVs(origIVs.begin(),
                                                 origIVs.begin() + tileCount);
  moveOpsAndReplaceIVs(tileLoop, elementLoop, newIVs, tiledOrigIVs, nOps,
                       rewriter);
 
  // Turn the fused loop into the tile group: scaled steps, gang only.
  rewriter.startOpModification(tileLoop);
  tileLoop.getStepMutable().clear();
  tileLoop.getStepMutable().append(tileLoopSteps);
  removeWorkerVectorFromLoop(tileLoop);
  rewriter.finalizeOpModification(tileLoop);
 
  return tileLoop;
}