[mlir][sparse][taco] Add support for float32.

[bixia]

Summary

Previously, we only support float64. We now support float32 and float64. When
constructing a tensor without providing a data type, the default is float32.

Fix the tests to data type consistency. All PyTACO application tests now use
float32 to match the default data type of TACO. Other tests may use float32 or
float64.

[aartbik]

You will have to rebase against head of main to get a clean presubmit run (I had the same problem today)

[bixia]

Rebase.

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

 # format) and store it as a compressed sparse fiber tensor. We use a small
 # tensor for the purpose of testing. To run the program using the data from
 # the real application, please download the data from:
 # http://frostt.io/tensors/nell-2/
 B = pt.read(os.path.join(_SCRIPT_PATH, "data/nell-2.tns"), csf)
 
 # These two lines have been modified from the original program to use static
 # data to support result comparison.
-C = pt.from_array(np.full((B.shape[1], 25), 1, dtype=np.float64))
-D = pt.from_array(np.full((B.shape[2], 25), 2, dtype=np.float64))
+C = pt.from_array(np.full((B.shape[1], 25), 1, dtype=np.float32))
+D = pt.from_array(np.full((B.shape[2], 25), 2, dtype=np.float32))
 
 # Declare the result to be a dense matrix.
 A = pt.tensor([B.shape[0], 25], rm)
 
 # Declare index vars.
 i, j, k, l = pt.get_index_vars(4)
 
 # Define the MTTKRP computation.

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

 sys.path.append(_SCRIPT_PATH)
 
 from tools import mlir_pytaco_api as pt
 from tools import testing_utils as utils
 
 i, j, k = pt.get_index_vars(3)
 
 # Set up dense matrices.
-A = pt.from_array(np.full((8, 8), 2.0))
-B = pt.from_array(np.full((8, 8), 3.0))
+A = pt.from_array(np.full((8, 8), 2.0, dtype=np.float32))
+B = pt.from_array(np.full((8, 8), 3.0, dtype=np.float32))
 
 # Set up sparse matrices.
 S = pt.tensor([8, 8], pt.format([pt.compressed, pt.compressed]))
 X = pt.tensor([8, 8], pt.format([pt.compressed, pt.compressed]))
 Y = pt.tensor([8, 8], pt.compressed)  # alternative syntax works too
 
 S.insert([0, 7], 42.0)
 

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

 # downloaded from here:
 # https://www.cise.ufl.edu/research/sparse/MM/Boeing/pwtk.tar.gz
 # In order to run the program using the matrix above, you can download the
 # matrix and replace this path to the actual path to the file.
 A = pt.read(os.path.join(_SCRIPT_PATH, "data/pwtk.mtx"), csr)
 
 # These two lines have been modified from the original program to use static
 # data to support result comparison.
-x = pt.from_array(np.full((A.shape[1],), 1, dtype=np.float64))
-z = pt.from_array(np.full((A.shape[0],), 2, dtype=np.float64))
+x = pt.from_array(np.full((A.shape[1],), 1, dtype=np.float32))
+z = pt.from_array(np.full((A.shape[0],), 2, dtype=np.float32))
 
 # Declare the result to be a dense vector
 y = pt.tensor([A.shape[0]], dv)
 
 # Declare index vars
 i, j = pt.get_index_vars(2)
 
 # Define the SpMV computation

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

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import numpy as np
 import os
 import sys
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
 from tools import mlir_pytaco_api as pt
 
 i, j = pt.get_index_vars(2)
 # Both tensors are true dense tensors.
 A = pt.from_array(np.full([2,3], 1, dtype=np.float64))
 B = pt.from_array(np.full([2,3], 2, dtype=np.float64))
 # Define the result tensor as a true dense tensor. The parameter is_dense=True
 # is an MLIR-PyTACO extension.
-C = pt.tensor([2, 3], is_dense=True)
+C = pt.tensor([2, 3], dtype=pt.float64, is_dense=True)
 
 C[i, j] = A[i, j] + B[i, j]
 
 # CHECK: [3. 3. 3. 3. 3. 3.]
 print(C.to_array().reshape(6))

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

