diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h index af64370a62dd7..92a33923ad57b 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h @@ -59,8 +59,10 @@ namespace sparse_tensor { /// Defines a strategy for loop ordering during sparse code generation. enum class LoopOrderingStrategy : unsigned { - kDefault, ///< Default strategy (eagerly selects last loop in topological - ///< sort). + kDefault, ///< Default strategy (eagerly selects last loop in topological + ///< sort). + kDenseOuter, ///< Prefer dense, then compressed, then singleton dimensions + ///< outermost. }; } // namespace sparse_tensor diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td index 75e77d67db1b3..0b8562e484f51 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td @@ -85,7 +85,9 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> { "mlir::sparse_tensor::LoopOrderingStrategy::kDefault", "Set the loop ordering strategy for sparse code generation", [{llvm::cl::values( clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDefault, "default", - "Default strategy (eagerly selects last loop in topological sort)"))}]>, + "Default strategy (eagerly selects last loop in topological sort)"), + clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDenseOuter, "dense-outer", + "Prefer dense, then compressed, then singleton dimensions outermost"))}]>, ]; } diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp index 73e0f3d2891d7..7ebe20e9674ce 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp @@ -80,6 +80,52 @@ 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. +/// 0 = Dense, 1 = Compressed, 2 = Singleton, 3 = Other/Unknown +static unsigned getLoopSparsityRank(unsigned loop, ArrayRef allTensors, + ArrayRef allMaps) { + unsigned bestRank = 3; // Default: most sparse (unknown/singleton-like) + + 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(expr)) { + if (dimExpr.getPosition() == loop) { + loopAccessesTensor = true; + break; + } + } + tensorDim++; + } + + if (loopAccessesTensor) { + const auto enc = getSparseTensorEncoding(tensor.getType()); + if (!enc) { + // Dense tensor - highest priority + return 0; + } 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 0; // Dense level + } else if (isCompressedLT(lvlType)) { + bestRank = std::min(bestRank, 1u); // Compressed level + } else if (isSingletonLT(lvlType)) { + bestRank = std::min(bestRank, 2u); // Singleton level + } + } + } + } + } + + return bestRank; +} + AffineMap IterationGraphSorter::topoSort() { // The sorted result will put the first Reduction iterator to the // latest possible position. @@ -107,10 +153,33 @@ AffineMap IterationGraphSorter::topoSort() { case sparse_tensor::LoopOrderingStrategy::kDefault: src = it.back(); break; + case sparse_tensor::LoopOrderingStrategy::kDenseOuter: { + // Prefer dense, then compressed, then singleton dimensions outermost. + // Create combined tensor and map lists for analysis. + SmallVector allTensors = ins; + allTensors.push_back(out); + SmallVector allMaps = loop2InsLvl; + allMaps.push_back(loop2OutLvl); + + // Find loop with best (lowest) sparsity rank. + unsigned bestLoop = it[0]; + unsigned bestRank = getLoopSparsityRank(bestLoop, allTensors, allMaps); + + for (auto candidateLoop : it) { + unsigned rank = getLoopSparsityRank(candidateLoop, allTensors, allMaps); + if (rank < bestRank || (rank == bestRank && candidateLoop < bestLoop)) { + bestLoop = candidateLoop; + bestRank = rank; + } + } + src = bestLoop; + break; + } } loopOrder.push_back(src); - it.pop_back(); + // 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) {