ENH: Simple COO format (#768)

Author: mtsokol

sparse/mlir_backend/_constructors.py

@@ -5,6 +5,7 @@
 import mlir.execution_engine
 import mlir.passmanager
 from mlir import ir
+from mlir import runtime as rt
 from mlir.dialects import arith, bufferization, func, sparse_tensor, tensor
 
 import numpy as np
@@ -13,7 +14,6 @@
 from ._common import fn_cache
 from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx
 from ._dtypes import DType, Index, asdtype
-from ._memref import make_memref_ctype, ranked_memref_from_np
 
 
 def _hold_self_ref_in_ret(fn):
@@ -108,7 +108,7 @@ def free_tensor(tensor_shaped):
     @classmethod
     def assemble(cls, module, arr: np.ndarray) -> ctypes.c_void_p:
         assert arr.ndim == 2
-        data = ranked_memref_from_np(arr.flatten())
+        data = rt.get_ranked_memref_descriptor(arr.flatten())
         out = ctypes.c_void_p()
         module.invoke(
             "assemble",
@@ -121,14 +121,14 @@ def assemble(cls, module, arr: np.ndarray) -> ctypes.c_void_p:
     def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p, dtype: type[DType]) -> np.ndarray:
         class Dense(ctypes.Structure):
             _fields_ = [
-                ("data", make_memref_ctype(dtype, 1)),
+                ("data", rt.make_nd_memref_descriptor(1, dtype.to_ctype())),
                 ("data_len", np.ctypeslib.c_intp),
                 ("shape_x", np.ctypeslib.c_intp),
                 ("shape_y", np.ctypeslib.c_intp),
             ]
 
             def to_np(self) -> np.ndarray:
-                data = self.data.to_numpy()[: self.data_len]
+                data = rt.ranked_memref_to_numpy([self.data])[: self.data_len]
                 return data.reshape((self.shape_x, self.shape_y))
 
         arr = Dense()
@@ -141,8 +141,107 @@ def to_np(self) -> np.ndarray:
 
 
 class COOFormat:
-    # TODO: implement
-    ...
+    @fn_cache
+    def get_module(shape: tuple[int], values_dtype: type[DType], index_dtype: type[DType]):
+        with ir.Location.unknown(ctx):
+            module = ir.Module.create()
+            values_dtype = values_dtype.get_mlir_type()
+            index_dtype = index_dtype.get_mlir_type()
+            index_width = getattr(index_dtype, "width", 0)
+            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])
+            encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, index_width, index_width)
+            coo_shaped = ir.RankedTensorType.get(list(shape), values_dtype, encoding)
+
+            tensor_1d_index = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], index_dtype)
+            tensor_2d_index = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size(), len(shape)], index_dtype)
+            tensor_1d_values = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], values_dtype)
+
+            with ir.InsertionPoint(module.body):
+
+                @func.FuncOp.from_py_func(tensor_1d_index, tensor_2d_index, tensor_1d_values)
+                def assemble(pos, index, values):
+                    return sparse_tensor.assemble(coo_shaped, (pos, index), values)
+
+                @func.FuncOp.from_py_func(coo_shaped)
+                def disassemble(tensor_shaped):
+                    nse = sparse_tensor.number_of_entries(tensor_shaped)
+                    pos = tensor.EmptyOp([arith.constant(ir.IndexType.get(), 2)], index_dtype)
+                    index = tensor.EmptyOp([nse, 2], index_dtype)
+                    values = tensor.EmptyOp([nse], values_dtype)
+                    pos, index, values, pos_len, index_len, values_len = sparse_tensor.disassemble(
+                        (tensor_1d_index, tensor_2d_index),
+                        tensor_1d_values,
+                        (index_dtype, index_dtype),
+                        index_dtype,
+                        tensor_shaped,
+                        (pos, index),
+                        values,
+                    )
+                    shape_consts = [arith.constant(index_dtype, s) for s in shape]
+                    return pos, index, values, pos_len, index_len, values_len, *shape_consts
+
+                @func.FuncOp.from_py_func(coo_shaped)
+                def free_tensor(tensor_shaped):
+                    bufferization.dealloc_tensor(tensor_shaped)
+
+            assemble.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+            disassemble.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+            free_tensor.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+            if DEBUG:
+                (CWD / "coo_module.mlir").write_text(str(module))
+            pm = mlir.passmanager.PassManager.parse("builtin.module(sparsifier{create-sparse-deallocs=1})")
+            pm.run(module.operation)
+            if DEBUG:
+                (CWD / "coo_module_opt.mlir").write_text(str(module))
+
+        module = mlir.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
+        return (module, coo_shaped)
+
+    @classmethod
+    def assemble(cls, module: ir.Module, arr: sps.coo_array) -> ctypes.c_void_p:
+        out = ctypes.c_void_p()
+        module.invoke(
+            "assemble",
+            ctypes.pointer(
+                ctypes.pointer(rt.get_ranked_memref_descriptor(np.array([0, arr.size], dtype=arr.coords[0].dtype)))
+            ),
+            ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np.stack(arr.coords, axis=1)))),
+            ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arr.data))),
+            ctypes.pointer(out),
+        )
+        return out
+
+    @classmethod
+    def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p, dtype: type[DType]) -> sps.coo_array:
+        class Coo(ctypes.Structure):
+            _fields_ = [
+                ("pos", rt.make_nd_memref_descriptor(1, Index.to_ctype())),
+                ("index", rt.make_nd_memref_descriptor(2, Index.to_ctype())),
+                ("values", rt.make_nd_memref_descriptor(1, dtype.to_ctype())),
+                ("pos_len", np.ctypeslib.c_intp),
+                ("index_len", np.ctypeslib.c_intp),
+                ("values_len", np.ctypeslib.c_intp),
+                ("shape_x", np.ctypeslib.c_intp),
+                ("shape_y", np.ctypeslib.c_intp),
+            ]
+
+            def to_sps(self) -> sps.coo_array:
+                pos = rt.ranked_memref_to_numpy([self.pos])[: self.pos_len]
+                index = rt.ranked_memref_to_numpy([self.index])[pos[0] : pos[1]]
+                values = rt.ranked_memref_to_numpy([self.values])[: self.values_len]
+                return sps.coo_array((values, index.T), shape=(self.shape_x, self.shape_y))
+
+        arr = Coo()
+        module.invoke(
+            "disassemble",
+            ctypes.pointer(ctypes.pointer(arr)),
+            ctypes.pointer(ptr),
+        )
+        return arr.to_sps()
 
 
 class CSRFormat:
