36#include <oneapi/mkl.hpp>
37#include <pybind11/pybind11.h>
38#include <sycl/sycl.hpp>
40#include "dpnp4pybind11.hpp"
43#include "ext/common.hpp"
46#include "utils/memory_overlap.hpp"
47#include "utils/type_dispatch.hpp"
49namespace ext_ns = ext::common;
50namespace py = pybind11;
51namespace td_ns = dpnp::tensor::type_dispatch;
53namespace dpnp::extensions::vm::py_internal
55template <
typename output_typesT,
typename contig_dispatchT>
56bool need_to_call_unary_ufunc(sycl::queue &exec_q,
57 const dpnp::tensor::usm_ndarray &src,
58 const dpnp::tensor::usm_ndarray &dst,
59 const output_typesT &output_type_vec,
60 const contig_dispatchT &contig_dispatch_vector)
63 int src_typenum = src.get_typenum();
64 int dst_typenum = dst.get_typenum();
66 auto array_types = td_ns::usm_ndarray_types();
67 int src_typeid = array_types.typenum_to_lookup_id(src_typenum);
68 int dst_typeid = array_types.typenum_to_lookup_id(dst_typenum);
71 int func_output_typeid = output_type_vec[src_typeid];
72 if (dst_typeid != func_output_typeid) {
77 if (!exec_q.get_device().has(sycl::aspect::fp64)) {
82 if (!dpnp::utils::queues_are_compatible(exec_q, {src, dst})) {
87 int dst_nd = dst.get_ndim();
88 if (dst_nd != src.get_ndim()) {
91 else if (dst_nd == 0) {
97 const py::ssize_t *src_shape = src.get_shape_raw();
98 const py::ssize_t *dst_shape = dst.get_shape_raw();
99 bool shapes_equal(
true);
100 size_t src_nelems(1);
102 for (
int i = 0; i < dst_nd; ++i) {
103 src_nelems *=
static_cast<size_t>(src_shape[i]);
104 shapes_equal = shapes_equal && (src_shape[i] == dst_shape[i]);
111 if (src_nelems == 0) {
116 auto dst_offsets = dst.get_minmax_offsets();
120 static_cast<size_t>(dst_offsets.second - dst_offsets.first);
121 if (range + 1 < src_nelems) {
127 auto const &overlap = dpnp::tensor::overlap::MemoryOverlap();
128 if (overlap(src, dst)) {
133 bool is_src_c_contig = src.is_c_contiguous();
134 bool is_dst_c_contig = dst.is_c_contiguous();
136 bool all_c_contig = (is_src_c_contig && is_dst_c_contig);
142 if (contig_dispatch_vector[src_typeid] ==
nullptr) {
148template <
typename output_typesT,
typename contig_dispatchT>
149bool need_to_call_unary_two_outputs_ufunc(
151 const dpnp::tensor::usm_ndarray &src,
152 const dpnp::tensor::usm_ndarray &dst1,
153 const dpnp::tensor::usm_ndarray &dst2,
154 const output_typesT &output_type_vec,
155 const contig_dispatchT &contig_dispatch_vector)
158 int src_typenum = src.get_typenum();
159 int dst1_typenum = dst1.get_typenum();
160 int dst2_typenum = dst2.get_typenum();
162 const auto &array_types = td_ns::usm_ndarray_types();
163 int src_typeid = array_types.typenum_to_lookup_id(src_typenum);
164 int dst1_typeid = array_types.typenum_to_lookup_id(dst1_typenum);
165 int dst2_typeid = array_types.typenum_to_lookup_id(dst2_typenum);
167 std::pair<int, int> func_output_typeids = output_type_vec[src_typeid];
170 if (dst1_typeid != func_output_typeids.first ||
171 dst2_typeid != func_output_typeids.second) {
176 if (!exec_q.get_device().has(sycl::aspect::fp64)) {
181 if (!dpnp::utils::queues_are_compatible(exec_q, {src, dst1, dst2})) {
186 int src_nd = src.get_ndim();
187 int dst1_nd = dst1.get_ndim();
188 int dst2_nd = dst2.get_ndim();
189 if (src_nd != dst1_nd || src_nd != dst2_nd) {
192 else if (dst1_nd == 0 || dst2_nd == 0) {
198 const py::ssize_t *src_shape = src.get_shape_raw();
199 const py::ssize_t *dst1_shape = dst1.get_shape_raw();
200 const py::ssize_t *dst2_shape = dst2.get_shape_raw();
201 bool shapes_equal(
true);
202 size_t src_nelems(1);
204 for (
int i = 0; i < src_nd; ++i) {
205 src_nelems *=
static_cast<std::size_t
>(src_shape[i]);
206 shapes_equal = shapes_equal && (src_shape[i] == dst1_shape[i]) &&
207 (src_shape[i] == dst2_shape[i]);
214 if (src_nelems == 0) {
219 auto dst1_offsets = dst1.get_minmax_offsets();
220 auto dst2_offsets = dst2.get_minmax_offsets();
224 static_cast<size_t>(dst1_offsets.second - dst1_offsets.first);
226 static_cast<size_t>(dst2_offsets.second - dst2_offsets.first);
227 if ((range1 + 1 < src_nelems) || (range2 + 1 < src_nelems)) {
233 auto const &overlap = dpnp::tensor::overlap::MemoryOverlap();
234 if (overlap(src, dst1) || overlap(src, dst2) || overlap(dst1, dst2)) {
239 bool is_src_c_contig = src.is_c_contiguous();
240 bool is_dst1_c_contig = dst1.is_c_contiguous();
241 bool is_dst2_c_contig = dst2.is_c_contiguous();
244 (is_src_c_contig && is_dst1_c_contig && is_dst2_c_contig);
250 if (contig_dispatch_vector[src_typeid] ==
nullptr) {
256template <
typename output_typesT,
typename contig_dispatchT>
257bool need_to_call_binary_ufunc(sycl::queue &exec_q,
258 const dpnp::tensor::usm_ndarray &src1,
259 const dpnp::tensor::usm_ndarray &src2,
260 const dpnp::tensor::usm_ndarray &dst,
261 const output_typesT &output_type_table,
262 const contig_dispatchT &contig_dispatch_table)
265 int src1_typenum = src1.get_typenum();
266 int src2_typenum = src2.get_typenum();
267 int dst_typenum = dst.get_typenum();
269 auto array_types = td_ns::usm_ndarray_types();
270 int src1_typeid = array_types.typenum_to_lookup_id(src1_typenum);
271 int src2_typeid = array_types.typenum_to_lookup_id(src2_typenum);
272 int dst_typeid = array_types.typenum_to_lookup_id(dst_typenum);
275 int output_typeid = output_type_table[src1_typeid][src2_typeid];
276 if (output_typeid != dst_typeid) {
281 if (src1_typeid != src2_typeid) {
286 if (!exec_q.get_device().has(sycl::aspect::fp64)) {
291 if (!dpnp::utils::queues_are_compatible(exec_q, {src1, src2, dst})) {
296 int dst_nd = dst.get_ndim();
297 if (dst_nd != src1.get_ndim() || dst_nd != src2.get_ndim()) {
300 else if (dst_nd == 0) {
306 const py::ssize_t *src1_shape = src1.get_shape_raw();
307 const py::ssize_t *src2_shape = src2.get_shape_raw();
308 const py::ssize_t *dst_shape = dst.get_shape_raw();
309 bool shapes_equal(
true);
310 size_t src_nelems(1);
312 for (
int i = 0; i < dst_nd; ++i) {
313 src_nelems *=
static_cast<size_t>(src1_shape[i]);
314 shapes_equal = shapes_equal && (src1_shape[i] == dst_shape[i] &&
315 src2_shape[i] == dst_shape[i]);
322 if (src_nelems == 0) {
327 auto dst_offsets = dst.get_minmax_offsets();
331 static_cast<size_t>(dst_offsets.second - dst_offsets.first);
332 if (range + 1 < src_nelems) {
338 auto const &overlap = dpnp::tensor::overlap::MemoryOverlap();
339 if (overlap(src1, dst) || overlap(src2, dst)) {
344 bool is_src1_c_contig = src1.is_c_contiguous();
345 bool is_src2_c_contig = src2.is_c_contiguous();
346 bool is_dst_c_contig = dst.is_c_contiguous();
349 (is_src1_c_contig && is_src2_c_contig && is_dst_c_contig);
355 if (contig_dispatch_table[src1_typeid] ==
nullptr) {
366#define MACRO_POPULATE_DISPATCH_VECTORS(__name__) \
367 template <typename fnT, typename T> \
368 struct ContigFactory \
372 if constexpr (std::is_same_v<typename OutputType<T>::value_type, \
377 return __name__##_contig_impl<T>; \
382 template <typename fnT, typename T> \
383 struct TypeMapFactory \
385 std::enable_if_t<std::is_same<fnT, int>::value, int> get() \
387 using rT = typename OutputType<T>::value_type; \
388 return td_ns::GetTypeid<rT>{}.get(); \
392 static void populate_dispatch_vectors(void) \
394 ext_ns::init_dispatch_vector<int, TypeMapFactory>( \
395 output_typeid_vector); \
396 ext_ns::init_dispatch_vector<unary_contig_impl_fn_ptr_t, \
397 ContigFactory>(contig_dispatch_vector); \
405#define MACRO_POPULATE_DISPATCH_2OUTS_VECTORS(__name__) \
406 template <typename fnT, typename T> \
407 struct ContigFactory \
411 if constexpr (std::is_same_v<typename OutputType<T>::value_type1, \
413 std::is_same_v<typename OutputType<T>::value_type2, \
419 fnT fn = __name__##_contig_impl<T>; \
425 template <typename fnT, typename T> \
426 struct TypeMapFactory \
428 std::enable_if_t<std::is_same<fnT, std::pair<int, int>>::value, \
429 std::pair<int, int>> \
432 using rT1 = typename OutputType<T>::value_type1; \
433 using rT2 = typename OutputType<T>::value_type2; \
434 return std::make_pair(td_ns::GetTypeid<rT1>{}.get(), \
435 td_ns::GetTypeid<rT2>{}.get()); \
439 static void populate_dispatch_vectors(void) \
441 ext_ns::init_dispatch_vector<std::pair<int, int>, TypeMapFactory>( \
442 output_typeid_vector); \
443 ext_ns::init_dispatch_vector<unary_two_outputs_contig_impl_fn_ptr_t, \
444 ContigFactory>(contig_dispatch_vector); \
452#define MACRO_POPULATE_DISPATCH_TABLES(__name__) \
453 template <typename fnT, typename T1, typename T2> \
454 struct ContigFactory \
458 if constexpr (std::is_same_v< \
459 typename OutputType<T1, T2>::value_type, \
464 return __name__##_contig_impl<T1, T2>; \
469 template <typename fnT, typename T1, typename T2> \
470 struct TypeMapFactory \
472 std::enable_if_t<std::is_same<fnT, int>::value, int> get() \
474 using rT = typename OutputType<T1, T2>::value_type; \
475 return td_ns::GetTypeid<rT>{}.get(); \
479 static void populate_dispatch_tables(void) \
481 ext_ns::init_dispatch_table<int, TypeMapFactory>( \
482 output_typeid_vector); \
483 ext_ns::init_dispatch_table<binary_contig_impl_fn_ptr_t, \
484 ContigFactory>(contig_dispatch_vector); \