[mlir][taco] Add a utility to create an MLIR sparse tensor from a file.

[bixia]

Summary

Move the functions that retrieve the supporting C library, compile an MLIR
module and build a JIT execution engine to mlir_pytaco_utils.

Add a function to create an MLIR sparse tensor from a file and return a pointer
to the MLIR sparse tensor as well as the shape of the sparse tensor.

Add unit tests.

[aartbik]

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

157

note that this can now be the following, using Wren's new pipeline

pipeline = f"builtin.module(sparse-compiler{{{opt}}})"

[bixia]

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

157

Will delay this change until https://reviews.llvm.org/D118658 is merged.

[bixia]

Rebase.

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

 import enum
 import numpy as np
 import functools
 import operator
 import os
 import threading
 
 # Import MLIR related modules.
-from mlir import all_passes_registration  # Register MLIR compiler passes.
-from mlir import execution_engine
 from mlir import ir
 from mlir import runtime
 from mlir.dialects import arith
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 from mlir.dialects import sparse_tensor
 from mlir.dialects.linalg.opdsl import lang
-from mlir.passmanager import PassManager
 
 from . import mlir_pytaco_utils as utils
 
 # TACO naming prefixes.
 _TACO_INDEX_PREFIX = "i"
 _TACO_TENSOR_PREFIX = "A"
 
 # Bitwidths for pointers and indices.
 _POINTER_BIT_WIDTH = 0
 _INDEX_BIT_WIDTH = 0
-# The name for the environment variable that provides the full path for the
-# supporting library.
-_SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
-# The default supporting library if the environment variable is not provided.
-_DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"
-# The JIT compiler optimization level.
-_OPT_LEVEL = 2
 # The entry point to the JIT compiled program.
 _ENTRY_NAME = "main"
 
 # Type aliases for type annotation.
 _BinaryOp = Callable[[Any, Any], Any]
 _ExprVisitor = Callable[..., None]
 _ExprInfoDict = Dict["IndexExpr", "_ExprInfo"]
 _LogicalOp = Callable[[bool, bool], bool]
dtype_to_irtype = {
       Type.INT32: ir.IntegerType.get_signless(32),
       Type.INT64: ir.IntegerType.get_signless(64),
       Type.FLOAT32: ir.F32Type.get(),
       Type.FLOAT64: ir.F64Type.get()
   }
   return dtype_to_irtype[dtype.kind]
 
 
