Add implementation of dpnp.heaviside function

dpnp/backend/extensions/ufunc/CMakeLists.txt

@@ -32,6 +32,7 @@ set(_elementwise_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/fmax.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/fmin.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/fmod.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/heaviside.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/logaddexp2.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/elementwise_functions/radians.cpp
 )

dpnp/backend/extensions/ufunc/elementwise_functions/common.cpp

@@ -32,6 +32,7 @@
 #include "fmax.hpp"
 #include "fmin.hpp"
 #include "fmod.hpp"
+#include "heaviside.hpp"
 #include "logaddexp2.hpp"
 #include "radians.hpp"
 
@@ -51,6 +52,7 @@ void init_elementwise_functions(py::module_ m)
     init_fmax(m);
     init_fmin(m);
     init_fmod(m);
+    init_heaviside(m);
     init_logaddexp2(m);
     init_radians(m);
 }

dpnp/backend/extensions/ufunc/elementwise_functions/heaviside.cpp

@@ -0,0 +1,145 @@
+//*****************************************************************************
+// Copyright (c) 2024, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// maxification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#include <sycl/sycl.hpp>
+
+#include "dpctl4pybind11.hpp"
+
+#include "heaviside.hpp"
+#include "kernels/elementwise_functions/heaviside.hpp"
+#include "populate.hpp"
+
+// include a local copy of elementwise common header from dpctl tensor:
+// dpctl/tensor/libtensor/source/elementwise_functions/elementwise_functions.hpp
+// TODO: replace by including dpctl header once available
+#include "../../elementwise_functions/elementwise_functions.hpp"
+
+// dpctl tensor headers
+#include "kernels/elementwise_functions/common.hpp"
+#include "kernels/elementwise_functions/maximum.hpp"
+#include "utils/type_dispatch.hpp"
+
+namespace dpnp::extensions::ufunc
+{
+namespace py = pybind11;
+namespace py_int = dpnp::extensions::py_internal;
+namespace td_ns = dpctl::tensor::type_dispatch;
+
+namespace impl
+{
+namespace ew_cmn_ns = dpctl::tensor::kernels::elementwise_common;
+
+template <typename T1, typename T2>
+struct OutputType
+{
+    using value_type = typename std::disjunction< // disjunction is C++17
+                                                  // feature, supported by DPC++
+        td_ns::BinaryTypeMapResultEntry<T1,
+                                        sycl::half,
+                                        T2,
+                                        sycl::half,
+                                        sycl::half>,
+        td_ns::BinaryTypeMapResultEntry<T1, float, T2, float, float>,
+        td_ns::BinaryTypeMapResultEntry<T1, double, T2, double, double>,
+        td_ns::DefaultResultEntry<void>>::result_type;
+};
+
+using dpnp::kernels::heaviside::HeavisideFunctor;
+
+template <typename argT1,
+          typename argT2,
+          typename resT,
+          unsigned int vec_sz = 4,
+          unsigned int n_vecs = 2,
+          bool enable_sg_loadstore = true>
+using ContigFunctor =
+    ew_cmn_ns::BinaryContigFunctor<argT1,
+                                   argT2,
+                                   resT,
+                                   HeavisideFunctor<argT1, argT2, resT>,
+                                   vec_sz,
+                                   n_vecs,
+                                   enable_sg_loadstore>;
+
+template <typename argT1, typename argT2, typename resT, typename IndexerT>
+using StridedFunctor =
+    ew_cmn_ns::BinaryStridedFunctor<argT1,
+                                    argT2,
+                                    resT,
+                                    IndexerT,
+                                    HeavisideFunctor<argT1, argT2, resT>>;
+
+using ew_cmn_ns::binary_contig_impl_fn_ptr_t;
+using ew_cmn_ns::binary_contig_matrix_contig_row_broadcast_impl_fn_ptr_t;
+using ew_cmn_ns::binary_contig_row_contig_matrix_broadcast_impl_fn_ptr_t;
+using ew_cmn_ns::binary_strided_impl_fn_ptr_t;
+
+static binary_contig_impl_fn_ptr_t
+    heaviside_contig_dispatch_table[td_ns::num_types][td_ns::num_types];
+static int heaviside_output_typeid_table[td_ns::num_types][td_ns::num_types];
+static binary_strided_impl_fn_ptr_t
+    heaviside_strided_dispatch_table[td_ns::num_types][td_ns::num_types];
+
+MACRO_POPULATE_DISPATCH_TABLES(heaviside);
+} // namespace impl
+
+void init_heaviside(py::module_ m)
+{
+    using arrayT = dpctl::tensor::usm_ndarray;
+    using event_vecT = std::vector<sycl::event>;
+    {
+        impl::populate_heaviside_dispatch_tables();
+        using impl::heaviside_contig_dispatch_table;
+        using impl::heaviside_output_typeid_table;
+        using impl::heaviside_strided_dispatch_table;
+
+        auto heaviside_pyapi = [&](const arrayT &src1, const arrayT &src2,
+                                   const arrayT &dst, sycl::queue &exec_q,
+                                   const event_vecT &depends = {}) {
+            return py_int::py_binary_ufunc(
+                src1, src2, dst, exec_q, depends, heaviside_output_typeid_table,
+                heaviside_contig_dispatch_table,
+                heaviside_strided_dispatch_table,
+                // no dedicated kernel for C-contig row with broadcasting
+                td_ns::NullPtrTable<
+                    impl::
+                        binary_contig_matrix_contig_row_broadcast_impl_fn_ptr_t>{},
+                td_ns::NullPtrTable<
+                    impl::
+                        binary_contig_row_contig_matrix_broadcast_impl_fn_ptr_t>{});
+        };
+        m.def("_heaviside", heaviside_pyapi, "", py::arg("src1"),
+              py::arg("src2"), py::arg("dst"), py::arg("sycl_queue"),
+              py::arg("depends") = py::list());
+
+        auto heaviside_result_type_pyapi = [&](const py::dtype &dtype1,
+                                               const py::dtype &dtype2) {
+            return py_int::py_binary_ufunc_result_type(
+                dtype1, dtype2, heaviside_output_typeid_table);
+        };
+        m.def("_heaviside_result_type", heaviside_result_type_pyapi);
+    }
+}
+} // namespace dpnp::extensions::ufunc

