[mlir][taco] Use sparse_tensor.out to write sparse tensors to files.

[bixia]

Summary

Add a Python method, output_sparse_tensor, to use sparse_tensor.out to write
a sparse tensor value to a file.

Modify the method that evaluates a tensor expression to return a pointer of the
MLIR sparse tensor for the result to delay the extraction of the coordinates and
non-zero values.

Implement the Tensor to_file method to evaluate the tensor assignment and write
the result to a file.

Add unit tests. Modify test golden files to reflect the change that TNS outputs
now have a comment line and two meta data lines.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/taco/test_MTTKRP.py

49–51

Why don't we simply use

filecmp.cmp()

as in test_SpMM.py?

and test for the exact contents of the golden file vs computed file.
We don't have any accuracy issues, it should be sorted the same way always,
and that we verify the metadata for correctness too.

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py

36

same

37

this seems the opposite of what I think is better

why not add the

; extended FROSTT format

to the golden file and we are done?

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMV.py

49–50

same

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

853–864

can we avoid the codeup for unpacked()?
perhaps simply by packing?

860

is this lexicographically sorted?

[bixia]

Add testing_utils for comparing results.

mlir/test/Integration/Dialect/SparseTensor/taco/test_MTTKRP.py

49–51

Per offline chat, I am adding a testing_utils for comparing result.

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py

36

Addressing this with testing_utils.

37

Addressing this with testing_utils.

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMV.py

49–50

Similar to the above.

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

853–864

Per offline chat, I add a TODO here and will fix this in the next PR.

860

The coordinates produced for a dense tensor are in lexicographically ordered. However, the values added by user through the Tensor.Insert method are not ordered.
Per offline conversation, I add a TODO here and will fix this in the next PR along with fixing the definition of "dense tensors".

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

860

OK, next revision sounds good.

mlir/test/Integration/Dialect/SparseTensor/taco/tools/testing_utils.py

10

I like this! All logic in once place, so we can later decide to test more/less/different!

21

note that technically there could be more than one comment line
of course, this is a very controlled environment, so okay for now

32

do we really need the float = 0.01 rtol? Can it be less?
That seems a pretty big difference, right?

[bixia]

mlir/test/Integration/Dialect/SparseTensor/taco/tools/testing_utils.py

32

right, changed to 0.0001

[bixia]

Change rtol from 0.001 to 0.00001.

mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_A.tns

+# See http://frostt.io/tensors/file-formats.html for FROSTT (.tns) format
+2 50
+2 25
 1 1 12
 1 2 12
 1 3 12
 1 4 12
 1 5 12
 1 6 12
 1 7 12
 1 8 12

mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_C.tns

+# See http://frostt.io/tensors/file-formats.html for FROSTT (.tns) format
+2 9
+3 3
 1.0 1.0 100.0
 1.0 2.0 107.0
 1.0 3.0 114.0
 2.0 1.0 201.0
 2.0 2.0 216.0
 2.0 3.0 231.0
 3.0 1.0 318.0
 3.0 2.0 342.0
 3.0 3.0 366.0

mlir/test/Integration/Dialect/SparseTensor/taco/data/gold_y.tns

 # See http://frostt.io/tensors/file-formats.html for FROSTT (.tns) format
+1 3
+3
 1 37102
 2 -20.4138
 3 804927

mlir/test/Integration/Dialect/SparseTensor/taco/test_MTTKRP.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import numpy as np
 import os
 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
 
 ###### This PyTACO part is taken from the TACO open-source project. ######
 # See http://tensor-compiler.org/docs/data_analytics/index.html.
 
 compressed = pt.compressed
 dense = pt.dense
 
 # Define formats for storing the sparse tensor and dense matrices.

 # 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)}")
+  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)}")

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMM.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import filecmp
 import numpy as np
 import os
 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.
 #
 # TODO: use sparse_tensor.out for output, so that C.tns becomes
 #       a file in extended FROSTT format
 #
 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 files True
