[Mlir-commits] [mlir] 8796d37 - [MLIR][SparseTensor] Added Sparse Outer Loop Ordering Strategy (#172198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jun 13 14:19:32 PDT 2026
Author: Govind Malasani
Date: 2026-06-13T14:19:27-07:00
New Revision: 8796d370aed0153ce884319c2365530e92b9bf5e
URL: https://github.com/llvm/llvm-project/commit/8796d370aed0153ce884319c2365530e92b9bf5e
DIFF: https://github.com/llvm/llvm-project/commit/8796d370aed0153ce884319c2365530e92b9bf5e.diff
LOG: [MLIR][SparseTensor] Added Sparse Outer Loop Ordering Strategy (#172198)
This PR builds upon the infrastructure set up for Sparse Tensor Loop
Ordering Heuristics (#154656) and the already existing Dense Outer loop
ordering strategy (#160168).
Added:
mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir
Modified:
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
Removed:
################################################################################
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 d0e265c2d72ad..b9cb492db9d22 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;
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..ec19b3cd84f6d
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering.mlir
@@ -0,0 +1,165 @@
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=default" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefixes=DEFAULT,DEFAULT-LOWERED
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=dense-outer" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefixes=DENSE,DENSE-LOWERED
+// RUN: mlir-opt %s --sparse-reinterpret-map="loop-ordering-strategy=sparse-outer" \
+// RUN: -sparsification --canonicalize | \
+// RUN: FileCheck %s --check-prefixes=SPARSE,SPARSE-LOWERED
+
+#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>
+}
+
+
+// DEFAULT-LOWERED-LABEL: func.func @sparse_loop_ordering_lowered
+// DEFAULT-LOWERED-DAG: %[[C0:.*]] = arith.constant 0 : index
+// DEFAULT-LOWERED-DAG: %[[C1:.*]] = arith.constant 1 : index
+// DEFAULT-LOWERED-DAG: %[[LVL_A:.*]] = sparse_tensor.lvl %arg0, %[[C0]]
+// DEFAULT-LOWERED-DAG: %[[POS_A:.*]] = sparse_tensor.positions %arg0 {level = 1 : index}
+// DEFAULT-LOWERED-DAG: %[[POS_B:.*]] = sparse_tensor.positions %arg1 {level = 0 : index}
+// DEFAULT-LOWERED-DAG: %[[LVL_B:.*]] = sparse_tensor.lvl %arg1, %{{.*}}
+// DEFAULT-LOWERED-DAG: %[[DIM_C0:.*]] = tensor.dim %arg2, %[[C0]]
+// DEFAULT-LOWERED-DAG: %[[DIM_C1:.*]] = tensor.dim %arg2, %[[C1]]
+// DEFAULT-LOWERED-DAG: %[[DIM_C2:.*]] = tensor.dim %arg2, %{{.*}}
+// DEFAULT-LOWERED: %[[B_COMPRESSED_START:.*]] = memref.load %[[POS_B]][%[[C0]]]
+// DEFAULT-LOWERED: %[[B_COMPRESSED_END:.*]] = memref.load %[[POS_B]][%[[C1]]]
+// DEFAULT-LOWERED: scf.for %[[B_COMPRESSED:.*]] = %[[B_COMPRESSED_START]] to %[[B_COMPRESSED_END]] step %[[C1]] {
+// DEFAULT-LOWERED: %[[B_COMPRESSED_PLUS_1:.*]] = arith.addi %[[B_COMPRESSED]], %[[C1]]
+// DEFAULT-LOWERED: scf.for %[[B_SINGLETON:.*]] = %[[B_COMPRESSED]] to %[[B_COMPRESSED_PLUS_1]] step %[[C1]] {
+// DEFAULT-LOWERED: scf.for %[[B_DENSE:.*]] = %[[C0]] to %[[LVL_B]] step %[[C1]] {
+// DEFAULT-LOWERED: scf.for %[[A_DENSE:.*]] = %[[C0]] to %[[LVL_A]] step %[[C1]] {
+// DEFAULT-LOWERED: %[[A_COMPRESSED_START:.*]] = memref.load %[[POS_A]][%[[A_DENSE]]]
+// DEFAULT-LOWERED: %[[A_DENSE_PLUS_1:.*]] = arith.addi %[[A_DENSE]], %[[C1]]
+// DEFAULT-LOWERED: %[[A_COMPRESSED_END:.*]] = memref.load %[[POS_A]][%[[A_DENSE_PLUS_1]]]
+// DEFAULT-LOWERED: scf.for %[[A_COMPRESSED:.*]] = %[[A_COMPRESSED_START]] to %[[A_COMPRESSED_END]] step %[[C1]] {
+// DEFAULT-LOWERED: %[[A_COMPRESSED_PLUS_1:.*]] = arith.addi %[[A_COMPRESSED]], %[[C1]]
+// DEFAULT-LOWERED: scf.for %[[A_SINGLETON:.*]] = %[[A_COMPRESSED]] to %[[A_COMPRESSED_PLUS_1]] step %[[C1]] {
+// DEFAULT-LOWERED: scf.for %[[C_DENSE_0:.*]] = %[[C0]] to %[[DIM_C0]] step %[[C1]] {
+// DEFAULT-LOWERED: scf.for %[[C_DENSE_1:.*]] = %[[C0]] to %[[DIM_C1]] step %[[C1]] {
+// DEFAULT-LOWERED: scf.for %[[C_DENSE_2:.*]] = %[[C0]] to %[[DIM_C2]] step %[[C1]] {
+
+// DENSE-LOWERED-LABEL: func.func @sparse_loop_ordering_lowered
+// DENSE-LOWERED-DAG: %[[C0:.*]] = arith.constant 0 : index
+// DENSE-LOWERED-DAG: %[[C1:.*]] = arith.constant 1 : index
+// DENSE-LOWERED-DAG: %[[LVL_A:.*]] = sparse_tensor.lvl %arg0, %[[C0]]
+// DENSE-LOWERED-DAG: %[[POS_A:.*]] = sparse_tensor.positions %arg0 {level = 1 : index}
+// DENSE-LOWERED-DAG: %[[POS_B:.*]] = sparse_tensor.positions %arg1 {level = 0 : index}
+// DENSE-LOWERED-DAG: %[[LVL_B:.*]] = sparse_tensor.lvl %arg1, %{{.*}}
+// DENSE-LOWERED-DAG: %[[DIM_C0:.*]] = tensor.dim %arg2, %[[C0]]
+// DENSE-LOWERED-DAG: %[[DIM_C1:.*]] = tensor.dim %arg2, %[[C1]]
+// DENSE-LOWERED-DAG: %[[DIM_C2:.*]] = tensor.dim %arg2, %{{.*}}
+// DENSE-LOWERED: scf.for %[[C_DENSE_0:.*]] = %[[C0]] to %[[DIM_C0]] step %[[C1]] {
+// DENSE-LOWERED: scf.for %[[C_DENSE_1:.*]] = %[[C0]] to %[[DIM_C1]] step %[[C1]] {
+// DENSE-LOWERED: scf.for %[[C_DENSE_2:.*]] = %[[C0]] to %[[DIM_C2]] step %[[C1]] {
+// DENSE-LOWERED: scf.for %[[A_DENSE:.*]] = %[[C0]] to %[[LVL_A]] step %[[C1]] iter_args
+// DENSE-LOWERED: %[[A_COMPRESSED_START:.*]] = memref.load %[[POS_A]][%[[A_DENSE]]]
+// DENSE-LOWERED: %[[A_DENSE_PLUS_1:.*]] = arith.addi %[[A_DENSE]], %[[C1]]
+// DENSE-LOWERED: %[[A_COMPRESSED_END:.*]] = memref.load %[[POS_A]][%[[A_DENSE_PLUS_1]]]
+// DENSE-LOWERED: scf.for %[[A_COMPRESSED:.*]] = %[[A_COMPRESSED_START]] to %[[A_COMPRESSED_END]] step %[[C1]] iter_args
+// DENSE-LOWERED: %[[B_COMPRESSED_START:.*]] = memref.load %[[POS_B]][%[[C0]]]
+// DENSE-LOWERED: %[[B_COMPRESSED_END:.*]] = memref.load %[[POS_B]][%[[C1]]]
+// DENSE-LOWERED: scf.for %[[B_COMPRESSED:.*]] = %[[B_COMPRESSED_START]] to %[[B_COMPRESSED_END]] step %[[C1]] iter_args
+// DENSE-LOWERED: %[[A_COMPRESSED_PLUS_1:.*]] = arith.addi %[[A_COMPRESSED]], %[[C1]]
+// DENSE-LOWERED: scf.for %[[A_SINGLETON:.*]] = %[[A_COMPRESSED]] to %[[A_COMPRESSED_PLUS_1]] step %[[C1]] iter_args
+// DENSE-LOWERED: %[[B_COMPRESSED_PLUS_1:.*]] = arith.addi %[[B_COMPRESSED]], %[[C1]]
+// DENSE-LOWERED: scf.for %[[B_SINGLETON:.*]] = %[[B_COMPRESSED]] to %[[B_COMPRESSED_PLUS_1]] step %[[C1]] iter_args
+// DENSE-LOWERED: scf.for %[[B_DENSE:.*]] = %[[C0]] to %[[LVL_B]] step %[[C1]] iter_args
+
+// SPARSE-LOWERED-LABEL: func.func @sparse_loop_ordering_lowered
+// SPARSE-LOWERED-DAG: %[[C0:.*]] = arith.constant 0 : index
+// SPARSE-LOWERED-DAG: %[[C1:.*]] = arith.constant 1 : index
+// SPARSE-LOWERED-DAG: %[[LVL_A:.*]] = sparse_tensor.lvl %arg0, %[[C0]]
+// SPARSE-LOWERED-DAG: %[[POS_A:.*]] = sparse_tensor.positions %arg0 {level = 1 : index}
+// SPARSE-LOWERED-DAG: %[[POS_B:.*]] = sparse_tensor.positions %arg1 {level = 0 : index}
+// SPARSE-LOWERED-DAG: %[[LVL_B:.*]] = sparse_tensor.lvl %arg1, %{{.*}}
+// SPARSE-LOWERED-DAG: %[[DIM_C0:.*]] = tensor.dim %arg2, %[[C0]]
+// SPARSE-LOWERED-DAG: %[[DIM_C1:.*]] = tensor.dim %arg2, %[[C1]]
+// SPARSE-LOWERED-DAG: %[[DIM_C2:.*]] = tensor.dim %arg2, %{{.*}}
+// SPARSE-LOWERED: %[[B_COMPRESSED_START:.*]] = memref.load %[[POS_B]][%[[C0]]]
+// SPARSE-LOWERED: %[[B_COMPRESSED_END:.*]] = memref.load %[[POS_B]][%[[C1]]]
+// SPARSE-LOWERED: scf.for %[[B_COMPRESSED:.*]] = %[[B_COMPRESSED_START]] to %[[B_COMPRESSED_END]] step %[[C1]] {
+// SPARSE-LOWERED: %[[B_COMPRESSED_PLUS_1:.*]] = arith.addi %[[B_COMPRESSED]], %[[C1]]
+// SPARSE-LOWERED: scf.for %[[B_SINGLETON:.*]] = %[[B_COMPRESSED]] to %[[B_COMPRESSED_PLUS_1]] step %[[C1]] {
+// SPARSE-LOWERED: scf.for %[[C_DENSE_0:.*]] = %[[C0]] to %[[DIM_C0]] step %[[C1]] {
+// SPARSE-LOWERED: scf.for %[[C_DENSE_1:.*]] = %[[C0]] to %[[DIM_C1]] step %[[C1]] {
+// SPARSE-LOWERED: scf.for %[[C_DENSE_2:.*]] = %[[C0]] to %[[DIM_C2]] step %[[C1]] {
+// SPARSE-LOWERED: scf.for %[[A_DENSE:.*]] = %[[C0]] to %[[LVL_A]] step %[[C1]] iter_args
+// SPARSE-LOWERED: %[[A_COMPRESSED_START:.*]] = memref.load %[[POS_A]][%[[A_DENSE]]]
+// SPARSE-LOWERED: %[[A_DENSE_PLUS_1:.*]] = arith.addi %[[A_DENSE]], %[[C1]]
+// SPARSE-LOWERED: %[[A_COMPRESSED_END:.*]] = memref.load %[[POS_A]][%[[A_DENSE_PLUS_1]]]
+// SPARSE-LOWERED: scf.for %[[A_COMPRESSED:.*]] = %[[A_COMPRESSED_START]] to %[[A_COMPRESSED_END]] step %[[C1]] iter_args
+// SPARSE-LOWERED: %[[A_COMPRESSED_PLUS_1:.*]] = arith.addi %[[A_COMPRESSED]], %[[C1]]
+// SPARSE-LOWERED: scf.for %[[A_SINGLETON:.*]] = %[[A_COMPRESSED]] to %[[A_COMPRESSED_PLUS_1]] step %[[C1]] iter_args
+// SPARSE-LOWERED: scf.for %[[B_DENSE:.*]] = %[[C0]] to %[[LVL_B]] step %[[C1]] iter_args
+
+func.func @sparse_loop_ordering_lowered(%A: tensor<?x?x?xf32, #X>,
+ %B: tensor<?x?x?xf32, #Y>,
+ %C: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+ %result = linalg.generic #trait
+ ins(%A, %B: tensor<?x?x?xf32, #X>, tensor<?x?x?xf32, #Y>)
+ outs(%C: tensor<?x?x?xf32>) {
+ ^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>
+ return %result : tensor<?x?x?xf32>
+}
More information about the Mlir-commits
mailing list