-def _compile_mlir(module: ir.Module) -> ir.Module:
-  """Compiles an MLIR module and returns the compiled module."""
-  # TODO: Replace this with a pipeline implemented for 
-  #   https://github.com/llvm/llvm-project/issues/51751.
-  pipeline = (
-      f"sparsification,"
-      f"sparse-tensor-conversion,"
-      f"builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
-      f"convert-scf-to-std,"
-      f"func-bufferize,"
-      f"tensor-constant-bufferize,"
-      f"builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),"
-      f"convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
-      f"lower-affine,"
-      f"convert-memref-to-llvm,"
-      f"convert-std-to-llvm,"
-      f"reconcile-unrealized-casts")
-  PassManager.parse(pipeline).run(module)
-  return module
-
-
[e-mail]_cache()
-def _get_support_lib_name() -> str:
-  """Returns the string for the supporting C shared library."""
-  return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
-
-
 def _ctype_pointer_from_array(array: np.ndarray) -> ctypes.pointer:
   """Returns the ctype pointer for the given numpy array."""
   return ctypes.pointer(
       ctypes.pointer(runtime.get_ranked_memref_descriptor(array)))
 
 
 class ModeFormat(enum.Enum):
   """The tensor dimension storage format class.
def ctype_pointer(self) -> ctypes.pointer:
     if self.is_dense():
       if self._dense_storage is None:
         self._dense_storage = np.zeros(self._shape, self._dtype.value)
       return _ctype_pointer_from_array(self._dense_storage)
 
     shape = np.array(self._shape, np.int64)
     indices = np.array(self._coords, np.int64)
     values = np.array(self._values, self._dtype.value)
-    ptr = utils.coo_tensor_to_sparse_tensor(_get_support_lib_name(), shape,
-                                            values, indices)
+    ptr = utils.coo_tensor_to_sparse_tensor(shape, values, indices)
     return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p))
 
   def get_coordinates_and_values(
       self) -> Tuple[List[Tuple[int, ...]], List[_AnyRuntimeType]]:
     """Returns the coordinates and values for the non-zero elements."""
     if not self.is_dense():
       return (self._coords, self._values)
 
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 = []
     self._visit(_gather_input_accesses_index_vars, (input_accesses,))
 
-    support_lib = _get_support_lib_name()
     # Build and compile the module to produce the execution engine.
     with ir.Context(), ir.Location.unknown():
       module = ir.Module.create()
       self._emit_assignment(module, dst, dst_indices, expr_to_info,
                             input_accesses)
-      compiled_module = _compile_mlir(module)
-
-      # We currently rely on an environment to pass in the full path of a
-      # supporting library for the execution engine.
-      engine = execution_engine.ExecutionEngine(
-          compiled_module, opt_level=_OPT_LEVEL, shared_libs=[support_lib])
+      engine = utils.compile_and_build_engine(module)
 
     # Gather the pointers for the input buffers.
     input_pointers = [a.tensor.ctype_pointer() for a in input_accesses]
     if dst.is_dense():
       # The pointer to receive dense output is the first argument to the
       # execution engine.
       arg_pointers = [dst.dense_dst_ctype_pointer()] + input_pointers
     else:
       # The pointer to receive sparse output is the last argument to the
       # execution engine. The pointer to receive a sparse tensor output is a
       # pointer to pointer of char.
       arg_pointers = input_pointers + [
           ctypes.pointer(ctypes.pointer(ctypes.c_char(0)))
       ]
 
     # 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(
-        support_lib,
         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))

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_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 process sparse tensor outputs.
 
-from typing import Tuple
+from typing import Sequence, Tuple
 import ctypes
 import functools
 import numpy as np
+import os
+
+# Import MLIR related modules.
+from mlir import all_passes_registration  # Register MLIR compiler passes.
+from mlir import execution_engine
+from mlir import ir
+from mlir import runtime
+from mlir.dialects import sparse_tensor
+from mlir.passmanager import PassManager
+
+# The name for the environment variable that provides the full path for the
+# supporting library.
+_SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
+# The default supporting library if the environment variable is not provided.
+_DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"
+
+# The JIT compiler optimization level.
+_OPT_LEVEL = 2
+# The entry point to the JIT compiled program.
+_ENTRY_NAME = "main"
 
 
 @functools.lru_cache()
-def _get_c_shared_lib(lib_name: str) -> ctypes.CDLL:
-  """Loads and returns the requested C shared library.
+def _get_support_lib_name() -> str:
+  """Gets the string name for the supporting C shared library."""
+  return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
 
-  Args:
-    lib_name: A string representing the C shared library.
+
[e-mail]_cache()
+def _get_c_shared_lib() -> ctypes.CDLL:
+  """Loads the supporting C shared library with the needed routines.
+
+  The name of the supporting C shared library is either provided by an
+  an environment variable or a default value.
 
   Returns:
-    The C shared library.
+    The supporting C shared library.
 
   Raises:
     OSError: If there is any problem in loading the shared library.
     ValueError: If the shared library doesn't contain the needed routines.
   """
   # This raises OSError exception if there is any problem in loading the shared
   # library.
