Merge pull request #776 from pydata/reshape-func

ENH: Implement `reshape` function

Author: mtsokol

sparse/mlir_backend/__init__.py

@@ -15,10 +15,12 @@
 )
 from ._ops import (
  add,
+ reshape,
 )
 
 __all__ = [
  "add",
  "asarray",
  "asdtype",
+ "reshape",
 ]

sparse/mlir_backend/_constructors.py

@@ -239,7 +239,7 @@ def from_sps(cls, arr: np.ndarray) -> "Dense":
 
  return dense_instance
 
- def to_sps(self, shape: tuple[int, ...]) -> sps.csr_array:
+ def to_sps(self, shape: tuple[int, ...]) -> np.ndarray:
  data = ranked_memref_to_numpy(self.data)
  return data.reshape(shape)
 

sparse/mlir_backend/_ops.py

@@ -5,10 +5,12 @@
 from mlir import ir
 from mlir.dialects import arith, func, linalg, sparse_tensor, tensor
 
+import numpy as np
+
 from ._common import fn_cache
-from ._constructors import Tensor
+from ._constructors import Tensor, numpy_to_ranked_memref
 from ._core import CWD, DEBUG, MLIR_C_RUNNER_UTILS, ctx, pm
-from ._dtypes import DType, FloatingDType
+from ._dtypes import DType, FloatingDType, Index
 
 
 @fn_cache
@@ -68,11 +70,35 @@ def add(a, b):
  return mlir.execution_engine.ExecutionEngine(module, opt_level=2, shared_libs=[MLIR_C_RUNNER_UTILS])
 
 
+@fn_cache
+def get_reshape_module(
+ a_tensor_type: ir.RankedTensorType,
+ shape_tensor_type: ir.RankedTensorType,
+ out_tensor_type: ir.RankedTensorType,
+) -> ir.Module:
+ with ir.Location.unknown(ctx):
+ module = ir.Module.create()
+
+ with ir.InsertionPoint(module.body):
+
+ @func.FuncOp.from_py_func(a_tensor_type, shape_tensor_type)
+ def reshape(a, shape):
+ return tensor.reshape(out_tensor_type, a, shape)
+
+ reshape.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+ if DEBUG:
+ (CWD / "reshape_module.mlir").write_text(str(module))
+ pm.run(module.operation)
+ if DEBUG:
+ (CWD / "reshape_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)
 
- # TODO: Add proper caching
  # TODO: Decide what will be the output tensor_type
  add_module = get_add_module(
  x1._obj.get_tensor_definition(x1.shape),
@@ -88,3 +114,24 @@ def add(x1: Tensor, x2: Tensor) -> Tensor:
  *x2._obj.to_module_arg(),
  )
  return Tensor(ret_obj, shape=out_tensor_type.shape)
+
+
+def reshape(x: Tensor, /, shape: tuple[int, ...]) -> Tensor:
+ ret_obj = x._format_class()
+ x_tensor_type = x._obj.get_tensor_definition(x.shape)
+ out_tensor_type = x._obj.get_tensor_definition(shape)
+
+ with ir.Location.unknown(ctx):
+ shape_tensor_type = ir.RankedTensorType.get([len(shape)], Index.get_mlir_type())
+
+ reshape_module = get_reshape_module(x_tensor_type, shape_tensor_type, out_tensor_type)
+
+ shape = np.array(shape)
+ reshape_module.invoke(
+ "reshape",
+ ctypes.pointer(ctypes.pointer(ret_obj)),
+ *x._obj.to_module_arg(),
+ ctypes.pointer(ctypes.pointer(numpy_to_ranked_memref(shape))),
+ )
+
+ return Tensor(ret_obj, shape=out_tensor_type.shape)

sparse/mlir_backend/tests/test_simple.py

@@ -75,6 +75,15 @@ def sampler_real_floating(size: tuple[int, ...]):
  raise NotImplementedError(f"{dtype=} not yet supported.")
 
 
+def get_exampe_csf_arrays(dtype: np.dtype) -> tuple:
+ pos_1 = np.array([0, 1, 3], dtype=np.int64)
+ crd_1 = np.array([1, 0, 1], dtype=np.int64)
+ pos_2 = np.array([0, 3, 5, 7], dtype=np.int64)
+ crd_2 = np.array([0, 1, 3, 0, 3, 0, 1], dtype=np.int64)
+ data = np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype)
+ return pos_1, crd_1, pos_2, crd_2, data
+
+
 @parametrize_dtypes
 @pytest.mark.parametrize("shape", [(100,), (10, 200), (5, 10, 20)])
 def test_dense_format(dtype, shape):
