[Mlir-commits] [mlir] [MLIR][SparseTensor] Dense Outer Loop Ordering Strategy (PR #160168)
Govind Malasani
llvmlistbot at llvm.org
Mon Oct 27 13:46:45 PDT 2025
https://github.com/gmalasan updated https://github.com/llvm/llvm-project/pull/160168
>From 6961a0d3c53a5bf56d874a244d9c777c6dda4112 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Sun, 21 Sep 2025 21:45:57 -0400
Subject: [PATCH 1/3] [MLIR][SparseTensor] Implemented denseOuter heuristic;
prefer dense loops to sparse/singleton loop
Conflicts to be resolved
---
.../Dialect/SparseTensor/Transforms/Passes.h | 2 +
.../Dialect/SparseTensor/Transforms/Passes.td | 4 +-
.../Transforms/Utils/IterationGraphSorter.cpp | 72 ++++++++++++++++++-
3 files changed, 76 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index af64370a62dd7..d03455ef1efd4 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -61,6 +61,8 @@ namespace sparse_tensor {
enum class LoopOrderingStrategy : unsigned {
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..17ee4be86c7fd 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -80,6 +80,53 @@ 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<Value> allTensors,
+ ArrayRef<AffineMap> 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<AffineDimExpr>(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 +154,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<Value> allTensors = ins;
+ allTensors.push_back(out);
+ SmallVector<AffineMap> 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) {
>From cc2e7864a74d9d18b8df5fe0e8a1f8dc2804b4c1 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Sun, 12 Oct 2025 18:40:11 -0700
Subject: [PATCH 2/3] [MLIR][SparseTensor] Format dense-outer heuristic
---
.../mlir/Dialect/SparseTensor/Transforms/Passes.h | 4 ++--
.../Transforms/Utils/IterationGraphSorter.cpp | 15 +++++++--------
2 files changed, 9 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index d03455ef1efd4..92a33923ad57b 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -59,8 +59,8 @@ 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.
};
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 17ee4be86c7fd..7ebe20e9674ce 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -83,11 +83,10 @@ inline static bool includesDenseOutput(SortMask mask) {
/// 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<Value> allTensors,
- ArrayRef<AffineMap> allMaps) {
+static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
+ ArrayRef<AffineMap> 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;
@@ -101,7 +100,7 @@ static unsigned getLoopSparsityRank(unsigned loop,
}
tensorDim++;
}
-
+
if (loopAccessesTensor) {
const auto enc = getSparseTensorEncoding(tensor.getType());
if (!enc) {
@@ -123,7 +122,7 @@ static unsigned getLoopSparsityRank(unsigned loop,
}
}
}
-
+
return bestRank;
}
@@ -161,11 +160,11 @@ AffineMap IterationGraphSorter::topoSort() {
allTensors.push_back(out);
SmallVector<AffineMap> 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)) {
>From bfe907bd5df5168405b36f66b79b1a6ae0ce2a08 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Mon, 27 Oct 2025 16:46:15 -0400
Subject: [PATCH 3/3] [MLIR][SparseTensor] Address PR feedback to fix style
---
.../Dialect/SparseTensor/Transforms/Passes.h | 7 +++----
.../Transforms/Utils/IterationGraphSorter.cpp | 16 ++++++++--------
2 files changed, 11 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 92a33923ad57b..419ecda80e9a5 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -58,11 +58,10 @@ enum class SparseEmitStrategy {
namespace sparse_tensor {
/// Defines a strategy for loop ordering during sparse code generation.
+/// See Passes.td for strategy descriptions.
enum class LoopOrderingStrategy : unsigned {
- kDefault, ///< Default strategy (eagerly selects last loop in topological
- ///< sort).
- kDenseOuter, ///< Prefer dense, then compressed, then singleton dimensions
- ///< outermost.
+ kDefault,
+ kDenseOuter,
};
} // namespace sparse_tensor
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 7ebe20e9674ce..c428988398f85 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -82,13 +82,13 @@ inline static bool includesDenseOutput(SortMask mask) {
/// 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
+/// 0 = Dense, 1 = Compressed, 2 = Singleton, 3 = Other/Unknown.
static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
ArrayRef<AffineMap> allMaps) {
- unsigned bestRank = 3; // Default: most sparse (unknown/singleton-like)
+ unsigned bestRank = 3; // Default: most sparse (unknown/singleton-like).
for (auto [tensor, map] : llvm::zip(allTensors, allMaps)) {
- // Check if this loop accesses this tensor
+ // Check if this loop accesses this tensor.
bool loopAccessesTensor = false;
unsigned tensorDim = 0;
for (AffineExpr expr : map.getResults()) {
@@ -104,19 +104,19 @@ static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
if (loopAccessesTensor) {
const auto enc = getSparseTensorEncoding(tensor.getType());
if (!enc) {
- // Dense tensor - highest priority
+ // Dense tensor - highest priority.
return 0;
} else {
- // Sparse tensor - check the level type for this dimension
+ // 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
+ return 0; // Dense level.
} else if (isCompressedLT(lvlType)) {
- bestRank = std::min(bestRank, 1u); // Compressed level
+ bestRank = std::min(bestRank, 1u); // Compressed level.
} else if (isSingletonLT(lvlType)) {
- bestRank = std::min(bestRank, 2u); // Singleton level
+ bestRank = std::min(bestRank, 2u); // Singleton level.
}
}
}
More information about the Mlir-commits
mailing list