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 --- /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 --- /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 --- /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 --- /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)}")