-  c_lib = ctypes.CDLL(lib_name)
+  c_lib = ctypes.CDLL(_get_support_lib_name())
 
   try:
     c_lib.convertToMLIRSparseTensor.restype = ctypes.c_void_p
   except Exception as e:
     raise ValueError("Missing function convertToMLIRSparseTensor from "
                      f"the supporting C shared library: {e} ") from e
 
   try:
     c_lib.convertFromMLIRSparseTensor.restype = ctypes.c_void_p
   except Exception as e:
     raise ValueError("Missing function convertFromMLIRSparseTensor from "
                      f"the C shared library: {e} ") from e
 
   return c_lib
 
 
 def sparse_tensor_to_coo_tensor(
-    lib_name: str,
     sparse_tensor: ctypes.c_void_p,
     dtype: np.dtype,
 ) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
   """Converts an MLIR sparse tensor to a COO-flavored format tensor.
 
   Args:
-     lib_name: A string for the supporting C shared library.
      sparse_tensor: A ctypes.c_void_p to the MLIR sparse tensor descriptor.
      dtype: The numpy data type for the tensor elements.
 
   Returns:
     A tuple that contains the following values for the COO-flavored format
     tensor:
     rank: An integer for the rank of the tensor.
     nse: An interger for the number of non-zero values in the tensor.
     shape: A 1D numpy array of integers, for the shape of the tensor.
     values: A 1D numpy array, for the non-zero values in the tensor.
     indices: A 2D numpy array of integers, representing the indices for the
       non-zero values in the tensor.
 
   Raises:
     OSError: If there is any problem in loading the shared library.
     ValueError: If the shared library doesn't contain the needed routines.
   """
-  c_lib = _get_c_shared_lib(lib_name)
+  c_lib = _get_c_shared_lib()
 
   rank = ctypes.c_ulonglong(0)
   nse = ctypes.c_ulonglong(0)
   shape = ctypes.POINTER(ctypes.c_ulonglong)()
   values = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))()
   indices = ctypes.POINTER(ctypes.c_ulonglong)()
   c_lib.convertFromMLIRSparseTensor(sparse_tensor, ctypes.byref(rank),
                                     ctypes.byref(nse), ctypes.byref(shape),
                                     ctypes.byref(values), ctypes.byref(indices))
 
   # Convert the returned values to the corresponding numpy types.
   shape = np.ctypeslib.as_array(shape, shape=[rank.value])
   values = np.ctypeslib.as_array(values, shape=[nse.value])
   indices = np.ctypeslib.as_array(indices, shape=[nse.value, rank.value])
-  return rank, nse, shape, values, indices
+  return rank.value, nse.value, shape, values, indices
 
 
-def coo_tensor_to_sparse_tensor(lib_name: str, np_shape: np.ndarray,
-                                np_values: np.ndarray,
+def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
                                 np_indices: np.ndarray) -> int:
   """Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.
 
   Args:
-     lib_name: A string for the supporting C shared library.
      np_shape: A 1D numpy array of integers, for the shape of the tensor.
      np_values: A 1D numpy array, for the non-zero values in the tensor.
      np_indices: A 2D numpy array of integers, representing the indices for the
        non-zero values in the tensor.
 
   Returns:
      An integer for the non-null ctypes.c_void_p to the MLIR sparse tensor
      descriptor.
 
   Raises:
     OSError: If there is any problem in loading the shared library.
     ValueError: If the shared library doesn't contain the needed routines.
   """
 
   rank = ctypes.c_ulonglong(len(np_shape))
   nse = ctypes.c_ulonglong(len(np_values))
   shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
   values = np_values.ctypes.data_as(
       ctypes.POINTER(np.ctypeslib.as_ctypes_type(np_values.dtype)))
   indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
 
-  c_lib = _get_c_shared_lib(lib_name)
+  c_lib = _get_c_shared_lib()
   ptr = c_lib.convertToMLIRSparseTensor(rank, nse, shape, values, indices)
   assert ptr is not None, "Problem with calling convertToMLIRSparseTensor"
   return ptr
