[Mlir-commits] [mlir] [MLIR][SparseTensor] Implemented LoopUseOuter and LoopUseInner heuristics for loop ordering (PR #232591)
Govind Malasani
llvmlistbot at llvm.org
Thu Oct 15 22:48:19 PDT 2026
https://github.com/gmalasan updated https://github.com/llvm/llvm-project/pull/232591
>From 1d5824016201fe45c61e0a7320c863d70325bc47 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Thu, 15 Oct 2026 20:37:34 -0700
Subject: [PATCH 1/3] [MLIR][SparseTensor] Implemented LoopUseOuter and
LoopUseInner heuristics for loop ordering
---
.../Dialect/SparseTensor/Transforms/Passes.h | 2 +
.../Dialect/SparseTensor/Transforms/Passes.td | 6 +-
.../Transforms/Utils/IterationGraphSorter.cpp | 59 +++++++++++++++-
.../Transforms/Utils/IterationGraphSorter.h | 4 ++
.../sparse_loop_ordering_inputs.mlir | 70 +++++++++++++++++++
5 files changed, 138 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 40b37dc05e92e6..058a567a4e8413 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -63,6 +63,8 @@ enum class LoopOrderingStrategy : unsigned {
kDefault,
kDenseOuter,
kSparseOuter,
+ kLoopUseOuter,
+ kLoopUseInner,
};
} // namespace sparse_tensor
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 0960254a314692..3be3dc5a836eb4 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -88,7 +88,11 @@ def SparseReinterpretMap : Pass<"sparse-reinterpret-map", "ModuleOp"> {
clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kDenseOuter, "dense-outer",
"Prefer dense, then compressed, then singleton dimensions outermost"),
clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kSparseOuter, "sparse-outer",
- "Prefer singleton, then compressed, then dense dimensions outermost"))}]>,
+ "Prefer singleton, then compressed, then dense dimensions outermost"),
+ clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kLoopUseOuter, "loop-use-outer",
+ "Prefer loops used by more input accesses outermost"),
+ clEnumValN(mlir::sparse_tensor::LoopOrderingStrategy::kLoopUseInner, "loop-use-inner",
+ "Prefer loops used by more input accesses innermost"))}]>,
];
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index a0180d228a36ad..959dfb622285fb 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -13,6 +13,7 @@
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/DenseSet.h"
using namespace mlir;
using namespace mlir::sparse_tensor;
@@ -135,6 +136,35 @@ static unsigned getLoopSparsityRank(unsigned loop, ArrayRef<Value> allTensors,
return minRank;
}
+/// Count distinct input accesses whose indexing expressions use each loop.
+static SmallVector<unsigned> getLoopUseCounts(linalg::GenericOp genericOp) {
+ SmallVector<unsigned> counts(genericOp.getNumLoops(), 0);
+ llvm::SmallDenseSet<std::pair<Value, AffineMap>> seen;
+ auto maps = genericOp.getIndexingMapsArray();
+
+ for (auto [idx, input] : llvm::enumerate(genericOp.getDpsInputs())) {
+ if (genericOp.getRegion().front().getArgument(idx).use_empty())
+ continue;
+
+ AffineMap map = maps[idx];
+ if (!seen.insert({input, map}).second)
+ continue;
+
+ llvm::SmallDenseSet<unsigned> accessedLoops;
+ for (AffineExpr expr : map.getResults()) {
+ // Count dependencies within expressions.
+ expr.walk([&](AffineExpr subExpr) {
+ if (auto dim = dyn_cast<AffineDimExpr>(subExpr))
+ accessedLoops.insert(dim.getPosition());
+ });
+ }
+ // Each access contributes once per loop.
+ for (unsigned loop : accessedLoops)
+ ++counts[loop];
+ }
+ return counts;
+}
+
AffineMap IterationGraphSorter::topoSort() {
// The sorted result will put the first Reduction iterator to the
// latest possible position.
@@ -162,6 +192,22 @@ AffineMap IterationGraphSorter::topoSort() {
case sparse_tensor::LoopOrderingStrategy::kDefault:
src = it.back();
break;
+ case sparse_tensor::LoopOrderingStrategy::kLoopUseOuter:
+ case sparse_tensor::LoopOrderingStrategy::kLoopUseInner: {
+ // Loop use counts must already be populated.
+ assert(loopUseCounts.size() == numLoops);
+
+ const bool preferMore = strategy == LoopOrderingStrategy::kLoopUseOuter;
+ src = it.front();
+ for (unsigned candidate : it) {
+ unsigned count = loopUseCounts[candidate];
+ unsigned bestCount = loopUseCounts[src];
+
+ if (preferMore ? count > bestCount : count < bestCount)
+ src = candidate;
+ }
+ break;
+ }
case sparse_tensor::LoopOrderingStrategy::kDenseOuter: {
// Prefer dense, then compressed, then singleton dimensions outermost.
// Create combined tensor and map lists for analysis.
@@ -250,14 +296,14 @@ IterationGraphSorter IterationGraphSorter::fromGenericOp(
genericOp.getIteratorTypesArray();
return IterationGraphSorter(std::move(ins), std::move(loopMap), out, outMap,
- std::move(iterTypes), strategy);
+ std::move(iterTypes), genericOp, strategy);
}
IterationGraphSorter::IterationGraphSorter(
SmallVector<Value> &&insArg, SmallVector<AffineMap> &&loop2InsLvlArg,
Value out, AffineMap loop2OutLvl,
SmallVector<utils::IteratorType> &&iterTypesArg,
- sparse_tensor::LoopOrderingStrategy strategy)
+ linalg::GenericOp genericOp, sparse_tensor::LoopOrderingStrategy strategy)
: ins(std::move(insArg)), loop2InsLvl(std::move(loop2InsLvlArg)), out(out),
loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypesArg)),
strategy(strategy) {
@@ -280,6 +326,15 @@ IterationGraphSorter::IterationGraphSorter(
return m.getNumResults() == 0;
}));
+ switch (strategy) {
+ case LoopOrderingStrategy::kLoopUseOuter:
+ case LoopOrderingStrategy::kLoopUseInner:
+ loopUseCounts = getLoopUseCounts(genericOp);
+ break;
+ default:
+ break;
+ }
+
itGraph.resize(getNumLoops(), std::vector<bool>(getNumLoops(), false));
inDegree.resize(getNumLoops());
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
index 35e58edeb2562c..820fc813944080 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.h
@@ -63,6 +63,7 @@ class IterationGraphSorter {
SmallVector<AffineMap> &&loop2InsLvlArg, Value out,
AffineMap loop2OutLvl,
SmallVector<utils::IteratorType> &&iterTypesArg,
+ linalg::GenericOp genericOp,
sparse_tensor::LoopOrderingStrategy strategy =
sparse_tensor::LoopOrderingStrategy::kDefault);
@@ -93,6 +94,9 @@ class IterationGraphSorter {
// Loop ordering strategy.
sparse_tensor::LoopOrderingStrategy strategy;
+
+ // Prepared once in the constructor for the loop-use strategies.
+ SmallVector<unsigned> loopUseCounts;
};
} // namespace sparse_tensor
diff --git a/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
new file mode 100644
index 00000000000000..1c0533d52d5b00
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
@@ -0,0 +1,70 @@
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=default" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,DEFAULT --enable-var-scope
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-outer" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,OUTER --enable-var-scope
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-inner" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,INNER --enable-var-scope
+
+// A[i] * B[j] * C[j]: i=1, j=2.
+// CHECK-LABEL: func.func @intersection(
+// CHECK: %[[A:.*]] = sparse_tensor.values %arg0
+// DEFAULT-NOT: scf.for
+// DEFAULT: scf.while
+// DEFAULT: scf.for %[[I:[^ ]+]] =
+// DEFAULT: memref.load %[[A]][%[[I]]]
+// OUTER-NOT: scf.for
+// OUTER: scf.while
+// OUTER: scf.for %[[I:[^ ]+]] =
+// OUTER: memref.load %[[A]][%[[I]]]
+// INNER-NOT: scf.while
+// INNER: scf.for %[[I:[^ ]+]] =
+// INNER: memref.load %[[A]][%[[I]]]
+// INNER: scf.while
+
+#S = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
+func.func @intersection(%a: tensor<64xi64, #S>, %b: tensor<64xi64, #S>, %c: tensor<64xi64, #S>, %init: tensor<i64>) -> tensor<i64> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(i,j)->(i)>, affine_map<(i,j)->(j)>, affine_map<(i,j)->(j)>, affine_map<(i,j)->()>],
+ iterator_types = ["reduction", "reduction"]
+ } ins(%a, %b, %c : tensor<64xi64, #S>, tensor<64xi64, #S>, tensor<64xi64, #S>) outs(%init : tensor<i64>) {
+ ^bb0(%x: i64, %y: i64, %z: i64, %acc: i64):
+ %p = arith.muli %x, %y : i64
+ %q = arith.muli %p, %z : i64
+ %s = arith.addi %q, %acc : i64
+ linalg.yield %s : i64
+ } -> tensor<i64>
+ return %r : tensor<i64>
+}
+
+// -----
+
+// A[2*i+1] * B[i] * C[j]: i=2, j=1
+// CHECK-LABEL: func.func @scaled(
+// CHECK: %[[B:.*]] = sparse_tensor.values %arg1
+// DEFAULT: scf.for %[[J:[^ ]+]] =
+// DEFAULT: tensor.extract %arg2[%[[J]]]
+// DEFAULT: scf.for %[[I:[^ ]+]] =
+// DEFAULT: memref.load %[[B]][%[[I]]]
+// OUTER: scf.for %[[I:[^ ]+]] =
+// OUTER: memref.load %[[B]][%[[I]]]
+// OUTER: scf.for %[[J:[^ ]+]] =
+// OUTER: tensor.extract %arg2[%[[J]]]
+// INNER: scf.for %[[J:[^ ]+]] =
+// INNER: tensor.extract %arg2[%[[J]]]
+// INNER: scf.for %[[I:[^ ]+]] =
+// INNER: memref.load %[[B]][%[[I]]]
+
+#S = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
+func.func @scaled(%a: tensor<16xi64>, %b: tensor<8xi64, #S>, %c: tensor<8xi64>, %init: tensor<i64>) -> tensor<i64> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(i,j)->(2*i+1)>, affine_map<(i,j)->(i)>, affine_map<(i,j)->(j)>, affine_map<(i,j)->()>],
+ iterator_types = ["reduction", "reduction"]
+ } ins(%a, %b, %c : tensor<16xi64>, tensor<8xi64, #S>, tensor<8xi64>) outs(%init : tensor<i64>) {
+ ^bb0(%x: i64, %y: i64, %z: i64, %acc: i64):
+ %p = arith.muli %x, %y : i64
+ %q = arith.muli %p, %z : i64
+ %s = arith.addi %q, %acc : i64
+ linalg.yield %s : i64
+ } -> tensor<i64>
+ return %r : tensor<i64>
+}
>From 48a1abd7b66cd5faf1026a746a3596656cc36149 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Thu, 15 Oct 2026 21:35:53 -0700
Subject: [PATCH 2/3] [MLIR][SparseTensor] Fix formatting
---
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 058a567a4e8413..5261d372ff7e19 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -66,7 +66,6 @@ enum class LoopOrderingStrategy : unsigned {
kLoopUseOuter,
kLoopUseInner,
};
-
} // namespace sparse_tensor
#define GEN_PASS_DECL
>From b655972ccf93172a942b7b1c1ae5049328614e90 Mon Sep 17 00:00:00 2001
From: gmalasan <145235389+gmalasan at users.noreply.github.com>
Date: Thu, 15 Oct 2026 22:45:39 -0700
Subject: [PATCH 3/3] [MLIR][SparseTensor] Fix brittle test to more explicitly
check affine maps and the lowered loop structure
---
.../sparse_loop_ordering_inputs.mlir | 77 ++++++++++++-------
1 file changed, 48 insertions(+), 29 deletions(-)
diff --git a/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
index 1c0533d52d5b00..34e6c41e70da93 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_loop_ordering_inputs.mlir
@@ -1,25 +1,38 @@
// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=default" \
-// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,DEFAULT --enable-var-scope
+// RUN: --mlir-print-local-scope | FileCheck %s --check-prefixes=CHECK,DEFAULT
// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-outer" \
-// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,OUTER --enable-var-scope
+// RUN: --mlir-print-local-scope | FileCheck %s --check-prefixes=CHECK,OUTER
// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-inner" \
-// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=CHECK,INNER --enable-var-scope
+// RUN: --mlir-print-local-scope | FileCheck %s --check-prefixes=CHECK,INNER
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=default" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=LOWER,DEFAULT-LOWER --enable-var-scope
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-outer" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=LOWER,OUTER-LOWER --enable-var-scope
+// RUN: mlir-opt %s --split-input-file --sparse-reinterpret-map="loop-ordering-strategy=loop-use-inner" \
+// RUN: --sparsification --canonicalize | FileCheck %s --check-prefixes=LOWER,INNER-LOWER --enable-var-scope
// A[i] * B[j] * C[j]: i=1, j=2.
// CHECK-LABEL: func.func @intersection(
-// CHECK: %[[A:.*]] = sparse_tensor.values %arg0
-// DEFAULT-NOT: scf.for
-// DEFAULT: scf.while
-// DEFAULT: scf.for %[[I:[^ ]+]] =
-// DEFAULT: memref.load %[[A]][%[[I]]]
-// OUTER-NOT: scf.for
-// OUTER: scf.while
-// OUTER: scf.for %[[I:[^ ]+]] =
-// OUTER: memref.load %[[A]][%[[I]]]
-// INNER-NOT: scf.while
-// INNER: scf.for %[[I:[^ ]+]] =
-// INNER: memref.load %[[A]][%[[I]]]
-// INNER: scf.while
+// CHECK: linalg.generic
+// DEFAULT-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> ()>]
+// OUTER-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> ()>]
+// INNER-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> ()>]
+// CHECK-SAME: sorted = true
+
+// LOWER-LABEL: func.func @intersection(
+// LOWER: %[[A:.*]] = sparse_tensor.values %arg0
+// DEFAULT-LOWER-NOT: scf.for
+// DEFAULT-LOWER: scf.while
+// DEFAULT-LOWER: scf.for %[[I:[^ ]+]] =
+// DEFAULT-LOWER: memref.load %[[A]][%[[I]]]
+// OUTER-LOWER-NOT: scf.for
+// OUTER-LOWER: scf.while
+// OUTER-LOWER: scf.for %[[I:[^ ]+]] =
+// OUTER-LOWER: memref.load %[[A]][%[[I]]]
+// INNER-LOWER-NOT: scf.while
+// INNER-LOWER: scf.for %[[I:[^ ]+]] =
+// INNER-LOWER: memref.load %[[A]][%[[I]]]
+// INNER-LOWER: scf.while
#S = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
func.func @intersection(%a: tensor<64xi64, #S>, %b: tensor<64xi64, #S>, %c: tensor<64xi64, #S>, %init: tensor<i64>) -> tensor<i64> {
@@ -40,19 +53,25 @@ func.func @intersection(%a: tensor<64xi64, #S>, %b: tensor<64xi64, #S>, %c: tens
// A[2*i+1] * B[i] * C[j]: i=2, j=1
// CHECK-LABEL: func.func @scaled(
-// CHECK: %[[B:.*]] = sparse_tensor.values %arg1
-// DEFAULT: scf.for %[[J:[^ ]+]] =
-// DEFAULT: tensor.extract %arg2[%[[J]]]
-// DEFAULT: scf.for %[[I:[^ ]+]] =
-// DEFAULT: memref.load %[[B]][%[[I]]]
-// OUTER: scf.for %[[I:[^ ]+]] =
-// OUTER: memref.load %[[B]][%[[I]]]
-// OUTER: scf.for %[[J:[^ ]+]] =
-// OUTER: tensor.extract %arg2[%[[J]]]
-// INNER: scf.for %[[J:[^ ]+]] =
-// INNER: tensor.extract %arg2[%[[J]]]
-// INNER: scf.for %[[I:[^ ]+]] =
-// INNER: memref.load %[[B]][%[[I]]]
+// CHECK: linalg.generic
+// DEFAULT-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1 * 2 + 1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> ()>]
+// OUTER-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0 * 2 + 1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> ()>]
+// INNER-SAME: indexing_maps = [affine_map<(d0, d1) -> (d1 * 2 + 1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> ()>]
+// CHECK-SAME: sorted = true
+
+// LOWER-LABEL: func.func @scaled(
+// LOWER-DAG: %[[C0:.*]] = arith.constant 0 : index
+// LOWER-DAG: %[[C8:.*]] = arith.constant 8 : index
+// LOWER-DAG: %[[B:.*]] = sparse_tensor.values %arg1
+// DEFAULT-LOWER: scf.for %{{[^ ]+}} = %[[C0]] to %[[C8]]
+// DEFAULT-LOWER: scf.for %[[I:[^ ]+]] =
+// DEFAULT-LOWER: memref.load %[[B]][%[[I]]]
+// OUTER-LOWER: scf.for %[[I:[^ ]+]] =
+// OUTER-LOWER: memref.load %[[B]][%[[I]]]
+// OUTER-LOWER: scf.for %{{[^ ]+}} = %[[C0]] to %[[C8]]
+// INNER-LOWER: scf.for %{{[^ ]+}} = %[[C0]] to %[[C8]]
+// INNER-LOWER: scf.for %[[I:[^ ]+]] =
+// INNER-LOWER: memref.load %[[B]][%[[I]]]
#S = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
func.func @scaled(%a: tensor<16xi64>, %b: tensor<8xi64, #S>, %c: tensor<8xi64>, %init: tensor<i64>) -> tensor<i64> {
More information about the Mlir-commits
mailing list