[mlir][sparse][pytaco] add PyTACO SpMM example

Also contains a few TODOs on future enhancements

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D118418

Author: aartbik

mlir/test/Integration/Dialect/SparseTensor/taco/data/A.mtx

@@ -0,0 +1,11 @@
+%%MatrixMarket matrix coordinate real general
+3 3 9
+1 1 1.0
+1 2 2.0
+1 3 4.0
+2 1 4.0
+2 2 5.0
+2 3 6.0
+3 1 7.0
+3 2 8.0
+3 3 9.0

mlir/test/Integration/Dialect/SparseTensor/taco/data/B.mtx

@@ -0,0 +1,11 @@
+%%MatrixMarket matrix coordinate real general
+3 3 9
+1 1 10.0
+1 2 11.0
+1 3 12.0
+2 1 13.0
+2 2 14.0
+2 3 15.0
+3 1 16.0
+3 2 17.0
+3 3 18.0

mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_C.tns

@@ -0,0 +1,9 @@
+1.0 1.0 100.0
+1.0 2.0 107.0
+1.0 3.0 114.0
+2.0 1.0 201.0
+2.0 2.0 216.0
+2.0 3.0 231.0
+3.0 1.0 318.0
+3.0 2.0 342.0
+3.0 3.0 366.0

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py

@@ -0,0 +1,41 @@
+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+import filecmp
+import numpy as np
+import os
+import sys
+import tempfile
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+
+from tools import mlir_pytaco_api as pt
+
+# Define the CSR format.
+#
+# TODO: accept "csr = pt.format([pt.dense, pt.compressed], [0, 1])"
+#
+csr = pt.format([pt.dense, pt.compressed])
+
+# Read matrices A and B from file, infer size of output matrix C.
+A = pt.read(os.path.join(_SCRIPT_PATH, "data/A.mtx"), csr)
+B = pt.read(os.path.join(_SCRIPT_PATH, "data/B.mtx"), csr)
+C = pt.tensor((A.shape[0], B.shape[1]), csr)
+
+# Define the kernel.
+i, j, k = pt.get_index_vars(3)
+C[i, j] = A[i, k] * B[k, j]
+
+# Force evaluation of the kernel by writing out C.
+#
+# TODO: use sparse_tensor.out for output, so that C.tns becomes
+#       a file in extended FROSTT format
+#
+with tempfile.TemporaryDirectory() as test_dir:
+  golden_file = os.path.join(_SCRIPT_PATH, "data/gold_C.tns")
+  out_file = os.path.join(test_dir, "C.tns")
+  pt.write(out_file, C)
+  #
+  # CHECK: Compare files True
+  #
+  print(f"Compare files {filecmp.cmp(golden_file, out_file)}")