ENH: MLIR backend POC (#755)

Author: mtsokol

.github/workflows/ci.yml

@@ -54,6 +54,31 @@ jobs:
         with:
           files: ./**/coverage*.xml
 
+  test_mlir:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout Repo
+      uses: actions/checkout@v4
+    - name: Setup Conda
+      uses: conda-incubator/setup-miniconda@v3
+      with:
+        python-version: '3.10'
+        channels: conda-forge
+        activate-environment: sparse-dev
+        miniforge-variant: Mambaforge
+        miniforge-version: latest
+        use-mamba: true
+    - name: Update Conda Environment
+      run: |
+        mamba env update -n sparse-dev -f ci/environment.yml
+        mamba run pip install '.[tests]'
+        mamba install conda-forge::mlir-python-bindings
+    - name: Build and run tests
+      shell: bash -l {0}
+      run: |
+        conda activate sparse-dev
+        SPARSE_BACKEND=MLIR pytest sparse/mlir_backend -v
+
   examples:
     runs-on: ubuntu-latest
     steps:
@@ -71,6 +96,7 @@ jobs:
     - name: Run examples
       run: |
         source ci/test_examples.sh
+
   notebooks:
     runs-on: ubuntu-latest
     steps:
@@ -87,6 +113,7 @@ jobs:
     - name: Run notebooks
       run: |
         source ci/test_notebooks.sh
+
   array_api_tests:
     strategy:
       matrix:
@@ -121,6 +148,7 @@ jobs:
       run: |
         cd ${GITHUB_WORKSPACE}/array-api-tests
         pytest array_api_tests -v -c pytest.ini -n 4 --max-examples=2 --derandomize --disable-deadline -o xfail_strict=True --xfails-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-xfails.txt --skips-file ${GITHUB_WORKSPACE}/ci/${{ matrix.backend }}-array-api-skips.txt
+
 on:
   # Trigger the workflow on push or pull request,
   # but only for the main branch

conftest.py

@@ -8,7 +8,7 @@ def add_doctest_modules(doctest_namespace):
     import numpy as np
 
     if sparse._BackendType.Numba != sparse._BACKEND:
-        pytest.skip()
+        pass  # TODO: pytest.skip() skips Finch and MLIR tests
 
     doctest_namespace["np"] = np
     doctest_namespace["sparse"] = sparse

pyproject.toml

@@ -46,7 +46,7 @@ tests = [
     "pre-commit",
     "scipy",
     "sparse[finch]",
-    "pytest-codspeed"
+    "pytest-codspeed",
 ]
 tox = ["sparse[tests]", "tox"]
 notebooks = ["sparse[tests]", "nbmake", "matplotlib"]

sparse/__init__.py

@@ -10,6 +10,7 @@
 class _BackendType(Enum):
     Numba = "Numba"
     Finch = "Finch"
+    MLIR = "MLIR"
 
 
 _ENV_VAR_NAME = "SPARSE_BACKEND"
@@ -40,6 +41,9 @@ class SparseFutureWarning(FutureWarning):
 if _BackendType.Finch == _BACKEND:
     from sparse.finch_backend import *  # noqa: F403
     from sparse.finch_backend import __all__
+elif _BackendType.MLIR == _BACKEND:
+    from sparse.mlir_backend import *  # noqa: F403
+    from sparse.mlir_backend import __all__
 else:
     from sparse.numba_backend import *  # noqa: F403
     from sparse.numba_backend import (  # noqa: F401

sparse/mlir_backend/__init__.py

@@ -0,0 +1,20 @@
+try:
+    import mlir  # noqa: F401
+except ModuleNotFoundError as e:
+    raise ImportError(
+        "MLIR Python bindings not installed. Run "
+        "`conda install conda-forge::mlir-python-bindings` "
+        "to enable MLIR backend."
+    ) from e
+
+from ._constructors import (
+    asarray,
+)
+from ._ops import (
+    add,
+)
+
+__all__ = [
+    "add",
+    "asarray",
+]

sparse/mlir_backend/_constructors.py

@@ -0,0 +1,257 @@
+import ctypes
+import ctypes.util
+
+import mlir.execution_engine
+import mlir.passmanager
+from mlir import ir
+from mlir.dialects import arith, bufferization, func, sparse_tensor, tensor
+
+import numpy as np
+import scipy.sparse as sps
+
+from ._core import DEBUG, MLIR_C_RUNNER_UTILS, SCRIPT_PATH, ctx
+from ._dtypes import DType, Float64, Index
+from ._memref import MemrefF64_1D, MemrefIdx_1D
+
+
+class Tensor:
+    def __init__(self, obj, module, tensor_type, disassemble_fn, values_dtype, index_dtype):
+        self.obj = obj
+        self.module = module
+        self.tensor_type = tensor_type
+        self.disassemble_fn = disassemble_fn
+        self.values_dtype = values_dtype
+        self.index_dtype = index_dtype
+
+    def __del__(self):
+        self.module.invoke("free_tensor", ctypes.pointer(self.obj))
+
+    def to_scipy_sparse(self):
+        """
+        Returns scipy.sparse or ndarray
+        """
+        return self.disassemble_fn(self.module, self.obj)
+
+
+class DenseFormat:
+    modules = {}
+
+    def get_module(shape: tuple[int], values_dtype: DType, index_dtype: DType):
+        with ir.Location.unknown(ctx):
+            module = ir.Module.create()
+            values_dtype = values_dtype.get()
+            index_dtype = index_dtype.get()
+            index_width = getattr(index_dtype, "width", 0)
+            levels = (sparse_tensor.LevelType.dense, sparse_tensor.LevelType.dense)
+            ordering = ir.AffineMap.get_permutation([0, 1])
+            encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, index_width, index_width)
+            dense_shaped = ir.RankedTensorType.get(list(shape), values_dtype, encoding)
+            tensor_1d = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], values_dtype)
+
+            with ir.InsertionPoint(module.body):
+
+                @func.FuncOp.from_py_func(tensor_1d)
+                def assemble(data):
+                    return sparse_tensor.assemble(dense_shaped, data, [])
+
+                @func.FuncOp.from_py_func(dense_shaped)
+                def disassemble(tensor_shaped):
+                    data = tensor.EmptyOp([arith.constant(ir.IndexType.get(), 0)], values_dtype)
+                    data, data_len = sparse_tensor.disassemble(
+                        tensor_1d,
+                        [],
+                        index_dtype,
+                        [],
+                        tensor_shaped,
+                        data,
+                        [],
+                    )
+                    shape_x = arith.constant(index_dtype, shape[0])
+                    shape_y = arith.constant(index_dtype, shape[1])
+                    return data, data_len, shape_x, shape_y
+
+                @func.FuncOp.from_py_func(dense_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:
+                (SCRIPT_PATH / "dense_module.mlir").write_text(str(module))
+            pm = mlir.passmanager.PassManager.parse("builtin.module(sparsifier{create-sparse-deallocs=1})")
+            pm.run(module.operation)
+            if DEBUG:
+                (SCRIPT_PATH / "dense_module_opt.mlir").write_text(str(module))
+
+        module = mlir.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
+        return (module, dense_shaped)
+
+    @classmethod
+    def assemble(cls, module, arr: np.ndarray) -> ctypes.c_void_p:
+        data = MemrefF64_1D.from_numpy(arr.flatten())
+        out = ctypes.c_void_p()
+        module.invoke(
+            "assemble",
+            ctypes.pointer(ctypes.pointer(data)),
+            ctypes.pointer(out),
+        )
+        return out
+
+    @classmethod
+    def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p) -> np.ndarray:
+        class Dense(ctypes.Structure):
+            _fields_ = [
+                ("data", MemrefF64_1D),
+                ("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]
+                return data.copy().reshape((self.shape_x, self.shape_y))
+
+        arr = Dense()
+        module.invoke(
+            "disassemble",
+            ctypes.pointer(ctypes.pointer(arr)),
+            ctypes.pointer(ptr),
+        )
+        return arr.to_np()
+
+
+class COOFormat:
+    modules = {}
+    # TODO: implement
+
+
+class CSRFormat:
+    modules = {}
+
+    def get_module(shape: tuple[int], values_dtype: DType, index_dtype: DType):
+        with ir.Location.unknown(ctx):
+            module = ir.Module.create()
+            values_dtype = values_dtype.get()
+            index_dtype = index_dtype.get()
+            index_width = getattr(index_dtype, "width", 0)
+            levels = (sparse_tensor.LevelType.dense, sparse_tensor.LevelType.compressed)
+            ordering = ir.AffineMap.get_permutation([0, 1])
+            encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, index_width, index_width)
+            csr_shaped = ir.RankedTensorType.get(list(shape), values_dtype, encoding)
+
+            tensor_1d_index = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], 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_1d_index, tensor_1d_values)
+                def assemble(pos, crd, data):
+                    return sparse_tensor.assemble(csr_shaped, data, (pos, crd))
+
+                @func.FuncOp.from_py_func(csr_shaped)
+                def disassemble(tensor_shaped):
+                    pos = tensor.EmptyOp([arith.constant(ir.IndexType.get(), 0)], index_dtype)
+                    crd = tensor.EmptyOp([arith.constant(ir.IndexType.get(), 0)], index_dtype)
+                    data = tensor.EmptyOp([arith.constant(ir.IndexType.get(), 0)], values_dtype)
+                    data, pos, crd, data_len, pos_len, crd_len = sparse_tensor.disassemble(
+                        tensor_1d_values,
+                        (tensor_1d_index, tensor_1d_index),
+                        index_dtype,
+                        (index_dtype, index_dtype),
+                        tensor_shaped,
+                        data,
+                        (pos, crd),
+                    )
+                    shape_x = arith.constant(index_dtype, shape[0])
+                    shape_y = arith.constant(index_dtype, shape[1])
+                    return data, pos, crd, data_len, pos_len, crd_len, shape_x, shape_y
+
+                @func.FuncOp.from_py_func(csr_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:
+                (SCRIPT_PATH / "scr_module.mlir").write_text(str(module))
+            pm = mlir.passmanager.PassManager.parse("builtin.module(sparsifier{create-sparse-deallocs=1})")
+            pm.run(module.operation)
+            if DEBUG:
+                (SCRIPT_PATH / "csr_module_opt.mlir").write_text(str(module))
+
+        module = mlir.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
+        return (module, csr_shaped)
+
+    @classmethod
+    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(MemrefIdx_1D.from_numpy(arr.indptr))),
+            ctypes.pointer(ctypes.pointer(MemrefIdx_1D.from_numpy(arr.indices))),
+            ctypes.pointer(ctypes.pointer(MemrefF64_1D.from_numpy(arr.data))),
+            ctypes.pointer(out),
+        )
+        return out
+
+    @classmethod
+    def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p) -> sps.csr_array:
+        class Csr(ctypes.Structure):
+            _fields_ = [
+                ("data", MemrefF64_1D),
+                ("pos", MemrefIdx_1D),
+                ("crd", MemrefIdx_1D),
+                ("data_len", np.ctypeslib.c_intp),
+                ("pos_len", np.ctypeslib.c_intp),
+                ("crd_len", np.ctypeslib.c_intp),
+                ("shape_x", np.ctypeslib.c_intp),
+                ("shape_y", np.ctypeslib.c_intp),
+            ]
+
+            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]
+                return sps.csr_array((data.copy(), crd.copy(), pos.copy()), shape=(self.shape_x, self.shape_y))
+
+        arr = Csr()
+        module.invoke(
+            "disassemble",
+            ctypes.pointer(ctypes.pointer(arr)),
+            ctypes.pointer(ptr),
+        )
+        return arr.to_sps()
+
+
+def _is_scipy_sparse_obj(x) -> bool:
+    return hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse")
+
+
+def _is_numpy_obj(x) -> bool:
+    return isinstance(x, np.ndarray)
+
+
+def asarray(obj) -> Tensor:
+    # TODO: discover obj's dtype
+    values_dtype = Float64
+    index_dtype = Index
+
+    # TODO: support other scipy formats
+    if _is_scipy_sparse_obj(obj):
+        format_class = CSRFormat
+    elif _is_numpy_obj(obj):
+        format_class = DenseFormat
+    else:
+        raise Exception(f"{type(obj)} not supported.")
+
+    # TODO: support proper caching
+    if hash(obj.shape) in format_class.modules:
+        module, tensor_type = format_class.modules[hash(obj.shape)]
+    else:
+        module, tensor_type = format_class.get_module(obj.shape, values_dtype, index_dtype)
+        format_class.modules[hash(obj.shape)] = module, tensor_type
+
+    assembled_obj = format_class.assemble(module, obj)
+    return Tensor(assembled_obj, module, tensor_type, format_class.disassemble, values_dtype, index_dtype)

sparse/mlir_backend/_core.py

@@ -0,0 +1,16 @@
+import ctypes
+import os
+import pathlib
+
+from mlir.ir import Context
+
+DEBUG = bool(int(os.environ.get("DEBUG", "0")))
+SCRIPT_PATH = pathlib.Path(__file__).parent
+
+MLIR_C_RUNNER_UTILS = ctypes.util.find_library("mlir_c_runner_utils")
+libc = ctypes.CDLL(ctypes.util.find_library("c"))
+libc.free.argtypes = [ctypes.c_void_p]
+libc.free.restype = None
+
+# TODO: remove global state
+ctx = Context()