[Mlir-commits] [mlir] 35ece3b - [mlir][sparse][pytaco] add PyTACO SpMM example

[Aart Bik]

Author: Aart Bik
Date: 2022-01-27T17:26:22-08:00
New Revision: 35ece3beaaae132c8784a1d5f6bb612ce380fe04

URL: https://github.com/llvm/llvm-project/commit/35ece3beaaae132c8784a1d5f6bb612ce380fe04
DIFF: https://github.com/llvm/llvm-project/commit/35ece3beaaae132c8784a1d5f6bb612ce380fe04.diff

LOG: [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

Added: 
    mlir/test/Integration/Dialect/SparseTensor/taco/data/A.mtx
    mlir/test/Integration/Dialect/SparseTensor/taco/data/B.mtx
    mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_C.tns
    mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py

Modified: 
    

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/data/A.mtx b/mlir/test/Integration/Dialect/SparseTensor/taco/data/A.mtx
new file mode 100644
index 0000000000000..6ea0893af6164
--- /dev/null
+++ b/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

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/data/B.mtx b/mlir/test/Integration/Dialect/SparseTensor/taco/data/B.mtx
new file mode 100644
index 0000000000000..9bb604d44c7c9
--- /dev/null
+++ b/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

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_C.tns b/mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_C.tns
new file mode 100644
index 0000000000000..e5c1ec14c4030
--- /dev/null
+++ b/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

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py
new file mode 100644
index 0000000000000..edee224fc8ce3
--- /dev/null
+++ b/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)}")