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
from tools import testing_utils as utils
# Define the CSR format.
csr = pt.format([pt.dense, pt.compressed], [0, 1])
# 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.
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 result True
#
print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")
# https://www.cise.ufl.edu/research/sparse/MM/Boeing/pwtk.tar.gz
# In order to run the program using the matrix above, you can download the
# matrix and replace this path to the actual path to the file.
A = pt.read(os.path.join(_SCRIPT_PATH, "data/pwtk.mtx"), csr)
# These two lines have been modified from the original program to use static
# data to support result comparison.
x = pt.from_array(np.full((A.shape[1], ), 1, dtype=np.float64))
z = pt.from_array(np.full((A.shape[0], ), 2, dtype=np.float64))
# Declare the result to be a dense vector
y = pt.tensor([A.shape[0]], dv)
# Declare index vars
i, j = pt.get_index_vars(2)
# Define the SpMV computation
y[i] = A[i, j] * x[j] + z[i]
##########################################################################
# Perform the SpMV computation and write the result to file
with tempfile.TemporaryDirectory() as test_dir:
golden_file = os.path.join(_SCRIPT_PATH, "data/gold_y.tns")
out_file = os.path.join(test_dir, "y.tns")
pt.write(out_file, y)
#
# CHECK: Compare result True
#
print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")
# sum(k, S[i, j] * A[i, k] * B[k, j])
# we only compute the intermediate dense matrix product that are actually
# needed to compute the result, with proper asymptotic complexity.
X[i, j] = S[i, j] * A[i, k] * B[k, j]
# Alternative way to define SDDMM kernel. Since this performs the reduction as
# sum(k, A[i, k] * B[k, j]) * S[i, j]
# the MLIR lowering results in two separate tensor index expressions that are
# fused prior to running the sparse compiler in order to guarantee proper
# asymptotic complexity.
Y[i, j] = A[i, k] * B[k, j] * S[i, j]
expected = """; extended FROSTT format
2 1
8 8
1 8 2016
"""
# Force evaluation of the kernels by writing out X and Y.
with tempfile.TemporaryDirectory() as test_dir:
x_file = os.path.join(test_dir, "X.tns")
y_file = os.path.join(test_dir, "Y.tns")
pt.write(x_file, X)
pt.write(y_file, Y)
#
# CHECK: Compare result True True
#
x_data = utils.file_as_string(x_file)
y_data = utils.file_as_string(y_file)
print(f"Compare result {x_data == expected} {y_data == expected}")
# tensor for the purpose of testing. To run the program using the data from
# the real application, please download the data from:
# http://frostt.io/tensors/nell-2/
B = pt.read(os.path.join(_SCRIPT_PATH, "data/nell-2.tns"), csf)
# These two lines have been modified from the original program to use static
# data to support result comparison.
C = pt.from_array(np.full((B.shape[1], 25), 1, dtype=np.float64))
D = pt.from_array(np.full((B.shape[2], 25), 2, dtype=np.float64))
# Declare the result to be a dense matrix.
A = pt.tensor([B.shape[0], 25], rm)
# Declare index vars.
i, j, k, l = pt.get_index_vars(4)
# Define the MTTKRP computation.
A[i, j] = B[i, k, l] * D[l, j] * C[k, j]
##########################################################################
# Perform the MTTKRP computation and write the result to file.
with tempfile.TemporaryDirectory() as test_dir:
golden_file = os.path.join(_SCRIPT_PATH, "data/gold_A.tns")
out_file = os.path.join(test_dir, "A.tns")
pt.write(out_file, A)
#
# CHECK: Compare result True
#
print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")
pt.format([pt.compressed, pt.compressed, pt.compressed]))
X = pt.tensor([8, 8, 8],
pt.format([pt.compressed, pt.compressed, pt.compressed]))
S.insert([0, 0, 0], 2.0)
S.insert([1, 1, 1], 3.0)
S.insert([4, 4, 4], 4.0)
S.insert([7, 7, 7], 5.0)
# TODO: make this work:
# X[i, j, k] = alpha[0] * S[i, j, k]
X[i, j, k] = S[i, j, k]
expected = """; extended FROSTT format
3 4
8 8 8
1 1 1 2
2 2 2 3
5 5 5 4
8 8 8 5
"""
# Force evaluation of the kernel by writing out X.
with tempfile.TemporaryDirectory() as test_dir:
x_file = os.path.join(test_dir, 'X.tns')
pt.write(x_file, X)
#
# CHECK: Compare result True
#
x_data = utils.file_as_string(x_file)
print(f'Compare result {x_data == expected}')
# tensor for the purpose of testing. To run the program using the data from
# the real application, please download the data from:
# http://frostt.io/tensors/nell-2/
B = pt.read(os.path.join(_SCRIPT_PATH, "data/nell-2.tns"), csf)
# These two lines have been modified from the original program to use static
# data to support result comparison.
C = pt.from_array(np.full((B.shape[1], 25), 1, dtype=np.float64))
D = pt.from_array(np.full((B.shape[2], 25), 2, dtype=np.float64))
# Declare the result to be a dense matrix.
A = pt.tensor([B.shape[0], 25], rm)
# Declare index vars.
i, j, k, l = pt.get_index_vars(4)
# Define the MTTKRP computation.
A[i, j] = B[i, k, l] * D[l, j] * C[k, j]
##########################################################################
# CHECK: Compare result True
# Perform the MTTKRP computation and write the result to file.
with tempfile.TemporaryDirectory() as test_dir:
actual_file = os.path.join(test_dir, "A.tns")
pt.write(actual_file, A)
actual = np.loadtxt(actual_file, np.float64)
expected = np.loadtxt(os.path.join(_SCRIPT_PATH, "data/gold_A.tns"),
np.float64)
print(f"Compare result {np.allclose(actual, expected, rtol=0.01)}")
# https://www.cise.ufl.edu/research/sparse/MM/Boeing/pwtk.tar.gz
# In order to run the program using the matrix above, you can download the
# matrix and replace this path to the actual path to the file.
A = pt.read(os.path.join(_SCRIPT_PATH, "data/pwtk.mtx"), csr)
# These two lines have been modified from the original program to use static
# data to support result comparison.
x = pt.from_array(np.full((A.shape[1], ), 1, dtype=np.float64))
z = pt.from_array(np.full((A.shape[0], ), 2, dtype=np.float64))
# Declare the result to be a dense vector
y = pt.tensor([A.shape[0]], dv)
# Declare index vars
i, j = pt.get_index_vars(2)
# Define the SpMV computation
y[i] = A[i, j] * x[j] + z[i]
##########################################################################
# CHECK: Compare result True
# Perform the SpMV computation and write the result to file
with tempfile.TemporaryDirectory() as test_dir:
actual_file = os.path.join(test_dir, "y.tns")
pt.write(actual_file, y)
actual = np.loadtxt(actual_file, np.float64)
expected = np.loadtxt(os.path.join(_SCRIPT_PATH, "data/gold_y.tns"),
np.float64)
print(f"Compare result {np.allclose(actual, expected, rtol=0.01)}")