diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h index cd626041834b1..728af841cc7b1 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h @@ -10,8 +10,8 @@ // //===----------------------------------------------------------------------===// -#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENENV_H_ -#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENENV_H_ +#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENENV_H_ +#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENENV_H_ #include "CodegenUtils.h" #include "LoopEmitter.h" @@ -206,4 +206,4 @@ class CodegenEnv { } // namespace sparse_tensor } // namespace mlir -#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENENV_H_ +#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENENV_H_ diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h index 57de437e2311c..8dc57e1b5479b 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h @@ -10,8 +10,8 @@ // //===----------------------------------------------------------------------===// -#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_ -#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_ +#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENUTILS_H_ +#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENUTILS_H_ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Complex/IR/Complex.h" @@ -434,4 +434,4 @@ inline bool isZeroRankedTensorOrScalar(Type type) { } // namespace sparse_tensor } // namespace mlir -#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_ +#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENUTILS_H_ diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp index b6011727f4127..588400f3dbaf0 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp @@ -1,4 +1,4 @@ -//===- LoopScheduler.cpp -------------------------------------------------===// +//===- 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. @@ -20,11 +20,10 @@ 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, returns the first DimExpr in the -/// expression. +/// 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 { public: explicit AffineDimFinder(ArrayRef itTypes) @@ -81,11 +80,9 @@ inline static bool includesDenseOutput(SortMask mask) { return includesAny(mask, SortMask::kIncludeDenseOutput); } -/// A helper to compute a topological sort. O(n^2) time complexity -/// as we use adj matrix for the graph. -/// The sorted result will put the first Reduction iterator to the -/// latest possible position. AffineMap IterationGraphSorter::topoSort() { + // The sorted result will put the first Reduction iterator to the + // latest possible position. std::vector redIt; // reduce iterator with 0 degree std::vector parIt; // parallel iterator with 0 degree const unsigned numLoops = getNumLoops(); @@ -170,6 +167,7 @@ AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) { // Reset the interation graph. for (auto &row : itGraph) std::fill(row.begin(), row.end(), false); + // Reset cached in-degree. std::fill(inDegree.begin(), inDegree.end(), 0); @@ -179,7 +177,6 @@ AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) { // Skip dense inputs when not requested. if ((!enc && !includesDenseInput(mask)) || in == ignored) continue; - addConstraints(in, map); } diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h index 52ee117029300..be94bb5dffde6 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h @@ -1,10 +1,17 @@ -//===- LoopScheduler.h -----------------------------------------*- C++ -*-===// +//===- 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/IR/AffineMap.h" @@ -21,7 +28,7 @@ class GenericOp; namespace sparse_tensor { -/// Iteration graph sorting. +/// Iteration graph sorting mask, enum class SortMask : unsigned { // The individual mask bits. kIncludeDenseOutput = 0x1, // b001 @@ -34,40 +41,52 @@ enum class SortMask : unsigned { class IterationGraphSorter { public: - // Constructs a scheduler from linalg.generic - // Maybe reuses the class to schedule foreach as well (to address - // non-permutation, e.g, traverse CSR in BSR order). + /// Factory method that construct an iteration graph sorter + /// for the given linalg.generic operation. static IterationGraphSorter fromGenericOp(linalg::GenericOp genericOp); - // Returns a permutation that represents the scheduled loop order. - // Note that the returned AffineMap could be null if the kernel can not be - // schedule due to cycles in the iteration graph. + /// 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 &&ins, SmallVector &&loop2InsLvl, Value out, AffineMap loop2OutLvl, SmallVector &&iterTypes); + // 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 ins; SmallVector loop2InsLvl; + // Output tensor and associated loop to level map. Value out; AffineMap loop2OutLvl; - // Loop type; + + // Loop itation types; SmallVector iterTypes; // Adjacent matrix that represents the iteration graph. std::vector> itGraph; + // InDegree used for topo sort. std::vector inDegree; }; } // namespace sparse_tensor } // namespace mlir + +#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_ITERATIONGRAPHSORTER_H_ diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h index fa8b0076f733b..78bb53e4483f6 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h @@ -6,8 +6,8 @@ // //===----------------------------------------------------------------------===// -#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORLOOPEMITTER_H_ -#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORLOOPEMITTER_H_ +#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_LOOPEMITTER_H_ +#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_LOOPEMITTER_H_ #include @@ -22,7 +22,7 @@ namespace sparse_tensor { // A compressed pair. using TensorLevel = unsigned; -//===----------------------------------------------------------------------===// +// // SparseTensorLoopEmiter class, manages sparse tensors and helps to // generate loop structure to (co)-iterate sparse tensors. // @@ -48,8 +48,7 @@ using TensorLevel = unsigned; // loopEmiter.exitCurrentLoop(); // exit k // loopEmiter.exitCurrentLoop(); // exit j // loopEmiter.exitCurrentLoop(); // exit i -//===----------------------------------------------------------------------===// - +// class LoopEmitter { public: /// Optional callback function to setup dense output tensors when @@ -705,4 +704,4 @@ class LoopEmitter { } // namespace sparse_tensor } // namespace mlir -#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORLOOPEMITTER_H_ +#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_LOOPEMITTER_H_ diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.h index 5c7d8aa4c9d96..3a61ec7a2236f 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.h @@ -10,8 +10,8 @@ // //===----------------------------------------------------------------------===// -#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORDESCRIPTOR_H_ -#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORDESCRIPTOR_H_ +#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_SPARSETENSORDESCRIPTOR_H_ +#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_SPARSETENSORDESCRIPTOR_H_ #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" #include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h" @@ -262,4 +262,4 @@ getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl &fields) { } // namespace sparse_tensor } // namespace mlir -#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSODESCRIPTOR_H_ +#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_SPARSETENSODESCRIPTOR_H_