@@ -207,9 +306,9 @@ def assemble(cls, module: ir.Module, arr: sps.csr_array) -> ctypes.c_void_p:
         out = ctypes.c_void_p()
         module.invoke(
             "assemble",
-            ctypes.pointer(ctypes.pointer(ranked_memref_from_np(arr.indptr))),
-            ctypes.pointer(ctypes.pointer(ranked_memref_from_np(arr.indices))),
-            ctypes.pointer(ctypes.pointer(ranked_memref_from_np(arr.data))),
+            ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arr.indptr))),
+            ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arr.indices))),
+            ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(arr.data))),
             ctypes.pointer(out),
         )
         return out
@@ -218,9 +317,9 @@ def assemble(cls, module: ir.Module, arr: sps.csr_array) -> ctypes.c_void_p:
     def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p, dtype: type[DType]) -> sps.csr_array:
         class Csr(ctypes.Structure):
             _fields_ = [
-                ("pos", make_memref_ctype(Index, 1)),
-                ("crd", make_memref_ctype(Index, 1)),
-                ("data", make_memref_ctype(dtype, 1)),
+                ("pos", rt.make_nd_memref_descriptor(1, Index.to_ctype())),
+                ("crd", rt.make_nd_memref_descriptor(1, Index.to_ctype())),
+                ("data", rt.make_nd_memref_descriptor(1, dtype.to_ctype())),
                 ("pos_len", np.ctypeslib.c_intp),
                 ("crd_len", np.ctypeslib.c_intp),
                 ("data_len", np.ctypeslib.c_intp),
@@ -229,9 +328,9 @@ class Csr(ctypes.Structure):
             ]
 
             def to_sps(self) -> sps.csr_array:
