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, l, m = pt.get_index_vars(5)
16
17# Set up scalar.
18alpha = pt.tensor(42.0)
19
20# Set up some sparse tensors with different dim annotations and ordering.
21S = pt.tensor([8, 8, 8],
22              pt.format([pt.compressed, pt.dense, pt.compressed], [1, 0, 2]))
23X = pt.tensor([8, 8, 8],
24              pt.format([pt.compressed, pt.compressed, pt.compressed],
25                        [1, 0, 2]))
26S.insert([0, 0, 0], 2.0)
27S.insert([1, 1, 1], 3.0)
28S.insert([4, 4, 4], 4.0)
29S.insert([7, 7, 7], 5.0)
30
31X[i, j, k] = alpha[0] * S[i, j, k]
32
33# Set up tensors with a dense last dimension. This results in a full
34# enveloping storage of all last "rows" with one or more nonzeros.
35T = pt.tensor([1, 2, 3, 4, 5],
36              pt.format([
37                  pt.compressed, pt.compressed, pt.compressed, pt.compressed,
38                  pt.dense
39              ]))
40Y = pt.tensor([1, 2, 3, 4, 5],
41              pt.format([
42                  pt.compressed, pt.compressed, pt.compressed, pt.compressed,
43                  pt.dense
44              ]))
45T.insert([0, 1, 2, 3, 4], -2.0)
46
47Y[i, j, k, l, m] = alpha[0] * T[i, j, k, l, m]
48
49# Set up a sparse tensor and dense tensor with different access.
50U = pt.tensor([2, 3], pt.format([pt.compressed, pt.compressed], [1, 0]))
51Z = pt.tensor([3, 2], pt.format([pt.dense, pt.dense]))
52U.insert([1, 2], 3.0)
53
54Z[i, j] = alpha[0] * U[j, i]
55
56x_expected = """; extended FROSTT format
573 4
588 8 8
591 1 1 84
602 2 2 126
615 5 5 168
628 8 8 210
63"""
64
65y_expected = """; extended FROSTT format
665 5
671 2 3 4 5
681 2 3 4 1 0
691 2 3 4 2 0
701 2 3 4 3 0
711 2 3 4 4 0
721 2 3 4 5 -84
73"""
74
75z_expected = """; extended FROSTT format
762 6
773 2
781 1 0
791 2 0
802 1 0
812 2 0
823 1 0
833 2 126
84"""
85
86# Force evaluation of the kernel by writing out X.
87with tempfile.TemporaryDirectory() as test_dir:
88  x_file = os.path.join(test_dir, 'X.tns')
89  pt.write(x_file, X)
90  y_file = os.path.join(test_dir, 'Y.tns')
91  pt.write(y_file, Y)
92  z_file = os.path.join(test_dir, 'Z.tns')
93  pt.write(z_file, Z)
94  #
95  # CHECK: Compare result True True True
96  #
97  x_data = utils.file_as_string(x_file)
98  y_data = utils.file_as_string(y_file)
99  z_data = utils.file_as_string(z_file)
100  print(
101      f'Compare result {x_data == x_expected} {y_data == y_expected} {z_data == z_expected}'
102  )
103