ENH: Support 3D COO and cross-rank reshape

Author: mtsokol

sparse/mlir_backend/_common.py

@@ -12,6 +12,10 @@ class MlirType(abc.ABC):
     def get_mlir_type(cls) -> ir.Type: ...
 
 
+class RefableList(list):
+    pass
+
+
 def fn_cache(f, maxsize: int | None = None):
     return functools.wraps(f)(functools.lru_cache(maxsize=maxsize)(f))
 

sparse/mlir_backend/_constructors.py

@@ -8,7 +8,7 @@
 import numpy as np
 import scipy.sparse as sps
 
-from ._common import _hold_self_ref_in_ret, _take_owneship, fn_cache
+from ._common import RefableList, _hold_self_ref_in_ret, _take_owneship, fn_cache
 from ._core import ctx, libc
 from ._dtypes import DType, asdtype
 
@@ -118,26 +118,31 @@ class Coo(ctypes.Structure):
         _index_dtype = index_dtype
 
         @classmethod
-        def from_sps(cls, arr: sps.coo_array) -> "Coo":
-            assert arr.has_canonical_format, "COO must have canonical format"
-            np_pos = np.array([0, arr.size], dtype=index_dtype.np_dtype)
-            np_coords = np.stack(arr.coords, axis=1, dtype=index_dtype.np_dtype)
+        def from_sps(cls, arr: sps.coo_array | np.ndarray) -> "Coo":
+            if isinstance(arr, sps.coo_array):
+                assert arr.has_canonical_format, "COO must have canonical format"
+                np_pos = np.array([0, arr.size], dtype=index_dtype.np_dtype)
+                np_coords = np.stack(arr.coords, axis=1, dtype=index_dtype.np_dtype)
+                np_data = arr.data
+            else:
+                assert len(arr) == 3, "COO must be comprised of three arrays"
+                np_pos, np_coords, np_data = arr
+
             pos = numpy_to_ranked_memref(np_pos)
             coords = numpy_to_ranked_memref(np_coords)
-            data = numpy_to_ranked_memref(arr.data)
-
+            data = numpy_to_ranked_memref(np_data)
             coo_instance = cls(pos=pos, coords=coords, data=data)
             _take_owneship(coo_instance, np_pos)
             _take_owneship(coo_instance, np_coords)
-            _take_owneship(coo_instance, arr)
+            _take_owneship(coo_instance, np_data)
 
             return coo_instance
 
-        def to_sps(self, shape: tuple[int, ...]) -> sps.coo_array:
+        def to_sps(self, shape: tuple[int, ...]) -> sps.coo_array | list[np.ndarray]:
             pos = ranked_memref_to_numpy(self.pos)
             coords = ranked_memref_to_numpy(self.coords)[pos[0] : pos[1]]
             data = ranked_memref_to_numpy(self.data)
-            return sps.coo_array((data, coords.T), shape=shape)
+            return sps.coo_array((data, coords.T), shape=shape) if len(shape) == 2 else RefableList([pos, coords, data])
 
         def to_module_arg(self) -> list:
             return [
@@ -159,8 +164,13 @@ def get_tensor_definition(cls, shape: tuple[int, ...]) -> ir.RankedTensorType:
                 compressed_lvl = sparse_tensor.EncodingAttr.build_level_type(
                     sparse_tensor.LevelFormat.compressed, [sparse_tensor.LevelProperty.non_unique]
                 )
-                levels = (compressed_lvl, sparse_tensor.LevelFormat.singleton)
-                ordering = ir.AffineMap.get_permutation([0, 1])
+                mid_singleton_lvls = [
+                    sparse_tensor.EncodingAttr.build_level_type(
+                        sparse_tensor.LevelFormat.singleton, [sparse_tensor.LevelProperty.non_unique]
+                    )
+                ] * (len(shape) - 2)
+                levels = (compressed_lvl, *mid_singleton_lvls, sparse_tensor.LevelFormat.singleton)
+                ordering = ir.AffineMap.get_permutation([*range(len(shape))])
                 encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, index_width, index_width)
                 return ir.RankedTensorType.get(list(shape), values_dtype, encoding)
 
@@ -191,10 +201,7 @@ def from_sps(cls, arrs: list[np.ndarray]) -> "Csf":
             return csf_instance
 
         def to_sps(self, shape: tuple[int, ...]) -> list[np.ndarray]:
-            class List(list):
-                pass
-
-            return List(ranked_memref_to_numpy(field) for field in self.get__fields_())
+            return RefableList(ranked_memref_to_numpy(field) for field in self.get__fields_())
 
         def to_module_arg(self) -> list:
             return [ctypes.pointer(ctypes.pointer(field)) for field in self.get__fields_()]
