text: //===- LoopAnalysis.cpp - Misc loop analysis routines //-------------------===//
  //
  // 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 implements miscellaneous loop analysis routines.
  //
  //===----------------------------------------------------------------------===//
  
  #include "mlir/Analysis/LoopAnalysis.h"
  
  #include "mlir/Analysis/AffineAnalysis.h"
  #include "mlir/Analysis/AffineStructures.h"
  #include "mlir/Analysis/NestedMatcher.h"
  #include "mlir/Dialect/Affine/IR/AffineOps.h"
  #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
  #include "mlir/Support/MathExtras.h"
  
  #include "llvm/ADT/DenseSet.h"
  #include "llvm/ADT/SmallString.h"
  #include <type_traits>
  
  using namespace mlir;
  
  /// Returns the trip count of the loop as an affine expression if the latter is
  /// expressible as an affine expression, and nullptr otherwise. The trip count
  /// expression is simplified before returning. This method only utilizes map
  /// composition to construct lower and upper bounds before computing the trip
  /// count expressions.
  void mlir::buildTripCountMapAndOperands(
      AffineForOp forOp, AffineMap *tripCountMap,
      SmallVectorImpl<Value> *tripCountOperands) {
    int64_t loopSpan;
  
    int64_t step = forOp.getStep();
    OpBuilder b(forOp.getOperation());
  
    if (forOp.hasConstantBounds()) {
      int64_t lb = forOp.getConstantLowerBound();
      int64_t ub = forOp.getConstantUpperBound();
      loopSpan = ub - lb;
      if (loopSpan < 0)
        loopSpan = 0;
      *tripCountMap = b.getConstantAffineMap(ceilDiv(loopSpan, step));
      tripCountOperands->clear();
      return;
    }
    auto lbMap = forOp.getLowerBoundMap();
    auto ubMap = forOp.getUpperBoundMap();
    if (lbMap.getNumResults() != 1) {
      *tripCountMap = AffineMap();
      return;
    }
  
    // Difference of each upper bound expression from the single lower bound
    // expression (divided by the step) provides the expressions for the trip
    // count map.
    AffineValueMap ubValueMap(ubMap, forOp.getUpperBoundOperands());
  
    SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),
                                           lbMap.getResult(0));
    auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),
                                     lbSplatExpr, b.getContext());
    AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());
  
    AffineValueMap tripCountValueMap;
    AffineValueMap::difference(ubValueMap, lbSplatValueMap, &tripCountValueMap);
    for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i)
      tripCountValueMap.setResult(i,
                                  tripCountValueMap.getResult(i).ceilDiv(step));
  
    *tripCountMap = tripCountValueMap.getAffineMap();
    tripCountOperands->assign(tripCountValueMap.getOperands().begin(),
                              tripCountValueMap.getOperands().end());
  }
  
  /// Returns the trip count of the loop if it's a constant, None otherwise. This
  /// method uses affine expression analysis (in turn using getTripCount) and is
  /// able to determine constant trip count in non-trivial cases.
  // FIXME(mlir-team): this is really relying on buildTripCountMapAndOperands;
  // being an analysis utility, it shouldn't. Replace with a version that just
  // works with analysis structures (FlatAffineConstraints) and thus doesn't
  // update the IR.
  Optional<uint64_t> mlir::getConstantTripCount(AffineForOp forOp) {
    SmallVector<Value, 4> operands;
    AffineMap map;
    buildTripCountMapAndOperands(forOp, &map, &operands);
  
    if (!map)
      return None;
  
    // Take the min if all trip counts are constant.
    Optional<uint64_t> tripCount;
    for (auto resultExpr : map.getResults()) {
      if (auto constExpr = resultExpr.dyn_cast<AffineConstantExpr>()) {
        if (tripCount.hasValue())
          tripCount = std::min(tripCount.getValue(),
                               static_cast<uint64_t>(constExpr.getValue()));
        else
          tripCount = constExpr.getValue();
      } else
        return None;
    }
    return tripCount;
  }
  
  /// Returns the greatest known integral divisor of the trip count. Affine
  /// expression analysis is used (indirectly through getTripCount), and
  /// this method is thus able to determine non-trivial divisors.
  uint64_t mlir::getLargestDivisorOfTripCount(AffineForOp forOp) {
    SmallVector<Value, 4> operands;
    AffineMap map;
    buildTripCountMapAndOperands(forOp, &map, &operands);
  
    if (!map)
      return 1;
  
    // The largest divisor of the trip count is the GCD of the individual largest
    // divisors.
    assert(map.getNumResults() >= 1 && "expected one or more results");
    Optional<uint64_t> gcd;
    for (auto resultExpr : map.getResults()) {
      uint64_t thisGcd;
      if (auto constExpr = resultExpr.dyn_cast<AffineConstantExpr>()) {
        uint64_t tripCount = constExpr.getValue();
        // 0 iteration loops (greatest divisor is 2^64 - 1).
        if (tripCount == 0)
          thisGcd = std::numeric_limits<uint64_t>::max();
        else
          // The greatest divisor is the trip count.
          thisGcd = tripCount;
      } else {
        // Trip count is not a known constant; return its largest known divisor.
        thisGcd = resultExpr.getLargestKnownDivisor();
      }
      if (gcd.hasValue())
        gcd = llvm::GreatestCommonDivisor64(gcd.getValue(), thisGcd);
      else
        gcd = thisGcd;
    }
    assert(gcd.hasValue() && "value expected per above logic");
    return gcd.getValue();
  }
  
  /// Given an induction variable `iv` of type AffineForOp and an access `index`
  /// of type index, returns `true` if `index` is independent of `iv` and
  /// false otherwise. The determination supports composition with at most one
  /// AffineApplyOp. The 'at most one AffineApplyOp' comes from the fact that
  /// the composition of AffineApplyOp needs to be canonicalized by construction
  /// to avoid writing code that composes arbitrary numbers of AffineApplyOps
  /// everywhere. To achieve this, at the very least, the compose-affine-apply
  /// pass must have been run.
  ///
  /// Prerequisites:
  ///   1. `iv` and `index` of the proper type;
  ///   2. at most one reachable AffineApplyOp from index;
  ///
  /// Returns false in cases with more than one AffineApplyOp, this is
  /// conservative.
  static bool isAccessIndexInvariant(Value iv, Value index) {
    assert(isForInductionVar(iv) && "iv must be a AffineForOp");
    assert(index.getType().isa<IndexType>() && "index must be of IndexType");
    SmallVector<Operation *, 4> affineApplyOps;
    getReachableAffineApplyOps({index}, affineApplyOps);
  
    if (affineApplyOps.empty()) {
      // Pointer equality test because of Value pointer semantics.
      return index != iv;
    }
  
    if (affineApplyOps.size() > 1) {
      affineApplyOps[0]->emitRemark(
          "CompositionAffineMapsPass must have been run: there should be at most "
          "one AffineApplyOp, returning false conservatively.");
      return false;
    }
  
    auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]);
    // We need yet another level of indirection because the `dim` index of the
    // access may not correspond to the `dim` index of composeOp.
    return !composeOp.getAffineValueMap().isFunctionOf(0, iv);
  }
  
  DenseSet<Value> mlir::getInvariantAccesses(Value iv, ArrayRef<Value> indices) {
    DenseSet<Value> res;
    for (unsigned idx = 0, n = indices.size(); idx < n; ++idx) {
      auto val = indices[idx];
      if (isAccessIndexInvariant(iv, val)) {
        res.insert(val);
      }
    }
    return res;
  }
  
  /// Given:
  ///   1. an induction variable `iv` of type AffineForOp;
  ///   2. a `memoryOp` of type const LoadOp& or const StoreOp&;
  /// determines whether `memoryOp` has a contiguous access along `iv`. Contiguous
  /// is defined as either invariant or varying only along a unique MemRef dim.
  /// Upon success, the unique MemRef dim is written in `memRefDim` (or -1 to
  /// convey the memRef access is invariant along `iv`).
  ///
  /// Prerequisites:
  ///   1. `memRefDim` ~= nullptr;
  ///   2. `iv` of the proper type;
  ///   3. the MemRef accessed by `memoryOp` has no layout map or at most an
  ///      identity layout map.
  ///
  /// Currently only supports no layoutMap or identity layoutMap in the MemRef.
  /// Returns false if the MemRef has a non-identity layoutMap or more than 1
  /// layoutMap. This is conservative.
  ///
  // TODO: check strides.
  template <typename LoadOrStoreOp>
  static bool isContiguousAccess(Value iv, LoadOrStoreOp memoryOp,
                                 int *memRefDim) {
    static_assert(
        llvm::is_one_of<LoadOrStoreOp, AffineLoadOp, AffineStoreOp>::value,
        "Must be called on either LoadOp or StoreOp");
    assert(memRefDim && "memRefDim == nullptr");
    auto memRefType = memoryOp.getMemRefType();
  
    auto layoutMap = memRefType.getAffineMaps();
    // TODO: remove dependence on Builder once we support non-identity layout map.
    Builder b(memoryOp.getContext());
    if (layoutMap.size() >= 2 ||
        (layoutMap.size() == 1 &&
         !(layoutMap[0] ==
           b.getMultiDimIdentityMap(layoutMap[0].getNumDims())))) {
      return memoryOp.emitError("NYI: non-trivial layoutMap"), false;
    }
  
    int uniqueVaryingIndexAlongIv = -1;
    auto accessMap = memoryOp.getAffineMap();
    SmallVector<Value, 4> mapOperands(memoryOp.getMapOperands());
    unsigned numDims = accessMap.getNumDims();
    for (unsigned i = 0, e = memRefType.getRank(); i < e; ++i) {
      // Gather map operands used result expr 'i' in 'exprOperands'.
      SmallVector<Value, 4> exprOperands;
      auto resultExpr = accessMap.getResult(i);
      resultExpr.walk([&](AffineExpr expr) {
        if (auto dimExpr = expr.dyn_cast<AffineDimExpr>())
          exprOperands.push_back(mapOperands[dimExpr.getPosition()]);
        else if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>())
          exprOperands.push_back(mapOperands[numDims + symExpr.getPosition()]);
      });
      // Check access invariance of each operand in 'exprOperands'.
      for (auto exprOperand : exprOperands) {
        if (!isAccessIndexInvariant(iv, exprOperand)) {
          if (uniqueVaryingIndexAlongIv != -1) {
            // 2+ varying indices -> do not vectorize along iv.
            return false;
          }
          uniqueVaryingIndexAlongIv = i;
        }
      }
    }
  
    if (uniqueVaryingIndexAlongIv == -1)
      *memRefDim = -1;
    else
      *memRefDim = memRefType.getRank() - (uniqueVaryingIndexAlongIv + 1);
    return true;
  }
  
  template <typename LoadOrStoreOp>
  static bool isVectorElement(LoadOrStoreOp memoryOp) {
    auto memRefType = memoryOp.getMemRefType();
    return memRefType.getElementType().template isa<VectorType>();
  }
  
  using VectorizableOpFun = std::function<bool(AffineForOp, Operation &)>;
  
  static bool
  isVectorizableLoopBodyWithOpCond(AffineForOp loop,
                                   VectorizableOpFun isVectorizableOp,
                                   NestedPattern &vectorTransferMatcher) {
    auto *forOp = loop.getOperation();
  
    // No vectorization across conditionals for now.
    auto conditionals = matcher::If();
    SmallVector<NestedMatch, 8> conditionalsMatched;
    conditionals.match(forOp, &conditionalsMatched);
    if (!conditionalsMatched.empty()) {
      return false;
    }
  
    // No vectorization across unknown regions.
    auto regions = matcher::Op([](Operation &op) -> bool {
      return op.getNumRegions() != 0 && !isa<AffineIfOp, AffineForOp>(op);
    });
    SmallVector<NestedMatch, 8> regionsMatched;
    regions.match(forOp, &regionsMatched);
    if (!regionsMatched.empty()) {
      return false;
    }
  
    SmallVector<NestedMatch, 8> vectorTransfersMatched;
    vectorTransferMatcher.match(forOp, &vectorTransfersMatched);
    if (!vectorTransfersMatched.empty()) {
      return false;
    }
  
    auto loadAndStores = matcher::Op(matcher::isLoadOrStore);
    SmallVector<NestedMatch, 8> loadAndStoresMatched;
    loadAndStores.match(forOp, &loadAndStoresMatched);
    for (auto ls : loadAndStoresMatched) {
      auto *op = ls.getMatchedOperation();
      auto load = dyn_cast<AffineLoadOp>(op);
      auto store = dyn_cast<AffineStoreOp>(op);
      // Only scalar types are considered vectorizable, all load/store must be
      // vectorizable for a loop to qualify as vectorizable.
      // TODO: ponder whether we want to be more general here.
      bool vector = load ? isVectorElement(load) : isVectorElement(store);
      if (vector) {
        return false;
      }
      if (isVectorizableOp && !isVectorizableOp(loop, *op)) {
        return false;
      }
    }
    return true;
  }
  
  bool mlir::isVectorizableLoopBody(AffineForOp loop, int *memRefDim,
                                    NestedPattern &vectorTransferMatcher) {
    *memRefDim = -1;
    VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) {
      auto load = dyn_cast<AffineLoadOp>(op);
      auto store = dyn_cast<AffineStoreOp>(op);
      int thisOpMemRefDim = -1;
      bool isContiguous = load ? isContiguousAccess(loop.getInductionVar(), load,
                                                    &thisOpMemRefDim)
                               : isContiguousAccess(loop.getInductionVar(), store,
                                                    &thisOpMemRefDim);
      if (thisOpMemRefDim != -1) {
        // If memory accesses vary across different dimensions then the loop is
        // not vectorizable.
        if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)
          return false;
        *memRefDim = thisOpMemRefDim;
      }
      return isContiguous;
    });
    return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher);
  }
  
  bool mlir::isVectorizableLoopBody(AffineForOp loop,
                                    NestedPattern &vectorTransferMatcher) {
    return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher);
  }
  
  /// Checks whether SSA dominance would be violated if a for op's body
  /// operations are shifted by the specified shifts. This method checks if a
  /// 'def' and all its uses have the same shift factor.
  // TODO: extend this to check for memory-based dependence violation when we have
  // the support.
  bool mlir::isOpwiseShiftValid(AffineForOp forOp, ArrayRef<uint64_t> shifts) {
    auto *forBody = forOp.getBody();
    assert(shifts.size() == forBody->getOperations().size());
  
    // Work backwards over the body of the block so that the shift of a use's
    // ancestor operation in the block gets recorded before it's looked up.
    DenseMap<Operation *, uint64_t> forBodyShift;
    for (auto it : llvm::enumerate(llvm::reverse(forBody->getOperations()))) {
      auto &op = it.value();
  
      // Get the index of the current operation, note that we are iterating in
      // reverse so we need to fix it up.
      size_t index = shifts.size() - it.index() - 1;
  
      // Remember the shift of this operation.
      uint64_t shift = shifts[index];
      forBodyShift.try_emplace(&op, shift);
  
      // Validate the results of this operation if it were to be shifted.
      for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) {
        Value result = op.getResult(i);
        for (auto *user : result.getUsers()) {
          // If an ancestor operation doesn't lie in the block of forOp,
          // there is no shift to check.
          if (auto *ancOp = forBody->findAncestorOpInBlock(*user)) {
            assert(forBodyShift.count(ancOp) > 0 && "ancestor expected in map");
            if (shift != forBodyShift[ancOp])
              return false;
          }
        }
      }
    }
    return true;
  }
  