[mlir][sparse][taco] Support complex types.

[bixia]

Summary

Support complex types of float and double. See the added test for an example.

[bixia]

Fixed comments.

[aartbik]

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

78

I have been using np.complex128/np.complex64 for these, since I like the bit-number at the end (which is also consistent with what we do for float). Do you mind changing that here?

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

44

same question on complex128, or is this defined by taco this way?

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

142

note that we now have

mlir.runtime.as_ctype()

and

mlir.runtime.C128/C64

so I don't think you need to do this "doubling" trick anymore

151

we have

mlir.runtime.to_numpy()

[aartbik]

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

75

Note that we can include np.float16 also now!

[bixia]

Use runtime.as_ctype and runtime.to_numpy.
replace csingle/cdouble with complex64/128.

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

142

use runtime.as_ctype and runtime.to_numpy

[aartbik]

few last naming nitpicks, but this looks solid!

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

14

complex64/128?

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

44

same question, did you keep this csingle/cdouble on purpose?
why not complex64/128?

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

96

np.complex64/128?

149

yes, nice!

[bixia]

Rename csingle/ccomplex to complex64/128.

[bixia]

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

75

Will add this and test in another PR.

mlir/test/Integration/Dialect/SparseTensor/taco/test_tensor_complex.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
+
+compressed = pt.compressed
+
+passed = 0
+all_types = [pt.complex64, pt.complex128]
+for t in all_types:
+  i, j = pt.get_index_vars(2)
+  A = pt.tensor([2, 3], dtype=t)
+  B = pt.tensor([2, 3], dtype=t)
+  C = pt.tensor([2, 3], compressed, dtype=t)
+  A.insert([0, 1], 10 + 20j)
+  A.insert([1, 2], 40 + 0.5j)
+  B.insert([0, 0], 20)
+  B.insert([1, 2], 30 + 15j)
+  C[i, j] = A[i, j] + B[i, j]
+
+  indices, values = C.get_coordinates_and_values()
+  passed += isinstance(values[0], t.value)
+  passed += np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
+  passed += np.allclose(values, [20, 10 + 20j, 70 + 15.5j])
+
+# CHECK: Number of passed: 6
+print("Number of passed:", passed)

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

class Type(enum.Enum):
   """The data types supported by TACO.
 
   We use numpy data types to implement the enum data types.
   """
   INT8 = np.int8
   INT16 = np.int16
   INT32 = np.int32
   INT64 = np.int64
   # numpy _ctype_from_dtype_scalar can't handle np.float16 yet.
   FLOAT32 = np.float32
   FLOAT64 = np.float64
+  COMPLEX64 = np.complex64
+  COMPLEX128 = np.complex128
 
 
 # All floating point type enums.
 _FLOAT_TYPES = (Type.FLOAT32, Type.FLOAT64)
 # All integral type enums.
 _INT_TYPES = (Type.INT8, Type.INT16, Type.INT32, Type.INT64)
