Add format generation from levels.

Author: hameerabbasi

pixi.toml

@@ -28,7 +28,7 @@ mkdocs-jupyter = "*"
 [feature.tests.tasks]
 test = "pytest --pyargs sparse -n auto"
 test-mlir = { cmd = "pytest --pyargs sparse/mlir_backend -vvv", env = { SPARSE_BACKEND = "MLIR" } }
-test-finch = { cmd = "pytest --pyargs sparse/tests  -n auto", env = { SPARSE_BACKEND = "Finch" }, depends-on = ["precompile"] }
+test-finch = { cmd = "pytest --pyargs sparse/tests -n auto -vvv", env = { SPARSE_BACKEND = "Finch", PYTHONFAULTHANDLER = "${HOME}/faulthandler.log" }, depends-on = ["precompile"] }
 
 [feature.tests.dependencies]
 pytest = ">=3.5"

sparse/mlir_backend/_common.py

@@ -4,8 +4,44 @@
 import weakref
 from dataclasses import dataclass
 
+import mlir.runtime as rt
 from mlir import ir
 
+import numpy as np
+
+from ._core import libc
+from ._dtypes import DType, asdtype
+
+
+def fn_cache(f, maxsize: int | None = None):
+    return functools.wraps(f)(functools.lru_cache(maxsize=maxsize)(f))
+
+
+@fn_cache
+def get_nd_memref_descr(rank: int, dtype: type[DType]) -> ctypes.Structure:
+    return rt.make_nd_memref_descriptor(rank, dtype.to_ctype())
+
+
+def numpy_to_ranked_memref(arr: np.ndarray) -> ctypes.Structure:
+    memref = rt.get_ranked_memref_descriptor(arr)
+    memref_descr = get_nd_memref_descr(arr.ndim, asdtype(arr.dtype))
+    # Required due to ctypes type checks
+    return memref_descr(
+        allocated=memref.allocated,
+        aligned=memref.aligned,
+        offset=memref.offset,
+        shape=memref.shape,
+        strides=memref.strides,
+    )
+
+
+def ranked_memref_to_numpy(ref: ctypes.Structure) -> np.ndarray:
+    return rt.ranked_memref_to_numpy([ref])
+
+
+def free_memref(obj: ctypes.Structure) -> None:
+    libc.free(ctypes.cast(obj.allocated, ctypes.c_void_p))
+
 
 class MlirType(abc.ABC):
     @classmethod
@@ -27,10 +63,6 @@ def __len__(self):
         return len(self.contents)
 
 
-def fn_cache(f, maxsize: int | None = None):
-    return functools.wraps(f)(functools.lru_cache(maxsize=maxsize)(f))
-
-
 def _hold_self_ref_in_ret(fn):
     @functools.wraps(fn)
     def wrapped(self, *a, **kw):

sparse/mlir_backend/_constructors.py

@@ -2,48 +2,30 @@
 from collections.abc import Iterable
 from typing import Any
 
-import mlir.runtime as rt
 from mlir import ir
 from mlir.dialects import sparse_tensor
 
 import numpy as np
 import scipy.sparse as sps
 
-from ._common import PackedArgumentTuple, _hold_self_ref_in_ret, _take_owneship, fn_cache
-from ._core import ctx, libc
+from ._common import (
+    PackedArgumentTuple,
+    _hold_self_ref_in_ret,
+    _take_owneship,
+    fn_cache,
+    free_memref,
+    get_nd_memref_descr,
+    numpy_to_ranked_memref,
+    ranked_memref_to_numpy,
+)
+from ._core import ctx
 from ._dtypes import DType, asdtype
 
 ###########
 # Memrefs #
 ###########
 
 
-@fn_cache
-def get_nd_memref_descr(rank: int, dtype: type[DType]) -> type:
-    return rt.make_nd_memref_descriptor(rank, dtype.to_ctype())
-
-
-def numpy_to_ranked_memref(arr: np.ndarray) -> ctypes.Structure:
-    memref = rt.get_ranked_memref_descriptor(arr)
-    memref_descr = get_nd_memref_descr(arr.ndim, asdtype(arr.dtype))
-    # Required due to ctypes type checks
-    return memref_descr(
-        allocated=memref.allocated,
-        aligned=memref.aligned,
-        offset=memref.offset,
-        shape=memref.shape,
-        strides=memref.strides,
-    )
-
-
-def ranked_memref_to_numpy(ref: ctypes.Structure) -> np.ndarray:
-    return rt.ranked_memref_to_numpy([ref])
-
-
-def free_memref(obj: ctypes.Structure) -> None:
-    libc.free(ctypes.cast(obj.allocated, ctypes.c_void_p))
-
-
 ###########
 # Formats #
 ###########

sparse/mlir_backend/_levels.py

