ENH: Implement broadcast_to function

Author: mtsokol

sparse/mlir_backend/__init__.py

@@ -15,11 +15,13 @@
 )
 from ._ops import (
     add,
+    broadcast_to,
     reshape,
 )
 
 __all__ = [
     "add",
+    "broadcast_to",
     "asarray",
     "asdtype",
     "reshape",

sparse/mlir_backend/_constructors.py

@@ -310,20 +310,20 @@ def __init__(
 
             if obj.format in ("csr", "csc"):
                 order = "r" if obj.format == "csr" else "c"
-                index_dtype = asdtype(obj.indptr.dtype)
-                self._format_class = get_csx_class(self._values_dtype, index_dtype, order)
+                self._index_dtype = asdtype(obj.indptr.dtype)
+                self._format_class = get_csx_class(self._values_dtype, self._index_dtype, order)
                 self._obj = self._format_class.from_sps(obj)
             elif obj.format == "coo":
-                index_dtype = asdtype(obj.coords[0].dtype)
-                self._format_class = get_coo_class(self._values_dtype, index_dtype)
+                self._index_dtype = asdtype(obj.coords[0].dtype)
+                self._format_class = get_coo_class(self._values_dtype, self._index_dtype)
                 self._obj = self._format_class.from_sps(obj)
             else:
                 raise Exception(f"{obj.format} SciPy format not supported.")
 
         elif _is_numpy_obj(obj):
             self._owns_memory = False
-            index_dtype = asdtype(np.intp)
-            self._format_class = get_dense_class(self._values_dtype, index_dtype)
+            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):
@@ -334,8 +334,8 @@ def __init__(
         elif format is not None:
             if format == "csf":
                 self._owns_memory = False
-                index_dtype = asdtype(np.intp)
-                self._format_class = get_csf_class(self._values_dtype, index_dtype)
+                self._index_dtype = asdtype(np.intp)
+                self._format_class = get_csf_class(self._values_dtype, self._index_dtype)
                 self._obj = self._format_class.from_sps(obj)
             else:
                 raise Exception(f"Format {format} not supported.")

sparse/mlir_backend/_ops.py

@@ -95,6 +95,32 @@ def reshape(a, shape):
     return mlir.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
 
 
+@fn_cache
+def get_broadcast_to_module(
+    in_tensor_type: ir.RankedTensorType,
+    out_tensor_type: ir.RankedTensorType,
+    dimensions: tuple[int, ...],
+) -> ir.Module:
+    with ir.Location.unknown(ctx):
+        module = ir.Module.create()
+
+        with ir.InsertionPoint(module.body):
+
+            @func.FuncOp.from_py_func(in_tensor_type)
+            def broadcast_to(in_tensor):
+                out = tensor.empty(out_tensor_type, [])
+                return linalg.broadcast(in_tensor, outs=[out], dimensions=dimensions)
+
+            broadcast_to.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.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
+
+
 def add(x1: Tensor, x2: Tensor) -> Tensor:
     ret_obj = x1._format_class()
     out_tensor_type = x1._obj.get_tensor_definition(x1.shape)
@@ -135,3 +161,32 @@ def reshape(x: Tensor, /, shape: tuple[int, ...]) -> Tensor:
     )
 
     return Tensor(ret_obj, shape=out_tensor_type.shape)
+
+
+def _infer_format_class(rank: int, values_dtype: type[DType], index_dtype: type[DType]) -> type[ctypes.Structure]:
+    from ._constructors import get_csf_class, get_csx_class, get_dense_class
+
+    if rank == 1:
+        return get_dense_class(values_dtype, index_dtype)
+    if rank == 2:
+        return get_csx_class(values_dtype, index_dtype, order="r")
+    if rank == 3:
+        return get_csf_class(values_dtype, index_dtype)
+    raise Exception(f"Rank not supported to infer format: {rank}")
+
+
+def broadcast_to(x: Tensor, /, shape: tuple[int, ...], dimensions: list[int]) -> Tensor:
+    x_tensor_type = x._obj.get_tensor_definition(x.shape)
+    format_class = _infer_format_class(len(shape), x._values_dtype, x._index_dtype)
+    out_tensor_type = format_class.get_tensor_definition(shape)
+    ret_obj = format_class()
+
+    broadcast_to_module = get_broadcast_to_module(x_tensor_type, out_tensor_type, tuple(dimensions))
+
+    broadcast_to_module.invoke(
+        "broadcast_to",
+        ctypes.pointer(ctypes.pointer(ret_obj)),
+        *x._obj.to_module_arg(),
+    )
+
+    return Tensor(ret_obj, shape=shape)

sparse/mlir_backend/tests/test_simple.py

@@ -217,8 +217,7 @@ def test_reshape(rng, dtype):
             arr = sps.random_array(
                 shape, density=DENSITY, format=format, dtype=dtype, random_state=rng, data_sampler=sampler
             )
-            if format == "coo":
-                arr.sum_duplicates()
+            arr.sum_duplicates()
 
             tensor = sparse.asarray(arr)
 
@@ -262,5 +261,55 @@ def test_reshape(rng, dtype):
             np.testing.assert_array_equal(actual, expected)
 
     # DENSE
-    # NOTE: dense reshape is probably broken in MLIR
+    # NOTE: dense reshape is probably broken in MLIR in 19.x branch
     # dense = np.arange(math.prod(SHAPE), dtype=dtype).reshape(SHAPE)
+
+
+@parametrize_dtypes
+def test_broadcast_to(dtype):
+    # CSR, CSC, COO
+    for shape, new_shape, dimensions, input_arr, expected_arrs in [
+        (
+            (3, 4),
+            (2, 3, 4),
+            [0],
+            np.array([[0, 1, 0, 3], [0, 0, 4, 5], [6, 7, 0, 0]]),
+            [
+                np.array([0, 3, 6]),
+                np.array([0, 1, 2, 0, 1, 2]),
+                np.array([0, 2, 4, 6, 8, 10, 12]),
+                np.array([1, 3, 2, 3, 0, 1, 1, 3, 2, 3, 0, 1]),
+                np.array([1.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.0, 3.0, 4.0, 5.0, 6.0, 7.0]),
+            ],
+        ),
+        (
+            (4, 2),
+            (4, 2, 2),
+            [1],
+            np.array([[0, 1], [0, 0], [2, 3], [4, 0]]),
+            [
+                np.array([0, 2, 2, 4, 6]),
+                np.array([0, 1, 0, 1, 0, 1]),
+                np.array([0, 1, 2, 4, 6, 7, 8]),
+                np.array([1, 1, 0, 1, 0, 1, 0, 0]),
+                np.array([1.0, 1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 4.0]),
+            ],
+        ),
+    ]:
+        for fn_format in [sps.csr_array, sps.csc_array, sps.coo_array]:
+            arr = fn_format(input_arr, shape=shape, dtype=dtype)
+            arr.sum_duplicates()
+            tensor = sparse.asarray(arr)
+            result = sparse.broadcast_to(tensor, new_shape, dimensions=dimensions).to_scipy_sparse()
+
+            for actual, expected in zip(result, expected_arrs, strict=False):
+                np.testing.assert_allclose(actual, expected)
+
+    # DENSE
+    np_arr = np.array([0, 0, 2, 3, 0, 1])
+    arr = np.asarray(np_arr, dtype=dtype)
+    tensor = sparse.asarray(arr)
+    result = sparse.broadcast_to(tensor, (3, 6), dimensions=[0]).to_scipy_sparse()
+
+    assert result.format == "csr"
+    np.testing.assert_allclose(result.todense(), np.repeat(np_arr[np.newaxis], 3, axis=0))