+
+
+def compile_and_build_engine(
+    module: ir.Module) -> execution_engine.ExecutionEngine:
+  """Compiles an MLIR module and builds a JIT execution engine.
+
+  Args:
+    module: The MLIR module.
+
+  Returns:
+    A JIT execution engine for the MLIR module.
+
+  """
+  pipeline = (
+      f"sparsification,"
+      f"sparse-tensor-conversion,"
+      f"builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),"
+      f"convert-scf-to-std,"
+      f"func-bufferize,"
+      f"tensor-constant-bufferize,"
+      f"builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),"
+      f"convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},"
+      f"lower-affine,"
+      f"convert-memref-to-llvm,"
+      f"convert-std-to-llvm,"
+      f"reconcile-unrealized-casts")
+  PassManager.parse(pipeline).run(module)
+  return execution_engine.ExecutionEngine(
+      module, opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])
+
+
+class _SparseTensorDescriptor(ctypes.Structure):
+  """A C structure for an MLIR sparse tensor."""
+  _fields_ = [
+      # A pointer for the MLIR sparse tensor storage.
+      ("storage", ctypes.POINTER(ctypes.c_ulonglong)),
+      # An MLIR MemRef descriptor for the shape of the sparse tensor.
+      ("shape", runtime.make_nd_memref_descriptor(1, ctypes.c_ulonglong)),
+  ]
+
+
+def _output_one_dim(dim: int, rank: int, shape: str) -> str:
+  """Produces the MLIR text code to output the size for the given dimension."""
+  return f"""
+  %c{dim} = arith.constant {dim} : index
+  %d{dim} = tensor.dim %t, %c{dim} : tensor<{shape}xf64, #enc>
+  memref.store %d{dim}, %b[%c{dim}] : memref<{rank}xindex>
+"""
+
+
+# TODO: With better support from MLIR, we may improve the current implementation
+# by doing the following:
+# (1) Use Python code to generate the kernel instead of doing MLIR text code
+#     stitching.
+# (2) Use scf.for instead of an unrolled loop to write out the dimension sizes
+#     when tensor.dim supports non-constant dimension value.
+def _get_create_sparse_tensor_kernel(
+    sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> str:
+  """Creates an MLIR text kernel to contruct a sparse tensor from a file.
+
+  The kernel returns a _SparseTensorDescriptor structure.
+  """
+  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))
+
+  # Get the MLIR text code to write the dimension sizes to the output buffer.
+  output_dims = "\n".join(
+      map(lambda d: _output_one_dim(d, rank, shape), range(rank)))
+
+  # Return the MLIR text kernel.
+  return f"""
+!Ptr = type !llvm.ptr<i8>
+#enc = #sparse_tensor.encoding<{{
+  dimLevelType = [ {sparsity} ]
+}}>
+func @{_ENTRY_NAME}(%filename: !Ptr) -> (tensor<{shape}xf64, #enc>, memref<{rank}xindex>)
+attributes {{ llvm.emit_c_interface }} {{
+  %t = sparse_tensor.new %filename : !Ptr to tensor<{shape}xf64, #enc>
+  %b = memref.alloc() : memref<{rank}xindex>
+  {output_dims}
+  return %t, %b : tensor<{shape}xf64, #enc>, memref<{rank}xindex>
+}}"""
+
+
+def create_sparse_tensor(
+    filename: str, sparsity: Sequence[sparse_tensor.DimLevelType]
+) -> Tuple[ctypes.c_void_p, np.ndarray]:
+  """Creates an MLIR sparse tensor from the input file.
+
+  Args:
+    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.
+
+  Returns:
+    A Tuple containing the following values:
+    storage: A ctypes.c_void_p for the MLIR sparse tensor storage.
+    shape: A 1D numpy array of integers, for the shape 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_create_sparse_tensor_kernel(sparsity)
+    module = ir.Module.parse(module)
+    engine = compile_and_build_engine(module)
+
+  # A sparse tensor descriptor to receive the kernel result.
+  c_tensor_desc = _SparseTensorDescriptor()
+  # Convert the filename to a byte stream.
+  c_filename = ctypes.c_char_p(bytes(filename, "utf-8"))
+
+  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

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

+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+from typing import Sequence
+import dataclasses
+import numpy as np
+import os
+import sys
+import tempfile
+
+from mlir.dialects import sparse_tensor
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco
+from tools import mlir_pytaco_utils as pytaco_utils
+
+# Define the aliases to shorten the code.
+_COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
+_DENSE = mlir_pytaco.ModeFormat.DENSE
+
+
+def _to_string(s: Sequence[int]) -> str:
+  """Converts a sequence of integer to a space separated value string."""
+  return " ".join(map(lambda e: str(e), s))
+
+
+def _add_one(s: Sequence[int]) -> Sequence[int]:
+  """Adds one to each element in the sequence of integer."""
+  return [i + 1 for i in s]
+
+
[e-mail](frozen=True)
+class _SparseTensorCOO:
+  """Values for a COO-flavored format sparse tensor.
+
+  Attributes:
+    rank: An integer rank for the tensor.
+    nse: An integer for the number of non-zero values.
+    shape: A sequence of integer for the dimension size.
+    values: A sequence of float for the non-zero values of the tensor.
+    indices: A sequence of coordinate, each coordinate is a sequence of integer.
+  """
+  rank: int
+  nse: int
+  shape: Sequence[int]
+  values: Sequence[float]
+  indices: Sequence[Sequence[int]]
+
+
+def _coo_values_to_tns_format(t: _SparseTensorCOO) -> str:
+  """Converts a sparse tensor COO-flavored values to TNS text format."""
+  # The coo_value_str contains one line for each (coordinate value) pair.
+  # Indices are 1-based in TNS text format but 0-based in MLIR.
+  coo_value_str = "\n".join(
+      map(lambda i: _to_string(_add_one(t.indices[i])) + " " + str(t.values[i]),
+          range(t.nse)))
+
+  # Returns the TNS text format representation for the tensor.
+  return f"""{t.rank} {t.nse}
+{_to_string(t.shape)}
+{coo_value_str}
+"""
+
+
+def _implement_read_tns_test(
+    t: _SparseTensorCOO,
+    sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> int:
+  tns_data = _coo_values_to_tns_format(t)
+
+  # Write sparse tensor data to a file.
+  with tempfile.TemporaryDirectory() as test_dir:
+    file_name = os.path.join(test_dir, "data.tns")
+    with open(file_name, "w") as file:
+      file.write(tns_data)
+
+    # Read the data from the file and construct an MLIR sparse tensor.
+    sparse_tensor, o_shape = pytaco_utils.create_sparse_tensor(
+        file_name, sparsity_codes)
+
+  passed = 0
+
+  # Verify the output shape for the tensor.
+  if np.allclose(o_shape, t.shape):
+    passed += 1
+
+  # Use the output MLIR sparse tensor pointer to retrieve the COO-flavored
+  # values and verify the values.
+  o_rank, o_nse, o_shape, o_values, o_indices = (
+      pytaco_utils.sparse_tensor_to_coo_tensor(sparse_tensor, np.float64))
+  if o_rank == t.rank and o_nse == t.nse and np.allclose(
+      o_shape, t.shape) and np.allclose(o_values, t.values) and np.allclose(
+          o_indices, t.indices):
+    passed += 1
+
+  return passed
+
+
+# A 2D sparse tensor data in COO-flavored format.
+_rank = 2
+_nse = 3
+_shape = [4, 5]
+_values = [3.0, 2.0, 4.0]
+_indices = [[0, 4], [1, 0], [3, 1]]
+
+_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
+_s = [_COMPRESSED, _COMPRESSED]
+# CHECK: PASSED 2D: 2
+print("PASSED 2D: ", _implement_read_tns_test(_t, _s))
+
+
+# A 3D sparse tensor data in COO-flavored format.
+_rank = 3
+_nse = 3
+_shape = [2, 5, 4]
+_values = [3.0, 2.0, 4.0]
+_indices = [[0, 4, 3], [1, 3, 0], [1, 3, 1]]
+
+_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
+_s = [_DENSE, _COMPRESSED, _COMPRESSED]
+# CHECK: PASSED 3D: 2
+print("PASSED 3D: ", _implement_read_tns_test(_t, _s))