+  # CHECK: Compare result True
   #
-  print(f"Compare files {filecmp.cmp(golden_file, out_file)}")
+  print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")

mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMV.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import numpy as np
 import os
 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
 
 ###### This PyTACO part is taken from the TACO open-source project. ######
 # See http://tensor-compiler.org/docs/scientific_computing/index.html.
 
 compressed = pt.compressed
 dense = pt.dense
 
 # Define formats for storing the sparse matrix and dense vectors.

 # 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)}")
+  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)}")

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

elif (isinstance(value_or_shape, tuple) or isinstance(
       # Create a tensor with the specified shape and format.
       self._shape = list(value_or_shape)
       self._init_format(fmt)
     else:
       raise ValueError("Invalid first argument. "
                        "Must be a tuple or list for a shape or a single value"
                        f"if initializing a scalar tensor: {value_or_shape}.")
 
+  def _set_packed_sparse_tensor(self, pointer: ctypes.c_void_p) -> None:
+    """Records the MLIR sparse tensor pointer."""
+    self._sparse_value_location = _SparseValueInfo._PACKED
+    self._packed_sparse_value = pointer
+
   def is_unpacked(self) -> bool:
     """Returns true if the tensor value is not packed as MLIR sparse tensor."""
     return (self._sparse_value_location == _SparseValueInfo._UNPACKED)
 
   def unpack(self) -> None:
     """Unpacks the MLIR sparse tensor representation."""
     if self.is_dense() or self.is_unpacked():
       return
) -> "Tensor":
 
     Returns:
       A tensor with the given non-zero values and storage format. The tensor
       value is stored as an MLIR sparse tensor.
     """
     sparse_tensor, shape = utils.create_sparse_tensor(filename,
                                                       fmt.format_pack.formats)
     tensor = Tensor(shape.tolist(), fmt)
-    tensor._sparse_value_location = _SparseValueInfo._PACKED
-    tensor._packed_sparse_value = sparse_tensor
+    tensor._set_packed_sparse_tensor(sparse_tensor)
 
     return tensor
 
+  def to_file(self, filename: str) -> None:
+    """Output the tensor value to a file.
+
+    This method evaluates any pending assignment to the tensor and outputs the
+    tensor value.
+
+    Args:
+      filename: A string file name.
+    """
+    self._sync_value()
+    if not self.is_unpacked():
+      utils.output_sparse_tensor(self._packed_sparse_value, filename,
+                                 self._format.format_pack.formats)
+      return
+
+    # TODO: Use MLIR code to output the value.
+    coords, values = self.get_coordinates_and_values()
+    assert len(coords) == len(values)
+    with open(filename, "w") as file:
+      # Output a comment line and the meta data.
+      file.write("; extended FROSTT format\n")
+      file.write(f"{self.order} {len(coords)}\n")
+      file.write(f"{' '.join(map(lambda i: str(i), self.shape))}\n")
+      # Output each (coordinate value) pair in a line.
+      for c, v in zip(coords, values):
+        # The coordinates are 1-based in the text file and 0-based in memory.
+        plus_one_to_str = lambda x: str(x + 1)
+        file.write(f"{' '.join(map(plus_one_to_str,c))} {v}\n")
+
   @property
   def dtype(self) -> DType:
     """Returns the data type for the Tensor."""
     return self._dtype
 
   @property
   def format(self) -> Format:
     """Returns the storage format for the Tensor."""
def evaluate(self) -> None:
     """Evaluates the assignment to the tensor."""
     result = self._assignment.expression.evaluate(self,
                                                   self._assignment.indices)
     self._assignment = None
     if self.is_dense():
       assert isinstance(result, np.ndarray)
       self._dense_storage = result
     else:
-      assert _all_instance_of(result, np.ndarray) and len(result) == 2
-      assert (result[0].ndim, result[1].ndim) == (1, 2)
-      (self._values, self._coords) = result
+      self._set_packed_sparse_tensor(result)
 
   def _sync_value(self) -> None:
     """Updates the tensor value by evaluating the pending assignment."""
     if self._assignment is not None:
       self.evaluate()
 
   def mlir_tensor_type(self) -> ir.RankedTensorType:
     """Returns the MLIR type for the tensor."""
with ir.InsertionPoint(module.body):
 
       linalg_funcop.func_op.attributes[
           "llvm.emit_c_interface"] = ir.UnitAttr.get()
 
   def evaluate(
       self,
       dst: Tensor,
       dst_indices: Tuple[IndexVar, ...],
-  ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+  ) -> Union[np.ndarray, ctypes.c_void_p]:
     """Evaluates tensor assignment dst[dst_indices] = expression.
 
     Args:
       dst: The destination tensor.
       dst_indices: The tuple of IndexVar used to access the destination tensor.
 
     Returns:
-      The result of the dense tensor represented in numpy ndarray or the sparse
-      tensor represented by two numpy ndarray for its non-zero values and
-      indices.
+      The result of the dense tensor represented in numpy ndarray or the pointer
+      to the MLIR sparse tensor.
 
     Raises:
       ValueError: If the expression is not proper or not supported.
     """
     expr_to_info = self._validate_and_collect_expr_info(dst, dst_indices)
 
     # Compute a list of input accesses.
     input_accesses = []
else:
       ]
 
     # Invoke the execution engine to run the module and return the result.
     engine.invoke(_ENTRY_NAME, *arg_pointers)
 
     if dst.is_dense():
       return runtime.ranked_memref_to_numpy(arg_pointers[0][0])
 
-    # Check and return the sparse tensor output.
-    rank, nse, shape, values, indices = utils.sparse_tensor_to_coo_tensor(
-        ctypes.cast(arg_pointers[-1][0], ctypes.c_void_p),
-        np.float64,
-    )
-    assert (np.equal(rank, dst.order)
-            and np.array_equal(shape, np.array(dst.shape)) and
-            np.equal(values.ndim, 1) and np.equal(values.shape[0], nse) and
-            np.equal(indices.ndim, 2) and np.equal(indices.shape[0], nse) and
-            np.equal(indices.shape[1], rank))
-    return (values, indices)
+    # Return the sparse tensor pointer.
+    return arg_pointers[-1][0]
 
 
 @dataclasses.dataclass(frozen=True)
 class Access(IndexExpr):
   """The tensor access class.
 
   We support the TACO API access class with an alias of this class.
 
def __post_init__(self) -> None:
     """
     if (not isinstance(self.indices, tuple) or
         not _all_instance_of(self.indices, IndexVar)):
       raise ValueError(f"Indices contain non IndexVar: {str(self.indices)}.")
     if self.tensor.order != len(self.indices):
       raise ValueError("Invalid indices for rank: "
                        f"str{self.tensor.order} != len({str(self.indices)}).")
 
+  def __repr__(self) -> str:
+    # The Tensor __repr__ method evaluates the pending assignment to the tensor.
+    # We want to define the __repr__ method here to avoid such evaluation of the
+    # tensor assignment.
+    indices_str = ", ".join(map(lambda i: i.name, self.indices))
+    return (f"Tensor({self.tensor.name}) " f"Indices({indices_str})")
+
   def _emit_expression(
       self,
       expr_to_opnd: Dict[IndexExpr, lang.OperandDef],
       expr_to_info: _ExprInfoDict,
   ) -> lang.ScalarExpression:
     """Emits a linalg dialect TensorUse expression for the tensor access."""
     assert self in expr_to_opnd
     dims = _mlir_dimensions_from_index_vars(self.indices)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_io.py

 DType = mlir_pytaco.DType
 Type = mlir_pytaco.Type
 
 # Constants used in the implementation.
 _MTX_FILENAME_SUFFIX = ".mtx"
 _TNS_FILENAME_SUFFIX = ".tns"
 
 
-def _write_tns(file: TextIO, tensor: Tensor) -> None:
-  """Outputs a tensor to a file using .tns format."""
-  coords, non_zeros = tensor.get_coordinates_and_values()
-  assert len(coords) == len(non_zeros)
-  # Output a coordinate and the corresponding value in a line.
-  for c, v in zip(coords, non_zeros):
-    # The coordinates are 1-based in the text file and 0-based in memory.
-    plus_one_to_str = lambda x: str(x + 1)
-    file.write(f"{' '.join(map(plus_one_to_str,c))} {v}\n")
-
-
 def read(filename: str, fmt: Format) -> Tensor:
   """Inputs a tensor from a given file.
 
   The name suffix of the file specifies the format of the input tensor. We
   currently only support .mtx format for support sparse tensors.
 
   Args:
     filename: A string input filename.
def write(filename: str, tensor: Tensor) -> None:
   """
   if (not isinstance(filename, str) or
       not filename.endswith(_TNS_FILENAME_SUFFIX)):
     raise ValueError("Expected string filename ends with"
                      f" {_TNS_FILENAME_SUFFIX}: {filename}.")
   if not isinstance(tensor, Tensor):
     raise ValueError(f"Expected a Tensor object: {tensor}.")
 
-  # TODO: combine the evaluation and the outputing into one step.
-  tensor._sync_value()
-  with open(filename, "w") as file:
-    return _write_tns(file, tensor)
+  tensor.to_file(filename)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

arg_pointers = [
       ctypes.byref(ctypes.pointer(c_tensor_desc)),
       ctypes.byref(c_filename)
   ]
 
   # Invoke the execution engine to run the module and return the result.
   engine.invoke(_ENTRY_NAME, *arg_pointers)
   shape = runtime.ranked_memref_to_numpy(ctypes.pointer(c_tensor_desc.shape))
   return c_tensor_desc.storage, shape
+
+
+# TODO: With better support from MLIR, we may improve the current implementation
+# by using Python code to generate the kernel instead of doing MLIR text code
+# stitching.
+def _get_output_sparse_tensor_kernel(
+    sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> str:
+  """Creates an MLIR text kernel to output a sparse tensor to a file.
+
+  The kernel returns void.
+  """
+  rank = len(sparsity_codes)
+
+  # Use ? to represent a dimension in the dynamic shape string representation.
+  shape = "x".join(map(lambda d: "?", range(rank)))
+
+  # Convert the encoded sparsity values to a string representation.
+  sparsity = ", ".join(
+      map(lambda s: '"compressed"' if s.value else '"dense"', sparsity_codes))
+
+  # Return the MLIR text kernel.
+  return f"""
+!Ptr = type !llvm.ptr<i8>
+#enc = #sparse_tensor.encoding<{{
+  dimLevelType = [ {sparsity} ]
+}}>
+func @{_ENTRY_NAME}(%t: tensor<{shape}xf64, #enc>, %filename: !Ptr)
+attributes {{ llvm.emit_c_interface }} {{
+  sparse_tensor.out %t, %filename : tensor<{shape}xf64, #enc>, !Ptr
+  std.return
+}}"""
+
+
+def output_sparse_tensor(
+    tensor: ctypes.c_void_p, filename: str,
+    sparsity: Sequence[sparse_tensor.DimLevelType]) -> None:
+  """Outputs an MLIR sparse tensor to the given file.
+
+  Args:
+    tensor: A C pointer to the MLIR sparse tensor.
+    filename: A string for the name of the file that contains the tensor data in
+      a COO-flavored format.
+    sparsity: A sequence of DimLevelType values, one for each dimension of the
+      tensor.
+
+  Raises:
+    OSError: If there is any problem in loading the supporting C shared library.
+    ValueError:  If the shared library doesn't contain the needed routine.
+  """
+  with ir.Context() as ctx, ir.Location.unknown():
+    module = _get_output_sparse_tensor_kernel(sparsity)
+    module = ir.Module.parse(module)
+    engine = compile_and_build_engine(module)
+
+  # Convert the filename to a byte stream.
+  c_filename = ctypes.c_char_p(bytes(filename, "utf-8"))
+
+  arg_pointers = [
+      ctypes.byref(ctypes.cast(tensor, ctypes.c_void_p)),
+      ctypes.byref(c_filename)
+  ]
+
+  # Invoke the execution engine to run the module and return the result.
+  engine.invoke(_ENTRY_NAME, *arg_pointers)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/testing_utils.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# This file contains the utilities to support testing.
+
+import numpy as np
+
+
+def compare_sparse_tns(expected: str, actual: str, rtol: float = 0.01) -> bool:
+  """Compares sparse tensor actual output file with expected output file.
+
+  This routine assumes the input files are in FROSTT format. See
+  http://frostt.io/tensors/file-formats.html for FROSTT (.tns) format.
+
+  It also assumes the first line in the output file is a comment line.
+
+  """
+  with open(actual, "r") as actual_f:
+    with open(expected, "r") as expected_f:
+      # Skip the first comment line.
+      _ = actual_f.readline()
+      _ = expected_f.readline()
+
+      # Compare the two lines of meta data
+      if actual_f.readline() != expected_f.readline() or actual_f.readline(
+      ) != expected_f.readline():
+        return FALSE
+
+  actual_data = np.loadtxt(actual, np.float64, skiprows=3)
+  expected_data = np.loadtxt(expected, np.float64, skiprows=3)
+  return np.allclose(actual_data, expected_data, rtol=rtol)

mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py

def test_read_tns():
   passed += (not a.is_unpacked())
   a.unpack()
   passed += (a.is_unpacked())
   coords, values = a.get_coordinates_and_values()
   passed += np.allclose(coords, [[0, 1], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)
+
+
+# CHECK-LABEL: test_write_unpacked_tns
+@_run
+def test_write_unpacked_tns():
+  a = mlir_pytaco.Tensor([2, 3])
+  a.insert([0, 1], 10)
+  a.insert([1, 2], 40)
+  a.insert([0, 0], 20)
+  with tempfile.TemporaryDirectory() as test_dir:
+    file_name = os.path.join(test_dir, "data.tns")
+    mlir_pytaco_io.write(file_name, a)
+    with open(file_name, "r") as file:
+      lines = file.readlines()
+  passed = 0
+  # Skip the comment line in the output.
+  if lines[1:] == ["2 3\n", "2 3\n", "1 2 10.0\n", "2 3 40.0\n", "1 1 20.0\n"]:
+    passed = 1
+  # CHECK: 1
+  print(passed)
+
+
+# CHECK-LABEL: test_write_packed_tns
+@_run
+def test_write_packed_tns():
+  a = mlir_pytaco.Tensor([2, 3])
+  a.insert([0, 1], 10)
+  a.insert([1, 2], 40)
+  a.insert([0, 0], 20)
+  b = mlir_pytaco.Tensor([2, 3])
+  i, j = mlir_pytaco.get_index_vars(2)
+  b[i, j] = a[i, j] + a[i, j]
+  with tempfile.TemporaryDirectory() as test_dir:
+    file_name = os.path.join(test_dir, "data.tns")
+    mlir_pytaco_io.write(file_name, b)
+    with open(file_name, "r") as file:
+      lines = file.readlines()
+  passed = 0
+  # Skip the comment line in the output.
+  if lines[1:] == ["2 3\n", "2 3\n", "1 1 40\n", "1 2 20\n", "2 3 80\n"]:
+    passed = 1
+  # CHECK: 1
+  print(passed)