diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d496836e..d9a40af7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,7 +18,7 @@ repos: exclude: ".ipynb" - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.6.3 + rev: v0.6.4 hooks: - id: ruff args: ["--fix"] diff --git a/sparse/mlir_backend/_constructors.py b/sparse/mlir_backend/_constructors.py index c772b246..b2c56ceb 100644 --- a/sparse/mlir_backend/_constructors.py +++ b/sparse/mlir_backend/_constructors.py @@ -1,19 +1,21 @@ +import abc import ctypes import functools +import typing import weakref 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 import scipy.sparse as sps from ._common import fn_cache -from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx +from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx, pm from ._dtypes import DType, Index, asdtype -from ._memref import make_memref_ctype, ranked_memref_from_np def _hold_self_ref_in_ret(fn): @@ -49,66 +51,133 @@ def to_scipy_sparse(self): """ Returns scipy.sparse or ndarray """ - return self.disassemble_fn(self.module, self.obj, self.values_dtype) + return self.disassemble_fn( + self.module, + self.obj, + self.tensor_type.shape, + self.values_dtype, + self.index_dtype, + ) + + +class BaseFormat(abc.ABC): + @classmethod + @abc.abstractmethod + def get_format_str(cls) -> str: + raise NotImplementedError + + @classmethod + @abc.abstractmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + raise NotImplementedError + + @classmethod + @abc.abstractmethod + def get_ordering(cls) -> ir.AffineMap: + raise NotImplementedError + + @classmethod + @abc.abstractmethod + def get_assemble_functions( + cls, + module: ir.Module, + tensor_shaped: ir.RankedTensorType, + values_dtype: type[DType], + index_dtype: type[DType], + ) -> tuple[typing.Callable, ...]: + raise NotImplementedError + + @classmethod + @abc.abstractmethod + def assemble(cls, module: ir.Module, arr: np.ndarray | sps.sparray) -> ctypes.c_void_p: + raise NotImplementedError + @classmethod + @abc.abstractmethod + def disassemble( + cls, + module: ir.Module, + ptr: ctypes.c_void_p, + shape: list[int], + values_dtype: type[DType], + index_dtype: type[DType], + ) -> np.ndarray | sps.sparray: + raise NotImplementedError -class DenseFormat: + @classmethod @fn_cache - def get_module(shape: tuple[int], values_dtype: DType, index_dtype: DType): + def get_module( + cls, + shape: tuple[int], + values_dtype: type[DType], + index_dtype: type[DType], + ) -> tuple[ir.Module, ir.RankedTensorType]: 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) - levels = (sparse_tensor.LevelFormat.dense, sparse_tensor.LevelFormat.dense) - ordering = ir.AffineMap.get_permutation([0, 1]) + levels = cls.get_levels() + ordering = cls.get_ordering() 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) + tensor_shaped = ir.RankedTensorType.get(list(shape), values_dtype, encoding) + + assemble, disassemble, free_tensor = cls.get_assemble_functions( + module, tensor_shaped, values_dtype, index_dtype + ) 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 / "dense_module.mlir").write_text(str(module)) - pm = mlir.passmanager.PassManager.parse("builtin.module(sparsifier{create-sparse-deallocs=1})") + (CWD / f"{cls.get_format_str()}.mlir").write_text(str(module)) pm.run(module.operation) if DEBUG: - (CWD / "dense_module_opt.mlir").write_text(str(module)) + (CWD / f"{cls.get_format_str()}_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) + return (module, tensor_shaped) + +class AbstractDenseFormat(BaseFormat): @classmethod - def assemble(cls, module, arr: np.ndarray) -> ctypes.c_void_p: - assert arr.ndim == 2 - data = ranked_memref_from_np(arr.flatten()) + def get_assemble_functions( + cls, + module: ir.Module, + tensor_shaped: ir.RankedTensorType, + values_dtype: type[DType], + index_dtype: type[DType], + ) -> tuple[typing.Callable, ...]: + 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(tensor_shaped, [], data) + + @func.FuncOp.from_py_func(tensor_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, + ) + return data, data_len + + @func.FuncOp.from_py_func(tensor_shaped) + def free_tensor(tensor_shaped): + bufferization.dealloc_tensor(tensor_shaped) + + return assemble, disassemble, free_tensor + + @classmethod + def assemble(cls, module: ir.Module, arr: np.ndarray) -> ctypes.c_void_p: + data = rt.get_ranked_memref_descriptor(arr) out = ctypes.c_void_p() module.invoke( "assemble", @@ -118,18 +187,23 @@ def assemble(cls, module, arr: np.ndarray) -> ctypes.c_void_p: return out @classmethod - def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p, dtype: type[DType]) -> np.ndarray: + def disassemble( + cls, + module: ir.Module, + ptr: ctypes.c_void_p, + shape: list[int], + values_dtype: type[DType], + index_dtype: type[DType], + ) -> np.ndarray: class Dense(ctypes.Structure): _fields_ = [ - ("data", make_memref_ctype(dtype, 1)), + ("data", rt.make_nd_memref_descriptor(1, values_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] - return data.reshape((self.shape_x, self.shape_y)) + data = rt.ranked_memref_to_numpy([self.data])[: self.data_len] + return data.reshape(shape) arr = Dense() module.invoke( @@ -140,99 +214,242 @@ def to_np(self) -> np.ndarray: return arr.to_np() -class COOFormat: - # TODO: implement - ... +class VectorFormat(AbstractDenseFormat): + @classmethod + def get_format_str(cls) -> str: + return "sparse_vector_format" + @classmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + return (sparse_tensor.LevelFormat.dense,) -class CSRFormat: - @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) - levels = (sparse_tensor.LevelFormat.dense, sparse_tensor.LevelFormat.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, (pos, crd), data) - - @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) - pos, crd, data, pos_len, crd_len, data_len = sparse_tensor.disassemble( - (tensor_1d_index, tensor_1d_index), - tensor_1d_values, - (index_dtype, index_dtype), - index_dtype, - tensor_shaped, - (pos, crd), - data, - ) - shape_x = arith.constant(index_dtype, shape[0]) - shape_y = arith.constant(index_dtype, shape[1]) - return pos, crd, data, pos_len, crd_len, data_len, shape_x, shape_y - - @func.FuncOp.from_py_func(csr_shaped) - def free_tensor(tensor_shaped): - bufferization.dealloc_tensor(tensor_shaped) + @classmethod + def get_ordering(cls) -> ir.AffineMap: + return ir.AffineMap.get_permutation([0]) - 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 / "csr_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 / "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) +class Dense2DFormat(AbstractDenseFormat): + @classmethod + def get_format_str(cls) -> str: + return "dense_2d_format" + + @classmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + return (sparse_tensor.LevelFormat.dense,) * 2 + + @classmethod + def get_ordering(cls) -> ir.AffineMap: + return ir.AffineMap.get_permutation([0, 1]) + + +class Dense3DFormat(AbstractDenseFormat): + @classmethod + def get_format_str(cls) -> str: + return "dense_3d_format" + + @classmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + return (sparse_tensor.LevelFormat.dense,) * 3 + + @classmethod + def get_ordering(cls) -> ir.AffineMap: + return ir.AffineMap.get_permutation([0, 2, 3]) + + +class COOFormat(BaseFormat): + @classmethod + def get_format_str(cls) -> str: + return "coo_format" + + @classmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + compressed_lvl = sparse_tensor.EncodingAttr.build_level_type( + sparse_tensor.LevelFormat.compressed, [sparse_tensor.LevelProperty.non_unique] + ) + return (compressed_lvl, sparse_tensor.LevelFormat.singleton) + + @classmethod + def get_ordering(cls) -> ir.AffineMap: + return ir.AffineMap.get_permutation([0, 1]) + + @classmethod + def get_assemble_functions( + cls, + module: ir.Module, + tensor_shaped: ir.RankedTensorType, + values_dtype: type[DType], + index_dtype: type[DType], + ) -> tuple[typing.Callable, ...]: + tensor_1d_index = tensor.RankedTensorType.get([ir.ShapedType.get_dynamic_size()], index_dtype) + tensor_2d_index = tensor.RankedTensorType.get( + [ir.ShapedType.get_dynamic_size(), tensor_shaped.rank], 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(tensor_shaped, (pos, index), values) + + @func.FuncOp.from_py_func(tensor_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, + ) + return pos, index, values, pos_len, index_len, values_len + + @func.FuncOp.from_py_func(tensor_shaped) + def free_tensor(tensor_shaped): + bufferization.dealloc_tensor(tensor_shaped) + + return assemble, disassemble, free_tensor + + @classmethod + def assemble(cls, module: ir.Module, arr: sps.coo_array) -> ctypes.c_void_p: + out = ctypes.c_void_p() + index_dtype = arr.coords[0].dtype + module.invoke( + "assemble", + ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np.array([0, arr.size], dtype=index_dtype)))), + ctypes.pointer( + ctypes.pointer(rt.get_ranked_memref_descriptor(np.stack(arr.coords, axis=1, dtype=index_dtype))) + ), + 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, + shape: list[int], + values_dtype: type[DType], + index_dtype: type[DType], + ) -> sps.coo_array: + class Coo(ctypes.Structure): + _fields_ = [ + ("pos", rt.make_nd_memref_descriptor(1, index_dtype.to_ctype())), + ("index", rt.make_nd_memref_descriptor(2, index_dtype.to_ctype())), + ("values", rt.make_nd_memref_descriptor(1, values_dtype.to_ctype())), + ("pos_len", np.ctypeslib.c_intp), + ("index_len", np.ctypeslib.c_intp), + ("values_len", 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=shape) + + arr = Coo() + module.invoke( + "disassemble", + ctypes.pointer(ctypes.pointer(arr)), + ctypes.pointer(ptr), + ) + return arr.to_sps() + + +class CSRFormat(BaseFormat): + @classmethod + def get_format_str(cls) -> str: + return "csr_format" + + @classmethod + def get_levels(cls) -> tuple[sparse_tensor.LevelFormat, ...]: + return (sparse_tensor.LevelFormat.dense, sparse_tensor.LevelFormat.compressed) + + @classmethod + def get_ordering(cls) -> ir.AffineMap: + return ir.AffineMap.get_permutation([0, 1]) + + @classmethod + def get_assemble_functions( + cls, + module: ir.Module, + tensor_shaped: ir.RankedTensorType, + values_dtype: type[DType], + index_dtype: type[DType], + ) -> tuple[typing.Callable, ...]: + 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(tensor_shaped, (pos, crd), data) + + @func.FuncOp.from_py_func(tensor_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) + pos, crd, data, pos_len, crd_len, data_len = sparse_tensor.disassemble( + (tensor_1d_index, tensor_1d_index), + tensor_1d_values, + (index_dtype, index_dtype), + index_dtype, + tensor_shaped, + (pos, crd), + data, + ) + return pos, crd, data, pos_len, crd_len, data_len + + @func.FuncOp.from_py_func(tensor_shaped) + def free_tensor(tensor_shaped): + bufferization.dealloc_tensor(tensor_shaped) + + return assemble, disassemble, free_tensor @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(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 @classmethod - def disassemble(cls, module: ir.Module, ptr: ctypes.c_void_p, dtype: type[DType]) -> sps.csr_array: + def disassemble( + cls, + module: ir.Module, + ptr: ctypes.c_void_p, + shape: list[int], + values_dtype: type[DType], + index_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_dtype.to_ctype())), + ("crd", rt.make_nd_memref_descriptor(1, index_dtype.to_ctype())), + ("data", rt.make_nd_memref_descriptor(1, values_dtype.to_ctype())), ("pos_len", np.ctypeslib.c_intp), ("crd_len", np.ctypeslib.c_intp), ("data_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, crd, pos), shape=(self.shape_x, self.shape_y)) + 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=shape) arr = Csr() module.invoke( @@ -257,11 +474,23 @@ 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 + index_dtype = asdtype(obj.indptr.dtype) + elif obj.format == "coo": + format_class = COOFormat + 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 + if obj.ndim == 1: + format_class = VectorFormat + elif obj.ndim == 2: + format_class = Dense2DFormat + elif obj.ndim == 3: + format_class = Dense3DFormat + else: + raise Exception(f"Rank {obj.ndim} of dense tensor not supported.") index_dtype = Index else: raise Exception(f"{type(obj)} not supported.") diff --git a/sparse/mlir_backend/_core.py b/sparse/mlir_backend/_core.py index 480914dd..3a02e66c 100644 --- a/sparse/mlir_backend/_core.py +++ b/sparse/mlir_backend/_core.py @@ -4,6 +4,7 @@ import pathlib from mlir.ir import Context +from mlir.passmanager import PassManager DEBUG = bool(int(os.environ.get("DEBUG", "0"))) CWD = pathlib.Path(".") @@ -15,3 +16,5 @@ # TODO: remove global state ctx = Context() + +pm = PassManager.parse("builtin.module(sparsifier{create-sparse-deallocs=1})", context=ctx) diff --git a/sparse/mlir_backend/_dtypes.py b/sparse/mlir_backend/_dtypes.py index 5ee5a461..2ab41401 100644 --- a/sparse/mlir_backend/_dtypes.py +++ b/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): ... diff --git a/sparse/mlir_backend/_memref.py b/sparse/mlir_backend/_memref.py deleted file mode 100644 index 387f3481..00000000 --- a/sparse/mlir_backend/_memref.py +++ /dev/null @@ -1,64 +0,0 @@ -import ctypes - -import numpy as np - -from ._common import fn_cache -from ._dtypes import DType, asdtype - - -def make_memref_ctype(dtype: type[DType], rank: int) -> type[ctypes.Structure]: - dtype = np.dtype(asdtype(dtype).np_dtype) - return _make_memref_ctype(dtype, rank) - - -@fn_cache -def _make_memref_ctype(dtype: np.dtype, rank: int) -> type[ctypes.Structure]: - ctype = np.ctypeslib.as_ctypes_type(dtype) - ptr_t = ctypes.POINTER(ctype) - - class MemrefType(ctypes.Structure): - _fields_ = [ - ("allocated", ctypes.c_void_p), - ("aligned", ptr_t), - ("offset", np.ctypeslib.c_intp), - ("shape", np.ctypeslib.c_intp * rank), - ("strides", np.ctypeslib.c_intp * rank), - ] - - @classmethod - def from_numpy(cls, arr: np.ndarray) -> "MemrefType": - if not arr.dtype == dtype: - raise TypeError(f"Expected {dtype=}, found {arr.dtype=}.") - if not all(s % arr.itemsize == 0 for s in arr.strides): - raise ValueError(f"Strides not item aligned: {arr.strides=}, {arr.itemsize=}") - - ptr = ctypes.cast(arr.ctypes.data, ctypes.c_void_p) - ptr_typed = ctypes.cast(ptr, ptr_t) - return cls( - allocated=ptr, - aligned=ptr_typed, - offset=0, - shape=arr.shape, - strides=tuple(s // arr.itemsize for s in arr.strides), - ) - - def to_numpy(self) -> np.ndarray: - if ctypes.cast(self.aligned, ctypes.c_void_p).value != ctypes.cast(self.allocated, ctypes.c_void_p).value: - raise RuntimeError("Encountered different values for `aligned` and `allocated`.") - shape = tuple(self.shape) - ptr = self.aligned - ret = np.ctypeslib.as_array(ptr, shape) - strides = tuple(s * ret.itemsize for s in self.strides) - if ret.strides != strides: - raise RuntimeError(f"Expected {ret.strides=} for {shape=}, got {strides=}.") - return ret - - def __hash__(self) -> int: - return hash(id(self)) - - return MemrefType - - -def ranked_memref_from_np(arr: np.ndarray) -> ctypes.Structure: - memref_type = _make_memref_ctype(arr.dtype, arr.ndim) - return memref_type.from_numpy(arr) diff --git a/sparse/mlir_backend/_ops.py b/sparse/mlir_backend/_ops.py index 4e678489..b9c7c28d 100644 --- a/sparse/mlir_backend/_ops.py +++ b/sparse/mlir_backend/_ops.py @@ -6,11 +6,16 @@ from mlir.dialects import arith, func, linalg, sparse_tensor, tensor from ._constructors import Tensor -from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx +from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx, pm from ._dtypes import DType, FloatingDType -def get_add_module(a_tensor_type, b_tensor_type, out_tensor_type, dtype: type[DType]): +def get_add_module( + a_tensor_type: ir.RankedTensorType, + b_tensor_type: ir.RankedTensorType, + out_tensor_type: ir.RankedTensorType, + dtype: type[DType], +) -> ir.Module: with ir.Location.unknown(ctx): module = ir.Module.create() # TODO: add support for complex dialect/dtypes @@ -53,9 +58,7 @@ def add(a, b): add.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get() if DEBUG: (CWD / "add_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 / "add_module_opt.mlir").write_text(str(module)) diff --git a/sparse/mlir_backend/tests/test_simple.py b/sparse/mlir_backend/tests/test_simple.py index 3b4cf11b..09729d70 100644 --- a/sparse/mlir_backend/tests/test_simple.py +++ b/sparse/mlir_backend/tests/test_simple.py @@ -1,3 +1,4 @@ +import math import typing import sparse @@ -71,16 +72,27 @@ def sampler_real_floating(size: tuple[int, ...]): raise NotImplementedError(f"{dtype=} not yet supported.") +@parametrize_dtypes +@pytest.mark.parametrize("shape", [(100,), (10, 20), (5, 10, 20)]) +def test_dense_format(dtype, shape): + data = np.arange(math.prod(shape), dtype=dtype) + tensor = sparse.asarray(data) + actual = tensor.to_scipy_sparse() + np.testing.assert_equal(actual, data) + + @parametrize_dtypes def test_constructors(rng, dtype): SHAPE = (10, 5) DENSITY = 0.5 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)) + c = np.arange(math.prod(SHAPE), dtype=dtype).reshape(SHAPE) + 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) @@ -88,6 +100,9 @@ def test_constructors(rng, dtype): 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 +130,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())