+# All complex type enums.
+_COMPLEX_TYPES = (Type.COMPLEX64, Type.COMPLEX128)
 # Type alias for any numpy type used to implement the runtime support for the
 # enum data types.
 _AnyRuntimeType = Union[np.int8, np.int16, np.int32, np.int64, np.float32,
-                        np.float64]
+                        np.float64, np.complex64, np.complex128]
 
 
 @dataclasses.dataclass(frozen=True)
 class DType:
   """The data type class.
 
   We support the TACO API dtype class with an alias of this class.
 
class DType:
   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
 
+  def is_complex(self) -> bool:
+    """Returns whether the data type represents a complex value."""
+    return self.kind in _COMPLEX_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: "i8",
       Type.INT16: "i16",
       Type.INT32: "i32",
       Type.INT64: "i64",
       Type.FLOAT32: "f32",
-      Type.FLOAT64: "f64"
+      Type.FLOAT64: "f64",
+      Type.COMPLEX64: "complex<f32>",
+      Type.COMPLEX128: "complex<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.int8: Type.INT8,
       np.int16: Type.INT16,
       np.int32: Type.INT32,
       np.int64: Type.INT64,
       np.float32: Type.FLOAT32,
-      np.float64: Type.FLOAT64
+      np.float64: Type.FLOAT64,
+      np.complex64: Type.COMPLEX64,
+      np.complex128: Type.COMPLEX128
   }
   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.INT8: ir.IntegerType.get_signless(8),
       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()
+      Type.FLOAT64: ir.F64Type.get(),
+      Type.COMPLEX64: ir.ComplexType.get(ir.F32Type.get()),
+      Type.COMPLEX128: ir.ComplexType.get(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)))
 
elif isinstance(fmt, Format):
         raise ValueError("Inconsistent shape and format: "
                          f"{self._shape}, {fmt}.")
       else:
         self._format = fmt
     else:
       raise ValueError(f"Invalid format argument: {fmt}.")
 
   def __init__(self,
-               value_or_shape: Optional[Union[List[int], Tuple[int, ...], float,
-                                              int]] = None,
+               value_or_shape: Optional[Union[List[int], Tuple[int, ...],
+                                              complex, float, int]] = None,
                fmt: Optional[Union[ModeFormat, List[ModeFormat],
                                    Format]] = None,
                dtype: Optional[DType] = None,
                name: Optional[str] = None,
                is_dense: bool = False):
     """The tensor constructor interface defined by TACO API.
 
     Args:
def __init__(self,
     fmt = fmt or ModeFormat.COMPRESSED
     # We currently use _coords and _values to host the sparse tensor value with
     # COO format, and _dense_storage to host the dense tensor value. We don't
     # support the conversion between the two storages.
     self._coords = []
     self._values = []
     self._stats = _Stats()
     if value_or_shape is None or isinstance(value_or_shape, int) or isinstance(
-        value_or_shape, float):
+        value_or_shape, float) or isinstance(value_or_shape, complex):
       # Create a scalar tensor and ignore the fmt parameter.
       self._shape = []
       self._format = _make_format([], [])
       if value_or_shape is not None:
         self._dense_storage = np.array(value_or_shape, dtype=self._dtype.value)
     elif (isinstance(value_or_shape, tuple) or isinstance(
         value_or_shape, list)) and _all_instance_of(value_or_shape, int):
       # Create a tensor with the specified shape and format.
class Tensor:
   def __repr__(self) -> str:
     self._sync_value()
     self.unpack()
     value_str = (f"{repr(self._dense_storage)})" if self.is_dense() else
                  f"{repr(self._coords)} {repr(self._values)})")
     return (f"Tensor(_name={repr(self._name)} "
             f"_dtype={repr(self._dtype)} : ") + value_str
 
-  def insert(self, coords: List[int], val: Union[float, int]) -> None:
+  def insert(self, coords: List[int], val: Union[complex, float, int]) -> None:
     """Inserts a value to the given coordinate.
 
     Args:
       coords: A list of integer coordinates. The length of the list must be the
         same as the rank of the tensor.
       val: A value being inserted. It is either an integral or a floating point
         value. This value will be converted to the data type of the tensor.
 
if not isinstance(coords, list):
       raise ValueError(f"Non list coordinate detected: {coords}.")
     if not _all_instance_of(coords, int):
       raise ValueError(f"Non integer coordinate detected: {coords}.")
     if (len(coords) != self.order or
         any([c < 0 or c >= self._shape[i] for i, c in enumerate(coords)])):
       raise ValueError("Invalid coordinate for rank: "
                        f"{self.order}, {coords}.")
 
-    if not isinstance(val, int) and not isinstance(val, float):
+    if not isinstance(val, int) and not isinstance(
+        val, float) and not isinstance(val, complex):
       raise ValueError(f"Value is neither int nor float: {val}.")
 
     self._coords.append(tuple(coords))
     self._values.append(self._dtype.value(val))
 
   def is_dense(self) -> bool:
     """Returns true if the tensor doesn't have sparsity annotation."""
     return self.order == 0 or self._format is None

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

 
 # Data type constants defined by PyTACO API.
 int8 = mlir_pytaco.DType(mlir_pytaco.Type.INT8)
 int16 = mlir_pytaco.DType(mlir_pytaco.Type.INT16)
 int32 = mlir_pytaco.DType(mlir_pytaco.Type.INT32)
 int64 = mlir_pytaco.DType(mlir_pytaco.Type.INT64)
 float32 = mlir_pytaco.DType(mlir_pytaco.Type.FLOAT32)
 float64 = mlir_pytaco.DType(mlir_pytaco.Type.FLOAT64)