@@ -176,11 +185,7 @@ def test_add(rng, dtype):
 @parametrize_dtypes
 def test_csf_format(dtype):
  SHAPE = (2, 2, 4)
- pos_1 = np.array([0, 1, 3], dtype=np.int64)
- crd_1 = np.array([1, 0, 1], dtype=np.int64)
- pos_2 = np.array([0, 3, 5, 7], dtype=np.int64)
- crd_2 = np.array([0, 1, 3, 0, 3, 0, 1], dtype=np.int64)
- data = np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype)
+ pos_1, crd_1, pos_2, crd_2, data = get_exampe_csf_arrays(dtype)
  csf = [pos_1, crd_1, pos_2, crd_2, data]
 
  csf_tensor = sparse.asarray(csf, shape=SHAPE, dtype=sparse.asdtype(dtype), format="csf")
@@ -192,3 +197,70 @@ def test_csf_format(dtype):
  csf_2 = [pos_1, crd_1, pos_2, crd_2, data * 2]
  for actual, expected in zip(res_tensor, csf_2, strict=False):
  np.testing.assert_array_equal(actual, expected)
+
+
+@parametrize_dtypes
+def test_reshape(rng, dtype):
+ DENSITY = 0.5
+ sampler = generate_sampler(dtype, rng)
+
+ # CSR, CSC, COO
+ for shape, new_shape in [((100, 50), (25, 200)), ((80, 1), (8, 10))]:
+ 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
+
+ arr = sps.random_array(
+ shape, density=DENSITY, format=format, dtype=dtype, random_state=rng, data_sampler=sampler
+ )
+ if format == "coo":
+ arr.sum_duplicates()
+
+ tensor = sparse.asarray(arr)
+
+ actual = sparse.reshape(tensor, shape=new_shape).to_scipy_sparse()
+ expected = arr.todense().reshape(new_shape)
+
+ np.testing.assert_array_equal(actual.todense(), expected)
+
+ # CSF
+ csf_shape = (2, 2, 4)
+ for shape, new_shape, expected_arrs in [
+ (
+ csf_shape,
+ (4, 4, 1),
+ [
+ np.array([0, 0, 3, 5, 7]),
+ np.array([0, 1, 3, 0, 3, 0, 1]),
+ np.array([0, 1, 2, 3, 4, 5, 6, 7]),
+ np.array([0, 0, 0, 0, 0, 0, 0]),
+ np.array([1, 2, 3, 4, 5, 6, 7]),
+ ],
+ ),
+ (
+ csf_shape,
+ (2, 1, 8),
+ [
+ np.array([0, 1, 2]),
+ np.array([0, 0]),
+ np.array([0, 3, 7]),
+ np.array([4, 5, 7, 0, 3, 4, 5]),
+ np.array([1, 2, 3, 4, 5, 6, 7]),
+ ],
+ ),
+ ]:
+ csf = get_exampe_csf_arrays(dtype)
+ csf_tensor = sparse.asarray(csf, shape=shape, dtype=sparse.asdtype(dtype), format="csf")
+
+ result = sparse.reshape(csf_tensor, shape=new_shape).to_scipy_sparse()
+
+ for actual, expected in zip(result, expected_arrs, strict=False):
+ np.testing.assert_array_equal(actual, expected)
+
+ # DENSE
+ # NOTE: dense reshape is probably broken in MLIR
+ # dense = np.arange(math.prod(SHAPE), dtype=dtype).reshape(SHAPE)