class DType:
   The following methods are defined by the TACO API:
     is_float: Returns whether the data type represents a floating point value.
     is_int:   Returns whether the data type represents an integral value.
 
   Attributes:
     kind: A Type enum representing the data type.
     value: The numpy data type for the TACO data type.
   """
-  kind: Type = Type.FLOAT64
+  kind: Type = Type.FLOAT32
 
   def is_float(self) -> bool:
     """Returns whether the data type represents a floating point value."""
     return self.kind in _FLOAT_TYPES
 
   def is_int(self) -> bool:
     """Returns whether the data type represents an integral value."""
     return self.kind in _INT_TYPES
 
   @property
   def value(self) -> _AnyRuntimeType:
     """Returns the numpy dtype for the data type."""
     return self.kind.value
 
 
+def _dtype_to_mlir_str(dtype: DType) -> str:
+  """Returns the MLIR string for the given dtype."""
+  dtype_to_str = {
+      Type.INT16: "i16",
+      Type.INT32: "i32",
+      Type.INT64: "i64",
+      Type.FLOAT32: "f32",
+      Type.FLOAT64: "f64"
+  }
+  return dtype_to_str[dtype.kind]
+
+
+def _nptype_to_taco_type(ty: np.dtype) -> DType:
+  """Returns the TACO type for the given numpy type."""
+  nptype_to_dtype = {
+      np.int16: Type.INT16,
+      np.int32: Type.INT32,
+      np.int64: Type.INT64,
+      np.float32: Type.FLOAT32,
+      np.float64: Type.FLOAT64
+  }
+  return DType(nptype_to_dtype[ty])
+
+
 def _mlir_type_from_taco_type(dtype: DType) -> ir.Type:
   """Returns the MLIR type corresponding to the given TACO type."""
   dtype_to_irtype = {
       Type.INT16: ir.IntegerType.get_signless(16),
       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 _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.
Args:
       is_dense: A boolean variable to indicate whether the tensor is a dense
         tensor without any sparsity annotation.
 
     Raises:
       ValueError: If there is any inconsistency among the input arguments.
     """
     # Take care of the argument default values common to both sparse tensors
     # and dense tensors.
-    dtype = dtype or DType(Type.FLOAT64)
+    dtype = dtype or DType(Type.FLOAT32)
     self._name = name or self._get_unique_name()
     self._assignment = None
     self._sparse_value_location = _SparseValueInfo._UNPACKED
     self._dense_storage = None
     self._dtype = dtype
 
     if is_dense:
       assert (fmt is None)
class Tensor:
   def unpack(self) -> None:
     """Unpacks the MLIR sparse tensor representation."""
     if self.is_dense() or self.is_unpacked():
       return
 
     # Use the output MLIR sparse tensor pointer to retrieve the COO-flavored
     # values and verify the values.
     rank, nse, shape, values, indices = utils.sparse_tensor_to_coo_tensor(
-        self._packed_sparse_value, np.float64)
+        self._packed_sparse_value, self._dtype.value)
     assert rank == self.order
     assert np.allclose(self.shape, shape)
     assert nse == len(values)
     self._coords = indices
     self._values = values
     self._sparse_value_location = _SparseValueInfo._UNPACKED
 
   def __repr__(self) -> str:
def to_array(self) -> np.ndarray:
     self._sync_value()
 
     return self._dense_storage
 
   @staticmethod
   def from_array(array: np.ndarray) -> "Tensor":
     """Returns a dense tensor with the value copied from the input array.
 
-    We currently only support the conversion of float64 numpy arrays to Tensor.
+    We currently only support the conversion of float32 and float64 numpy arrays
+    to Tensor.
 
     Args:
       array: The numpy array that provides the data type, shape and value for
         the tensor.
 
     Returns:
       A Tensor object.
 
     Raises:
-      ValueError if the data type of the numpy array is not float64.
+      ValueError if the data type of the numpy array is not supported.
     """
-    if array.dtype != np.float64:
-      raise ValueError(f"Expected float64 value type: {array.dtype}.")
-    tensor = Tensor(array.shape, is_dense=True)
+    if array.dtype != np.float32 and array.dtype != np.float64:
+      raise ValueError(f"Expected floating point value type: {array.dtype}.")
+    tensor = Tensor(
+        array.shape,
+        dtype=_nptype_to_taco_type(array.dtype.type),
+        is_dense=True)
     tensor._dense_storage = np.copy(array)
     return tensor
 
   @staticmethod
   def from_coo(
       coordinates: List[Tuple[int, ...]],
       values: List[_AnyRuntimeType],
       fmt: Format,
) -> "Tensor":
     rank = fmt.rank()
     assert all(len(c) == rank and _all_instance_of(c, int) for c in coordinates)
 
     # Find the maximum coordinate value for each dimension.
     max_coordinate = list(map(max, zip(*coordinates)))
     # The size of each dimension is one more that such a maximum coordinate
     # value.
     shape = [c + 1 for c in max_coordinate]