+complex64 = mlir_pytaco.DType(mlir_pytaco.Type.COMPLEX64)
+complex128 = mlir_pytaco.DType(mlir_pytaco.Type.COMPLEX128)
 
 # Storage format constants defined by the PyTACO API. In PyTACO, each storage
 # format constant has two aliasing names.
 compressed = mlir_pytaco.ModeFormat.COMPRESSED
 Compressed = mlir_pytaco.ModeFormat.COMPRESSED
 dense = mlir_pytaco.ModeFormat.DENSE
 Dense = mlir_pytaco.ModeFormat.DENSE

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

support_types = [(np.int8, c_lib.convertToMLIRSparseTensorI8,
                       c_lib.convertFromMLIRSparseTensorI16),
                      (np.int32, c_lib.convertToMLIRSparseTensorI32,
                       c_lib.convertFromMLIRSparseTensorI32),
                      (np.int64, c_lib.convertToMLIRSparseTensorI64,
                       c_lib.convertFromMLIRSparseTensorI64),
                      (np.float32, c_lib.convertToMLIRSparseTensorF32,
                       c_lib.convertFromMLIRSparseTensorF32),
                      (np.float64, c_lib.convertToMLIRSparseTensorF64,
-                      c_lib.convertFromMLIRSparseTensorF64)]
+                      c_lib.convertFromMLIRSparseTensorF64),
+                     (np.complex64, c_lib.convertToMLIRSparseTensorC32,
+                      c_lib.convertFromMLIRSparseTensorC32),
+                     (np.complex128, c_lib.convertToMLIRSparseTensorC64,
+                      c_lib.convertFromMLIRSparseTensorC64)]
   except Exception as e:
     raise ValueError(f"Missing supporting function: {e}") from e
   for i, info in enumerate(support_types):
     _record_support_funcs(info[0], info[1], info[2], type_to_funcs)
 
   def get_support_funcs(ty: np.dtype):
     funcs = type_to_funcs[ty]
     assert funcs is not None
) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
   Raises:
     OSError: If there is any problem in loading the shared library.
     ValueError: If the shared library doesn't contain the needed routines.
   """
   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))()
+
+  values = ctypes.POINTER(runtime.as_ctype(np.dtype(dtype)))()
   indices = ctypes.POINTER(ctypes.c_ulonglong)()
   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])
+  values = runtime.to_numpy(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,
                                 np_indices: np.ndarray, np_perm: np.ndarray,
                                 np_sparse: np.ndarray) -> int:
   """Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.
 
   Args:
Raises:
     ValueError: If the shared library doesn't contain the needed routines.
   """
 
   r = len(np_shape)
   rank = ctypes.c_ulonglong(r)
   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)))
+      ctypes.POINTER(runtime.as_ctype(np.dtype(np_values.dtype))))
   indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
 
   perm = np_perm.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
   sparse = np_sparse.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
 
   convert_to = _get_support_func_locator()(np_values.dtype.type)[0]
   ptr = convert_to(rank, nse, shape, values, indices, perm, sparse)
   assert ptr is not None, "Problem with calling convertToMLIRSparseTensorF64"