text: //===- LoopTiling.cpp --- Loop tiling pass ------------------------------*-===//
  //
  // Copyright 2019 The MLIR Authors.
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //   http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  // =============================================================================
  //
  // This file implements a pass to tile loop nests.
  //
  //===----------------------------------------------------------------------===//
  
  #include "mlir/AffineOps/AffineOps.h"
  #include "mlir/Analysis/AffineAnalysis.h"
  #include "mlir/Analysis/AffineStructures.h"
  #include "mlir/Analysis/LoopAnalysis.h"
  #include "mlir/Analysis/Utils.h"
  #include "mlir/IR/Builders.h"
  #include "mlir/Pass/Pass.h"
  #include "mlir/Transforms/LoopUtils.h"
  #include "mlir/Transforms/Passes.h"
  #include "mlir/Transforms/Utils.h"
  #include "llvm/Support/CommandLine.h"
  #include "llvm/Support/Debug.h"
  #include <iomanip>
  #include <sstream>
  
  using namespace mlir;
  
  #define DEBUG_TYPE "loop-tile"
  
  static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
  
  static llvm::cl::opt<unsigned long long>
      clCacheSizeKiB("tile-cache-size",
                     llvm::cl::desc("Set size of cache to tile for in KiB"),
                     llvm::cl::cat(clOptionsCategory));
  
  // Tile size to use for all loops (overrides -tile-sizes if provided).
  static llvm::cl::opt<unsigned>
      clTileSize("tile-size", llvm::cl::desc("Use this tile size for all loops"),
                 llvm::cl::cat(clOptionsCategory));
  
  // List of tile sizes. If any of them aren't provided, they are filled with
  // clTileSize / kDefaultTileSize.
  static llvm::cl::list<unsigned> clTileSizes(
      "tile-sizes",
      llvm::cl::desc(
          "List of tile sizes for each perfect nest (overridden by -tile-size)"),
      llvm::cl::ZeroOrMore, llvm::cl::cat(clOptionsCategory));
  
  namespace {
  
  /// A pass to perform loop tiling on all suitable loop nests of a Function.
  struct LoopTiling : public FunctionPass<LoopTiling> {
    explicit LoopTiling(uint64_t cacheSizeBytes = kDefaultCacheMemCapacity,
                        bool avoidMaxMinBounds = true)
        : cacheSizeBytes(cacheSizeBytes), avoidMaxMinBounds(avoidMaxMinBounds) {}
  
    void runOnFunction() override;
    void getTileSizes(ArrayRef<AffineForOp> band,
                      SmallVectorImpl<unsigned> *tileSizes);
  
    // Default tile size if nothing is provided.
    constexpr static unsigned kDefaultTileSize = 4;
    constexpr static uint64_t kDefaultCacheMemCapacity = 512 * 1024UL;
  
    // Capacity of the cache to tile for.
    uint64_t cacheSizeBytes;
    // If true, tile sizes are set to avoid max/min in bounds if possible.
    bool avoidMaxMinBounds;
  };
  
  } // end anonymous namespace
  
  /// Creates a pass to perform loop tiling on all suitable loop nests of a
  /// Function.
  FunctionPassBase *mlir::createLoopTilingPass(uint64_t cacheSizeBytes) {
    return new LoopTiling(cacheSizeBytes);
  }
  
  // Move the loop body of AffineForOp 'src' from 'src' into the specified
  // location in destination's body, ignoring the terminator.
  static inline void moveLoopBody(AffineForOp src, AffineForOp dest,
                                  Block::iterator loc) {
    auto &insts = src.getBody()->getOperations();
    dest.getBody()->getOperations().splice(loc, insts, insts.begin(),
                                           std::prev(insts.end()));
  }
  
  // Move the loop body of AffineForOp 'src' from 'src' to the start of dest's
  // body.
  static inline void moveLoopBody(AffineForOp src, AffineForOp dest) {
    moveLoopBody(src, dest, dest.getBody()->begin());
  }
  
  /// Constructs and sets new loop bounds after tiling for the case of
  /// hyper-rectangular index sets, where the bounds of one dimension do not
  /// depend on other dimensions. Bounds of each dimension can thus be treated
  /// independently, and deriving the new bounds is much simpler and faster
  /// than for the case of tiling arbitrary polyhedral shapes.
  static void
  constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
                                  MutableArrayRef<AffineForOp> newLoops,
                                  ArrayRef<unsigned> tileSizes) {
    assert(!origLoops.empty());
    assert(origLoops.size() == tileSizes.size());
  
    FuncBuilder b(origLoops[0].getOperation());
    unsigned width = origLoops.size();
  
    // Bounds for tile space loops.
    for (unsigned i = 0; i < width; i++) {
      auto lbOperands = origLoops[i].getLowerBoundOperands();
      auto ubOperands = origLoops[i].getUpperBoundOperands();
      SmallVector<Value *, 4> newLbOperands(lbOperands);
      SmallVector<Value *, 4> newUbOperands(ubOperands);
      newLoops[i].setLowerBound(newLbOperands, origLoops[i].getLowerBoundMap());
      newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap());
      newLoops[i].setStep(tileSizes[i]);
    }
    // Bounds for intra-tile loops.
    for (unsigned i = 0; i < width; i++) {
      int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]);
      auto mayBeConstantCount = getConstantTripCount(origLoops[i]);
      // The lower bound is just the tile-space loop.
      AffineMap lbMap = b.getDimIdentityMap();
      newLoops[width + i].setLowerBound(
          /*operands=*/newLoops[i].getInductionVar(), lbMap);
  
      // Set the upper bound.
      if (mayBeConstantCount.hasValue() &&
          mayBeConstantCount.getValue() < tileSizes[i]) {
        // Trip count is less than tile size; upper bound is the trip count.
        auto ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue());
        newLoops[width + i].setUpperBoundMap(ubMap);
      } else if (largestDiv % tileSizes[i] != 0) {
        // Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
        // Construct the upper bound map; the operands are the original operands
        // with 'i' (tile-space loop) appended to it. The new upper bound map is
        // the original one with an additional expression i + tileSize appended.
        SmallVector<Value *, 4> ubOperands(origLoops[i].getUpperBoundOperands());
        ubOperands.push_back(newLoops[i].getInductionVar());
  
        auto origUbMap = origLoops[i].getUpperBoundMap();
        SmallVector<AffineExpr, 4> boundExprs;
        boundExprs.reserve(1 + origUbMap.getNumResults());
        auto dim = b.getAffineDimExpr(origUbMap.getNumInputs());
        // The new upper bound map is the original one with an additional
        // expression i + tileSize appended.
        boundExprs.push_back(dim + tileSizes[i]);
        boundExprs.append(origUbMap.getResults().begin(),
                          origUbMap.getResults().end());
        auto ubMap =
            b.getAffineMap(origUbMap.getNumInputs() + 1, 0, boundExprs, {});
        newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
      } else {
        // No need of the min expression.
        auto dim = b.getAffineDimExpr(0);
        auto ubMap = b.getAffineMap(1, 0, dim + tileSizes[i], {});
        newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
      }
    }
  }
  
  /// Tiles the specified band of perfectly nested loops creating tile-space loops
  /// and intra-tile loops. A band is a contiguous set of loops.
  //  TODO(bondhugula): handle non hyper-rectangular spaces.
  LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band,
                                  ArrayRef<unsigned> tileSizes) {
    assert(!band.empty());
    assert(band.size() == tileSizes.size() && "Incorrect number of tile sizes");
  
    // Check if the supplied for op's are all successively nested.
    for (unsigned i = 1, e = band.size(); i < e; i++) {
      assert(band[i].getOperation()->getParentOp() == band[i - 1].getOperation());
    }
  
    auto origLoops = band;
  
    AffineForOp rootAffineForOp = origLoops[0];
    auto loc = rootAffineForOp.getLoc();
    // Note that width is at least one since band isn't empty.
    unsigned width = band.size();
  
    SmallVector<AffineForOp, 12> newLoops(2 * width);
    AffineForOp innermostPointLoop;
  
    // The outermost among the loops as we add more..
    auto *topLoop = rootAffineForOp.getOperation();
  
    // Add intra-tile (or point) loops.
    for (unsigned i = 0; i < width; i++) {
      FuncBuilder b(topLoop);
      // Loop bounds will be set later.
      auto pointLoop = b.create<AffineForOp>(loc, 0, 0);
      pointLoop.getBody()->getOperations().splice(
          pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
          topLoop);
      newLoops[2 * width - 1 - i] = pointLoop;
      topLoop = pointLoop.getOperation();
      if (i == 0)
        innermostPointLoop = pointLoop;
    }
  
    // Add tile space loops;
    for (unsigned i = width; i < 2 * width; i++) {
      FuncBuilder b(topLoop);
      // Loop bounds will be set later.
      auto tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
      tileSpaceLoop.getBody()->getOperations().splice(
          tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
          topLoop);
      newLoops[2 * width - i - 1] = tileSpaceLoop;
      topLoop = tileSpaceLoop.getOperation();
    }
  
    // Move the loop body of the original nest to the new one.
    moveLoopBody(origLoops[origLoops.size() - 1], innermostPointLoop);
  
    SmallVector<Value *, 8> origLoopIVs;
    extractForInductionVars(band, &origLoopIVs);
    SmallVector<Optional<Value *>, 6> ids(origLoopIVs.begin(), origLoopIVs.end());
    FlatAffineConstraints cst;
    getIndexSet(band, &cst);
  
    if (!cst.isHyperRectangular(0, width)) {
      rootAffineForOp.emitError("tiled code generation unimplemented for the "
                                "non-hyperrectangular case");
      return failure();
    }
  
    constructTiledIndexSetHyperRect(origLoops, newLoops, tileSizes);
    // In this case, the point loop IVs just replace the original ones.
    for (unsigned i = 0; i < width; i++) {
      origLoopIVs[i]->replaceAllUsesWith(newLoops[i + width].getInductionVar());
    }
  
    // Erase the old loop nest.
    rootAffineForOp.erase();
  
    return success();
  }
  
  // Identify valid and profitable bands of loops to tile. This is currently just
  // a temporary placeholder to test the mechanics of tiled code generation.
  // Returns all maximal outermost perfect loop nests to tile.
  static void getTileableBands(Function &f,
                               std::vector<SmallVector<AffineForOp, 6>> *bands) {
    // Get maximal perfect nest of 'affine.for' insts starting from root
    // (inclusive).
    auto getMaximalPerfectLoopNest = [&](AffineForOp root) {
      SmallVector<AffineForOp, 6> band;
      AffineForOp currInst = root;
      do {
        band.push_back(currInst);
      } while (currInst.getBody()->getOperations().size() == 2 &&
               (currInst = currInst.getBody()->front().dyn_cast<AffineForOp>()));
      bands->push_back(band);
    };
  
    for (auto &block : f)
      for (auto &op : block)
        if (auto forOp = op.dyn_cast<AffineForOp>())
          getMaximalPerfectLoopNest(forOp);
  }
  
  // Reduce each tile size to the largest divisor of the corresponding trip count
  // (if the trip count is known).
  static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band,
                                           SmallVectorImpl<unsigned> *tileSizes) {
    assert(band.size() == tileSizes->size() && "invalid tile size count");
    for (unsigned i = 0, e = band.size(); i < e; i++) {
      unsigned &tSizeAdjusted = (*tileSizes)[i];
      auto mayConst = getConstantTripCount(band[i]);
      if (!mayConst.hasValue())
        continue;
      // Adjust the tile size to largest factor of the trip count less than
      // tSize.
      uint64_t constTripCount = mayConst.getValue();
      if (tSizeAdjusted > constTripCount / 2)
        tSizeAdjusted = constTripCount / 2;
      while (constTripCount % tSizeAdjusted != 0)
        tSizeAdjusted--;
    }
  }
  
  // Returns tile sizes to use. Checks CL options; if none are specified, sets it
  // based on a simple model that looks at the memory footprint and determines
  // tile sizes assuming identity accesses / 1:1 tile size proportional footprint
  // along each of the dimensions being tiled.
  // TODO(mlir-team): evolve this model. Tile size determination is a large area
  // to play with in general.
  void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
                                SmallVectorImpl<unsigned> *tileSizes) {
    if (band.empty())
      return;
  
    tileSizes->resize(band.size());
  
    // Use clTileSize for all loops if specified.
    if (clTileSize.getNumOccurrences() > 0) {
      std::fill(tileSizes->begin(), tileSizes->end(), clTileSize);
      return;
    }
  
    // Use clTileSizes and fill them with default tile size if it's short.
    if (!clTileSizes.empty()) {
      std::fill(tileSizes->begin(), tileSizes->end(),
                LoopTiling::kDefaultTileSize);
      std::copy(clTileSizes.begin(),
                clTileSizes.begin() + std::min(clTileSizes.size(), band.size()),
                tileSizes->begin());
      return;
    }
  
    // The first loop in the band.
    auto rootForOp = band[0];
    (void)rootForOp;
  
    // Obtain memory footprint and set tile sizes so that a tile fits in
    // the cache size. This is an approximation with the assumption that the
    // footprint increases with the tile size linearly in that dimension (i.e.,
    // assumes one-to-one access function).
    auto fp = getMemoryFootprintBytes(band[0], 0);
    if (!fp.hasValue()) {
      // Fill with default tile sizes if footprint is unknown.
      std::fill(tileSizes->begin(), tileSizes->end(),
                LoopTiling::kDefaultTileSize);
      if (avoidMaxMinBounds)
        adjustToDivisorsOfTripCounts(band, tileSizes);
      LLVM_DEBUG(
          rootForOp.emitWarning("memory footprint unknown: using default tile "
                                "sizes adjusted to trip count divisors"));
      return;
    }
  
    // Check how many times larger the cache size is when compared to footprint.
    uint64_t excessFactor = llvm::divideCeil(fp.getValue(), cacheSizeBytes);
    if (excessFactor <= 1) {
      // No need of any tiling - set tile size to 1.
      std::fill(tileSizes->begin(), tileSizes->end(), 1);
      return;
    }
  
    // Divide all loops equally in an attempt to reduce footprint.
    // TODO(bondhugula): this is approximate. Ideally, obtain reuse factor /
    // profitability along each dimension and weight tile sizes based on that as
    // one possible approach. Or compute a polynomial in tile sizes and solve for
    // it.
  
    // For an n-d tilable band, compute n^th root of the excess.
    unsigned tSize =
        static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size())));
    // We'll keep a running product to determine the last tile size better.
    unsigned cumulProductOfTileSizes = 1;
    for (unsigned i = 0, e = band.size(); i < e; i++) {
      if (i < e - 1)
        (*tileSizes)[i] = tSize;
      else
        // Set last tile size to cover the balance.
        (*tileSizes)[i] = std::max(
            1U, static_cast<unsigned>(excessFactor / cumulProductOfTileSizes));
      cumulProductOfTileSizes *= (*tileSizes)[i];
    }
    if (avoidMaxMinBounds)
      adjustToDivisorsOfTripCounts(band, tileSizes);
  }
  
  void LoopTiling::runOnFunction() {
    // Override cache size if provided on command line.
    if (clCacheSizeKiB.getNumOccurrences() > 0)
      cacheSizeBytes = clCacheSizeKiB * 1024;
  
    // Bands of loops to tile.
    std::vector<SmallVector<AffineForOp, 6>> bands;
    getTileableBands(getFunction(), &bands);
  
    for (auto &band : bands) {
      // Set up tile sizes; fill missing tile sizes at the end with default tile
      // size or clTileSize if one was provided.
      SmallVector<unsigned, 6> tileSizes;
      getTileSizes(band, &tileSizes);
      if (llvm::DebugFlag) {
        std::stringstream msg;
        msg << "using tile sizes [";
        for (auto tSize : tileSizes)
          msg << tSize << " ";
        msg << "]\n";
        auto rootForOp = band[0];
        rootForOp.emitNote(msg.str());
      }
      if (failed(tileCodeGen(band, tileSizes)))
        return signalPassFailure();
    }
  }
  
  constexpr unsigned LoopTiling::kDefaultTileSize;
  constexpr uint64_t LoopTiling::kDefaultCacheMemCapacity;
  
  static PassRegistration<LoopTiling> pass("loop-tile", "Tile loop nests");
  