-                pos = self.pos.to_numpy()[: self.pos_len]
-                crd = self.crd.to_numpy()[: self.crd_len]
-                data = self.data.to_numpy()[: self.data_len]
+                pos = rt.ranked_memref_to_numpy([self.pos])[: self.pos_len]
+                crd = rt.ranked_memref_to_numpy([self.crd])[: self.crd_len]
+                data = rt.ranked_memref_to_numpy([self.data])[: self.data_len]
                 return sps.csr_array((data, crd, pos), shape=(self.shape_x, self.shape_y))
 
         arr = Csr()
@@ -257,9 +356,16 @@ def asarray(obj) -> Tensor:
 
     # TODO: support other scipy formats
     if _is_scipy_sparse_obj(obj):
-        format_class = CSRFormat
-        # This can be int32 or int64
-        index_dtype = asdtype(obj.indptr.dtype)
+        if obj.format == "csr":
+            format_class = CSRFormat
+            # This can be int32 or int64
+            index_dtype = asdtype(obj.indptr.dtype)
+        elif obj.format == "coo":
+            format_class = COOFormat
+            # This can be int32 or int64
+            index_dtype = asdtype(obj.coords[0].dtype)
+        else:
+            raise Exception(f"{obj.format} SciPy format not supported.")
     elif _is_numpy_obj(obj):
         format_class = DenseFormat
         index_dtype = Index

sparse/mlir_backend/_dtypes.py

@@ -50,6 +50,10 @@ class DType(MlirType):
     np_dtype: np.dtype
     bit_width: int
 
+    @classmethod
+    def to_ctype(cls):
+        return np.ctypeslib.as_ctypes_type(cls.np_dtype)
+
 
 class FloatingDType(DType): ...
 

sparse/mlir_backend/_memref.py

@@ -78,16 +78,21 @@ def test_constructors(rng, dtype):
     sampler = generate_sampler(dtype, rng)
     a = sps.random_array(SHAPE, density=DENSITY, format="csr", dtype=dtype, random_state=rng, data_sampler=sampler)
     c = np.arange(50, dtype=dtype).reshape((10, 5))
+    d = sps.random_array(SHAPE, density=DENSITY, format="coo", dtype=dtype, random_state=rng, data_sampler=sampler)
 
     a_tensor = sparse.asarray(a)
     c_tensor = sparse.asarray(c)
+    d_tensor = sparse.asarray(d)
 
     a_retured = a_tensor.to_scipy_sparse()
     assert_csr_equal(a, a_retured)
 
     c_returned = c_tensor.to_scipy_sparse()
     np.testing.assert_equal(c, c_returned)
 
+    d_returned = d_tensor.to_scipy_sparse()
+    np.testing.assert_equal(d.todense(), d_returned.todense())
+
 
 @parametrize_dtypes
 def test_add(rng, dtype):
@@ -115,3 +120,10 @@ def test_add(rng, dtype):
     expected = a + c
     assert isinstance(actual, np.ndarray)
     np.testing.assert_array_equal(actual, expected)
+
+    # TODO: Blocked by https://github.com/llvm/llvm-project/issues/107477
+    # d = sps.random_array(SHAPE, density=DENSITY, format="coo", dtype=dtype, random_state=rng)
+    # d_tensor = sparse.asarray(d)
+    # actual = sparse.add(b_tensor, d_tensor).to_scipy_sparse()
+    # expected = b + d
+    # np.testing.assert_array_equal(actual.todense(), expected.todense())