ENH: Implement broadcast_to function (#782)

Author: mtsokol

sparse/mlir_backend/__init__.py

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

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)
@@ -152,3 +178,20 @@ def reshape(x: Tensor, /, shape: tuple[int, ...]) -> Tensor:
  )
 
  return Tensor(ret_obj, shape=out_tensor_type.shape)
+
+
+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

@@ -289,5 +289,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))