[Mlir-commits] [mlir] [MLIR][SparseTensor] Added Sparse Outer Loop Ordering Strategy (PR #172198)
Govind Malasani
llvmlistbot at llvm.org
Sun Dec 28 23:48:05 PST 2025
https://github.com/gmalasan updated https://github.com/llvm/llvm-project/pull/172198
>From 5ea32b12d040c19f5cc28bf26d90880a4b4fdcec Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Sat, 13 Dec 2025 23:00:45 -0800
Subject: [PATCH 1/2] [MLIR][SparseTensor] Added SparseOuter loop ordering
heuristic
---
.../Dialect/SparseTensor/Transforms/Passes.h | 1 +
.../Dialect/SparseTensor/Transforms/Passes.td | 4 +-
.../Transforms/Utils/IterationGraphSorter.cpp | 54 +++++++++++++++----
3 files changed, 47 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 419ecda80e9a5..40b37dc05e92e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -62,6 +62,7 @@ namespace sparse_tensor {
enum class LoopOrderingStrategy : unsigned {
kDefault,
kDenseOuter,
+ kSparseOuter,
};
} // namespace sparse_tensor
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 0b8562e484f51..7ad54a0d2218d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -87,7 +87,9 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> {
clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDefault, "default",
"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"))}]>,
+ "Prefer dense, then compressed, then singleton dimensions outermost"),
+ clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kSparseOuter, "sparse-outer",
+ "Prefer singleton, then compressed, then dense dimensions outermost"))}]>,
];
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 99048034b4f0c..a0180d228a36a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -82,11 +82,19 @@ 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.
+/// 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) {
- // Start with highest rank.
- unsigned minRank = 3;
+ 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.
@@ -105,19 +113,19 @@ static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
if (loopAccessesTensor) {
const auto enc = getSparseTensorEncoding(tensor.getType());
if (!enc) {
- // Dense tensor - lowest rank.
- return 0;
+ // 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 0; // Dense level.
+ return denseRank; // Dense level.
} else if (isCompressedLT(lvlType)) {
- minRank = std::min(minRank, 1u); // Compressed level.
+ minRank = std::min(minRank, compressedRank); // Compressed level.
} else if (isSingletonLT(lvlType)) {
- minRank = std::min(minRank, 2u); // Singleton level.
+ minRank = std::min(minRank, singletonRank); // Singleton level.
}
}
}
@@ -164,10 +172,34 @@ AffineMap IterationGraphSorter::topoSort() {
// Find loop with minimum (lowest) sparsity rank.
unsigned minLoop = it[0];
- unsigned minRank = getLoopSparsityRank(minLoop, allTensors, allMaps);
+ 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);
+ unsigned rank =
+ getLoopSparsityRank(candidateLoop, allTensors, allMaps, false);
if (rank < minRank || (rank == minRank && candidateLoop < minLoop)) {
minLoop = candidateLoop;
minRank = rank;
>From 4463c59d17157dc4effe292d7a4a430d2d78c3d3 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Sun, 28 Dec 2025 23:47:51 -0800
Subject: [PATCH 2/2] [MLIR][SparseTensor] Add check test for default,
dense-outer, sparse-outer heuristics
---
.../SparseTensor/sparse_loop_ordering.mlir | 69 +++++++++++++++++++
1 file changed, 69 insertions(+)
create mode 100644 mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir
diff --git a/mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir
new file mode 100644
index 0000000000000..709d12566b657
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir
@@ -0,0 +1,69 @@
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=default" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefix=DEFAULT
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=dense-outer" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefix=DENSE
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=sparse-outer" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefix=SPARSE
+
+#X = #sparse_tensor.encoding<{
+ map = (d0, d1, d2) -> (
+ d0 : dense,
+ d1 : compressed,
+ d2 : singleton)
+}>
+
+#Y = #sparse_tensor.encoding<{
+ map = (d0, d1, d2) -> (
+ d0 : compressed,
+ d1 : singleton,
+ d2 : dense)
+}>
+
+#trait = {
+ indexing_maps = [
+ affine_map<(i,j,k,l,m,n,o,p,q) -> (l,m,n)>,
+ affine_map<(i,j,k,l,m,n,o,p,q) -> (o,p,q)>,
+ affine_map<(i,j,k,l,m,n,o,p,q) -> (i,j,k)>
+ ],
+ iterator_types = ["parallel", "parallel", "parallel",
+ "parallel", "parallel", "parallel",
+ "parallel", "parallel", "parallel"]
+}
+
+// DEFAULT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d3, d4, d5)>
+// DEFAULT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2)>
+// DEFAULT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d6, d7, d8)>
+// DEFAULT-LABEL: func.func @sparse_loop_ordering
+// DEFAULT: linalg.generic
+// DEFAULT-SAME: sorted = true
+
+// DENSE: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d3, d6)>
+// DENSE: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d4, d7, d8)>
+// DENSE: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5)>
+// DENSE-LABEL: func.func @sparse_loop_ordering
+// DENSE: linalg.generic
+// DENSE-SAME: sorted = true
+
+// SPARSE: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d5, d6, d7)>
+// SPARSE: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d8)>
+// SPARSE: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d3, d4)>
+// SPARSE-LABEL: func.func @sparse_loop_ordering
+// SPARSE: linalg.generic
+// SPARSE-SAME: sorted = true
+
+func.func @sparse_loop_ordering(%A: tensor<?x?x?xf32, #X>,
+ %B: tensor<?x?x?xf32, #Y>,
+ %C: tensor<?x?x?xf32, #X>) -> tensor<?x?x?xf32, #X> {
+ %result = linalg.generic #trait
+ ins(%A, %B: tensor<?x?x?xf32, #X>, tensor<?x?x?xf32, #Y>)
+ outs(%C: tensor<?x?x?xf32, #X>) {
+ ^bb(%a: f32, %b: f32, %c: f32):
+ %ab = arith.mulf %a, %b : f32
+ %sum = arith.addf %c, %ab : f32
+ linalg.yield %sum : f32
+ } -> tensor<?x?x?xf32, #X>
+ return %result : tensor<?x?x?xf32, #X>
+}
More information about the Mlir-commits
mailing list