[Mlir-commits] [mlir] [mlir][sparse] Fix mixed sparse-dense iterator lowering (PR #219753)
Kohei Yamaguchi
llvmlistbot at llvm.org
Sat Oct 31 10:20:41 PDT 2026
https://github.com/sott0n updated https://github.com/llvm/llvm-project/pull/219753
>From bf7793f02179204f7de51b44a2172a8b0b704197 Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Sun, 30 Aug 2026 08:09:34 +0900
Subject: [PATCH 1/3] [mlir][sparse] Fix mixed sparse-dense iterator lowering
---
.../Transforms/SparseIterationToScf.cpp | 16 +++--
.../Transforms/Sparsification.cpp | 67 ++++++++++++-------
.../sparse_kernels_to_iterator.mlir | 55 +++++++++++++++
3 files changed, 107 insertions(+), 31 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
index 09ca7474ccbbe8..6e793d96c33e51 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
@@ -228,9 +228,11 @@ class SparseIterateOpConverter : public OpConversionPattern<IterateOp> {
LogicalResult
matchAndRewrite(IterateOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- if (!op.getCrdUsedLvls().empty())
+ I64BitSet crdUsedLvls = op.getCrdUsedLvls();
+ if (crdUsedLvls.count() > 1 ||
+ (!crdUsedLvls.empty() && !crdUsedLvls[op.getSpaceDim() - 1]))
return rewriter.notifyMatchFailure(
- op, "non-empty coordinates list not implemented.");
+ op, "coordinates other than the final level are not implemented");
Location loc = op.getLoc();
@@ -254,13 +256,15 @@ class SparseIterateOpConverter : public OpConversionPattern<IterateOp> {
Block *block = rewriter.applySignatureConversion(
op.getBody(), signatureConversion, getTypeConverter());
+ bool usesCrd = !crdUsedLvls.empty();
ValueRange ret = genLoopWithIterator(
rewriter, loc, it.get(), ivs,
- [block](PatternRewriter &rewriter, Location loc, Region &loopBody,
- SparseIterator *it, ValueRange reduc) -> SmallVector<Value> {
+ [block, usesCrd](PatternRewriter &rewriter, Location loc,
+ Region &loopBody, SparseIterator *it,
+ ValueRange reduc) -> SmallVector<Value> {
SmallVector<Value> blockArgs(reduc);
- // TODO: Also appends coordinates if used.
- // blockArgs.push_back(it->deref(rewriter, loc));
+ if (usesCrd)
+ blockArgs.push_back(it->deref(rewriter, loc));
llvm::append_range(blockArgs, it->getCursor());
Block *dstBlock = &loopBody.getBlocks().front();
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 860ea68029c3a4..9b307b31de8214 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -1008,7 +1008,7 @@ static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
static bool getAllTidLvlsInLatPoints(
CodegenEnv &env, LatPointId li, LoopId curr,
- llvm::function_ref<void(TensorLevel, AffineExpr)> callback) {
+ llvm::function_ref<void(TensorLevel, AffineExpr, bool)> callback) {
const BitVector &simple = env.lat(li).simple;
const TensorId outTid = env.merger().getOutTensorID();
const std::optional<Level> outLvl = env.merger().getLvl(outTid, curr);
@@ -1020,7 +1020,8 @@ static bool getAllTidLvlsInLatPoints(
LevelType lt, bool isIdxReduc) {
if (simple[b]) {
if (isIdxReduc) {
- callback(env.makeTensorLevel(tid, *lvl), nullptr);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr,
+ /*isLoopCond=*/true);
numloopCond++;
return;
}
@@ -1044,10 +1045,12 @@ static bool getAllTidLvlsInLatPoints(
}
}
hasNonUnique = !isUniqueLT(lt) || hasNonUnique;
- callback(env.makeTensorLevel(tid, *lvl), nullptr);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr,
+ /*isLoopCond=*/true);
numloopCond++;
} else if (lt.hasDenseSemantic() || isIdxReduc) {
- callback(env.makeTensorLevel(tid, *lvl), nullptr);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr,
+ /*isLoopCond=*/false);
} else {
assert(isUndefLT(lt));
linalg::GenericOp op = env.op();
@@ -1082,7 +1085,8 @@ static bool getAllTidLvlsInLatPoints(
// level. We need to generate the address according to the
// affine expression. This is also the best place we can do it
// to avoid putting it inside inner loops.
- callback(env.makeTensorLevel(tid, l), exp);
+ callback(env.makeTensorLevel(tid, l), exp,
+ /*isLoopCond=*/false);
}
}
}
@@ -1097,14 +1101,16 @@ static bool getAllTidLvlsInLatPoints(
// TODO: we should avoid introducing corner cases for all-dense sparse
// tensors.
if (stt.hasEncoding() && stt.isAllDense())
- callback(env.makeTensorLevel(outTid, *outLvl), nullptr);
+ callback(env.makeTensorLevel(outTid, *outLvl), nullptr,
+ /*isLoopCond=*/false);
}
if (numloopCond == 0) {
// Corner cases where the loop bound is defined by a *unused* operand, in
// this case, we just generate a dense "fake" loop by iterating over the
// synthetic tensor.
- callback(env.makeTensorLevel(env.merger().getSynTensorID(), curr), nullptr);
+ callback(env.makeTensorLevel(env.merger().getSynTensorID(), curr), nullptr,
+ /*isLoopCond=*/true);
numloopCond++;
}
// If we just need to one loop conditions and the conditions is not imposed on
@@ -1129,14 +1135,16 @@ static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
const LatPointId l0 = env.set(lts)[0];
SmallVector<TensorLevel> tidLvls;
- getAllTidLvlsInLatPoints(env, l0, curr, [&](TensorLevel tl, AffineExpr) {
- // TODO: remove this! The same tensor level might be added for multiple
- // times due to the special handling for all-dense "sparse" output tensor
- // (see L1038).
- if (llvm::is_contained(tidLvls, tl))
- return;
- tidLvls.emplace_back(tl);
- });
+ getAllTidLvlsInLatPoints(env, l0, curr,
+ [&](TensorLevel tl, AffineExpr, bool) {
+ // TODO: remove this! The same tensor level might
+ // be added for multiple times due to the special
+ // handling for all-dense "sparse" output tensor
+ // (see L1038).
+ if (llvm::is_contained(tidLvls, tl))
+ return;
+ tidLvls.emplace_back(tl);
+ });
env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tidLvls);
@@ -1190,14 +1198,18 @@ static void genInitConstantDenseAddress(CodegenEnv &env,
static bool translateBitsToTidLvlPairs(
CodegenEnv &env, LatPointId li, LoopId curr,
SmallVectorImpl<TensorLevel> &tidLvls,
+ SmallVectorImpl<TensorLevel> &condTidLvls,
SmallVectorImpl<std::pair<TensorLevel, AffineExpr>> &affineTidLvls) {
- return getAllTidLvlsInLatPoints(env, li, curr,
- [&](TensorLevel tl, AffineExpr exp) {
- if (exp)
- affineTidLvls.emplace_back(tl, exp);
- else
- tidLvls.emplace_back(tl);
- });
+ return getAllTidLvlsInLatPoints(
+ env, li, curr, [&](TensorLevel tl, AffineExpr exp, bool isLoopCond) {
+ if (exp)
+ affineTidLvls.emplace_back(tl, exp);
+ else {
+ tidLvls.emplace_back(tl);
+ if (isLoopCond)
+ condTidLvls.emplace_back(tl);
+ }
+ });
}
/// Starts a single loop in current sequence.
@@ -1209,15 +1221,20 @@ static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
// after fully migration.
// The set of tensors + lvls to generate loops on
SmallVector<TensorLevel> tidLvls;
+ SmallVector<TensorLevel> condTidLvls;
// The set of dense tensors with non-trivial affine expression that just
// becomes invariant and the address are generated at the current level.
SmallVector<std::pair<TensorLevel, AffineExpr>> affineTidLvls;
- bool isSingleCond =
- translateBitsToTidLvlPairs(env, li, curr, tidLvls, affineTidLvls);
+ bool isSingleCond = translateBitsToTidLvlPairs(env, li, curr, tidLvls,
+ condTidLvls, affineTidLvls);
// Emit the for/while-loop control.
- Operation *loop = genLoop(env, builder, curr, numCases, needsUniv, tidLvls);
+ ArrayRef<TensorLevel> loopTidLvls = env.generatingSparseIterator()
+ ? ArrayRef(condTidLvls)
+ : ArrayRef(tidLvls);
+ Operation *loop =
+ genLoop(env, builder, curr, numCases, needsUniv, loopTidLvls);
Location loc = env.op().getLoc();
for (auto [tidLvl, exp] : affineTidLvls) {
env.emitter().locateLvlAtAffineAddress(builder, loc, tidLvl, exp);
diff --git a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
index 4ca13514498905..0ad5ca575fd5cb 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
@@ -16,6 +16,10 @@
map = (d0) -> (d0 : compressed)
}>
+#CSR = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : dense, d1 : compressed)
+}>
+
// CHECK-LABEL: func.func @sqsum(
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
@@ -167,3 +171,54 @@ func.func @add(%arg0: tensor<10xi32, #VEC>, %arg1: tensor<10xi32, #VEC>) -> tens
} -> tensor<10xi32>
return %0 : tensor<10xi32>
}
+
+// ITER-LABEL: func.func @matvec(
+// ITER: %[[VECTOR:.*]] = bufferization.to_buffer
+// ITER: %[[ROW_SPACE:.*]] = sparse_tensor.extract_iteration_space %{{.*}} lvls = 0
+// ITER: sparse_tensor.iterate %[[ROW_ITER:.*]] in %[[ROW_SPACE]] at(%[[ROW:[a-zA-Z0-9_]+]])
+// ITER: %[[COL_SPACE:.*]] = sparse_tensor.extract_iteration_space %{{.*}} at %[[ROW_ITER]] lvls = 1
+// ITER: sparse_tensor.iterate %[[COL_ITER:.*]] in %[[COL_SPACE]] at(%[[COL:[a-zA-Z0-9_]+]])
+// ITER: %[[MATRIX_VALUE:.*]] = sparse_tensor.extract_value %{{.*}} at %[[COL_ITER]]
+// ITER: %[[VECTOR_VALUE:.*]] = memref.load %[[VECTOR]][%[[COL]]]
+// ITER: %[[PRODUCT:.*]] = arith.mulf %[[MATRIX_VALUE]], %[[VECTOR_VALUE]] : f64
+// ITER: arith.addf %{{.*}}, %[[PRODUCT]] : f64
+
+// CHECK-LABEL: func.func @matvec(
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[VECTOR:.*]] = bufferization.to_buffer
+// CHECK: %[[OUTPUT:.*]] = bufferization.to_buffer
+// CHECK: scf.for %[[ROW:.*]] = %[[C0]] to %{{.*}} step %[[C1]] {
+// CHECK: %[[SUM:.*]] = memref.load %[[OUTPUT]][%[[ROW]]]
+// CHECK: %[[COL_LO:.*]] = memref.load %{{.*}}[%[[ROW]]]
+// CHECK: %[[ROW_END:.*]] = arith.addi %[[ROW]], %[[C1]] : index
+// CHECK: %[[COL_HI:.*]] = memref.load %{{.*}}[%[[ROW_END]]]
+// CHECK: %[[RESULT:.*]] = scf.for %[[POS:.*]] = %[[COL_LO]] to %[[COL_HI]] step %[[C1]] iter_args(%[[ACC:.*]] = %[[SUM]]) -> (f64) {
+// CHECK: %[[COL:.*]] = memref.load %{{.*}}[%[[POS]]]
+// CHECK: %[[MATRIX_VALUE:.*]] = memref.load %{{.*}}[%[[POS]]]
+// CHECK: %[[VECTOR_VALUE:.*]] = memref.load %[[VECTOR]][%[[COL]]]
+// CHECK: %[[PRODUCT:.*]] = arith.mulf %[[MATRIX_VALUE]], %[[VECTOR_VALUE]] : f64
+// CHECK: %[[NEXT:.*]] = arith.addf %[[ACC]], %[[PRODUCT]] : f64
+// CHECK: scf.yield %[[NEXT]] : f64
+// CHECK: }
+// CHECK: memref.store %[[RESULT]], %[[OUTPUT]][%[[ROW]]]
+// CHECK: }
+func.func @matvec(%matrix: tensor<4x8xf64, #CSR>,
+ %vector: tensor<8xf64>,
+ %output: tensor<4xf64>) -> tensor<4xf64> {
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(i, j) -> (i, j)>,
+ affine_map<(i, j) -> (j)>,
+ affine_map<(i, j) -> (i)>
+ ],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%matrix, %vector : tensor<4x8xf64, #CSR>, tensor<8xf64>)
+ outs(%output : tensor<4xf64>) {
+ ^bb0(%matrixValue: f64, %vectorValue: f64, %sum: f64):
+ %product = arith.mulf %matrixValue, %vectorValue : f64
+ %next = arith.addf %sum, %product : f64
+ linalg.yield %next : f64
+ } -> tensor<4xf64>
+ return %result : tensor<4xf64>
+}
>From c8cde7c380de70ac1a2b307f4671a77939a8b1c7 Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Sun, 1 Nov 2026 00:48:12 +0900
Subject: [PATCH 2/3] Fix mixed sparse iterator loop handling
---
.../Transforms/Sparsification.cpp | 69 +++++------
.../Transforms/Utils/LoopEmitter.cpp | 107 +++++++++++++++---
.../Transforms/Utils/LoopEmitter.h | 9 +-
.../sparse_kernels_to_iterator.mlir | 41 +++++++
4 files changed, 165 insertions(+), 61 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 9b307b31de8214..d4c662c4b9c72f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -1008,7 +1008,7 @@ static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
static bool getAllTidLvlsInLatPoints(
CodegenEnv &env, LatPointId li, LoopId curr,
- llvm::function_ref<void(TensorLevel, AffineExpr, bool)> callback) {
+ llvm::function_ref<void(TensorLevel, AffineExpr)> callback) {
const BitVector &simple = env.lat(li).simple;
const TensorId outTid = env.merger().getOutTensorID();
const std::optional<Level> outLvl = env.merger().getLvl(outTid, curr);
@@ -1020,8 +1020,7 @@ static bool getAllTidLvlsInLatPoints(
LevelType lt, bool isIdxReduc) {
if (simple[b]) {
if (isIdxReduc) {
- callback(env.makeTensorLevel(tid, *lvl), nullptr,
- /*isLoopCond=*/true);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr);
numloopCond++;
return;
}
@@ -1045,12 +1044,10 @@ static bool getAllTidLvlsInLatPoints(
}
}
hasNonUnique = !isUniqueLT(lt) || hasNonUnique;
- callback(env.makeTensorLevel(tid, *lvl), nullptr,
- /*isLoopCond=*/true);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr);
numloopCond++;
} else if (lt.hasDenseSemantic() || isIdxReduc) {
- callback(env.makeTensorLevel(tid, *lvl), nullptr,
- /*isLoopCond=*/false);
+ callback(env.makeTensorLevel(tid, *lvl), nullptr);
} else {
assert(isUndefLT(lt));
linalg::GenericOp op = env.op();
@@ -1085,8 +1082,7 @@ static bool getAllTidLvlsInLatPoints(
// level. We need to generate the address according to the
// affine expression. This is also the best place we can do it
// to avoid putting it inside inner loops.
- callback(env.makeTensorLevel(tid, l), exp,
- /*isLoopCond=*/false);
+ callback(env.makeTensorLevel(tid, l), exp);
}
}
}
@@ -1101,16 +1097,14 @@ static bool getAllTidLvlsInLatPoints(
// TODO: we should avoid introducing corner cases for all-dense sparse
// tensors.
if (stt.hasEncoding() && stt.isAllDense())
- callback(env.makeTensorLevel(outTid, *outLvl), nullptr,
- /*isLoopCond=*/false);
+ callback(env.makeTensorLevel(outTid, *outLvl), nullptr);
}
if (numloopCond == 0) {
// Corner cases where the loop bound is defined by a *unused* operand, in
// this case, we just generate a dense "fake" loop by iterating over the
// synthetic tensor.
- callback(env.makeTensorLevel(env.merger().getSynTensorID(), curr), nullptr,
- /*isLoopCond=*/true);
+ callback(env.makeTensorLevel(env.merger().getSynTensorID(), curr), nullptr);
numloopCond++;
}
// If we just need to one loop conditions and the conditions is not imposed on
@@ -1135,16 +1129,14 @@ static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
const LatPointId l0 = env.set(lts)[0];
SmallVector<TensorLevel> tidLvls;
- getAllTidLvlsInLatPoints(env, l0, curr,
- [&](TensorLevel tl, AffineExpr, bool) {
- // TODO: remove this! The same tensor level might
- // be added for multiple times due to the special
- // handling for all-dense "sparse" output tensor
- // (see L1038).
- if (llvm::is_contained(tidLvls, tl))
- return;
- tidLvls.emplace_back(tl);
- });
+ getAllTidLvlsInLatPoints(env, l0, curr, [&](TensorLevel tl, AffineExpr) {
+ // TODO: remove this! The same tensor level might be added for multiple
+ // times due to the special handling for all-dense "sparse" output tensor
+ // (see L1038).
+ if (llvm::is_contained(tidLvls, tl))
+ return;
+ tidLvls.emplace_back(tl);
+ });
env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tidLvls);
@@ -1198,18 +1190,14 @@ static void genInitConstantDenseAddress(CodegenEnv &env,
static bool translateBitsToTidLvlPairs(
CodegenEnv &env, LatPointId li, LoopId curr,
SmallVectorImpl<TensorLevel> &tidLvls,
- SmallVectorImpl<TensorLevel> &condTidLvls,
SmallVectorImpl<std::pair<TensorLevel, AffineExpr>> &affineTidLvls) {
- return getAllTidLvlsInLatPoints(
- env, li, curr, [&](TensorLevel tl, AffineExpr exp, bool isLoopCond) {
- if (exp)
- affineTidLvls.emplace_back(tl, exp);
- else {
- tidLvls.emplace_back(tl);
- if (isLoopCond)
- condTidLvls.emplace_back(tl);
- }
- });
+ return getAllTidLvlsInLatPoints(env, li, curr,
+ [&](TensorLevel tl, AffineExpr exp) {
+ if (exp)
+ affineTidLvls.emplace_back(tl, exp);
+ else
+ tidLvls.emplace_back(tl);
+ });
}
/// Starts a single loop in current sequence.
@@ -1221,20 +1209,15 @@ static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
// after fully migration.
// The set of tensors + lvls to generate loops on
SmallVector<TensorLevel> tidLvls;
- SmallVector<TensorLevel> condTidLvls;
// The set of dense tensors with non-trivial affine expression that just
// becomes invariant and the address are generated at the current level.
SmallVector<std::pair<TensorLevel, AffineExpr>> affineTidLvls;
- bool isSingleCond = translateBitsToTidLvlPairs(env, li, curr, tidLvls,
- condTidLvls, affineTidLvls);
+ bool isSingleCond =
+ translateBitsToTidLvlPairs(env, li, curr, tidLvls, affineTidLvls);
// Emit the for/while-loop control.
- ArrayRef<TensorLevel> loopTidLvls = env.generatingSparseIterator()
- ? ArrayRef(condTidLvls)
- : ArrayRef(tidLvls);
- Operation *loop =
- genLoop(env, builder, curr, numCases, needsUniv, loopTidLvls);
+ Operation *loop = genLoop(env, builder, curr, numCases, needsUniv, tidLvls);
Location loc = env.op().getLoc();
for (auto [tidLvl, exp] : affineTidLvls) {
env.emitter().locateLvlAtAffineAddress(builder, loc, tidLvl, exp);
@@ -1315,7 +1298,7 @@ static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp,
const LatPointId li = env.set(lts)[0];
auto [loop, isSingleCond] =
startLoop(env, rewriter, curr, li, lsize, needsUniv);
- assert(isSingleCond == llvm::isa<IterateOp>(loop));
+ isSingleCond = llvm::isa<IterateOp, scf::ForOp>(loop);
// We cannot change this to `for (const LatPointId li : env.set(lts))`
// because the loop body causes data-movement which invalidates
// the iterator.
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
index 94d6e640da4367..1746adb90927f9 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
@@ -275,22 +275,24 @@ void LoopEmitter::initializeLoopEmit(
}
}
- // The sparse iterator values will only be available after the loop is
- // constructed.
- if (emitStrategy == SparseEmitStrategy::kSparseIterator)
- return;
-
// For every synthetic tensor, set the high bound by calling the callback.
if (synSetter) {
TensorId synId = getSynTensorId();
for (unsigned i = 0, e = loopHighs.size(); i < e; i++) {
Value sz = loopHighs[i] = synSetter(builder, loc, i);
+ if (emitStrategy == SparseEmitStrategy::kSparseIterator)
+ continue;
auto [stl, it] = makeSynLevelAndIterator(sz, synId, i, emitStrategy);
lvls[synId][i] = std::move(stl);
iters[synId][i].emplace_back(std::move(it));
}
}
+ // The sparse iterator values will only be available after the loop is
+ // constructed.
+ if (emitStrategy == SparseEmitStrategy::kSparseIterator)
+ return;
+
// For every manifest tensor:
// * For every level:
// * get the positions and coordinates buffers
@@ -536,6 +538,38 @@ bool LoopEmitter::shouldIteratedByForLoop(ArrayRef<SparseIterator *> spIters) {
return true;
}
+void LoopEmitter::linkSparseIteratorDenseLvls(OpBuilder &builder, Location loc,
+ ArrayRef<TensorLevel> tidLvls,
+ Value crd) {
+ // Dense levels do not drive coiteration, but their iterator values are
+ // needed when extracting a child iteration space.
+ for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
+ SparseTensorType stt = getSparseTensorType(tensors[tid]);
+ assert(stt.getLvlType(lvl).hasDenseSemantic());
+
+ Value pos = crd;
+ if (lvl != 0) {
+ Value parentIter = spIterVals[tid][lvl - 1];
+ assert(parentIter && "expected parent iterator for encoded dense level");
+ IteratorType parentTp = cast<IteratorType>(parentIter.getType());
+ SmallVector<Type> cursorTps(parentTp.isUnique() ? 1 : 2,
+ builder.getIndexType());
+ auto parentCursor = UnrealizedConversionCastOp::create(
+ builder, loc, cursorTps, parentIter);
+ Value parentPos = parentCursor.getResults().back();
+ Value lvlSize = LvlOp::create(builder, loc, tensors[tid], lvl);
+ pos = ADDI(MULI(parentPos, lvlSize), crd);
+ }
+
+ IteratorType itTp = IterSpaceType::get(builder.getContext(),
+ stt.getEncoding(), lvl, lvl + 1)
+ .getIteratorType();
+ spIterVals[tid][lvl] =
+ UnrealizedConversionCastOp::create(builder, loc, itTp, pos)
+ .getResult(0);
+ }
+}
+
Region *LoopEmitter::enterCurrentCoIterationCase(OpBuilder &builder,
Location loc,
I64BitSet caseBit,
@@ -586,6 +620,9 @@ Region *LoopEmitter::enterCurrentCoIterationCase(OpBuilder &builder,
}
// Must have consumed all iterator SSA values.
assert(iters.empty());
+ linkSparseIteratorDenseLvls(builder, loc,
+ loopStack.back().auxiliaryDenseTidLvls,
+ loopStack.back().iv);
return &caseRegion;
}
@@ -595,10 +632,43 @@ Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
bool needsUniv) {
// TODO: Argument `numCases` only used when generating iterator-based sparse
// loops. Simplify the code upon feature complete.
- // TODO: handle coiteration with sparse iterator.
if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
- if (tidLvls.size() == 1) {
- auto [tid, lvl] = unpackTensorLevel(tidLvls.front());
+ SmallVector<TensorLevel> iteratedTidLvls;
+ SmallVector<TensorLevel> auxiliaryDenseTidLvls;
+ for (auto [tidLvl, tidAndLvl] :
+ llvm::zip_equal(tidLvls, unpackTensorLevelRange(tidLvls))) {
+ auto [tid, lvl] = tidAndLvl;
+ if (isSynTensor(tid))
+ continue;
+ SparseTensorType stt = getSparseTensorType(tensors[tid]);
+ if (!stt.hasEncoding())
+ continue;
+ if (stt.getLvlType(lvl).hasDenseSemantic())
+ auxiliaryDenseTidLvls.push_back(tidLvl);
+ else
+ iteratedTidLvls.push_back(tidLvl);
+ }
+
+ if (iteratedTidLvls.empty() && !auxiliaryDenseTidLvls.empty()) {
+ iteratedTidLvls.push_back(auxiliaryDenseTidLvls.front());
+ auxiliaryDenseTidLvls.erase(auxiliaryDenseTidLvls.begin());
+ }
+
+ // Unencoded dense and synthetic levels are represented by an SCF loop.
+ if (iteratedTidLvls.empty()) {
+ Value hi = loopHighs[loopStack.size()];
+ assert(hi && "expected a loop bound for dense iteration");
+ scf::ForOp forOp =
+ scf::ForOp::create(builder, loc, C_IDX(0), hi, C_IDX(1), reduc);
+ llvm::copy(forOp.getRegionIterArgs(), reduc.begin());
+ builder.setInsertionPointToStart(forOp.getBody());
+ loopStack.emplace_back(tidLvls, forOp, builder.getInsertionBlock(),
+ forOp.getInductionVar(), loopTag);
+ return forOp;
+ }
+
+ if (iteratedTidLvls.size() == 1) {
+ auto [tid, lvl] = unpackTensorLevel(iteratedTidLvls.front());
Value t = tensors[tid];
// Extract and iterate over the iteration space.
@@ -615,14 +685,17 @@ Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
llvm::copy(iterOp.getRegionIterArgs(), reduc.begin());
// Set the insertion point to loop body.
builder.setInsertionPointToStart(iterOp.getBody());
- loopStack.emplace_back(tidLvls, iterOp, builder.getInsertionBlock(),
- iterOp.getCrds().front(), loopTag);
+ linkSparseIteratorDenseLvls(builder, loc, auxiliaryDenseTidLvls,
+ iterOp.getCrds().front());
+ loopStack.emplace_back(
+ iteratedTidLvls, iterOp, builder.getInsertionBlock(),
+ iterOp.getCrds().front(), loopTag, auxiliaryDenseTidLvls);
return iterOp;
}
// CoIteration Loops.
SmallVector<Value> spaces;
- for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
+ for (auto [tid, lvl] : unpackTensorLevelRange(iteratedTidLvls)) {
Value t = tensors[tid];
ExtractIterSpaceOp extractSpaceOp =
lvl == 0 ? ExtractIterSpaceOp::create(builder, loc, t)
@@ -633,8 +706,10 @@ Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
auto coIterOp = CoIterateOp::create(builder, loc, spaces, reduc, numCases);
// The CoIterationOp does not have insertion block nor induction variable.
// TODO: the `struct LoopInfo` should be simplied after full migration.
- loopStack.emplace_back(tidLvls, coIterOp, /*insertion block*/ nullptr,
- /*induction variable*/ nullptr, loopTag);
+ loopStack.emplace_back(iteratedTidLvls, coIterOp,
+ /*insertion block*/ nullptr,
+ /*induction variable*/ nullptr, loopTag,
+ auxiliaryDenseTidLvls);
return coIterOp;
}
@@ -725,8 +800,7 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
MutableArrayRef<Value> reduc) {
const LoopInfo &loopInfo = loopStack.back();
- if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
- auto iterateOp = llvm::cast<IterateOp>(loopInfo.loop);
+ if (auto iterateOp = llvm::dyn_cast<IterateOp>(loopInfo.loop)) {
assert(reduc.size() == iterateOp.getNumResults());
sparse_tensor::YieldOp::create(rewriter, loc, reduc);
// Exit the loop.
@@ -869,7 +943,8 @@ void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
// Clean up the values, it would help use to discover potential bug at a
// earlier stage (instead of silently using a wrong value).
const LoopInfo &loopInfo = loopStack.back();
- if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
+ if (emitStrategy == SparseEmitStrategy::kSparseIterator &&
+ !isa<scf::ForOp>(loopInfo.loop)) {
Operation *p = loopInfo.loop;
if (isa<IterateOp>(p))
sparse_tensor::YieldOp::create(rewriter, loc, reduc);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
index a47062987b4b6c..879dc03b057e21 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
@@ -253,8 +253,10 @@ class LoopEmitter {
// the set of tensors levels that the loop is iterating over.
struct LoopInfo final {
LoopInfo(ArrayRef<TensorLevel> tidLvls, Operation *loop, Block *userBlock,
- Value iv, StringAttr loopTag)
- : tidLvls(tidLvls), loop(loop), userCodeBlock(userBlock), iv(iv) {
+ Value iv, StringAttr loopTag,
+ ArrayRef<TensorLevel> auxiliaryDenseTidLvls = {})
+ : tidLvls(tidLvls), auxiliaryDenseTidLvls(auxiliaryDenseTidLvls),
+ loop(loop), userCodeBlock(userBlock), iv(iv) {
// Attached a special tag to loop emitter generated loop.
if (loopTag)
loop->setDiscardableAttr(LoopEmitter::getLoopEmitterLoopAttrName(),
@@ -265,6 +267,7 @@ class LoopEmitter {
// required for levels with non-tivial index expressions, which is
// maintained by the sliceDrivenInfo array below.
const llvm::SmallVector<TensorLevel> tidLvls;
+ const llvm::SmallVector<TensorLevel> auxiliaryDenseTidLvls;
Operation *loop; // the loop operation
Block *const userCodeBlock; // the block holding users' generated code.
Value iv; // the induction variable for the loop
@@ -273,6 +276,8 @@ class LoopEmitter {
void categorizeIterators(ArrayRef<TensorLevel> tidLvls,
SmallVectorImpl<SparseIterator *> &raIters,
SmallVectorImpl<SparseIterator *> &spIters);
+ void linkSparseIteratorDenseLvls(OpBuilder &builder, Location loc,
+ ArrayRef<TensorLevel> tidLvls, Value crd);
///
/// LoopEmitter internal helper functions.
///
diff --git a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
index 0ad5ca575fd5cb..8585b9a4a2f06e 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
@@ -222,3 +222,44 @@ func.func @matvec(%matrix: tensor<4x8xf64, #CSR>,
} -> tensor<4xf64>
return %result : tensor<4xf64>
}
+
+// ITER-LABEL: func.func @mixed_sparse_encoded_dense_parent(
+// ITER: sparse_tensor.iterate %{{.*}} in %{{.*}} at(%[[ROW:.*]])
+// ITER: %[[DENSE_ITER:.*]] = builtin.unrealized_conversion_cast %[[ROW]]
+// ITER: sparse_tensor.extract_iteration_space %{{.*}} at %[[DENSE_ITER]] lvls = 1
+func.func @mixed_sparse_encoded_dense_parent(
+ %rows: tensor<4xf64, #VEC>, %matrix: tensor<4x8xf64, #CSR>,
+ %output: tensor<4x8xf64>) -> tensor<4x8xf64> {
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(i, j) -> (i)>,
+ affine_map<(i, j) -> (i, j)>,
+ affine_map<(i, j) -> (i, j)>
+ ],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%rows, %matrix : tensor<4xf64, #VEC>, tensor<4x8xf64, #CSR>)
+ outs(%output : tensor<4x8xf64>) {
+ ^bb0(%rowValue: f64, %matrixValue: f64, %out: f64):
+ %product = arith.mulf %rowValue, %matrixValue : f64
+ linalg.yield %product : f64
+ } -> tensor<4x8xf64>
+ return %result : tensor<4x8xf64>
+}
+
+// ITER-LABEL: func.func @sparse_broadcast(
+// ITER: sparse_tensor.iterate
+// ITER: scf.for
+func.func @sparse_broadcast(%input: tensor<4xf64, #VEC>,
+ %output: tensor<4x8xf64>) -> tensor<4x8xf64> {
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(i, j) -> (i)>,
+ affine_map<(i, j) -> (i, j)>
+ ],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%input : tensor<4xf64, #VEC>) outs(%output : tensor<4x8xf64>) {
+ ^bb0(%inputValue: f64, %outputValue: f64):
+ linalg.yield %inputValue : f64
+ } -> tensor<4x8xf64>
+ return %result : tensor<4x8xf64>
+}
>From ea07ef9187372712e2da88cff13f52714f0f324e Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Sun, 1 Nov 2026 02:19:18 +0900
Subject: [PATCH 3/3] Test sparse iterator SCF lowering
---
.../sparse_kernels_to_iterator.mlir | 20 +++++++++++++++++++
1 file changed, 20 insertions(+)
diff --git a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
index 8585b9a4a2f06e..df2776d65f4663 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_kernels_to_iterator.mlir
@@ -227,6 +227,17 @@ func.func @matvec(%matrix: tensor<4x8xf64, #CSR>,
// ITER: sparse_tensor.iterate %{{.*}} in %{{.*}} at(%[[ROW:.*]])
// ITER: %[[DENSE_ITER:.*]] = builtin.unrealized_conversion_cast %[[ROW]]
// ITER: sparse_tensor.extract_iteration_space %{{.*}} at %[[DENSE_ITER]] lvls = 1
+
+// CHECK-LABEL: func.func @mixed_sparse_encoded_dense_parent(
+// CHECK: %[[MIXED_C1:.*]] = arith.constant 1 : index
+// CHECK: scf.for %[[MIXED_ROW_POS:.*]] = %{{.*}} to %{{.*}} step %[[MIXED_C1]] {
+// CHECK: %[[MIXED_ROW:.*]] = memref.load %{{.*}}[%[[MIXED_ROW_POS]]]
+// CHECK: %[[MIXED_COL_LO:.*]] = memref.load %{{.*}}[%[[MIXED_ROW]]]
+// CHECK: %[[MIXED_ROW_END:.*]] = arith.addi %[[MIXED_ROW]], %[[MIXED_C1]] : index
+// CHECK: %[[MIXED_COL_HI:.*]] = memref.load %{{.*}}[%[[MIXED_ROW_END]]]
+// CHECK: scf.for %[[MIXED_COL_POS:.*]] = %[[MIXED_COL_LO]] to %[[MIXED_COL_HI]] step %[[MIXED_C1]] {
+// CHECK: %[[MIXED_COL:.*]] = memref.load %{{.*}}[%[[MIXED_COL_POS]]]
+// CHECK: memref.store %{{.*}}, %{{.*}}[%[[MIXED_ROW]], %[[MIXED_COL]]]
func.func @mixed_sparse_encoded_dense_parent(
%rows: tensor<4xf64, #VEC>, %matrix: tensor<4x8xf64, #CSR>,
%output: tensor<4x8xf64>) -> tensor<4x8xf64> {
@@ -249,6 +260,15 @@ func.func @mixed_sparse_encoded_dense_parent(
// ITER-LABEL: func.func @sparse_broadcast(
// ITER: sparse_tensor.iterate
// ITER: scf.for
+
+// CHECK-LABEL: func.func @sparse_broadcast(
+// CHECK: %[[BROADCAST_C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BROADCAST_C8:.*]] = arith.constant 8 : index
+// CHECK: scf.for %[[BROADCAST_ROW_POS:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK: %[[BROADCAST_ROW:.*]] = memref.load %{{.*}}[%[[BROADCAST_ROW_POS]]]
+// CHECK: %[[BROADCAST_VALUE:.*]] = memref.load %{{.*}}[%[[BROADCAST_ROW_POS]]]
+// CHECK: scf.for %[[BROADCAST_COL:.*]] = %[[BROADCAST_C0]] to %[[BROADCAST_C8]] step %{{.*}} {
+// CHECK: memref.store %[[BROADCAST_VALUE]], %{{.*}}[%[[BROADCAST_ROW]], %[[BROADCAST_COL]]]
func.func @sparse_broadcast(%input: tensor<4xf64, #VEC>,
%output: tensor<4x8xf64>) -> tensor<4x8xf64> {
%result = linalg.generic {
More information about the Mlir-commits
mailing list