From 8dbf87eaac929c8097187614c6d9694d4179ad5b Mon Sep 17 00:00:00 2001
From: Avimitin <dev@avimit.in>
Date: Sun, 25 Jun 2023 15:24:40 +0800
Subject: [PATCH] [example][MLIRSparseTensor] add
 sparse_tensor.number_of_entries

---
 examples/MLIRSparseTensor/makefile            | 13 ++++++++++
 .../sparse-tensor-number-of-entries.mlir      | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 examples/MLIRSparseTensor/sparse-tensor-number-of-entries.mlir

diff --git a/examples/MLIRSparseTensor/makefile b/examples/MLIRSparseTensor/makefile
index d575059c6..c6b5564d5 100644
--- a/examples/MLIRSparseTensor/makefile
+++ b/examples/MLIRSparseTensor/makefile
@@ -57,3 +57,16 @@ sparse-tensor-dump-run:
 		--sparse-compiler="enable-runtime-library=true" | \
 	${MLIR_CPU_RUNNER} -e main -O0 --entry-point-result=void \
 		--shared-libs=${MLIR_RUNNER_UTILS},${MLIR_C_RUNNER_UTILS}
+
+sparse-tensor-number-of-entries-lower:
+	@${MLIR_OPT} ./sparse-tensor-number-of-entries.mlir \
+		--sparse-compiler="enable-runtime-library=false" -o log.mlir
+sparse-tensor-number-of-entries-translate:
+	@${MLIR_OPT} ./sparse-tensor-number-of-entries.mlir \
+		--sparse-compiler="enable-runtime-library=false" | \
+	${MLIR_TRANSLATE} --mlir-to-llvmir -o log.ll
+sparse-tensor-number-of-entries-run:
+	@${MLIR_OPT} ./sparse-tensor-number-of-entries.mlir \
+		--sparse-compiler="enable-runtime-library=false" | \
+	${MLIR_CPU_RUNNER} -e main -O0 --entry-point-result=void \
+		--shared-libs=${MLIR_RUNNER_UTILS},${MLIR_C_RUNNER_UTILS}
diff --git a/examples/MLIRSparseTensor/sparse-tensor-number-of-entries.mlir b/examples/MLIRSparseTensor/sparse-tensor-number-of-entries.mlir
new file mode 100644
index 000000000..5709bf4ed
--- /dev/null
+++ b/examples/MLIRSparseTensor/sparse-tensor-number-of-entries.mlir
@@ -0,0 +1,24 @@
+// This is an example of operation `sparse_tensor.number_of_entries` in MLIR SparseTensor Dialect.
+// The `sparse_tensor.number_of_entries` operation accpet a tensor with sparse attribute, then return the number
+// of the non-zero entry in the given sparse tensor.
+
+#SparseMatrix = #sparse_tensor.encoding<{
+  dimLevelType = ["compressed", "compressed"]
+}>
+
+func.func @main() {
+  %d0 = arith.constant dense<[
+      [1.0, 0.0, 0.0, 0.0],
+      [0.0, 0.0, 2.0, 0.0],
+      [0.0, 0.0, 0.0, 0.0],
+      [0.0, 0.0, 0.0, 3.0]
+    ]> : tensor<4x4xf64>
+
+  %s0 = sparse_tensor.convert %d0 : tensor<4x4xf64> to tensor<4x4xf64, #SparseMatrix>
+  %n0 = sparse_tensor.number_of_entries %s0 : tensor<4x4xf64, #SparseMatrix>
+
+  // This should print 3
+  vector.print %n0 : index
+
+  return
+}