Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve performance of dpnp.matmul and dpnp.dot with out keyword #1694

Merged
merged 6 commits into from
Feb 8, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions dpnp/dpnp_iface.py
Original file line number Diff line number Diff line change
Expand Up @@ -484,17 +484,20 @@ def get_result_array(a, out=None, casting="safe"):
if out is None:
return a
else:
dpnp.check_supported_arrays_type(out)
if out.shape != a.shape:
raise ValueError(
f"Output array of shape {a.shape} is needed, got {out.shape}."
)
elif isinstance(out, dpt.usm_ndarray):
out = dpnp_array._create_from_usm_ndarray(out)
if a is out:
return out
else:
dpnp.check_supported_arrays_type(out)
if out.shape != a.shape:
raise ValueError(
f"Output array of shape {a.shape} is needed, got {out.shape}."
)
elif isinstance(out, dpt.usm_ndarray):
out = dpnp_array._create_from_usm_ndarray(out)

dpnp.copyto(out, a, casting=casting)
dpnp.copyto(out, a, casting=casting)

return out
return out


def get_usm_ndarray(a):
Expand Down
54 changes: 41 additions & 13 deletions dpnp/dpnp_utils/dpnp_utils_linearalgebra.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,42 @@
__all__ = ["dpnp_dot", "dpnp_matmul"]


def _create_result_array(x1, x2, out, shape, dtype, usm_type, sycl_queue):
"""
Create the result array.

If `out` is not ``None`` and its features match the specified `shape`, `dtype,
`usm_type`, and `sycl_queue` and it is C-contiguous or F-contiguous and
does not have any memory overlap with `x1` and `x2`, `out` itself is returned.
If these conditions are not statisfied, an empty array is returned with the
specified `shape`, `dtype, `usm_type`, and `sycl_queue`.
"""

if out is not None:
dpnp.check_supported_arrays_type(out)
vtavana marked this conversation as resolved.
Show resolved Hide resolved
x1_usm = dpnp.get_usm_ndarray(x1)
x2_usm = dpnp.get_usm_ndarray(x2)
out_usm = dpnp.get_usm_ndarray(out)

if (
out.dtype == dtype
and out.shape == shape
and out.usm_type == usm_type
and out.sycl_queue == sycl_queue
and (out.flags.c_contiguous or out.flags.f_contiguous)
and not ti._array_overlap(x1_usm, out_usm)
and not ti._array_overlap(x2_usm, out_usm)
):
return out

return dpnp.empty(
shape,
dtype=dtype,
usm_type=usm_type,
sycl_queue=sycl_queue,
)


def _copy_array(x, dep_events, host_events, contig_copy=False, dtype=None):
"""
Creating a copy of input array if needed.
Expand Down Expand Up @@ -212,14 +248,9 @@ def dpnp_dot(a, b, /, out=None):
a, b, dtype=None, casting="no", sycl_queue=exec_q
)

# create result array
result = dpnp.empty(
(),
dtype=dot_dtype,
usm_type=res_usm_type,
sycl_queue=exec_q,
result = _create_result_array(
a, b, out, (), dot_dtype, res_usm_type, exec_q
)

# input arrays should have the proper data type
dep_events_list = []
host_tasks_list = []
Expand Down Expand Up @@ -361,13 +392,10 @@ def dpnp_matmul(
x2_shape = x2.shape
res_shape = tuple(tmp_shape) + (x1_shape[-2], x2_shape[-1])

# calculate results
result = dpnp.empty(
res_shape,
dtype=gemm_dtype,
usm_type=res_usm_type,
sycl_queue=exec_q,
result = _create_result_array(
x1, x2, out, res_shape, gemm_dtype, res_usm_type, exec_q
)
# calculate result
if result.size == 0:
pass
elif x1.size == 0 or x2.size == 0:
Expand Down
Loading