IterationGraphSorter.h

Go to the documentation of this file.

//===- IterationGraphSorter.h -----------------------------------*- C++ -*-===//
//
// 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 header file defines the iteration graph sorter (top-sort scheduling).
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_
 
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/IR/AffineMap.h"
 
namespace mlir {
 
// Forward declarations.
class Value;
namespace utils {
enum class IteratorType : uint32_t;
} // namespace utils
namespace linalg {
class GenericOp;
} // namespace linalg
 
namespace sparse_tensor {
 
/// Iteration graph sorting mask,
enum class SortMask : unsigned {
  // The individual mask bits.
  kIncludeDenseOutput = 0x1, // b001
  kIncludeDenseInput = 0x2,  // b010
  // The subsets of mask bits.
  kIncludeAll = 0x7,   // b111
  kIncludeDense = 0x3, // b011
  kSparseOnly = 0x0,   // b000
};
 
class IterationGraphSorter {
public:
  /// Factory method that constructs an iteration graph sorter
  /// for the given linalg.generic operation with a specific loop ordering
  /// strategy.
  static IterationGraphSorter
  fromGenericOp(linalg::GenericOp genericOp,
                sparse_tensor::LoopOrderingStrategy strategy);
 
  /// Returns a permutation that represents the scheduled loop order.
  /// Note that the returned AffineMap could be null if the kernel
  /// cannot be scheduled due to cyclic iteration graph.
  [[nodiscard]] AffineMap sort(SortMask mask, Value ignored = nullptr);
 
  /// Returns the number of loops in the iteration graph.
  unsigned getNumLoops() const { return loop2OutLvl.getNumDims(); }
 
private:
  // Private constructor.
  IterationGraphSorter(SmallVector<Value> &&insArg,
                       SmallVector<AffineMap> &&loop2InsLvlArg, Value out,
                       AffineMap loop2OutLvl,
                       SmallVector<utils::IteratorType> &&iterTypesArg,
                       sparse_tensor::LoopOrderingStrategy strategy =
                           sparse_tensor::LoopOrderingStrategy::kDefault);
 
  // Adds all the constraints in the given loop to level map.
  void addConstraints(Value t, AffineMap loop2LvlMap);
 
  /// A helper to compute a topological sort. The method has an
  /// O(n^2) time complexity since we use an adjacency matrix
  /// representation for the iteration graph.
  AffineMap topoSort();
 
  // Input tensors and associated loop to level maps.
  SmallVector<Value> ins;
  SmallVector<AffineMap> loop2InsLvl;
 
  // Output tensor and associated loop to level map.
  Value out;
  AffineMap loop2OutLvl;
 
  // Loop itation types;
  SmallVector<utils::IteratorType> iterTypes;
 
  // Adjacency matrix that represents the iteration graph.
  std::vector<std::vector<bool>> itGraph;
 
  // InDegree used for topo sort.
  std::vector<unsigned> inDegree;
 
  // Loop ordering strategy.
  sparse_tensor::LoopOrderingStrategy strategy;
};
 
} // namespace sparse_tensor
} // namespace mlir
 
#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_