[Mlir-commits] [mlir] [mlir][sparse] Fix number_of_entries for loose-compressed tensors (PR #224167)
Kohei Yamaguchi
llvmlistbot at llvm.org
Mon Oct 5 20:02:57 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/3] [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 38558777d537c7a..31774a8c5f23fda 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 0e88d31dae8e84a..da2f1651a61e158 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 42a6c4d77dbdd81..545b76e8b46e47f 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 000000000000000..32bc2a1eb54145f
--- /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 000000000000000..3ec9387cdf8278f
--- /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 000000000000000..c544249d9db67b3
--- /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 000000000000000..ac8517e4ebe1384
--- /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/3] [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 545b76e8b46e47f..d242e5277661f38 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;
>From 7bb33c964e0ac79cad33165a49bfacd1a17efd2b Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Tue, 6 Oct 2026 11:35:26 +0900
Subject: [PATCH 3/3] [mlir][sparse] Separate output entry count from NSE
---
.../SparseTensor/IR/SparseTensorOps.td | 11 +--
.../Transforms/SparseTensorCodegen.cpp | 9 +--
.../Transforms/SparseTensorConversion.cpp | 8 +-
.../Transforms/SparseTensorRewriting.cpp | 51 +++++--------
.../SparseTensor/rewriting_for_codegen.mlir | 10 ++-
.../sparse_number_of_entries.mlir | 46 ------------
.../sparse_number_of_entries_invalid.mlir | 14 ----
.../CPU/sparse_number_of_entries.mlir | 73 +++----------------
.../SparseTensor/CPU/sparse_out_loose.mlir | 2 +-
9 files changed, 45 insertions(+), 179 deletions(-)
delete mode 100644 mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
delete mode 100644 mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index e41e9dec8c19ac1..85baffed37640f4 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -384,12 +384,13 @@ def SparseTensor_ToValuesOp : SparseTensor_Op<"values",
}
def SparseTensor_NumberOfEntriesOp : SparseTensor_Op<"number_of_entries", [Pure]> {
- let summary = "Returns the number of entries that are stored in the tensor.";
+ let summary = "Returns the used extent of the values storage.";
let description = [{
- Returns the number of entries that are stored in the given sparse tensor.
- Note that this is typically the number of nonzero elements in the tensor,
- but since explicit zeros may appear in the storage formats, the more
- accurate nomenclature is used.
+ Returns the used extent of the values storage for the given sparse tensor.
+ This count includes explicit zeros and may include storage-only slots, such
+ as holes in a loose-compressed level. Therefore, it need not equal either
+ the number of entries produced by sparse iteration or the number of
+ mathematically nonzero elements.
Example:
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 31774a8c5f23fda..89779ad5232c282 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1302,14 +1302,7 @@ 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.
+ // Query memSizes for the used extent of the values storage.
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 da2f1651a61e158..18755010fb2d388 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -589,13 +589,7 @@ class SparseNumberOfEntriesConverter
matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
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.
+ // Query the used extent of the values storage.
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 d242e5277661f38..35a11b8b19c7e65 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1358,34 +1358,6 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
}
};
-/// 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;
-
- 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:
@@ -1566,9 +1538,21 @@ struct OutRewriter : public OpRewritePattern<OutOp> {
LogicalResult matchAndRewrite(OutOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
- // Calculate NNZ.
Value src = op.getTensor();
- Value nnz = NumberOfEntriesOp::create(rewriter, loc, src);
+
+ // The writer metadata precedes the records, so count exactly the entries
+ // that sparse iteration will emit before creating the writer.
+ Value zero = constantIndex(rewriter, loc, 0);
+ Value one = constantIndex(rewriter, loc, 1);
+ auto count = ForeachOp::create(
+ rewriter, loc, src, 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);
+ });
+ Value numEntries = count.getResult(0);
// Allocate a temporary buffer for storing dimension-sizes/coordinates.
const auto srcTp = getSparseTensorType(src);
@@ -1593,7 +1577,8 @@ struct OutRewriter : public OpRewritePattern<OutOp> {
.getResult(0);
Value rankValue = constantIndex(rewriter, loc, dimRank);
createFuncCall(rewriter, loc, "outSparseTensorWriterMetaData", {},
- {writer, rankValue, nnz, dimSizes}, EmitCInterface::On);
+ {writer, rankValue, numEntries, dimSizes},
+ EmitCInterface::On);
Value dimCoords = dimSizes; // Reuse the dimSizes buffer for dimCoords.
Type eltTp = srcTp.getElementType();
@@ -1648,8 +1633,8 @@ void mlir::populateLowerSparseOpsToForeachPatterns(RewritePatternSet &patterns,
ReshapeRewriter<tensor::CollapseShapeOp>,
Sparse2SparseReshapeRewriter<tensor::ExpandShapeOp>,
Sparse2SparseReshapeRewriter<tensor::CollapseShapeOp>,
- SparseTensorDimOpRewriter, TensorReshapeRewriter,
- NumberOfEntriesRewriter, OutRewriter>(patterns.getContext());
+ SparseTensorDimOpRewriter, TensorReshapeRewriter, OutRewriter>(
+ patterns.getContext());
if (enableConvert)
patterns.add<DirectConvertRewriter>(patterns.getContext());
diff --git a/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir b/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir
index 11fc502822539b0..c5046e6afa1469c 100644
--- a/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir
+++ b/mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir
@@ -52,12 +52,18 @@ func.func @sparse_new_coo(%arg0: !llvm.ptr) -> tensor<?x?xf32, #COO> {
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
// CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
-// CHECK: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
+// CHECK: %[[COUNT:.*]] = scf.for {{.*}} iter_args({{.*}} = %[[C0]]) -> (index) {
+// CHECK: %[[INNER:.*]] = scf.for {{.*}} iter_args(%[[ACC:.*]] = {{.*}}) -> (index) {
+// CHECK: %[[NEXT:.*]] = arith.addi %[[ACC]], %[[C1]] : index
+// CHECK: scf.yield %[[NEXT]] : index
+// CHECK: }
+// CHECK: scf.yield %[[INNER]] : index
+// CHECK: }
// CHECK: %[[DS:.*]] = memref.alloca(%[[C2]]) : memref<?xindex>
// CHECK: memref.store %[[C10]], %[[DS]]{{\[}}%[[C0]]] : memref<?xindex>
// CHECK: memref.store %[[C20]], %[[DS]]{{\[}}%[[C1]]] : memref<?xindex>
// CHECK: %[[W:.*]] = call @createSparseTensorWriter(%[[B]])
-// CHECK: call @outSparseTensorWriterMetaData(%[[W]], %[[C2]], %[[NNZ]], %[[DS]])
+// CHECK: call @outSparseTensorWriterMetaData(%[[W]], %[[C2]], %[[COUNT]], %[[DS]])
// CHECK: %[[V:.*]] = memref.alloca() : memref<f32>
// CHECK: scf.for %{{.*}} = %[[C0]] to %[[C10]] step %[[C1]] {
// CHECK: scf.for {{.*}} {
diff --git a/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir b/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
deleted file mode 100644
index 32bc2a1eb54145f..000000000000000
--- a/mlir/test/Dialect/SparseTensor/sparse_number_of_entries.mlir
+++ /dev/null
@@ -1,46 +0,0 @@
-// 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
deleted file mode 100644
index 3ec9387cdf8278f..000000000000000
--- a/mlir/test/Dialect/SparseTensor/sparse_number_of_entries_invalid.mlir
+++ /dev/null
@@ -1,14 +0,0 @@
-// 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
index c544249d9db67b3..d22fedf3b981b6d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_number_of_entries.mlir
@@ -25,20 +25,15 @@
#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>) {
+func.func @main() {
+ %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>
+
%n = sparse_tensor.number_of_entries %s : tensor<2x4xf64, #Loose>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
@@ -48,8 +43,13 @@ func.func @check(%pos: tensor<4xindex>, %crd: tensor<6xindex>,
%next = arith.addi %acc, %c1 : index
sparse_tensor.yield %next : index
}
+ // The NSE is the used extent of the values storage, including holes, while
+ // sparse iteration only visits the three entries stored in those intervals.
+ // CHECK: 6
+ // CHECK-NEXT: 3
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.
@@ -57,56 +57,3 @@ func.func @check(%pos: tensor<4xindex>, %crd: tensor<6xindex>,
}
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
index ac8517e4ebe1384..fd25cf6ea97048c 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_loose.mlir
@@ -18,7 +18,7 @@
// DEFINE: %{env} =
//--------------------------------------------------------------------------------------------------
-// Write to stdout so that FileCheck verifies both the NSE header and entries.
+// Write to stdout so that FileCheck verifies the entry-count header and entries.
// REDEFINE: %{env} = TENSOR0=
// RUN: %{compile} | env %{env} %{run} | FileCheck %s
// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
More information about the Mlir-commits
mailing list