Skip to content

Commit 93cdddf

Browse files
antonwolfyantonwolfy
committed
Add mkl_vm::arg() extnesion to be used in dpnp.angle
1 parent 9792685 commit 93cdddf

File tree

6 files changed

+179
-0
lines changed

6 files changed

+179
-0
lines changed

‎dpnp/backend/extensions/vm/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ if(NOT _use_onemkl_interfaces)
2929
${CMAKE_CURRENT_SOURCE_DIR}/acos.cpp
3030
${CMAKE_CURRENT_SOURCE_DIR}/acosh.cpp
3131
${CMAKE_CURRENT_SOURCE_DIR}/add.cpp
32+
${CMAKE_CURRENT_SOURCE_DIR}/arg.cpp
3233
${CMAKE_CURRENT_SOURCE_DIR}/asin.cpp
3334
${CMAKE_CURRENT_SOURCE_DIR}/asinh.cpp
3435
${CMAKE_CURRENT_SOURCE_DIR}/atan.cpp

‎dpnp/backend/extensions/vm/arg.cpp‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
//*****************************************************************************
2+
// Copyright (c) 2025, Intel Corporation
3+
// All rights reserved.
4+
//
5+
// Redistribution and use in source and binary forms, with or without
6+
// modification, are permitted provided that the following conditions are met:
7+
// - Redistributions of source code must retain the above copyright notice,
8+
// this list of conditions and the following disclaimer.
9+
// - Redistributions in binary form must reproduce the above copyright notice,
10+
// this list of conditions and the following disclaimer in the documentation
11+
// and/or other materials provided with the distribution.
12+
//
13+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
14+
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
15+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
16+
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
17+
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
18+
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
19+
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
20+
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
21+
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
22+
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
23+
// THE POSSIBILITY OF SUCH DAMAGE.
24+
//*****************************************************************************
25+
26+
#include <oneapi/mkl.hpp>
27+
#include <sycl/sycl.hpp>
28+
29+
#include "dpctl4pybind11.hpp"
30+
31+
#include "arg.hpp"
32+
#include "common.hpp"
33+
34+
// include a local copy of elementwise common header from dpctl tensor:
35+
// dpctl/tensor/libtensor/source/elementwise_functions/elementwise_functions.hpp
36+
// TODO: replace by including dpctl header once available
37+
#include "../elementwise_functions/elementwise_functions.hpp"
38+
39+
// dpctl tensor headers
40+
#include "kernels/elementwise_functions/common.hpp"
41+
#include "utils/type_dispatch.hpp"
42+
#include "utils/type_utils.hpp"
43+
44+
namespace dpnp::extensions::vm
45+
{
46+
namespace py = pybind11;
47+
namespace py_int = dpnp::extensions::py_internal;
48+
namespace td_ns = dpctl::tensor::type_dispatch;
49+
50+
namespace impl
51+
{
52+
namespace ew_cmn_ns = dpctl::tensor::kernels::elementwise_common;
53+
namespace mkl_vm = oneapi::mkl::vm; // OneMKL namespace with VM functions
54+
namespace tu_ns = dpctl::tensor::type_utils;
55+
56+
/**
57+
* @brief A factory to define pairs of supported types for which
58+
* MKL VM library provides support in oneapi::mkl::vm::arg<T> function.
59+
*
60+
* @tparam T Type of input vector `a` and of result vector `y`.
61+
*/
62+
template <typename T>
63+
struct OutputType
64+
{
65+
using value_type = typename std::disjunction<
66+
td_ns::TypeMapResultEntry<T, std::complex<double>, double>,
67+
td_ns::TypeMapResultEntry<T, std::complex<float>, float>,
68+
td_ns::DefaultResultEntry<void>>::result_type;
69+
};
70+
71+
template <typename T>
72+
static sycl::event arg_contig_impl(sycl::queue &exec_q,
73+
std::size_t in_n,
74+
const char *in_a,
75+
char *out_y,
76+
const std::vector<sycl::event> &depends)
77+
{
78+
tu_ns::validate_type_for_device<T>(exec_q);
79+
80+
std::int64_t n = static_cast<std::int64_t>(in_n);
81+
const T *a = reinterpret_cast<const T *>(in_a);
82+
83+
using resTy = typename OutputType<T>::value_type;
84+
resTy *y = reinterpret_cast<resTy *>(out_y);
85+
86+
return mkl_vm::arg(exec_q,
87+
n, // number of elements to be calculated
88+
a, // pointer `a` containing input vector of size n
89+
y, // pointer `y` to the output vector of size n
90+
depends);
91+
}
92+
93+
using ew_cmn_ns::unary_contig_impl_fn_ptr_t;
94+
using ew_cmn_ns::unary_strided_impl_fn_ptr_t;
95+
96+
static int output_typeid_vector[td_ns::num_types];
97+
static unary_contig_impl_fn_ptr_t contig_dispatch_vector[td_ns::num_types];
98+
99+
MACRO_POPULATE_DISPATCH_VECTORS(arg);
100+
} // namespace impl
101+
102+
void init_arg(py::module_ m)
103+
{
104+
using arrayT = dpctl::tensor::usm_ndarray;
105+
using event_vecT = std::vector<sycl::event>;
106+
107+
impl::populate_dispatch_vectors();
108+
using impl::contig_dispatch_vector;
109+
using impl::output_typeid_vector;
110+
111+
auto arg_pyapi = [&](sycl::queue &exec_q, const arrayT &src,
112+
const arrayT &dst, const event_vecT &depends = {}) {
113+
return py_int::py_unary_ufunc(
114+
src, dst, exec_q, depends, output_typeid_vector,
115+
contig_dispatch_vector,
116+
// no support of strided implementation in OneMKL
117+
td_ns::NullPtrVector<impl::unary_strided_impl_fn_ptr_t>{});
118+
};
119+
m.def("_arg", arg_pyapi,
120+
"Call `arg` function from OneMKL VM library to compute "
121+
"the inverse tangent of vector elements",
122+
py::arg("sycl_queue"), py::arg("src"), py::arg("dst"),
123+
py::arg("depends") = py::list());
124+
125+
auto arg_need_to_call_pyapi = [&](sycl::queue &exec_q, const arrayT &src,
126+
const arrayT &dst) {
127+
return py_internal::need_to_call_unary_ufunc(
128+
exec_q, src, dst, output_typeid_vector, contig_dispatch_vector);
129+
};
130+
m.def("_mkl_arg_to_call", arg_need_to_call_pyapi,
131+
"Check input arguments to answer if `arg` function from "
132+
"OneMKL VM library can be used",
133+
py::arg("sycl_queue"), py::arg("src"), py::arg("dst"));
134+
}
135+
} // namespace dpnp::extensions::vm

