IterationGraphSorter.cpp

Go to the documentation of this file.

//===- IterationGraphSorter.cpp -------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "IterationGraphSorter.h"
 
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/BuiltinTypes.h"
#include "llvm/ADT/DenseSet.h"
 
using namespace mlir;
using namespace mlir::sparse_tensor;
 
namespace {
 
/// A helper class that visits an affine expression and tries to find
/// an AffineDimExpr to which the corresponding iterator from a GenericOp
/// matches the desired iterator type. If there is no matched iterator
/// type, the method returns the first DimExpr in the expression.
class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
public:
  explicit AffineDimFinder(ArrayRef<utils::IteratorType> itTypes)
      : iterTypes(itTypes) {}
 
  /// Overrides the visit method from AffineExprVisitor.
  void visitDimExpr(AffineDimExpr expr) {
    if (pickedDim == nullptr || pickIterType == iterTypes[expr.getPosition()])
      pickedDim = expr;
  }
 
  /// Sets the desired iterator type that we want to pick.
  void setPickedIterType(utils::IteratorType iterType) {
    pickIterType = iterType;
  }
 
  /// Gets the desired AffineDimExpr.
  AffineDimExpr getDimExpr() const {
    return llvm::cast<AffineDimExpr>(pickedDim);
  }
 
  /// Walks the graph in post order to find dim expr.
  void walkPostOrder(AffineExpr expr) {
    pickedDim = nullptr;
    AffineExprVisitor<AffineDimFinder>::walkPostOrder(expr);
  }
 
private:
  /// The picked AffineDimExpr after visit.
  AffineExpr pickedDim;
  /// The iterator type that we want.
  utils::IteratorType pickIterType;
  /// The mapping between levels and iterator types.
  ArrayRef<utils::IteratorType> iterTypes;
};
 
/// Flattens an affine expression into a list of AffineDimExprs.
struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> {
  // Overrides method from AffineExprVisitor.
  void visitDimExpr(AffineDimExpr expr) { dims.push_back(expr); }
  SmallVector<AffineDimExpr> dims;
};
 
} // namespace
 
inline static bool includesAny(SortMask mask1, SortMask mask2) {
  return static_cast<unsigned>(mask1) & static_cast<unsigned>(mask2);
}
 
inline static bool includesDenseInput(SortMask mask) {
  return includesAny(mask, SortMask::kIncludeDenseInput);
}
 
inline static bool includesDenseOutput(SortMask mask) {
  return includesAny(mask, SortMask::kIncludeDenseOutput);
}
 
/// Returns a sparsity rank for loop ordering: lower values indicate
/// dimensions that should be placed in outer loops.
/// When preferDenseOuter is true the ranking is
///   0 = Dense, 1 = Compressed, 2 = Singleton, 3 = Other/Unknown.
/// Otherwise
///   0 = Singleton, 1 = Compressed, 2 = Dense, 3 = Other/Unknown.
static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
                                    ArrayRef<AffineMap> allMaps,
                                    bool preferDenseOuter) {
  const unsigned denseRank = preferDenseOuter ? 0 : 2;
  const unsigned singletonRank = preferDenseOuter ? 2 : 0;
  const unsigned compressedRank = 1;
  const unsigned unknownRank = 3;
 
  unsigned minRank = unknownRank;
 
  for (auto [tensor, map] : llvm::zip(allTensors, allMaps)) {
    // Check if this loop accesses this tensor.
    bool loopAccessesTensor = false;
    unsigned tensorDim = 0;
    for (AffineExpr expr : map.getResults()) {
      if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
        if (dimExpr.getPosition() == loop) {
          loopAccessesTensor = true;
          break;
        }
      }
      tensorDim++;
    }
 
    if (loopAccessesTensor) {
      const auto enc = getSparseTensorEncoding(tensor.getType());
      if (!enc) {
        // Dense tensor.
        return denseRank;
      } else {
        // Sparse tensor - check the level type for this dimension.
        auto lvlTypes = enc.getLvlTypes();
        if (tensorDim < lvlTypes.size()) {
          auto lvlType = lvlTypes[tensorDim];
          if (isDenseLT(lvlType)) {
            return denseRank; // Dense level.
          } else if (isCompressedLT(lvlType)) {
            minRank = std::min(minRank, compressedRank); // Compressed level.
          } else if (isSingletonLT(lvlType)) {
            minRank = std::min(minRank, singletonRank); // Singleton level.
          }
        }
      }
    }
  }
 
  return minRank;
}
 
