1# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
2
3from typing import Sequence
4import dataclasses
5import numpy as np
6import os
7import sys
8import tempfile
9
10from mlir.dialects import sparse_tensor
11
12_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
13sys.path.append(_SCRIPT_PATH)
14from tools import mlir_pytaco
15from tools import mlir_pytaco_utils as pytaco_utils
16
17# Define the aliases to shorten the code.
18_COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
19_DENSE = mlir_pytaco.ModeFormat.DENSE
20
21
22def _to_string(s: Sequence[int]) -> str:
23  """Converts a sequence of integer to a space separated value string."""
24  return " ".join(map(lambda e: str(e), s))
25
26
27def _add_one(s: Sequence[int]) -> Sequence[int]:
28  """Adds one to each element in the sequence of integer."""
29  return [i + 1 for i in s]
30
31
32@dataclasses.dataclass(frozen=True)
33class _SparseTensorCOO:
34  """Values for a COO-flavored format sparse tensor.
35
36  Attributes:
37    rank: An integer rank for the tensor.
38    nse: An integer for the number of non-zero values.
39    shape: A sequence of integer for the dimension size.
40    values: A sequence of float for the non-zero values of the tensor.
41    indices: A sequence of coordinate, each coordinate is a sequence of integer.
42  """
43  rank: int
44  nse: int
45  shape: Sequence[int]
46  values: Sequence[float]
47  indices: Sequence[Sequence[int]]
48
49
50def _coo_values_to_tns_format(t: _SparseTensorCOO) -> str:
51  """Converts a sparse tensor COO-flavored values to TNS text format."""
52  # The coo_value_str contains one line for each (coordinate value) pair.
53  # Indices are 1-based in TNS text format but 0-based in MLIR.
54  coo_value_str = "\n".join(
55      map(lambda i: _to_string(_add_one(t.indices[i])) + " " + str(t.values[i]),
56          range(t.nse)))
57
58  # Returns the TNS text format representation for the tensor.
59  return f"""{t.rank} {t.nse}
60{_to_string(t.shape)}
61{coo_value_str}
62"""
63
64
65def _implement_read_tns_test(
66    t: _SparseTensorCOO,
67    sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> int:
68  tns_data = _coo_values_to_tns_format(t)
69
70  # Write sparse tensor data to a file.
71  with tempfile.TemporaryDirectory() as test_dir:
72    file_name = os.path.join(test_dir, "data.tns")
73    with open(file_name, "w") as file:
74      file.write(tns_data)
75
76    # Read the data from the file and construct an MLIR sparse tensor.
77    sparse_tensor, o_shape = pytaco_utils.create_sparse_tensor(
78        file_name, sparsity_codes, "f64")
79
80  passed = 0
81
82  # Verify the output shape for the tensor.
83  if np.array_equal(o_shape, t.shape):
84    passed += 1
85
86  # Use the output MLIR sparse tensor pointer to retrieve the COO-flavored
87  # values and verify the values.
88  o_rank, o_nse, o_shape, o_values, o_indices = (
89      pytaco_utils.sparse_tensor_to_coo_tensor(sparse_tensor, np.float64))
90  if o_rank == t.rank and o_nse == t.nse and np.array_equal(
91      o_shape, t.shape) and np.allclose(o_values, t.values) and np.array_equal(
92          o_indices, t.indices):
93    passed += 1
94
95  return passed
96
97
98# A 2D sparse tensor data in COO-flavored format.
99_rank = 2
100_nse = 3
101_shape = [4, 5]
102_values = [3.0, 2.0, 4.0]
103_indices = [[0, 4], [1, 0], [3, 1]]
104
105_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
106_s = [_COMPRESSED, _COMPRESSED]
107# CHECK: PASSED 2D: 2
108print("PASSED 2D: ", _implement_read_tns_test(_t, _s))
109
110
111# A 3D sparse tensor data in COO-flavored format.
112_rank = 3
113_nse = 3
114_shape = [2, 5, 4]
115_values = [3.0, 2.0, 4.0]
116_indices = [[0, 4, 3], [1, 3, 0], [1, 3, 1]]
117
118_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
119_s = [_DENSE, _COMPRESSED, _COMPRESSED]
120# CHECK: PASSED 3D: 2
121print("PASSED 3D: ", _implement_read_tns_test(_t, _s))
122