‎dpnp/backend/extensions/vm/arg.hpp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
//*****************************************************************************
2+
// Copyright (c) 2025, Intel Corporation
3+
// All rights reserved.
4+
//
5+
// Redistribution and use in source and binary forms, with or without
6+
// modification, are permitted provided that the following conditions are met:
7+
// - Redistributions of source code must retain the above copyright notice,
8+
// this list of conditions and the following disclaimer.
9+
// - Redistributions in binary form must reproduce the above copyright notice,
10+
// this list of conditions and the following disclaimer in the documentation
11+
// and/or other materials provided with the distribution.
12+
//
13+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
14+
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
15+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
16+
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
17+
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
18+
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
19+
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
20+
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
21+
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
22+
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
23+
// THE POSSIBILITY OF SUCH DAMAGE.
24+
//*****************************************************************************
25+
26+
#pragma once
27+
28+
#include <pybind11/pybind11.h>
29+
30+
namespace py = pybind11;
31+
32+
namespace dpnp::extensions::vm
33+
{
34+
void init_arg(py::module_ m);
35+
} // namespace dpnp::extensions::vm

‎dpnp/backend/extensions/vm/vm_py.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "acos.hpp"
3333
#include "acosh.hpp"
3434
#include "add.hpp"
35+
#include "arg.hpp"
3536
#include "asin.hpp"
3637
#include "asinh.hpp"
3738
#include "atan.hpp"
@@ -80,6 +81,7 @@ PYBIND11_MODULE(_vm_impl, m)
8081
vm_ns::init_acos(m);
8182
vm_ns::init_acosh(m);
8283
vm_ns::init_add(m);
84+
vm_ns::init_arg(m);
8385
vm_ns::init_asin(m);
8486
vm_ns::init_asinh(m);
8587
vm_ns::init_atan(m);

‎dpnp/dpnp_algo/dpnp_elementwise_common.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -497,12 +497,16 @@ def __init__(
497497
result_type_resolver_fn,
498498
unary_dp_impl_fn,
499499
docs,
500+
mkl_fn_to_call=None,
501+
mkl_impl_fn=None,
500502
):
501503
super().__init__(
502504
name,
503505
result_type_resolver_fn,
504506
unary_dp_impl_fn,
505507
docs,
508+
mkl_fn_to_call=mkl_fn_to_call,
509+
mkl_impl_fn=mkl_impl_fn,
506510
)
507511

508512
def __call__(self, x, deg=False, out=None, order="K"):

‎dpnp/dpnp_iface_mathematical.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,8 @@ def _process_ediff1d_args(arg, arg_name, ary_dtype, ary_sycl_queue, usm_type):
561561
ti._angle_result_type,
562562
ti._angle,
563563
_ANGLE_DOCSTRING,
564+
mkl_fn_to_call="_mkl_arg_to_call",
565+
mkl_impl_fn="_arg",
564566
)
565567

566568

0 commit comments

Comments
 (0)