@@ -0,0 +1,165 @@
+import ctypes
+import dataclasses
+import enum
+import itertools
+import re
+import typing
+
+import mlir.runtime as rt
+from mlir import ir
+from mlir.dialects import sparse_tensor
+
+import numpy as np
+
+from ._common import (
+    PackedArgumentTuple,
+    _take_owneship,
+    fn_cache,
+    numpy_to_ranked_memref,
+    ranked_memref_to_numpy,
+)
+from ._dtypes import DType, asdtype
+
+_CAMEL_TO_SNAKE = [re.compile("(.)([A-Z][a-z]+)"), re.compile("([a-z0-9])([A-Z])")]
+
+
+def _camel_to_snake(name: str) -> str:
+    for exp in _CAMEL_TO_SNAKE:
+        name = exp.sub(r"\1_\2", name)
+
+    return name.lower()
+
+
+@fn_cache
+def get_nd_memref_descr(rank: int, dtype: type[DType]) -> type:
+    return rt.make_nd_memref_descriptor(rank, dtype.to_ctype())
+
+
+class LevelProperties(enum.Flag):
+    NonOrdered = enum.auto()
+    NonUnique = enum.auto()
+
+    def build(self) -> list[sparse_tensor.LevelProperty]:
+        return [getattr(sparse_tensor.LevelProperty, _camel_to_snake(p.name)) for p in type(self) if p in self]
+
+
+class LevelFormat(enum.Enum):
+    Dense = "dense"
+    Compressed = "compressed"
+    Singleton = "singleton"
+
+    def build(self) -> sparse_tensor.LevelFormat:
+        return getattr(sparse_tensor.LevelFormat, self.value)
+
+
+@dataclasses.dataclass(eq=True, frozen=True, kw_only=True)
+class Level:
+    format: LevelFormat
+    properties: LevelProperties = LevelProperties(0)
+
+    def build(self):
+        sparse_tensor.EncodingAttr.build_level_type(self.format.build(), self.properties.build())
+
+
+@dataclasses.dataclass(kw_only=True)
+class StorageFormat:
+    levels: tuple[Level, ...]
+    order: typing.Literal["C", "F"] | tuple[int, ...]
+    pos_width: int
+    crd_width: int
+    dtype: type[DType]
+
+    @property
+    def storage_rank(self) -> int:
+        return len(self.levels)
+
+    @property
+    def rank(self) -> int:
+        return self.storage_rank
+
+    def __post_init__(self):
+        rank = self.storage_rank
+        self.dtype = asdtype(self.dtype)
+        if self.order == "C":
+            self.order = tuple(range(rank))
+            return
+
+        if self.order == "F":
+            self.order = tuple(reversed(range(rank)))
+            return
+
+        if sorted(self.order) != list(range(rank)):
+            raise ValueError(f"`sorted(self.order) != list(range(rank))`, {self.order=}, {rank=}.")
+
+        self.order = tuple(self.order)
+
+    @fn_cache
+    def get_mlir_type(self, *, shape: tuple[int, ...]) -> ir.RankedTensorType:
+        if len(shape) != self.rank:
+            raise ValueError(f"`len(shape) != self.rank`, {shape=}, {self.rank=}")
+        mlir_levels = [level.build() for level in self.levels]
+        mlir_order = list(self.order)
+        mlir_reverse_order = [0] * self.rank
+        for i, r in enumerate(mlir_order):
+            mlir_reverse_order[r] = i
+
+        dtype = self.dtype.get_mlir_type()
+        encoding = sparse_tensor.EncodingAttr.get(
+            mlir_levels, mlir_order, mlir_reverse_order, self.pos_width, self.crd_width
+        )
+        return ir.RankedTensorType.get(list(shape), dtype, encoding)
+
+    @fn_cache
+    def get_ctypes_type(self):
+        ptr_dtype = asdtype(getattr(np, f"uint{self.pos_width}"))
+        idx_dtype = asdtype(getattr(np, f"uint{self.crd_width}"))
+
+        def get_fields():
+            fields = []
+            compressed_counter = 0
+            for level, next_level in itertools.zip_longest(self.levels, self.levels[1:]):
+                if LevelFormat.Compressed == level.format:
+                    compressed_counter += 1
+                    fields.append((f"pointers_to_{compressed_counter}", get_nd_memref_descr(1, ptr_dtype)))
+                    if next_level is not None and LevelFormat.Singleton == next_level.format:
+                        fields.append((f"indices_{compressed_counter}", get_nd_memref_descr(2, idx_dtype)))
+                    else:
+                        fields.append((f"indices_{compressed_counter}", get_nd_memref_descr(1, idx_dtype)))
+
+            fields.append(("values", get_nd_memref_descr(1, self.dtype.np_dtype)))
+            return fields
+
+        storage_format = self
+
+        class Format(ctypes.Structure):
+            _fields_ = get_fields()
+
+            def get_mlir_type(self, *, shape: tuple[int, ...]):
+                return self.get_storage_format().get_mlir_type(shape=shape)
+
+            def to_module_arg(self) -> list:
+                return [ctypes.pointer(ctypes.pointer(f) for f in self.get__fields_())]
+
+            def get__fields_(self) -> list:
+                return [getattr(self, field[0]) for field in self._fields_]
+
+            def to_constituent_arrays(self) -> PackedArgumentTuple:
+                return PackedArgumentTuple(tuple(ranked_memref_to_numpy(field) for field in self.get__fields_()))
+
+            def get_storage_format(self) -> StorageFormat:
+                return storage_format
+
+            @classmethod
+            def from_constituent_arrays(cls, arrs: list[np.ndarray]) -> "Format":
+                inst = cls(*[numpy_to_ranked_memref(arr) for arr in arrs])
+                for arr in arrs:
+                    _take_owneship(inst, arr)
+                return inst
+
+        return Format
+
+    def __hash__(self):
+        return hash(id(self))
+
+    def __eq__(self, value):
+        return self is value