-    tensor = Tensor(shape, fmt)
+    tensor = Tensor(shape, fmt, dtype=dtype)
     tensor._coords = coordinates
     tensor._values = values
 
     return tensor
 
   @staticmethod
   def from_file(
       filename: str,
       fmt: Format,
       dtype: DType,
   ) -> "Tensor":
     """Constructs a sparse tensor using the COO-flavored values from a file.
 
     Args:
       filename: A string for the name of the file that contains the sparse
         tensor data.
       fmt: The tensor storage format.
       dtype: The tensor element data type.
 
     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)
+                                                      fmt.format_pack.formats,
+                                                      _dtype_to_mlir_str(dtype))
+    tensor = Tensor(shape.tolist(), fmt, dtype=dtype)
     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
if self.is_dense():
       raise ValueError("Writing dense tensors without sparsity annotation to "
                        "file is not supported.")
 
     if self.is_unpacked():
       raise ValueError("Writing unpacked sparse tensors to file is not "
                        "supported.")
 
     utils.output_sparse_tensor(self._packed_sparse_value, filename,
-                               self._format.format_pack.formats)
+                               self._format.format_pack.formats,
+                               _dtype_to_mlir_str(self._dtype))
 
   @property
   def dtype(self) -> DType:
     """Returns the data type for the Tensor."""
     return self._dtype
 
   @property
   def format(self) -> Format:

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 read(filename: str, fmt: Format) -> Tensor:
+def read(filename: str, fmt: Format,
+         dtype: DType = DType(Type.FLOAT32)) -> 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.
     fmt: The storage format of the tensor.
+    dtype: The data type, default to float32.
 
   Raises:
     ValueError: If filename doesn't end with .mtx or .tns, or fmt is not an
     instance of Format or fmt is not a sparse tensor.
   """
   if (not isinstance(filename, str) or
       (not filename.endswith(_MTX_FILENAME_SUFFIX) and
        not filename.endswith(_TNS_FILENAME_SUFFIX))):
     raise ValueError("Expected string filename ends with "
                      f"{_MTX_FILENAME_SUFFIX} or {_TNS_FILENAME_SUFFIX}: "
                      f"{filename}.")
 
-  return Tensor.from_file(filename, fmt, DType(Type.FLOAT64))
+  return Tensor.from_file(filename, fmt, dtype)
 
 
 def write(filename: str, tensor: Tensor) -> None:
   """Outputs a tensor to a given file.
 
   The name suffix of the file specifies the format of the output. We currently
   only support .tns format.
 

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 Sequence, Tuple
+from typing import Callable, Dict, 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
 
+# Type aliases for type annotation.
+_SupportFunc = Callable[..., None]
+_SupportFuncLocator = Callable[[np.dtype], Tuple[_SupportFunc, _SupportFunc]]
+
 # 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_support_lib_name() -> str:
   """Gets the string name for the supporting C shared library."""
   return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
 
 
+def _record_support_funcs(
+    ty: np.dtype, to_func: _SupportFunc, from_func: _SupportFunc,
+    ty_to_funcs: Dict[np.dtype, Tuple[_SupportFunc, _SupportFunc]]) -> None:
+  """Records the two supporting functions for a given data type."""
+  to_func.restype = ctypes.c_void_p
+  from_func.restype = ctypes.c_void_p
+  ty_to_funcs[ty] = (to_func, from_func)
+
+
 @functools.lru_cache()
-def _get_c_shared_lib() -> ctypes.CDLL:
-  """Loads the supporting C shared library with the needed routines.
+def _get_support_func_locator() -> _SupportFuncLocator:
+  """Constructs a function to locate the supporting functions for a data type.
+
+  Loads the supporting C shared library with the needed routines. Constructs a
+  dictionary from the supported data types to the routines for the data types,
+  and then a function to look up the dictionary for a given data type.
 
   The name of the supporting C shared library is either provided by an
   an environment variable or a default value.
 
   Returns:
-    The supporting C shared library.
+    The function to look up the supporting functions for a given data type.
 
   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(_get_support_lib_name())
 
+  type_to_funcs = {}
   try:
-    c_lib.convertToMLIRSparseTensorF64.restype = ctypes.c_void_p
+    _record_support_funcs(np.float32, c_lib.convertToMLIRSparseTensorF32,
+                          c_lib.convertFromMLIRSparseTensorF32, type_to_funcs)
   except Exception as e:
-    raise ValueError("Missing function convertToMLIRSparseTensorF64 from "
-                     f"the supporting C shared library: {e} ") from e
+    raise ValueError(f"Missing supporting function: {e}") from e
 
   try:
-    c_lib.convertFromMLIRSparseTensorF64.restype = ctypes.c_void_p
+    _record_support_funcs(np.float64, c_lib.convertToMLIRSparseTensorF64,
+                          c_lib.convertFromMLIRSparseTensorF64, type_to_funcs)
   except Exception as e:
-    raise ValueError("Missing function convertFromMLIRSparseTensorF64 from "
-                     f"the C shared library: {e} ") from e
+    raise ValueError(f"Missing supporting function: {e}") from e
+
+  def get_support_funcs(ty: np.dtype):
+    funcs = type_to_funcs[ty]
+    assert funcs is not None
+    return funcs
 
-  return c_lib
+  return get_support_funcs
 
 
 def sparse_tensor_to_coo_tensor(
     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.
 
Returns:
     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()
-
+  convert_from = _get_support_func_locator()(dtype)[1]
   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.convertFromMLIRSparseTensorF64(sparse_tensor, ctypes.byref(rank),
-                                       ctypes.byref(nse), ctypes.byref(shape),
-                                       ctypes.byref(values),
-                                       ctypes.byref(indices))
+  convert_from(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.value, nse.value, shape, values, indices
 
 
def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
 
   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()
-  ptr = c_lib.convertToMLIRSparseTensorF64(rank, nse, shape, values, indices)
+  convert_to = _get_support_func_locator()(np_values.dtype.type)[0]
+  ptr = convert_to(rank, nse, shape, values, indices)
   assert ptr is not None, "Problem with calling convertToMLIRSparseTensorF64"
   return ptr
 
 
 def compile_and_build_engine(
     module: ir.Module) -> execution_engine.ExecutionEngine:
   """Compiles an MLIR module and builds a JIT execution engine.
 