dpnp/backend/extensions/ufunc/elementwise_functions/heaviside.hpp

@@ -0,0 +1,35 @@
+//*****************************************************************************
+// Copyright (c) 2024, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace dpnp::extensions::ufunc
+{
+void init_heaviside(py::module_ m);
+} // namespace dpnp::extensions::ufunc

dpnp/backend/kernels/elementwise_functions/heaviside.hpp

@@ -0,0 +1,57 @@
+//*****************************************************************************
+// Copyright (c) 2024, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#pragma once
+
+#include <sycl/sycl.hpp>
+
+// dpctl tensor headers
+#include "utils/math_utils.hpp"
+#include "utils/type_utils.hpp"
+
+namespace dpnp::kernels::heaviside
+{
+namespace mu_ns = dpctl::tensor::math_utils;
+namespace tu_ns = dpctl::tensor::type_utils;
+
+template <typename argT1, typename argT2, typename resT>
+struct HeavisideFunctor
+{
+    using supports_sg_loadstore = std::negation<
+        std::disjunction<tu_ns::is_complex<argT1>, tu_ns::is_complex<argT2>>>;
+    using supports_vec = typename std::false_type;
+
+    resT operator()(const argT1 &in1, const argT2 &in2) const
+    {
+        if (std::isnan(in1)) {
+            return in1;
+        }
+        else if (in1 == 0) {
+            return in2;
+        }
+        return resT(in1 > 0);
+    }
+};
+} // namespace dpnp::kernels::heaviside

dpnp/dpnp_iface_mathematical.py

@@ -100,6 +100,7 @@
     "fmin",
     "fmod",
     "gradient",
+    "heaviside",
     "imag",
     "maximum",
     "minimum",
@@ -2179,6 +2180,62 @@ def gradient(f, *varargs, axis=None, edge_order=1):
     return tuple(outvals)
 
 
