ENH: Update MLIR backend to LLVM 20.dev

Author: mtsokol

sparse/mlir_backend/_constructors.py

@@ -108,11 +108,11 @@ def get_tensor_definition(cls, shape: tuple[int, ...]) -> ir.RankedTensorType:
 
 
 @fn_cache
-def get_coo_class(values_dtype: type[DType], index_dtype: type[DType]) -> type[ctypes.Structure]:
+def get_coo_class(values_dtype: type[DType], index_dtype: type[DType], *, rank: int = 2) -> type[ctypes.Structure]:
     class Coo(ctypes.Structure):
         _fields_ = [
             ("pos", get_nd_memref_descr(1, index_dtype)),
-            ("coords", get_nd_memref_descr(2, index_dtype)),
+            *[(f"coords_{i}", get_nd_memref_descr(1, index_dtype)) for i in range(rank)],
             ("data", get_nd_memref_descr(1, values_dtype)),
         ]
         dtype = values_dtype
@@ -124,42 +124,46 @@ def from_sps(cls, arr: sps.coo_array | Iterable[np.ndarray]) -> "Coo":
                 if not arr.has_canonical_format:
                     raise Exception("COO must have canonical format")
                 np_pos = np.array([0, arr.size], dtype=index_dtype.np_dtype)
-                np_coords = np.stack(arr.coords, axis=1, dtype=index_dtype.np_dtype)
+                np_coords = [np.array(coord, dtype=index_dtype.np_dtype) for coord in arr.coords]
                 np_data = arr.data
             else:
                 if len(arr) != 3:
                     raise Exception("COO must be comprised of three arrays")
                 np_pos, np_coords, np_data = arr
 
             pos = numpy_to_ranked_memref(np_pos)
-            coords = numpy_to_ranked_memref(np_coords)
+            coords = [numpy_to_ranked_memref(coord) for coord in np_coords]
             data = numpy_to_ranked_memref(np_data)
-            coo_instance = cls(pos=pos, coords=coords, data=data)
+            coo_instance = cls(pos, *(coords + [data]))
             _take_owneship(coo_instance, np_pos)
-            _take_owneship(coo_instance, np_coords)
+            for coord in np_coords:
+                _take_owneship(coo_instance, coord)
             _take_owneship(coo_instance, np_data)
 
             return coo_instance
 
         def to_sps(self, shape: tuple[int, ...]) -> sps.coo_array | list[np.ndarray]:
             pos = ranked_memref_to_numpy(self.pos)
-            coords = ranked_memref_to_numpy(self.coords)[pos[0] : pos[1]]
+            coords = [ranked_memref_to_numpy(coord) for coord in self.get_coord_list()]
             data = ranked_memref_to_numpy(self.data)
             return (
-                sps.coo_array((data, coords.T), shape=shape)
+                sps.coo_array((data, np.stack(coords, axis=0, dtype=index_dtype.np_dtype)), shape=shape)
                 if len(shape) == 2
                 else PackedArgumentTuple((pos, coords, data))
             )
 
         def to_module_arg(self) -> list:
             return [
                 ctypes.pointer(ctypes.pointer(self.pos)),
-                ctypes.pointer(ctypes.pointer(self.coords)),
+                *[ctypes.pointer(ctypes.pointer(coord)) for coord in self.get_coord_list()],
                 ctypes.pointer(ctypes.pointer(self.data)),
             ]
 
         def get__fields_(self) -> list:
-            return [self.pos, self.coords, self.data]
+            return [self.pos, *self.get_coord_list(), self.data]
+
+        def get_coord_list(self) -> list:
+            return [getattr(self, f"coords_{i}") for i in range(rank)]
 
         @classmethod
         @fn_cache
@@ -173,10 +177,14 @@ def get_tensor_definition(cls, shape: tuple[int, ...]) -> ir.RankedTensorType:
                 )
                 mid_singleton_lvls = [
                     sparse_tensor.EncodingAttr.build_level_type(
-                        sparse_tensor.LevelFormat.singleton, [sparse_tensor.LevelProperty.non_unique]
+                        sparse_tensor.LevelFormat.singleton,
+                        [sparse_tensor.LevelProperty.non_unique, sparse_tensor.LevelProperty.soa],
                     )
                 ] * (len(shape) - 2)
