[Mlir-commits] [mlir] [mlir][sparse] Fix number_of_entries for loose-compressed tensors (PR #224167)
Kohei Yamaguchi
llvmlistbot at llvm.org
Mon Sep 28 17:11:33 PDT 2026
https://github.com/sott0n updated https://github.com/llvm/llvm-project/pull/224167
>From d4cc208fba041e0ee13f2499670f1f584d0fce87 Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Thu, 17 Sep 2026 05:07:39 +0900
Subject: [PATCH 1/2] [mlir][sparse] Fix number_of_entries for loose-compressed
tensors
---
.../Transforms/SparseTensorCodegen.cpp | 9 +-
.../Transforms/SparseTensorConversion.cpp | 8 +-
.../Transforms/SparseTensorRewriting.cpp | 33 +++++-
.../sparse_number_of_entries.mlir | 46 +++++++
.../sparse_number_of_entries_invalid.mlir | 14 +++
.../CPU/sparse_number_of_entries.mlir | 112 ++++++++++++++++++
.../SparseTensor/CPU/sparse_out_loose.mlir | 60 ++++++++++
7 files changed, 277 insertions(+), 5 deletions(-)
create mode 100644 mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
create mode 100644 mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir
create mode 100644 mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir
create mode 100644 mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 38558777d537c..31774a8c5f23f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1302,9 +1302,14 @@ class SparseNumberOfEntriesConverter
LogicalResult
matchAndRewrite(NumberOfEntriesOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
+ // Loose-compressed tensors may have holes in their values buffer. Their
+ // entry counts must be lowered to foreach before converting storage.
+ if (llvm::any_of(getSparseTensorType(op.getTensor()).getLvlTypes(),
+ isLooseCompressedLT))
+ return rewriter.notifyMatchFailure(
+ op, "loose-compressed entry count must be lowered to foreach");
+
// Query memSizes for the actually stored values.
- // FIXME: the nse value computed in this way might be wrong when there is
- // any "loose_compressed" level.
auto desc = getDescriptorFromTensorTuple(adaptor.getTensor(),
op.getTensor().getType());
rewriter.replaceOp(op, desc.getValMemSize(rewriter, op.getLoc()));
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 0e88d31dae8e8..da2f1651a61e1 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -588,8 +588,14 @@ class SparseNumberOfEntriesConverter
LogicalResult
matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- // Query values array size for the actually stored values size.
auto stt = getSparseTensorType(op.getTensor());
+ // Loose-compressed tensors may have holes in their values buffer. Their
+ // entry counts must be lowered to foreach before converting storage.
+ if (llvm::any_of(stt.getLvlTypes(), isLooseCompressedLT))
+ return rewriter.notifyMatchFailure(
+ op, "loose-compressed entry count must be lowered to foreach");
+
+ // Query values array size for the actually stored values size.
auto vals = genValuesCall(rewriter, op.getLoc(), stt, adaptor.getTensor());
auto zero = constantIndex(rewriter, op.getLoc(), 0);
rewriter.replaceOpWithNewOp<memref::DimOp>(op, vals, zero);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 42a6c4d77dbdd..545b76e8b46e4 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1358,6 +1358,35 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
}
};
+/// Count entries in loose-compressed tensors by traversing their stored
+/// elements. The values buffer may contain holes, so its size is not the
+/// number of entries. Keep the constant-time buffer-size query for formats
+/// without loose-compressed levels.
+struct NumberOfEntriesRewriter : public OpRewritePattern<NumberOfEntriesOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(NumberOfEntriesOp op,
+ PatternRewriter &rewriter) const override {
+ if (!llvm::any_of(getSparseTensorType(op.getTensor()).getLvlTypes(),
+ isLooseCompressedLT))
+ return failure();
+
+ Location loc = op.getLoc();
+ Value zero = constantIndex(rewriter, loc, 0);
+ Value one = constantIndex(rewriter, loc, 1);
+ auto count =
+ ForeachOp::create(rewriter, loc, op.getTensor(), ValueRange{zero},
+ [one](OpBuilder &builder, Location loc, ValueRange,
+ Value, ValueRange iterArgs) {
+ Value next = arith::AddIOp::create(
+ builder, loc, iterArgs.front(), one);
+ sparse_tensor::YieldOp::create(builder, loc, next);
+ });
+ rewriter.replaceOp(op, count.getResults());
+ return success();
+ }
+};
+
/// Sparse rewriting rule for the foreach operator.
struct ForeachRewriter : public OpRewritePattern<ForeachOp> {
public:
@@ -1620,8 +1649,8 @@ void mlir::populateLowerSparseOpsToForeachPatterns(RewritePatternSet &patterns,
ReshapeRewriter<tensor::CollapseShapeOp>,
Sparse2SparseReshapeRewriter<tensor::ExpandShapeOp>,
Sparse2SparseReshapeRewriter<tensor::CollapseShapeOp>,
- SparseTensorDimOpRewriter, TensorReshapeRewriter, OutRewriter>(
- patterns.getContext());
+ SparseTensorDimOpRewriter, TensorReshapeRewriter,
+ NumberOfEntriesRewriter, OutRewriter>(patterns.getContext());
if (enableConvert)
patterns.add<DirectConvertRewriter>(patterns.getContext());
diff --git a/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir b/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
new file mode 100644
index 0000000000000..32bc2a1eb5414
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-opt %s --lower-sparse-ops-to-foreach --canonicalize | FileCheck %s
+// RUN: mlir-opt %s --lower-sparse-ops-to-foreach='enable-runtime-library=false' --canonicalize | FileCheck %s
+
+#Loose = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j : loose_compressed)
+}>
+#LooseDense = #sparse_tensor.encoding<{
+ map = (i, j, k) -> (i : dense, j : loose_compressed, k : dense)
+}>
+#CSR = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @loose(
+// CHECK-SAME: %[[T:.*]]: tensor
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1 : index
+// CHECK: %[[COUNT:.*]] = sparse_tensor.foreach in %[[T]] init(%[[ZERO]])
+// CHECK: ^bb0({{.*}}, %[[ACC:.*]]: index):
+// CHECK: %[[NEXT:.*]] = arith.addi %[[ACC]], %[[ONE]] : index
+// CHECK: sparse_tensor.yield %[[NEXT]] : index
+// CHECK: return %[[COUNT]] : index
+func.func @loose(%t: tensor<?x?xf64, #Loose>) -> index {
+ %n = sparse_tensor.number_of_entries %t : tensor<?x?xf64, #Loose>
+ return %n : index
+}
+
+// Check all levels, not only the innermost one.
+// CHECK-LABEL: func.func @loose_dense(
+// CHECK: sparse_tensor.foreach
+// CHECK-NOT: sparse_tensor.number_of_entries
+// CHECK: return
+func.func @loose_dense(%t: tensor<?x?x2xf64, #LooseDense>) -> index {
+ %n = sparse_tensor.number_of_entries %t : tensor<?x?x2xf64, #LooseDense>
+ return %n : index
+}
+
+// Formats with contiguous storage retain their constant-time query.
+// CHECK-LABEL: func.func @compressed(
+// CHECK-NOT: sparse_tensor.foreach
+// CHECK: %[[N:.*]] = sparse_tensor.number_of_entries
+// CHECK: return %[[N]] : index
+func.func @compressed(%t: tensor<?x?xf64, #CSR>) -> index {
+ %n = sparse_tensor.number_of_entries %t : tensor<?x?xf64, #CSR>
+ return %n : index
+}
diff --git a/mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir b/mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir
new file mode 100644
index 0000000000000..3ec9387cdf827
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s --sparse-tensor-codegen --verify-diagnostics
+// RUN: mlir-opt %s --sparse-tensor-conversion --verify-diagnostics
+
+// Without lower-sparse-ops-to-foreach and lower-sparse-foreach-to-scf,
+// storage conversion must not silently return the values buffer length.
+#Loose = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j : loose_compressed)
+}>
+
+func.func @loose(%t: tensor<?x?xf64, #Loose>) -> index {
+ // expected-error at +1 {{failed to legalize operation 'sparse_tensor.number_of_entries'}}
+ %n = sparse_tensor.number_of_entries %t : tensor<?x?xf64, #Loose>
+ return %n : index
+}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir
new file mode 100644
index 0000000000000..c544249d9db67
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir
@@ -0,0 +1,112 @@
+//--------------------------------------------------------------------------------------------------
+// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
+//
+// Set-up that's shared across all tests in this directory. In principle, this
+// config could be moved to lit.local.cfg. However, there are downstream users that
+// do not use these LIT config files. Hence why this is kept inline.
+//
+// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
+// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
+// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
+// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
+// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
+// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs}
+// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}
+//
+// DEFINE: %{env} =
+//--------------------------------------------------------------------------------------------------
+
+// RUN: %{compile} | %{run} | FileCheck %s
+// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
+// RUN: %{compile} | %{run} | FileCheck %s
+
+#Loose = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j : loose_compressed)
+}>
+#LooseDense = #sparse_tensor.encoding<{
+ map = (i, j, k) -> (i : dense, j : loose_compressed, k : dense),
+ posWidth = 32,
+ crdWidth = 32
+}>
+#LooseCompressed = #sparse_tensor.encoding<{
+ map = (i, j, k) -> (i : dense, j : loose_compressed, k : compressed)
+}>
+
+func.func @check(%pos: tensor<4xindex>, %crd: tensor<6xindex>,
+ %val: tensor<6xf64>) {
+ %s = sparse_tensor.assemble (%pos, %crd), %val :
+ (tensor<4xindex>, tensor<6xindex>), tensor<6xf64>
+ to tensor<2x4xf64, #Loose>
+ %n = sparse_tensor.number_of_entries %s : tensor<2x4xf64, #Loose>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %count = sparse_tensor.foreach in %s init(%c0) :
+ tensor<2x4xf64, #Loose>, index -> index do {
+ ^bb0(%i: index, %j: index, %v: f64, %acc: index):
+ %next = arith.addi %acc, %c1 : index
+ sparse_tensor.yield %next : index
+ }
+ vector.print %n : index
+ vector.print %count : index
+ %has_runtime = sparse_tensor.has_runtime_library
+ scf.if %has_runtime {
+ // Assemble copies the buffers only on the runtime-library path.
+ bufferization.dealloc_tensor %s : tensor<2x4xf64, #Loose>
+ }
+ return
+}
+
+func.func @main() {
+ %crd = arith.constant dense<[0, 1, 2, 0, 0, 3]> : tensor<6xindex>
+ %val = arith.constant dense<[0.0, 10.0, 20.0, 999.0, 0.0, 20.0]> : tensor<6xf64>
+ %contiguous = arith.constant dense<[0, 1, 1, 3]> : tensor<4xindex>
+ %gaps = arith.constant dense<[1, 2, 4, 6]> : tensor<4xindex>
+ %empty = arith.constant dense<[1, 1, 6, 6]> : tensor<4xindex>
+
+ // Contiguous intervals ignore unused trailing capacity.
+ // CHECK: 3
+ // CHECK-NEXT: 3
+ call @check(%contiguous, %crd, %val) : (tensor<4xindex>, tensor<6xindex>, tensor<6xf64>) -> ()
+ // Ignore leading/interior holes, but include the explicit zero at slot 4.
+ // CHECK-NEXT: 3
+ // CHECK-NEXT: 3
+ call @check(%gaps, %crd, %val) : (tensor<4xindex>, tensor<6xindex>, tensor<6xf64>) -> ()
+ // Nonzero positions do not imply any stored entries.
+ // CHECK-NEXT: 0
+ // CHECK-NEXT: 0
+ call @check(%empty, %crd, %val) : (tensor<4xindex>, tensor<6xindex>, tensor<6xf64>) -> ()
+
+ // A non-innermost loose level followed by a dense level. Every stored
+ // coordinate has two values, including explicitly stored zeros.
+ %pos32 = arith.constant dense<[1, 2, 4, 6]> : tensor<4xi32>
+ %crd32 = arith.constant dense<[0, 1, 2, 0, 0, 3]> : tensor<6xi32>
+ %dense_values = arith.constant dense<0.0> : tensor<12xf64>
+ %dense = sparse_tensor.assemble (%pos32, %crd32), %dense_values :
+ (tensor<4xi32>, tensor<6xi32>), tensor<12xf64>
+ to tensor<2x4x2xf64, #LooseDense>
+ %dense_n = sparse_tensor.number_of_entries %dense : tensor<2x4x2xf64, #LooseDense>
+ // CHECK-NEXT: 6
+ vector.print %dense_n : index
+
+ // Only the compressed children of reachable loose coordinates count.
+ // Slots 1, 4 and 5 have respectively 2, 2 and 1 children. Summing the
+ // lengths of all compressed intervals would incorrectly count 8 entries.
+ %child_pos = arith.constant dense<[0, 1, 3, 4, 5, 7, 8]> : tensor<7xindex>
+ %child_crd = arith.constant dense<[0, 0, 1, 0, 1, 0, 1, 1]> : tensor<8xindex>
+ %child_val = arith.constant dense<0.0> : tensor<8xf64>
+ %compressed = sparse_tensor.assemble (%gaps, %crd, %child_pos, %child_crd), %child_val :
+ (tensor<4xindex>, tensor<6xindex>, tensor<7xindex>, tensor<8xindex>), tensor<8xf64>
+ to tensor<2x4x2xf64, #LooseCompressed>
+ %compressed_n = sparse_tensor.number_of_entries %compressed : tensor<2x4x2xf64, #LooseCompressed>
+ // CHECK-NEXT: 5
+ vector.print %compressed_n : index
+
+ %has_runtime = sparse_tensor.has_runtime_library
+ scf.if %has_runtime {
+ bufferization.dealloc_tensor %dense : tensor<2x4x2xf64, #LooseDense>
+ bufferization.dealloc_tensor %compressed : tensor<2x4x2xf64, #LooseCompressed>
+ }
+ return
+}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir
new file mode 100644
index 0000000000000..ac8517e4ebe13
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir
@@ -0,0 +1,60 @@
+//--------------------------------------------------------------------------------------------------
+// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
+//
+// Set-up that's shared across all tests in this directory. In principle, this
+// config could be moved to lit.local.cfg. However, there are downstream users that
+// do not use these LIT config files. Hence why this is kept inline.
+//
+// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
+// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
+// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
+// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
+// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
+// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs}
+// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}
+//
+// DEFINE: %{env} =
+//--------------------------------------------------------------------------------------------------
+
+// Write to stdout so that FileCheck verifies both the NSE header and entries.
+// REDEFINE: %{env} = TENSOR0=
+// RUN: %{compile} | env %{env} %{run} | FileCheck %s
+// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
+// RUN: %{compile} | env %{env} %{run} | FileCheck %s
+
+!Filename = !llvm.ptr
+
+#Loose = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : dense, j : loose_compressed)
+}>
+
+func.func private @getTensorFilename(index) -> !Filename
+
+func.func @main() {
+ %c0 = arith.constant 0 : index
+ %filename = call @getTensorFilename(%c0) : (index) -> !Filename
+ %pos = arith.constant dense<[1, 2, 4, 6]> : tensor<4xindex>
+ %crd = arith.constant dense<[0, 1, 2, 0, 0, 3]> : tensor<6xindex>
+ %val = arith.constant dense<[111.0, 10.0, 222.0, 333.0, 0.0, 20.0]> : tensor<6xf64>
+ %s = sparse_tensor.assemble (%pos, %crd), %val :
+ (tensor<4xindex>, tensor<6xindex>), tensor<6xf64>
+ to tensor<2x4xf64, #Loose>
+
+ // CHECK: # extended FROSTT format
+ // CHECK-NEXT: 2 3
+ // CHECK-NEXT: 2 4
+ // CHECK-NEXT: 1 2 10
+ // CHECK-NEXT: 2 1 0
+ // CHECK-NEXT: 2 4 20
+ // CHECK-NOT: {{.}}
+ sparse_tensor.out %s, %filename : tensor<2x4xf64, #Loose>, !Filename
+
+ %has_runtime = sparse_tensor.has_runtime_library
+ scf.if %has_runtime {
+ // Assemble copies the buffers only on the runtime-library path.
+ bufferization.dealloc_tensor %s : tensor<2x4xf64, #Loose>
+ }
+ return
+}
>From d65556b36d143683acb4102c97f019499d6ac895 Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Tue, 29 Sep 2026 09:11:13 +0900
Subject: [PATCH 2/2] [mlir][sparse] Clarify number_of_entries traversal cost
---
.../SparseTensor/Transforms/SparseTensorRewriting.cpp | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 545b76e8b46e4..d242e5277661f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1358,10 +1358,9 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
}
};
-/// Count entries in loose-compressed tensors by traversing their stored
-/// elements. The values buffer may contain holes, so its size is not the
-/// number of entries. Keep the constant-time buffer-size query for formats
-/// without loose-compressed levels.
+/// Count entries in loose-compressed tensors with an O(n) traversal of their
+/// stored elements instead of a constant-time buffer-size query. This is
+/// necessary because the values buffer may contain holes.
struct NumberOfEntriesRewriter : public OpRewritePattern<NumberOfEntriesOp> {
using OpRewritePattern::OpRewritePattern;
More information about the Mlir-commits
mailing list