/// Count distinct input accesses whose indexing expressions use each loop.
static SmallVector<unsigned> getLoopUseCounts(linalg::GenericOp genericOp) {
  SmallVector<unsigned> counts(genericOp.getNumLoops(), 0);
  llvm::SmallDenseSet<std::pair<Value, AffineMap>> seen;
  auto maps = genericOp.getIndexingMapsArray();
 
  for (auto [idx, input] : llvm::enumerate(genericOp.getDpsInputs())) {
    if (genericOp.getRegion().front().getArgument(idx).use_empty())
      continue;
 
    AffineMap map = maps[idx];
    if (!seen.insert({input, map}).second)
      continue;
 
    llvm::SmallDenseSet<unsigned> accessedLoops;
    for (AffineExpr expr : map.getResults()) {
      // Count dependencies within expressions.
      expr.walk([&](AffineExpr subExpr) {
        if (auto dim = dyn_cast<AffineDimExpr>(subExpr))
          accessedLoops.insert(dim.getPosition());
      });
    }
    // Each access contributes once per loop.
    for (unsigned loop : accessedLoops)
      ++counts[loop];
  }
  return counts;
}
 
AffineMap IterationGraphSorter::topoSort() {
  // The sorted result will put the first Reduction iterator to the
  // latest possible position.
  std::vector<unsigned> redIt; // reduce iterator with 0 degree
  std::vector<unsigned> parIt; // parallel iterator with 0 degree
  const unsigned numLoops = getNumLoops();
  for (unsigned i = 0; i < numLoops; i++) {
    if (inDegree[i] == 0) {
      if (iterTypes[i] == utils::IteratorType::reduction)
        redIt.push_back(i);
      else
        parIt.push_back(i);
    }
  }
 
  SmallVector<unsigned> loopOrder;
  while (!redIt.empty() || !parIt.empty()) {
    // We always prefer a parallel loop over a reduction loop because putting
    // a reduction loop early might make the loop sequence inadmissible.
    auto &it = !parIt.empty() ? parIt : redIt;
 
    // Select loop based on strategy.
    unsigned src;
    switch (strategy) {
    case sparse_tensor::LoopOrderingStrategy::kDefault:
      src = it.back();
      break;
    case sparse_tensor::LoopOrderingStrategy::kLoopUseOuter:
    case sparse_tensor::LoopOrderingStrategy::kLoopUseInner: {
      // Loop use counts must already be populated.
      assert(loopUseCounts.size() == numLoops);
 
      const bool preferMore = strategy == LoopOrderingStrategy::kLoopUseOuter;
      src = it.front();
      for (unsigned candidate : it) {
        unsigned count = loopUseCounts[candidate];
        unsigned bestCount = loopUseCounts[src];
 
        if (preferMore ? count > bestCount : count < bestCount)
          src = candidate;
      }
      break;
    }
    case sparse_tensor::LoopOrderingStrategy::kDenseOuter: {
      // Prefer dense, then compressed, then singleton dimensions outermost.
      // Create combined tensor and map lists for analysis.
      SmallVector<Value> allTensors = ins;
      allTensors.push_back(out);
      SmallVector<AffineMap> allMaps = loop2InsLvl;
      allMaps.push_back(loop2OutLvl);
 
      // Find loop with minimum (lowest) sparsity rank.
      unsigned minLoop = it[0];
      unsigned minRank =
          getLoopSparsityRank(minLoop, allTensors, allMaps, true);
 
      for (auto candidateLoop : it) {
        unsigned rank =
            getLoopSparsityRank(candidateLoop, allTensors, allMaps, true);
        if (rank < minRank || (rank == minRank && candidateLoop < minLoop)) {
          minLoop = candidateLoop;
          minRank = rank;
        }
      }
      src = minLoop;
      break;
    }
    case sparse_tensor::LoopOrderingStrategy::kSparseOuter: {
      // Prefer singleton, then compressed, then dense dimensions outermost.
      SmallVector<Value> allTensors = ins;
      allTensors.push_back(out);
      SmallVector<AffineMap> allMaps = loop2InsLvl;
      allMaps.push_back(loop2OutLvl);
 
      unsigned minLoop = it[0];
      unsigned minRank =
          getLoopSparsityRank(minLoop, allTensors, allMaps, false);
 
      for (auto candidateLoop : it) {
        unsigned rank =
            getLoopSparsityRank(candidateLoop, allTensors, allMaps, false);
        if (rank < minRank || (rank == minRank && candidateLoop < minLoop)) {
          minLoop = candidateLoop;
          minRank = rank;
        }
      }
      src = minLoop;
      break;
    }
    }
 
    loopOrder.push_back(src);
    // Remove the selected loop from the worklist.
    it.erase(std::find(it.begin(), it.end(), src));
    // Update in-degree, and push 0-degree node into worklist.
    for (unsigned dst = 0; dst < numLoops; dst++) {
      if (itGraph[src][dst] && --inDegree[dst] == 0) {
        if (iterTypes[dst] == utils::IteratorType::reduction)
          redIt.push_back(dst);
        else
          parIt.push_back(dst);
      }
    }
  }
 
  // Return the topological sort on success.
  if (loopOrder.size() == numLoops)
    return AffineMap::getPermutationMap(loopOrder, out.getContext());
 
  // Cycle detected.
  return AffineMap();
}
 