-                levels = (compressed_lvl, *mid_singleton_lvls, sparse_tensor.LevelFormat.singleton)
+                last_singleton_lvl = sparse_tensor.EncodingAttr.build_level_type(
+                    sparse_tensor.LevelFormat.singleton, [sparse_tensor.LevelProperty.soa]
+                )
+                levels = (compressed_lvl, *mid_singleton_lvls, last_singleton_lvl)
                 ordering = ir.AffineMap.get_permutation([*range(len(shape))])
                 encoding = sparse_tensor.EncodingAttr.get(levels, ordering, ordering, index_width, index_width)
                 return ir.RankedTensorType.get(list(shape), values_dtype, encoding)
@@ -320,6 +328,7 @@ def __init__(
         self._values_dtype = dtype if dtype is not None else asdtype(obj.dtype)
 
         if _is_scipy_sparse_obj(obj):
+            self.format = obj.format
             self._owns_memory = False
 
             if obj.format in ("csr", "csc"):
@@ -335,22 +344,26 @@ def __init__(
                 raise Exception(f"{obj.format} SciPy format not supported.")
 
         elif _is_numpy_obj(obj):
+            self.format = "dense"
             self._owns_memory = False
             self._index_dtype = asdtype(np.intp)
             self._format_class = get_dense_class(self._values_dtype, self._index_dtype)
             self._obj = self._format_class.from_sps(obj)
 
         elif _is_mlir_obj(obj):
+            self.format = "custom"
             self._owns_memory = True
             self._format_class = type(obj)
             self._obj = obj
 
         elif format is not None:
+            self.format = format
             if format in ["csf", "coo"]:
                 fn_format_class = get_csf_class if format == "csf" else get_coo_class
+                kwargs = {} if format == "csf" else {"rank": len(self.shape)}
                 self._owns_memory = False
                 self._index_dtype = asdtype(np.intp)
-                self._format_class = fn_format_class(self._values_dtype, self._index_dtype)
+                self._format_class = fn_format_class(self._values_dtype, self._index_dtype, **kwargs)
                 self._obj = self._format_class.from_sps(obj)
 
             else:

sparse/mlir_backend/_ops.py

@@ -32,7 +32,7 @@ def get_add_module(
 
             @func.FuncOp.from_py_func(a_tensor_type, b_tensor_type)
             def add(a, b):
-                out = tensor.empty(out_tensor_type, [])
+                out = tensor.empty(out_tensor_type.shape, dtype, encoding=out_tensor_type.encoding)
                 generic_op = linalg.GenericOp(
                     [out_tensor_type],
                     [a, b],
@@ -108,7 +108,9 @@ def get_broadcast_to_module(
 
             @func.FuncOp.from_py_func(in_tensor_type)
             def broadcast_to(in_tensor):
-                out = tensor.empty(out_tensor_type, [])
+                out = tensor.empty(
+                    out_tensor_type.shape, out_tensor_type.element_type, encoding=out_tensor_type.encoding
+                )
                 return linalg.broadcast(in_tensor, outs=[out], dimensions=dimensions)
 
             broadcast_to.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
@@ -156,7 +158,7 @@ def _infer_format_class(rank: int, values_dtype: type[DType], index_dtype: type[
 
 def reshape(x: Tensor, /, shape: tuple[int, ...]) -> Tensor:
     x_tensor_type = x._obj.get_tensor_definition(x.shape)
-    if len(x.shape) == len(shape):
+    if len(x.shape) == len(shape) or x.format == "dense":
         out_tensor_type = x._obj.get_tensor_definition(shape)
         ret_obj = x._format_class()
     else:

sparse/mlir_backend/tests/test_simple.py

@@ -166,20 +166,18 @@ def test_add(rng, dtype):
     assert isinstance(actual, np.ndarray)
     np.testing.assert_array_equal(actual, expected)
 
-    # NOTE: Fixed in https://github.com/llvm/llvm-project/pull/108615
-    # actual = sparse.add(c_tensor, c_tensor).to_scipy_sparse()
-    # expected = c + c
-    # assert isinstance(actual, np.ndarray)
-    # np.testing.assert_array_equal(actual, expected)
+    actual = sparse.add(dense_tensor, dense_tensor).to_scipy_sparse()
+    expected = dense + dense
+    assert isinstance(actual, np.ndarray)
+    np.testing.assert_array_equal(actual, expected)
 
     actual = sparse.add(csr_2_tensor, coo_tensor).to_scipy_sparse()
     expected = csr_2 + coo
     np.testing.assert_array_equal(actual.todense(), expected.todense())
 
-    # NOTE: https://discourse.llvm.org/t/passmanager-fails-on-simple-coo-addition-example/81247
-    # actual = sparse.add(d_tensor, d_tensor).to_scipy_sparse()
-    # expected = d + d
-    # np.testing.assert_array_equal(actual.todense(), expected.todense())
+    actual = sparse.add(coo_tensor, coo_tensor).to_scipy_sparse()
+    expected = coo + coo
+    np.testing.assert_array_equal(actual.todense(), expected.todense())
 
 
 @parametrize_dtypes
@@ -203,7 +201,7 @@ def test_csf_format(dtype):
 def test_coo_3d_format(dtype):
     SHAPE = (2, 2, 4)
     pos = np.array([0, 7])
-    crd = np.array([[0, 1, 0, 0, 1, 1, 0], [1, 3, 1, 0, 0, 1, 0], [3, 1, 1, 0, 1, 1, 1]])
+    crd = [np.array([0, 1, 0, 0, 1, 1, 0]), np.array([1, 3, 1, 0, 0, 1, 0]), np.array([3, 1, 1, 0, 1, 1, 1])]
     data = np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype)
     coo = [pos, crd, data]
 
@@ -212,11 +210,10 @@ def test_coo_3d_format(dtype):
     for actual, expected in zip(result, coo, strict=False):
         np.testing.assert_array_equal(actual, expected)
 
-    # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
-    # res_tensor = sparse.add(coo_tensor, coo_tensor).to_scipy_sparse()
-    # coo_2 = [pos, crd, data * 2]
-    # for actual, expected in zip(res_tensor, coo_2, strict=False):
-    #     np.testing.assert_array_equal(actual, expected)
+    res_tensor = sparse.add(coo_tensor, coo_tensor).to_scipy_sparse()
+    coo_2 = [pos, crd, data * 2]
+    for actual, expected in zip(res_tensor, coo_2, strict=False):
+        np.testing.assert_array_equal(actual, expected)
 
 
 @parametrize_dtypes
@@ -232,9 +229,6 @@ def test_reshape(rng, dtype):
         ((80, 1), (80,)),
     ]:
         for format in ["csr", "csc", "coo"]:
-            if format == "coo":
-                # NOTE: Blocked by https://github.com/llvm/llvm-project/pull/109135
-                continue
             if format == "csc":
                 # NOTE: Blocked by https://github.com/llvm/llvm-project/issues/109641
                 continue
@@ -289,8 +283,18 @@ def test_reshape(rng, dtype):
             np.testing.assert_array_equal(actual, expected)
 
     # DENSE
-    # NOTE: dense reshape is probably broken in MLIR in 19.x branch
-    # dense = np.arange(math.prod(SHAPE), dtype=dtype).reshape(SHAPE)
+    for shape, new_shape in [
+        ((100, 50), (25, 200)),
+        ((100, 50), (10, 500, 1)),
+        ((80, 1), (8, 10)),
+        ((80, 1), (80,)),
+    ]:
+        dense = np.arange(math.prod(shape), dtype=dtype).reshape(shape)
+        dense_tensor = sparse.asarray(dense)
+        actual = sparse.reshape(dense_tensor, shape=new_shape).to_scipy_sparse()
+        expected = dense.reshape(new_shape)
+
+        np.testing.assert_array_equal(actual, expected)
 
 
 @parametrize_dtypes