@@ -310,20 +317,20 @@ def __init__(
 
             if obj.format in ("csr", "csc"):
                 order = "r" if obj.format == "csr" else "c"
-                index_dtype = asdtype(obj.indptr.dtype)
-                self._format_class = get_csx_class(self._values_dtype, index_dtype, order)
+                self._index_dtype = asdtype(obj.indptr.dtype)
+                self._format_class = get_csx_class(self._values_dtype, self._index_dtype, order)
                 self._obj = self._format_class.from_sps(obj)
             elif obj.format == "coo":
-                index_dtype = asdtype(obj.coords[0].dtype)
-                self._format_class = get_coo_class(self._values_dtype, index_dtype)
+                self._index_dtype = asdtype(obj.coords[0].dtype)
+                self._format_class = get_coo_class(self._values_dtype, self._index_dtype)
                 self._obj = self._format_class.from_sps(obj)
             else:
                 raise Exception(f"{obj.format} SciPy format not supported.")
 
         elif _is_numpy_obj(obj):
             self._owns_memory = False
-            index_dtype = asdtype(np.intp)
-            self._format_class = get_dense_class(self._values_dtype, index_dtype)
+            self._index_dtype = asdtype(np.intp)
+            self._format_class = get_dense_class(self._values_dtype, self._index_dtype)
             self._obj = self._format_class.from_sps(obj)
 
         elif _is_mlir_obj(obj):
@@ -332,11 +339,13 @@ def __init__(
             self._obj = obj
 
         elif format is not None:
-            if format == "csf":
+            if format in ["csf", "coo"]:
+                fn_format_class = get_csf_class if format == "csf" else get_coo_class
                 self._owns_memory = False
-                index_dtype = asdtype(np.intp)
-                self._format_class = get_csf_class(self._values_dtype, index_dtype)
+                self._index_dtype = asdtype(np.intp)
+                self._format_class = fn_format_class(self._values_dtype, self._index_dtype)
                 self._obj = self._format_class.from_sps(obj)
+
             else:
                 raise Exception(f"Format {format} not supported.")
 

sparse/mlir_backend/_ops.py

@@ -116,10 +116,27 @@ def add(x1: Tensor, x2: Tensor) -> Tensor:
     return Tensor(ret_obj, shape=out_tensor_type.shape)
 
 
+def _infer_format_class(rank: int, values_dtype: type[DType], index_dtype: type[DType]) -> type[ctypes.Structure]:
+    from ._constructors import get_csf_class, get_csx_class, get_dense_class
+
+    if rank == 1:
+        return get_dense_class(values_dtype, index_dtype)
+    if rank == 2:
+        return get_csx_class(values_dtype, index_dtype, order="r")
+    if rank == 3:
+        return get_csf_class(values_dtype, index_dtype)
+    raise Exception(f"Rank not supported to infer format: {rank}")
+
+
 def reshape(x: Tensor, /, shape: tuple[int, ...]) -> Tensor:
-    ret_obj = x._format_class()
     x_tensor_type = x._obj.get_tensor_definition(x.shape)
-    out_tensor_type = x._obj.get_tensor_definition(shape)
+    if len(x.shape) == len(shape):
+        out_tensor_type = x._obj.get_tensor_definition(shape)
+        ret_obj = x._format_class()
+    else:
+        format_class = _infer_format_class(len(shape), x._values_dtype, x._index_dtype)
+        out_tensor_type = format_class.get_tensor_definition(shape)
+        ret_obj = format_class()
 
     with ir.Location.unknown(ctx):
         shape_tensor_type = ir.RankedTensorType.get([len(shape)], Index.get_mlir_type())

sparse/mlir_backend/tests/test_simple.py

@@ -199,13 +199,38 @@ def test_csf_format(dtype):
         np.testing.assert_array_equal(actual, expected)
 
 
+@parametrize_dtypes
+def test_coo_3d_format(dtype):
+    SHAPE = (2, 2, 4)
+    pos = np.array([0, 7])
+    crd = np.array([[0, 1, 0, 0, 1, 1, 0], [1, 3, 1, 0, 0, 1, 0], [3, 1, 1, 0, 1, 1, 1]])
+    data = np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype)
+    coo = [pos, crd, data]
+
+    coo_tensor = sparse.asarray(coo, shape=SHAPE, dtype=sparse.asdtype(dtype), format="coo")
+    result = coo_tensor.to_scipy_sparse()
+    for actual, expected in zip(result, coo, strict=False):
+        np.testing.assert_array_equal(actual, expected)
+
+    # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
+    # res_tensor = sparse.add(coo_tensor, coo_tensor).to_scipy_sparse()
+    # coo_2 = [pos, crd, data * 2]
+    # for actual, expected in zip(res_tensor, coo_2, strict=False):
+    #     np.testing.assert_array_equal(actual, expected)
+
+
 @parametrize_dtypes
 def test_reshape(rng, dtype):
     DENSITY = 0.5
     sampler = generate_sampler(dtype, rng)
 
     # CSR, CSC, COO
-    for shape, new_shape in [((100, 50), (25, 200)), ((80, 1), (8, 10))]:
+    for shape, new_shape in [
+        ((100, 50), (25, 200)),
+        ((100, 50), (10, 500, 1)),
+        ((80, 1), (8, 10)),
+        ((80, 1), (80,)),
+    ]:
         for format in ["csr", "csc", "coo"]:
             if format == "coo":
                 # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
@@ -217,15 +242,17 @@ def test_reshape(rng, dtype):
             arr = sps.random_array(
                 shape, density=DENSITY, format=format, dtype=dtype, random_state=rng, data_sampler=sampler
             )
-            if format == "coo":
-                arr.sum_duplicates()
-
+            arr.sum_duplicates()
             tensor = sparse.asarray(arr)
 
             actual = sparse.reshape(tensor, shape=new_shape).to_scipy_sparse()
+            if isinstance(actual, list):
+                continue  # skip checking CSF output
+            if not isinstance(actual, np.ndarray):
+                actual = actual.todense()
             expected = arr.todense().reshape(new_shape)
 
-            np.testing.assert_array_equal(actual.todense(), expected)
+            np.testing.assert_array_equal(actual, expected)
 
     # CSF
     csf_shape = (2, 2, 4)