1# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
2
3import filecmp
4import numpy as np
5import os
6import sys
7import tempfile
8
9_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
10sys.path.append(_SCRIPT_PATH)
11
12from tools import mlir_pytaco_api as pt
13from tools import testing_utils as utils
14
15# Define the CSR format.
16csr = pt.format([pt.dense, pt.compressed], [0, 1])
17
18# Read matrices A and B from file, infer size of output matrix C.
19A = pt.read(os.path.join(_SCRIPT_PATH, "data/A.mtx"), csr)
20B = pt.read(os.path.join(_SCRIPT_PATH, "data/B.mtx"), csr)
21C = pt.tensor([A.shape[0], B.shape[1]], csr)
22
23# Define the kernel.
24i, j, k = pt.get_index_vars(3)
25C[i, j] = A[i, k] * B[k, j]
26
27# Force evaluation of the kernel by writing out C.
28with tempfile.TemporaryDirectory() as test_dir:
29  golden_file = os.path.join(_SCRIPT_PATH, "data/gold_C.tns")
30  out_file = os.path.join(test_dir, "C.tns")
31  pt.write(out_file, C)
32  #
33  # CHECK: Compare result True
34  #
35  print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")
36