[examples][MLIRSparseTensor] Add example for sparse_tensor.new (#166)

* [MLIRSparseTensor] Add example for sparse_tensor.new

This commit adds an example for sparse_tensor.new operation to use the matrix market format.

Signed-off-by: Avimitin <dev@avimit.in>

examples/MLIRSparseTensor/.gitignore

@@ -0,0 +1,3 @@
+log.mlir
+log.ll
+log.s

examples/MLIRSparseTensor/data/sa.mtx

@@ -0,0 +1,12 @@
+%%MatrixMarket matrix coordinate integer general
+%
+% This ASCII file represent a sparse 4x8 matrix with 4 nonzeros in MatrixMarket format.
+% <https://math.nist.gov/MatrixMarket/formats.html>
+%
+% Indice are 1-based, so (1,1) is the first element.
+%
+4 8 4
+1 2 2
+3 2 2
+1 6 6
+4 8 1

examples/MLIRSparseTensor/data/sb.mtx

@@ -0,0 +1,15 @@
+%%MatrixMarket matrix coordinate integer general
+%
+% This ASCII file represent a sparse 8x4 matrix with 7 nonzeros in MatrixMarket format.
+% <https://math.nist.gov/MatrixMarket/formats.html>
+%
+% Indice are 1-based, so (1,1) is the first element.
+%
+8 4 7
+1 4 1
+2 3 2
+3 2 3
+4 1 4
+6 2 5
+7 3 6
+8 3 7

examples/MLIRSparseTensor/makefile

@@ -1,5 +1,33 @@
 MLIR_OPT := ../../llvm/build/bin/mlir-opt
+MLIR_CPU_RUNNER := ../../llvm/build/bin/mlir-cpu-runner
+MLIR_TRANSLATE := ../../llvm/build/bin/mlir-translate
+ifeq ($(shell uname),Linux)
+MLIR_RUNNER_UTILS := ../../llvm/build/lib/libmlir_runner_utils.so
+MLIR_C_RUNNER_UTILS := ../../llvm/build/lib/libmlir_c_runner_utils.so
+MTRIPLE := x86_64-unknown-linux-gnu
+else ifeq ($(shell uname),Darwin)
+MLIR_RUNNER_UTILS := ../../llvm/build/lib/libmlir_runner_utils.dylib
+MLIR_C_RUNNER_UTILS := ../../llvm/build/lib/libmlir_c_runner_utils.dylib
+MTRIPLE := x86_64-apple-darwin
+endif
+
+SPARSE_MATRIX_A := ./data/sa.mtx
+SPARSE_MATRIX_B := ./data/sb.mtx
 
 sparse-tensor-fuse-tensor-cast-lower:
 	@${MLIR_OPT} ./sparse-tensor-fuse-tensor-cast.mlir\
 		--pre-sparsification-rewrite -o ./log.mlir
+
+sparse-tensor-new-lower:
+	@${MLIR_OPT} ./sparse-tensor-new.mlir \
+		--sparse-compiler="enable-runtime-library=true" -o log.mlir
+sparse-tensor-new-translate:
+	@${MLIR_OPT} ./sparse-tensor-new.mlir \
+		--sparse-compiler="enable-runtime-library=true" | \
+	${MLIR_TRANSLATE} --mlir-to-llvmir -o log.ll
+sparse-tensor-new-run:
+	@${MLIR_OPT} ./sparse-tensor-new.mlir \
+		--sparse-compiler="enable-runtime-library=true" | \
+	TENSOR0=${SPARSE_MATRIX_A} TENSOR1=${SPARSE_MATRIX_B} \
+	${MLIR_CPU_RUNNER} -e main -O0 --entry-point-result=void \
+		--shared-libs=${MLIR_RUNNER_UTILS},${MLIR_C_RUNNER_UTILS}

examples/MLIRSparseTensor/sparse-tensor-new.mlir

@@ -0,0 +1,46 @@
+!Filename = !llvm.ptr<i8>
+
+#SparseMatrix = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ]
+}>
+
+// Get tensor filename from env $TENSOR+index
+func.func private @getTensorFilename(index) -> (!Filename)
+func.func private @printMemrefI32(memref<*xi32>) -> ()
+
+//
+// Computes C = A x B with all matrices dense.
+//
+func.func @matmul( %A: tensor<4x8xi32, #SparseMatrix>, %B: tensor<8x4xi32, #SparseMatrix>) -> tensor<4x4xi32, #SparseMatrix> {
+  %C = bufferization.alloc_tensor() : tensor<4x4xi32, #SparseMatrix>
+  %D = linalg.matmul
+    ins(%A, %B: tensor<4x8xi32, #SparseMatrix>, tensor<8x4xi32, #SparseMatrix>)
+    outs(%C: tensor<4x4xi32, #SparseMatrix>) -> tensor<4x4xi32, #SparseMatrix>
+  return %D: tensor<4x4xi32, #SparseMatrix>
+}
+
+func.func @main() {
+  // Get tensor filename from env variable TENSOR0
+  %c0 = arith.constant 0 : index
+  %file0 = call @getTensorFilename(%c0) : (index) -> (!Filename)
+  // Get tensor filename from env variable TENSOR1
+  %c1 = arith.constant 1 : index
+  %file1 = call @getTensorFilename(%c1) : (index) -> (!Filename)
+
+  // Read tensor from file and convert them to tensor<m x n x i32, #attr>
+  %sa = sparse_tensor.new %file0 : !Filename to tensor<4x8xi32, #SparseMatrix>
+  %sb = sparse_tensor.new %file1 : !Filename to tensor<8x4xi32, #SparseMatrix>
+
+  // Call the kernel
+  %1 = call @matmul(%sa, %sb)
+       : (tensor<4x8xi32, #SparseMatrix>, tensor<8x4xi32, #SparseMatrix>)
+          -> tensor<4x4xi32, #SparseMatrix>
+
+  // Extract numerical values to an array
+  %val1 = sparse_tensor.values %1 : tensor<4x4xi32, #SparseMatrix> to memref<?xi32>
+  // Cast the array to a printable unrank memref
+  %val_buf_1 = memref.cast %val1 : memref<?xi32> to memref<*xi32>
+  call @printMemrefI32(%val_buf_1) : (memref<*xi32>) -> ()
+
+  return
+}

examples/lit.cfg.py

@@ -46,7 +46,8 @@
     'MLIRGPU',
     'MLIRPDL',
     'MLIRPython',
-    'MLIRSCF', 
+    'MLIRSCF',
+    'MLIRSparseTensor',
     'MLIRTOSA',
     'Pooling',
     'RISCVBuddyExt',