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
15i, j, k = pt.get_index_vars(3)
16
17# Set up dense matrices.
18A = pt.from_array(np.full((8, 8), 2.0, dtype=np.float32))
19B = pt.from_array(np.full((8, 8), 3.0, dtype=np.float32))
20
21# Set up sparse matrices.
22S = pt.tensor([8, 8], pt.format([pt.compressed, pt.compressed]))
23X = pt.tensor([8, 8], pt.format([pt.compressed, pt.compressed]))
24Y = pt.tensor([8, 8], pt.compressed)  # alternative syntax works too
25
26S.insert([0, 7], 42.0)
27
28# Define the SDDMM kernel. Since this performs the reduction as
29#   sum(k, S[i, j] * A[i, k] * B[k, j])
30# we only compute the intermediate dense matrix product that are actually
31# needed to compute the result, with proper asymptotic complexity.
32X[i, j] = S[i, j] * A[i, k] * B[k, j]
33
34# Alternative way to define SDDMM kernel. Since this performs the reduction as
35#   sum(k, A[i, k] * B[k, j]) * S[i, j]
36# the MLIR lowering results in two separate tensor index expressions that are
37# fused prior to running the sparse compiler in order to guarantee proper
38# asymptotic complexity.
39Y[i, j] = A[i, k] * B[k, j] * S[i, j]
40
41expected = """; extended FROSTT format
422 1
438 8
441 8 2016
45"""
46
47# Force evaluation of the kernels by writing out X and Y.
48with tempfile.TemporaryDirectory() as test_dir:
49  x_file = os.path.join(test_dir, "X.tns")
50  y_file = os.path.join(test_dir, "Y.tns")
51  pt.write(x_file, X)
52  pt.write(y_file, Y)
53  #
54  # CHECK: Compare result True True
55  #
56  x_data = utils.file_as_string(x_file)
57  y_data = utils.file_as_string(y_file)
58  print(f"Compare result {x_data == expected} {y_data == expected}")
59