+_HEAVISIDE_DOCSTRING = """
+Compute the Heaviside step function.
+
+The Heaviside step function is defined as::
+
+                          0   if x1 < 0
+    heaviside(x1, x2) =  x2   if x1 == 0
+                          1   if x1 > 0
+
+where `x2` is often taken to be 0.5, but 0 and 1 are also sometimes used.
+
+Parameters
+----------
+x1 : {dpnp.ndarray, usm_ndarray, scalar}
+    Input values.
+    Both inputs `x1` and `x2` can not be scalars at the same time.
+x2 : {dpnp.ndarray, usm_ndarray, scalar}
+    The value of the function when `x1` is ``0``.
+    Both inputs `x1` and `x2` can not be scalars at the same time.
+out : {None, dpnp.ndarray, usm_ndarray}, optional
+    Output array to populate.
+    Array must have the correct shape and the expected data type.
+    Default: ``None``.
+order : {"C", "F", "A", "K"}, optional
+    Memory layout of the newly output array, if parameter `out` is ``None``.
+    Default: ``"K"``.
+
+Returns
+-------
+out : dpnp.ndarray
+    The output array, element-wise Heaviside step function of `x1`.
+
+Limitations
+-----------
+Parameters `where` and `subok` are supported with their default values.
+Keyword argument `kwargs` is currently unsupported.
+Otherwise ``NotImplementedError`` exception will be raised.
+
+Examples
+--------
+>>> import dpnp as np
+>>> a = np.array([-1.5, 0, 2.0])
+>>> np.heaviside(a, 0.5)
+array([0. , 0.5, 1. ])
+>>> np.heaviside(a, 1)
+array([0., 1., 1.])
+"""
+
+heaviside = DPNPBinaryFunc(
+    "heaviside",
+    ufi._heaviside_result_type,
+    ufi._heaviside,
+    _HEAVISIDE_DOCSTRING,
+)
+
+
 _IMAG_DOCSTRING = """
 Computes imaginary part of each element `x_i` for input array `x`.
 

tests/skipped_tests.tbl

@@ -19,8 +19,6 @@ tests/test_umath.py::test_umaths[('frexp', 'f')]
 tests/test_umath.py::test_umaths[('frexp', 'd')]
 tests/test_umath.py::test_umaths[('gcd', 'ii')]
 tests/test_umath.py::test_umaths[('gcd', 'll')]
-tests/test_umath.py::test_umaths[('heaviside', 'ff')]
-tests/test_umath.py::test_umaths[('heaviside', 'dd')]
 tests/test_umath.py::test_umaths[('lcm', 'ii')]
 tests/test_umath.py::test_umaths[('lcm', 'll')]
 tests/test_umath.py::test_umaths[('ldexp', 'fi')]
@@ -193,8 +191,6 @@ tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_fx
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_x
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_size1
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_to_nan
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_heaviside
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_heaviside_nan_inf
 
 tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsBeta_param_0_{a_shape=(), b_shape=(), shape=(4, 3, 2)}::test_beta
 tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsBeta_param_1_{a_shape=(), b_shape=(), shape=(3, 2)}::test_beta

tests/skipped_tests_gpu.tbl

@@ -26,8 +26,6 @@ tests/test_umath.py::test_umaths[('frexp', 'f')]
 tests/test_umath.py::test_umaths[('frexp', 'd')]
 tests/test_umath.py::test_umaths[('gcd', 'ii')]
 tests/test_umath.py::test_umaths[('gcd', 'll')]
-tests/test_umath.py::test_umaths[('heaviside', 'ff')]
-tests/test_umath.py::test_umaths[('heaviside', 'dd')]
 tests/test_umath.py::test_umaths[('lcm', 'ii')]
 tests/test_umath.py::test_umaths[('lcm', 'll')]
 tests/test_umath.py::test_umaths[('ldexp', 'fi')]
@@ -247,8 +245,6 @@ tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_fx
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_x
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_size1
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_inf_to_nan
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_heaviside
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_heaviside_nan_inf
 
 tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsBeta_param_0_{a_shape=(), b_shape=(), shape=(4, 3, 2)}::test_beta
 tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsBeta_param_1_{a_shape=(), b_shape=(), shape=(3, 2)}::test_beta