class _SparseTensorDescriptor(ctypes.Structure):
   _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:
+def _output_one_dim(dim: int, rank: int, shape: str, type: 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>
+  %d{dim} = tensor.dim %t, %c{dim} : tensor<{shape}x{type}, #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:
+    sparsity_codes: Sequence[sparse_tensor.DimLevelType], type: str) -> 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)))
+      map(lambda d: _output_one_dim(d, rank, shape, type), 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>)
+func @{_ENTRY_NAME}(%filename: !Ptr) -> (tensor<{shape}x{type}, #enc>, memref<{rank}xindex>)
 attributes {{ llvm.emit_c_interface }} {{
-  %t = sparse_tensor.new %filename : !Ptr to tensor<{shape}xf64, #enc>
+  %t = sparse_tensor.new %filename : !Ptr to tensor<{shape}x{type}, #enc>
   %b = memref.alloc() : memref<{rank}xindex>
   {output_dims}
-  return %t, %b : tensor<{shape}xf64, #enc>, memref<{rank}xindex>
+  return %t, %b : tensor<{shape}x{type}, #enc>, memref<{rank}xindex>
 }}"""
 
 
-def create_sparse_tensor(
-    filename: str, sparsity: Sequence[sparse_tensor.DimLevelType]
-) -> Tuple[ctypes.c_void_p, np.ndarray]:
+def create_sparse_tensor(filename: str,
+                         sparsity: Sequence[sparse_tensor.DimLevelType],
+                         type: str) -> 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 = _get_create_sparse_tensor_kernel(sparsity, type)
     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
 
 
 # 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:
+    sparsity_codes: Sequence[sparse_tensor.DimLevelType], type: str) -> 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)
+func @{_ENTRY_NAME}(%t: tensor<{shape}x{type}, #enc>, %filename: !Ptr)
 attributes {{ llvm.emit_c_interface }} {{
-  sparse_tensor.out %t, %filename : tensor<{shape}xf64, #enc>, !Ptr
+  sparse_tensor.out %t, %filename : tensor<{shape}x{type}, #enc>, !Ptr
   std.return
 }}"""
 
 
-def output_sparse_tensor(
-    tensor: ctypes.c_void_p, filename: str,
-    sparsity: Sequence[sparse_tensor.DimLevelType]) -> None:
+def output_sparse_tensor(tensor: ctypes.c_void_p, filename: str,
+                         sparsity: Sequence[sparse_tensor.DimLevelType],
+                         type: str) -> 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.
+    type: The MLIR string for the data type.
 
   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 = _get_output_sparse_tensor_kernel(sparsity, type)
     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/unit_test_tensor_utils.py

def _implement_read_tns_test(
   # 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)
+        file_name, sparsity_codes, "f64")
 
   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