IterationGraphSorter IterationGraphSorter::fromGenericOp(
    linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy) {
  // Must be a demapped sparse kernel.
  assert(!hasAnyNonIdentityOperandsOrResults(genericOp) &&
         hasAnySparseOperandOrResult(genericOp) &&
         genericOp.getNumDpsInits() == 1);
 
  SmallVector<AffineMap> loopMap = genericOp.getIndexingMapsArray();
  SmallVector<Value> ins = genericOp.getDpsInputs();
 
  AffineMap outMap = loopMap.back();
  loopMap.pop_back();
 
  Value out = genericOp.getDpsInitOperand(0)->get();
  SmallVector<utils::IteratorType> iterTypes =
      genericOp.getIteratorTypesArray();
 
  return IterationGraphSorter(std::move(ins), std::move(loopMap), out, outMap,
                              std::move(iterTypes), genericOp, strategy);
}
 
IterationGraphSorter::IterationGraphSorter(
    SmallVector<Value> &&insArg, SmallVector<AffineMap> &&loop2InsLvlArg,
    Value out, AffineMap loop2OutLvl,
    SmallVector<utils::IteratorType> &&iterTypesArg,
    linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy)
    : ins(std::move(insArg)), loop2InsLvl(std::move(loop2InsLvlArg)), out(out),
      loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypesArg)),
      strategy(strategy) {
  // One map per tensor.
  assert(loop2InsLvl.size() == ins.size());
  // All the affine maps have the same number of dimensions (loops).
  assert(llvm::all_equal(llvm::map_range(
      loop2InsLvl, [](AffineMap m) { return m.getNumDims(); })));
  // The number of results of the map should match the rank of the tensor.
  assert(llvm::all_of(llvm::zip(loop2InsLvl, ins), [](auto mvPair) {
    auto [m, v] = mvPair;
 
    // For ranked types the rank must match.
    // Simply return true for UnrankedTensorType
    if (auto shapedType = llvm::dyn_cast<ShapedType>(v.getType())) {
      return !shapedType.hasRank() ||
             (m.getNumResults() == shapedType.getRank());
    }
    // Non-shaped (scalar) types behave like rank-0.
    return m.getNumResults() == 0;
  }));
 
  switch (strategy) {
  case LoopOrderingStrategy::kLoopUseOuter:
  case LoopOrderingStrategy::kLoopUseInner:
    loopUseCounts = getLoopUseCounts(genericOp);
    break;
  default:
    break;
  }
 
  itGraph.resize(getNumLoops(), std::vector<bool>(getNumLoops(), false));
  inDegree.resize(getNumLoops());
}
 
AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) {
  // Reset the adjacency matrix that represents the iteration graph.
  for (auto &row : itGraph)
    llvm::fill(row, false);
 
  // Reset in-degree.
  llvm::fill(inDegree, 0);
 
  // Add the constraints for the loop to level map.
  for (auto [in, map] : llvm::zip(ins, loop2InsLvl)) {
    // Get map and encoding.
    const auto enc = getSparseTensorEncoding(in.getType());
    // Skip dense inputs when not requested.
    if ((!enc && !includesDenseInput(mask)) || in == ignored)
      continue;
    addConstraints(in, map);
  }
 
  // Add the constraints for the output map.
  const auto enc = getSparseTensorEncoding(out.getType());
  if ((enc || includesDenseOutput(mask)) && out != ignored)
    addConstraints(out, loop2OutLvl);
 
  // Return the topological sort (empty for cyclic).
  return topoSort();
}
 
void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) {
  auto addIterOrdering = [this](unsigned f, unsigned t) {
    if (!itGraph[f][t] && f != t) {
      itGraph[f][t] = true;
      inDegree[t]++;
    }
  };
 
  // Set up a reduction finder.
  AffineDimFinder finder(iterTypes);
  finder.setPickedIterType(utils::IteratorType::reduction);
 
  // To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
  // we require there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
  // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
  const Level lvlRank = loop2LvlMap.getNumResults();
  for (Level lvl = 1; lvl < lvlRank; lvl++) {
    const AffineExpr fa = loop2LvlMap.getResult(lvl - 1);
    const AffineExpr ta = loop2LvlMap.getResult(lvl);
 
    if (llvm::isa<AffineDimExpr>(fa) || llvm::isa<AffineDimExpr>(ta)) {
      // Special case when at least one loop2LvlExp is a simple AffineDimExpr
      // (say, d0) and we require d0 > {d1, d2, ...} or {d1, d2, ...} > d0
      AffineDimCollector fCollector;
      fCollector.walkPostOrder(fa);
      AffineDimCollector tCollector;
      tCollector.walkPostOrder(ta);
 
      for (auto fd : fCollector.dims) {
        for (auto td : tCollector.dims) {
          const unsigned f = fd.getPosition();
          const unsigned t = td.getPosition();
          addIterOrdering(f, t);
        }
      }
      continue;
    }
 
    // When both loop2LvlExpr is compound, we pick an abitrary reduction loop
    // from lhs and rhs and use them as d_x and d_y.
    finder.walkPostOrder(fa);
    const AffineDimExpr fexp = finder.getDimExpr();
    const unsigned fldx = fexp.getPosition();
 
    finder.walkPostOrder(ta);
    const AffineDimExpr texp = finder.getDimExpr();
    const unsigned tldx = texp.getPosition();
 
    // d_x > d_y
    addIterOrdering(fldx, tldx);
 
    AffineDimCollector fCollector;
    fCollector.walkPostOrder(fa);
    AffineDimCollector tCollector;
    tCollector.walkPostOrder(ta);
 
    // Make sure dx and dy is the last.
    for (auto fd : fCollector.dims) {
      const unsigned f = fd.getPosition();
      addIterOrdering(f, fldx);
    }
    for (auto td : tCollector.dims) {
      const unsigned t = td.getPosition();
      addIterOrdering(t, tldx);
    }
    // {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
    // This is to ensure that the affine expressions are reduced in sparse
    // tensor level ordering.
    for (auto fd : fCollector.dims) {
      const unsigned f = fd.getPosition();
      if (f == fldx) // skip d_x
        continue;
      for (auto td : tCollector.dims) {
        const unsigned t = td.getPosition();
        if (t == tldx) // skip d_y
          continue;
        addIterOrdering(f, t);
      }
    }
  }
}