Tests passing after rebase.

Author: hameerabbasi

pixi.toml

@@ -60,7 +60,9 @@ PYTHONFAULTHANDLER = "${HOME}/faulthandler.log"
 
 [feature.mlir.dependencies]
 scipy = ">=0.19"
-mlir-python-bindings = "19.*"
+
+[feature.mlir.target.osx-arm64.pypi-dependencies]
+finch-mlir = ">=0.0.2"
 
 [feature.mlir.activation.env]
 SPARSE_BACKEND = "MLIR"

sparse/mlir_backend/_array.py

@@ -41,5 +41,10 @@ def copy(self) -> "Array":
         arrs = tuple(arr.copy() for arr in self.get_constituent_arrays())
         return from_constituent_arrays(format=self.format, arrays=arrs, shape=self.shape)
 
+    def asformat(self, format: StorageFormat) -> "Array":
+        from ._ops import asformat
+
+        return asformat(self, format=format)
+
     def get_constituent_arrays(self) -> tuple[np.ndarray, ...]:
         return self._storage.get_constituent_arrays()

sparse/mlir_backend/_ops.py

@@ -11,7 +11,7 @@
 from ._common import as_shape, fn_cache
 from ._core import CWD, DEBUG, OPT_LEVEL, SHARED_LIBS, ctx, pm
 from ._dtypes import DType, IeeeComplexFloatingDType, IeeeRealFloatingDType, IntegerDType
-from .levels import _determine_format
+from .levels import StorageFormat, _determine_format
 
 
 @fn_cache
@@ -135,7 +135,31 @@ def broadcast_to(in_tensor):
     return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
 
 
-def add(x1: Array, x2: Array) -> Array:
+@fn_cache
+def get_convert_module(
+    in_tensor_type: ir.RankedTensorType,
+    out_tensor_type: ir.RankedTensorType,
+):
+    with ir.Location.unknown(ctx):
+        module = ir.Module.create()
+
+        with ir.InsertionPoint(module.body):
+
+            @func.FuncOp.from_py_func(in_tensor_type)
+            def convert(in_tensor):
+                return sparse_tensor.convert(out_tensor_type, in_tensor)
+
+            convert.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+            if DEBUG:
+                (CWD / "broadcast_to_module.mlir").write_text(str(module))
+            pm.run(module.operation)
+            if DEBUG:
+                (CWD / "broadcast_to_module_opt.mlir").write_text(str(module))
+
+    return mlir_finch.execution_engine.ExecutionEngine(module, opt_level=OPT_LEVEL, shared_libs=SHARED_LIBS)
+
+
+def add(x1: Array, x2: Array, /) -> Array:
     ret_storage_format = _determine_format(x1.format, x2.format, dtype=x1.dtype, union=True)
     ret_storage = ret_storage_format._get_ctypes_type(owns_memory=True)()
     out_tensor_type = ret_storage_format._get_mlir_type(shape=np.broadcast_shapes(x1.shape, x2.shape))
@@ -156,6 +180,24 @@ def add(x1: Array, x2: Array) -> Array:
     return Array(storage=ret_storage, shape=tuple(out_tensor_type.shape))
 
 
+def asformat(x: Array, /, format: StorageFormat) -> Array:
+    out_tensor_type = format._get_mlir_type(shape=x.shape)
+    ret_storage = format._get_ctypes_type(owns_memory=True)()
+
+    convert_module = get_convert_module(
+        x._get_mlir_type(),
+        out_tensor_type,
+    )
+
+    convert_module.invoke(
+        "convert",
+        ctypes.pointer(ctypes.pointer(ret_storage)),
+        *x._to_module_arg(),
+    )
+
+    return Array(storage=ret_storage, shape=x.shape)
+
+
 def reshape(x: Array, /, shape: tuple[int, ...]) -> Array:
     from ._conversions import _from_numpy
 

sparse/mlir_backend/tests/test_simple.py

@@ -185,7 +185,9 @@ def test_add(rng, dtype):
     expected = csr_2 + coo
     assert_csx_equal(expected, actual)
 
-    actual = sparse.to_scipy(sparse.add(coo_tensor, coo_tensor))
+    # This ends up being DCSR, not COO
+    actual_tensor = sparse.add(coo_tensor, coo_tensor)
+    actual = sparse.to_scipy(actual_tensor.asformat(coo_tensor.format))
     expected = coo + coo
     np.testing.assert_array_equal(actual.todense(), expected.todense())
 
@@ -247,7 +249,7 @@ def test_coo_3d_format(dtype):
     for actual, expected in zip(result, carrs, strict=True):
         np.testing.assert_array_equal(actual, expected)
 
-    result_arrays = sparse.add(coo_array, coo_array).get_constituent_arrays()
+    result_arrays = sparse.add(coo_array, coo_array).asformat(coo_array.format).get_constituent_arrays()
     constituent_arrays = (pos, *crd, data * 2)
     for actual, expected in zip(result_arrays, constituent_arrays, strict=False):
         np.testing.assert_array_equal(actual, expected)