diff --git a/dpnp/backend/extensions/indexing/CMakeLists.txt b/dpnp/backend/extensions/indexing/CMakeLists.txt index 827a830e43a1..8f3fdfc53dde 100644 --- a/dpnp/backend/extensions/indexing/CMakeLists.txt +++ b/dpnp/backend/extensions/indexing/CMakeLists.txt @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # ***************************************************************************** -# Copyright (c) 2025, Intel Corporation +# Copyright (c) 2026, Intel Corporation # All rights reserved. # # Redistribution and use in source and binary forms, with or without @@ -30,6 +30,7 @@ set(python_module_name _indexing_impl) set(_module_src ${CMAKE_CURRENT_SOURCE_DIR}/choose.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/putmask.cpp ${CMAKE_CURRENT_SOURCE_DIR}/indexing_py.cpp ) diff --git a/dpnp/backend/extensions/indexing/indexing_py.cpp b/dpnp/backend/extensions/indexing/indexing_py.cpp index a2d0b2efd512..695725196646 100644 --- a/dpnp/backend/extensions/indexing/indexing_py.cpp +++ b/dpnp/backend/extensions/indexing/indexing_py.cpp @@ -1,5 +1,5 @@ //***************************************************************************** -// Copyright (c) 2025, Intel Corporation +// Copyright (c) 2026, Intel Corporation // All rights reserved. // // Redistribution and use in source and binary forms, with or without @@ -33,8 +33,10 @@ #include #include "choose.hpp" +#include "putmask.hpp" PYBIND11_MODULE(_indexing_impl, m) { dpnp::extensions::indexing::init_choose(m); + dpnp::extensions::indexing::init_putmask(m); } diff --git a/dpnp/backend/extensions/indexing/putmask.cpp b/dpnp/backend/extensions/indexing/putmask.cpp new file mode 100644 index 000000000000..86dce6ad7fa7 --- /dev/null +++ b/dpnp/backend/extensions/indexing/putmask.cpp @@ -0,0 +1,346 @@ +//***************************************************************************** +// Copyright (c) 2026, 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. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// 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 +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include "dpnp4pybind11.hpp" + +#include "kernels/indexing/putmask.hpp" + +// dpnp tensor headers +#include "utils/offset_utils.hpp" +#include "utils/output_validation.hpp" +#include "utils/type_dispatch.hpp" + +// utils extension headers +#include "ext/common.hpp" +#include "ext/validation_utils.hpp" + +namespace py = pybind11; +namespace td_ns = dpnp::tensor::type_dispatch; + +using dpnp::tensor::usm_ndarray; + +using ext::common::dtype_from_typenum; +using ext::validation::array_names; +using ext::validation::check_has_dtype; +using ext::validation::check_num_dims; +using ext::validation::check_queue; +using ext::validation::check_same_dtype; +using ext::validation::check_same_size; +using ext::validation::check_writable; + +namespace dpnp::extensions::indexing +{ +using ext::common::init_dispatch_vector; + +typedef sycl::event (*putmask_strided_fn_ptr_t)( + sycl::queue &, + const int, // nd + const std::size_t, // nelems + const py::ssize_t *, // shape_strides + char *, // dst + py::ssize_t, // dst_offset + const char *, // mask + py::ssize_t, // mask_offset + const char *, // values + const std::size_t, // values_size + const std::vector &); + +template +sycl::event putmask_strided_call(sycl::queue &q, + const int nd, + const std::size_t nelems, + const py::ssize_t *shape_strides, + char *dst_p, + py::ssize_t dst_offset, + const char *mask_p, + py::ssize_t mask_offset, + const char *values_p, + const std::size_t values_size, + const std::vector &depends) +{ + return dpnp::kernels::putmask::putmask_strided_impl( + q, nd, nelems, shape_strides, dst_p, dst_offset, mask_p, mask_offset, + values_p, values_size, depends); +} + +typedef sycl::event (*putmask_contig_fn_ptr_t)( + sycl::queue &, + const std::size_t, // nelems + char *, // dst + const char *, // mask + const char *, // values + const std::size_t, // values_size + const std::vector &); + +template +sycl::event putmask_contig_call(sycl::queue &q, + const std::size_t nelems, + char *dst_p, + const char *mask_p, + const char *values_p, + const std::size_t values_size, + const std::vector &depends) +{ + return dpnp::kernels::putmask::putmask_contig_impl( + q, nelems, dst_p, mask_p, values_p, values_size, depends); +} + +putmask_strided_fn_ptr_t putmask_strided_dispatch_vector[td_ns::num_types]; +putmask_contig_fn_ptr_t putmask_contig_dispatch_vector[td_ns::num_types]; + +std::pair + py_putmask(const usm_ndarray &dst, + const usm_ndarray &mask, + const usm_ndarray &values, + sycl::queue &exec_q, + const std::vector &depends = {}) +{ + array_names names = {{&dst, "dst"}, {&mask, "mask"}, {&values, "values"}}; + + check_same_dtype(&dst, &values, names); + check_has_dtype(&mask, td_ns::typenum_t::BOOL, names); + + check_same_size({&dst, &mask}, names); + const int nd = dst.get_ndim(); + check_num_dims({&mask}, nd, names); + + check_queue({&dst, &mask, &values}, names, exec_q); + check_writable({&dst}, names); + + auto types = td_ns::usm_ndarray_types(); + // dst_typeid == values_typeid (check_same_dtype(&dst, &values, names)) + int dst_values_typeid = types.typenum_to_lookup_id(dst.get_typenum()); + + const py::ssize_t *dst_shape = dst.get_shape_raw(); + const py::ssize_t *mask_shape = mask.get_shape_raw(); + bool shapes_equal(true); + std::size_t nelems(1); + + for (int i = 0; i < std::max(nd, 1); ++i) { + const py::ssize_t d = (nd == 0 ? 1 : dst_shape[i]); + const py::ssize_t m = (nd == 0 ? 1 : mask_shape[i]); + nelems *= static_cast(d); + shapes_equal = shapes_equal && (d == m); + } + if (!shapes_equal) { + throw py::value_error("`mask` and `dst` shapes must match"); + } + + // if nelems is zero, return + if (nelems == 0) { + return {sycl::event(), sycl::event()}; + } + + dpnp::tensor::validation::AmpleMemory::throw_if_not_ample(dst, nelems); + + char *dst_p = dst.get_data(); + const char *mask_p = mask.get_data(); + const char *values_p = values.get_data(); + const std::size_t values_size = values.get_size(); + + // the contig kernel cycles `values` by the memory-linear index, which + // matches numpy's C-order `values.flat` only for C-contiguous data + const bool all_c_contig = dst.is_c_contiguous() && mask.is_c_contiguous() && + values.is_c_contiguous(); + + if (all_c_contig) { + auto contig_fn = putmask_contig_dispatch_vector[dst_values_typeid]; + + if (contig_fn == nullptr) { + py::dtype dst_values_dtype_py = + dtype_from_typenum(dst_values_typeid); + throw std::runtime_error( + "Contiguous implementation is missing for " + + std::string(py::str(dst_values_dtype_py)) + " data type"); + } + + auto comp_ev = contig_fn(exec_q, nelems, dst_p, mask_p, values_p, + values_size, depends); + sycl::event ht_ev = dpnp::utils::keep_args_alive( + exec_q, {dst, mask, values}, {comp_ev}); + + return std::make_pair(ht_ev, comp_ev); + } + + // strided path: the iteration space is intentionally not simplified, so + // the kernel's linear index stays equal to the C-order flat index used to + // cycle `values` (simplify_iteration_space may reorder axes and break it) + const auto &dst_strides = dst.get_strides_vector(); + const auto &mask_strides = mask.get_strides_vector(); + + using shT = std::vector; + shT common_shape; + shT s_dst_strides; + shT s_mask_strides; + + int eff_nd = nd; + if (nd == 0) { + // scalar arrays: single-element 1D iteration + eff_nd = 1; + common_shape = {1}; + s_dst_strides = {0}; + s_mask_strides = {0}; + } + else { + common_shape.assign(dst_shape, dst_shape + nd); + s_dst_strides = dst_strides; + s_mask_strides = mask_strides; + } + + // trivial offsets: shape and strides are passed without simplification + constexpr py::ssize_t dst_off = 0; + constexpr py::ssize_t mask_off = 0; + + auto strided_fn = putmask_strided_dispatch_vector[dst_values_typeid]; + if (strided_fn == nullptr) { + py::dtype dt = dtype_from_typenum(dst_values_typeid); + throw std::runtime_error("Strided implementation is missing for " + + std::string(py::str(dt)) + " data type"); + } + + using dpnp::tensor::offset_utils::device_allocate_and_pack; + + std::vector host_tasks; + host_tasks.reserve(2); + + auto pack = device_allocate_and_pack( + exec_q, host_tasks, common_shape, s_dst_strides, s_mask_strides); + + auto shape_strides_owner = std::move(std::get<0>(pack)); + const py::ssize_t *shape_strides_dev = shape_strides_owner.get(); + const sycl::event &cpy_ev = std::get<2>(pack); + + std::vector all_deps = depends; + all_deps.push_back(cpy_ev); + + sycl::event comp_ev = + strided_fn(exec_q, eff_nd, nelems, shape_strides_dev, dst_p, dst_off, + mask_p, mask_off, values_p, values_size, all_deps); + + sycl::event cleanup_ev = dpnp::tensor::alloc_utils::async_smart_free( + exec_q, {comp_ev}, shape_strides_owner); + host_tasks.push_back(cleanup_ev); + + sycl::event ht_ev = + dpnp::utils::keep_args_alive(exec_q, {dst, mask, values}, host_tasks); + + return std::make_pair(ht_ev, comp_ev); +} + +/** + * @brief A factory to define pairs of supported types for which + * putmask function is available. + * + * @tparam T Type of input vector `dst` and `values` and of result vector `dst`. + */ +template +struct PutMaskOutputType +{ + using value_type = typename std::disjunction< + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry, + td_ns::TypeMapResultEntry>, + td_ns::TypeMapResultEntry>, + td_ns::DefaultResultEntry>::result_type; +}; + +template +struct PutMaskStridedFactory +{ + fnT get() + { + if constexpr (std::is_same_v::value_type, + void>) { + return nullptr; + } + else { + return putmask_strided_call; + } + } +}; + +template +struct PutMaskContigFactory +{ + fnT get() + { + if constexpr (std::is_same_v::value_type, + void>) { + return nullptr; + } + else { + return putmask_contig_call; + } + } +}; + +static void populate_putmask_dispatch_vectors() +{ + init_dispatch_vector( + putmask_strided_dispatch_vector); + init_dispatch_vector( + putmask_contig_dispatch_vector); +} + +void init_putmask(py::module_ &m) +{ + populate_putmask_dispatch_vectors(); + + m.def("_putmask", &py_putmask, "", py::arg("dst"), py::arg("mask"), + py::arg("values"), py::arg("sycl_queue"), + py::arg("depends") = py::list()); + + return; +} + +} // namespace dpnp::extensions::indexing diff --git a/dpnp/backend/extensions/indexing/putmask.hpp b/dpnp/backend/extensions/indexing/putmask.hpp new file mode 100644 index 000000000000..6c65231edc6d --- /dev/null +++ b/dpnp/backend/extensions/indexing/putmask.hpp @@ -0,0 +1,38 @@ +//***************************************************************************** +// Copyright (c) 2026, 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. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// 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 + +namespace py = pybind11; + +namespace dpnp::extensions::indexing +{ +void init_putmask(py::module_ &m); +} // namespace dpnp::extensions::indexing diff --git a/dpnp/backend/kernels/indexing/putmask.hpp b/dpnp/backend/kernels/indexing/putmask.hpp new file mode 100644 index 000000000000..954e18c91eef --- /dev/null +++ b/dpnp/backend/kernels/indexing/putmask.hpp @@ -0,0 +1,303 @@ +//***************************************************************************** +// Copyright (c) 2026, 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. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// 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 +#include +#include + +#include + +// dpnp tensor headers +#include "kernels/alignment.hpp" +#include "kernels/dpnp_tensor_types.hpp" +#include "utils/offset_utils.hpp" +#include "utils/sycl_utils.hpp" +#include "utils/type_utils.hpp" + +namespace dpnp::kernels::putmask +{ +using dpnp::tensor::ssize_t; + +template +struct PutMaskStridedFunctor +{ +private: + T *dst_ = nullptr; + const std::uint8_t *mask_u8_ = nullptr; + const T *vals_ = nullptr; + const TwoOffsets_IndexerT indexer_; + const std::size_t values_size_ = 0; + +public: + PutMaskStridedFunctor(T *dst, + const bool *mask, + const T *vals, + const TwoOffsets_IndexerT &indexer, + std::size_t values_size) + : dst_(dst), mask_u8_(reinterpret_cast(mask)), + vals_(vals), indexer_(indexer), values_size_(values_size) + { + } + + void operator()(sycl::id<1> wid) const + { + // empty `values` is a no-op (also guards the read below) + if (values_size_ == 0) { + return; + } + + const std::size_t lin = wid.get(0); + auto offset = indexer_(static_cast(lin)); + + const dpnp::tensor::ssize_t dst_off = offset.get_first_offset(); + const dpnp::tensor::ssize_t mask_off = offset.get_second_offset(); + + if (mask_u8_[mask_off]) { + const std::size_t vlin = lin % values_size_; + dst_[dst_off] = vals_[vlin]; + } + } +}; + +template +struct PutMaskContigFunctor +{ +private: + T *dst_ = nullptr; + const std::uint8_t *mask_u8_ = nullptr; + const T *values_ = nullptr; + std::size_t nelems_ = 0; + std::size_t val_size_ = 0; + +public: + PutMaskContigFunctor(T *dst, + const bool *mask, + const T *values, + std::size_t nelems, + std::size_t val_size) + : dst_(dst), mask_u8_(reinterpret_cast(mask)), + values_(values), nelems_(nelems), val_size_(val_size) + { + } + + void operator()(sycl::nd_item<1> ndit) const + { + if (val_size_ == 0 || nelems_ == 0) { + return; + } + + constexpr std::uint8_t elems_per_wi = n_vecs * vec_sz; + /* Each work-item processes vec_sz elements, contiguous in memory */ + /* NOTE: work-group size must be divisible by sub-group size */ + + using dpnp::tensor::type_utils::is_complex_v; + if constexpr (enable_sg_loadstore && !is_complex_v) { + auto sg = ndit.get_sub_group(); + const std::uint32_t sgSize = sg.get_max_local_range()[0]; + const std::size_t lane_id = sg.get_local_id()[0]; + + const std::size_t base = + elems_per_wi * (ndit.get_group(0) * ndit.get_local_range(0) + + sg.get_group_id()[0] * sgSize); + + const bool values_no_repeat = (val_size_ >= nelems_); + + if (base + elems_per_wi * sgSize <= nelems_) { + using dpnp::tensor::sycl_utils::sub_group_load; + using dpnp::tensor::sycl_utils::sub_group_store; + +#pragma unroll + for (std::uint8_t it = 0; it < elems_per_wi; it += vec_sz) { + const std::size_t offset = base + it * sgSize; + + auto dst_multi_ptr = sycl::address_space_cast< + sycl::access::address_space::global_space, + sycl::access::decorated::yes>(&dst_[offset]); + auto mask_multi_ptr = sycl::address_space_cast< + sycl::access::address_space::global_space, + sycl::access::decorated::yes>(&mask_u8_[offset]); + + const sycl::vec dst_vec = + sub_group_load(sg, dst_multi_ptr); + const sycl::vec mask_vec = + sub_group_load(sg, mask_multi_ptr); + + sycl::vec val_vec; + + if (values_no_repeat) { + auto values_multi_ptr = sycl::address_space_cast< + sycl::access::address_space::global_space, + sycl::access::decorated::yes>(&values_[offset]); + + val_vec = sub_group_load(sg, values_multi_ptr); + } + else { + const std::size_t idx = offset + lane_id; +#pragma unroll + for (std::uint8_t k = 0; k < vec_sz; ++k) { + const std::size_t g = + idx + static_cast(k) * sgSize; + val_vec[k] = values_[g % val_size_]; + } + } + + sycl::vec out_vec; +#pragma unroll + for (std::uint8_t vec_id = 0; vec_id < vec_sz; ++vec_id) { + out_vec[vec_id] = + (mask_vec[vec_id] != static_cast(0)) + ? val_vec[vec_id] + : dst_vec[vec_id]; + } + + sub_group_store(sg, out_vec, dst_multi_ptr); + } + } + else { + const std::size_t lane_id = sg.get_local_id()[0]; + for (std::size_t k = base + lane_id; k < nelems_; k += sgSize) { + if (mask_u8_[k]) { + const std::size_t v = + values_no_repeat ? k : (k % val_size_); + dst_[k] = values_[v]; + } + } + } + } + else { + const std::size_t gid = ndit.get_global_linear_id(); + const std::size_t gws = ndit.get_global_range(0); + + const bool values_no_repeat = (val_size_ >= nelems_); + for (std::size_t offset = gid; offset < nelems_; offset += gws) { + if (mask_u8_[offset]) { + const std::size_t v = + values_no_repeat ? offset : (offset % val_size_); + dst_[offset] = values_[v]; + } + } + } + } +}; + +template +sycl::event putmask_strided_impl(sycl::queue &exec_q, + const int nd, + std::size_t nelems, + const dpnp::tensor::ssize_t *shape_strides, + char *dst_cp, + const dpnp::tensor::ssize_t dst_offset, + const char *mask_cp, + const dpnp::tensor::ssize_t mask_offset, + const char *values_cp, + std::size_t values_size, + const std::vector &depends = {}) +{ + dpnp::tensor::type_utils::validate_type_for_device(exec_q); + + T *dst_tp = reinterpret_cast(dst_cp); + const bool *mask_tp = reinterpret_cast(mask_cp); + const T *vals_tp = reinterpret_cast(values_cp); + + using IndexerT = dpnp::tensor::offset_utils::TwoOffsets_StridedIndexer; + const IndexerT indexer{nd, dst_offset, mask_offset, shape_strides}; + + return exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(depends); + + using PutMaskFunc = PutMaskStridedFunctor; + cgh.parallel_for( + sycl::range<1>(nelems), + PutMaskFunc(dst_tp, mask_tp, vals_tp, indexer, values_size)); + }); +} + +template +sycl::event putmask_contig_impl(sycl::queue &exec_q, + std::size_t nelems, + char *dst_cp, + const char *mask_cp, + const char *values_cp, + std::size_t values_size, + const std::vector &depends = {}) +{ + T *dst_tp = reinterpret_cast(dst_cp); + const bool *mask_tp = reinterpret_cast(mask_cp); + const T *values_tp = reinterpret_cast(values_cp); + + constexpr std::uint8_t elems_per_wi = n_vecs * vec_sz; + const std::size_t n_work_items_needed = nelems / elems_per_wi; + const std::size_t empirical_threshold = std::size_t(1) << 21; + const std::size_t lws = (n_work_items_needed <= empirical_threshold) + ? std::size_t(128) + : std::size_t(256); + + const std::size_t n_groups = + ((nelems + lws * elems_per_wi - 1) / (lws * elems_per_wi)); + const auto gws_range = sycl::range<1>(n_groups * lws); + const auto lws_range = sycl::range<1>(lws); + + using dpnp::tensor::kernels::alignment_utils::is_aligned; + using dpnp::tensor::kernels::alignment_utils::required_alignment; + + const bool aligned = is_aligned(dst_tp) && + is_aligned(mask_tp) && + is_aligned(values_tp); + + sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(depends); + + if (aligned) { + constexpr bool enable_sg = true; + using PutMaskFunc = + PutMaskContigFunctor; + + cgh.parallel_for( + sycl::nd_range<1>(gws_range, lws_range), + PutMaskFunc(dst_tp, mask_tp, values_tp, nelems, values_size)); + } + else { + constexpr bool enable_sg = false; + using PutMaskFunc = + PutMaskContigFunctor; + + cgh.parallel_for( + sycl::nd_range<1>(gws_range, lws_range), + PutMaskFunc(dst_tp, mask_tp, values_tp, nelems, values_size)); + } + }); + + return comp_ev; +} + +} // namespace dpnp::kernels::putmask diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 2f778fc5d15c..8a34a581c528 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -1,5 +1,5 @@ # ***************************************************************************** -# Copyright (c) 2016, Intel Corporation +# Copyright (c) 2026, Intel Corporation # All rights reserved. # # Redistribution and use in source and binary forms, with or without @@ -56,12 +56,8 @@ import dpnp.tensor as dpt import dpnp.tensor._tensor_impl as ti -# pylint: disable=no-name-in-module -from .dpnp_algo import ( - dpnp_putmask, -) from .dpnp_array import dpnp_array -from .dpnp_utils import call_origin, get_usm_allocations +from .dpnp_utils import get_usm_allocations from .exceptions import ExecutionPlacementError from .tensor._copy_utils import _nonzero_impl from .tensor._indexing_functions import _get_indexing_mode @@ -1807,30 +1803,102 @@ def put_along_axis(a, ind, values, axis, mode="wrap"): dpt.put_along_axis(usm_a, usm_ind, usm_vals, axis=axis, mode=mode) -def putmask(x1, mask, values): +def putmask(a, /, mask, values): """ Changes elements of an array based on conditional and input values. + Sets ``a.flat[n] = values[n]`` for each ``n`` where ``mask.flat[n]`` is + ``True``. + + If `values` is not the same size as `a` and `mask` then it will + repeat. This gives behavior different from ``a[mask] = values``. + For full documentation refer to :obj:`numpy.putmask`. + Parameters + ---------- + a : {dpnp.ndarray, usm_ndarray} + Target array. + mask : {dpnp.ndarray, usm_ndarray} + Boolean mask array. It has to be the same shape as `a`. + values : {dpnp.ndarray, usm_ndarray, scalar} + Values to put into `a` where `mask` is ``True``. If `values` is smaller + than `a` it will be repeated. + Limitations ----------- - Input arrays ``arr``, ``mask`` and ``values`` are supported - as :obj:`dpnp.ndarray`. + Input array ``a`` and ``mask`` are expected to have the same shape (unlike + :obj:`numpy.putmask`, which only requires the same size). + + See Also + -------- + :obj:`dpnp.place` : Change elements of an array based on conditional and + input values. + :obj:`dpnp.put` : Replaces specified elements of an array with given values. + :obj:`dpnp.take` : Take elements from an array along an axis. + :obj:`dpnp.copyto` : Copies values from one array to another. + + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(6).reshape(2, 3) + >>> np.putmask(x, x>2, x**2) + >>> x + array([[ 0, 1, 2], + [ 9, 16, 25]]) + + If `values` is smaller than `x1` it is repeated: + + >>> x = np.arange(5) + >>> np.putmask(x, x>1, np.array([-33, -44])) + >>> x + array([ 0, 1, -33, -44, -33]) """ - x1_desc = dpnp.get_dpnp_descriptor( - x1, copy_when_strides=False, copy_when_nondefault_queue=False - ) - mask_desc = dpnp.get_dpnp_descriptor(mask, copy_when_nondefault_queue=False) - values_desc = dpnp.get_dpnp_descriptor( - values, copy_when_nondefault_queue=False - ) - if x1_desc and mask_desc and values_desc: - return dpnp_putmask(x1_desc, mask_desc, values_desc) + dpnp.check_supported_arrays_type(a, mask) + dpnp.check_supported_arrays_type(values, scalar_type=True, all_scalars=True) + + if not a.shape == mask.shape: + raise ValueError("mask and data must be the same size") + + mask = dpnp.astype(mask, dpnp.bool, copy=False) + + if dpnp.isscalar(values): + a[mask] = values - return call_origin(numpy.putmask, x1, mask, values, dpnp_inplace=True) + elif not dpnp.can_cast(values.dtype, a.dtype): + raise TypeError( + f"Cannot cast array data from {values.dtype} to {a.dtype} " + "according to the rule 'safe'" + ) + + elif a.shape == values.shape: + a[mask] = values[mask] + + else: + # numpy putmask cycles values by the C-order flat index of the + # destination (values.flat[c % N]), independent of memory layout, so + # values must always be flattened in C-order. + values_1d = values.ravel(order="C") + if a.dtype != values_1d.dtype: + values_1d = dpnp.astype( + values_1d, a.dtype, casting="safe", copy=False + ) + _, exec_q = get_usm_allocations([a, mask, values_1d]) + + _manager = dpu.SequentialOrderManager[exec_q] + dep_evs = _manager.submitted_events + + h_ev, putmask_ev = indexing_ext._putmask( + a.get_array(), + mask.get_array(), + values_1d.get_array(), + exec_q, + dep_evs, + ) + _manager.add_event_pair(h_ev, putmask_ev) def ravel_multi_index(multi_index, dims, mode="raise", order="C"): diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py index 84bf62d03562..d6cabb09c857 100644 --- a/dpnp/tests/test_indexing.py +++ b/dpnp/tests/test_indexing.py @@ -1169,126 +1169,136 @@ def test_indices(dimension, dtype, sparse): assert_array_equal(Xnp, X) -@pytest.mark.parametrize( - "mask", - [ - [[True, False], [False, True]], - [[False, True], [True, False]], - [[False, False], [True, True]], - ], - ids=[ - "[[True, False], [False, True]]", - "[[False, True], [True, False]]", - "[[False, False], [True, True]]", - ], -) -@pytest.mark.parametrize( - "arr", - [[[0, 0], [0, 0]], [[1, 2], [1, 2]], [[1, 2], [3, 4]]], - ids=["[[0, 0], [0, 0]]", "[[1, 2], [1, 2]]", "[[1, 2], [3, 4]]"], -) -def test_putmask1(arr, mask): - a = numpy.array(arr) - ia = dpnp.array(a) - m = numpy.array(mask) - im = dpnp.array(m) - v = numpy.array([100, 200]) - iv = dpnp.array(v) - numpy.putmask(a, m, v) - dpnp.putmask(ia, im, iv) - assert_array_equal(a, ia) +class TestPutMask: + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("shape", [(7,), (2, 3), (4, 3, 2)]) + def test_same_shape_values(self, dt, shape): + a = generate_random_numpy_array(shape, dtype=dt) + mask = generate_random_numpy_array(shape, dtype=dpnp.bool) + vals = generate_random_numpy_array(shape, dtype=dt) + ia, imask, ivals = dpnp.array(a), dpnp.array(mask), dpnp.array(vals) + numpy.putmask(a, mask, vals) + dpnp.putmask(ia, imask, ivals) + assert_array_equal(ia, a) -@pytest.mark.parametrize( - "vals", - [ - [100, 200], - [100, 200, 300, 400, 500, 600], - [100, 200, 300, 400, 500, 600, 800, 900], - ], - ids=[ - "[100, 200]", - "[100, 200, 300, 400, 500, 600]", - "[100, 200, 300, 400, 500, 600, 800, 900]", - ], -) -@pytest.mark.parametrize( - "mask", - [ + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("shape", [(7,), (2, 3), (4, 3, 2)]) + @pytest.mark.parametrize("n_vals", [1, 3, 15]) + def test_broadcast_values(self, dt, order, shape, n_vals): + a = generate_random_numpy_array(shape, dtype=dt, order=order) + mask = generate_random_numpy_array(shape, dtype=dpnp.bool, order=order) + vals = generate_random_numpy_array((n_vals,), dtype=dt) + ia = dpnp.array(a, order=order) + imask = dpnp.array(mask, order=order) + ivals = dpnp.array(vals) + + numpy.putmask(a, mask, vals) + dpnp.putmask(ia, imask, ivals) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize( + "slice_spec", [ - [[True, False], [False, True]], - [[False, True], [True, False]], - [[False, False], [True, True]], - ] - ], - ids=[ - "[[[True, False], [False, True]], [[False, True], [True, False]], [[False, False], [True, True]]]" - ], -) -@pytest.mark.parametrize( - "arr", - [[[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]], - ids=["[[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]"], -) -def test_putmask2(arr, mask, vals): - a = numpy.array(arr) - ia = dpnp.array(a) - m = numpy.array(mask) - im = dpnp.array(m) - v = numpy.array(vals) - iv = dpnp.array(v) - numpy.putmask(a, m, v) - dpnp.putmask(ia, im, iv) - assert_array_equal(a, ia) + (slice(None), slice(None, None, 2)), + (slice(None, None, 2), slice(None)), + (slice(None, None, 2), slice(None, None, 2)), + ], + ) + def test_strided(self, dt, slice_spec): + a = generate_random_numpy_array((4, 6), dtype=dt) + ia = dpnp.array(a) + a, ia = a[slice_spec], ia[slice_spec] + mask = generate_random_numpy_array(a.shape, dtype=dpnp.bool) + vals = generate_random_numpy_array((5,), dtype=dt) + imask, ivals = dpnp.array(mask), dpnp.array(vals) + numpy.putmask(a, mask, vals) + dpnp.putmask(ia, imask, ivals) + assert_array_equal(ia, a) -@pytest.mark.parametrize( - "vals", - [ - [100, 200], - [100, 200, 300, 400, 500, 600], - [100, 200, 300, 400, 500, 600, 800, 900], - ], - ids=[ - "[100, 200]", - "[100, 200, 300, 400, 500, 600]", - "[100, 200, 300, 400, 500, 600, 800, 900]", - ], -) -@pytest.mark.parametrize( - "mask", - [ - [ - [[[False, False], [True, True]], [[True, True], [True, True]]], - [[[False, False], [True, True]], [[False, False], [False, False]]], - ] - ], - ids=[ - "[[[[False, False], [True, True]], [[True, True], [True, True]]], [[[False, False], [True, True]], [[False, False], [False, False]]]]" - ], -) -@pytest.mark.parametrize( - "arr", - [ - [ - [[[1, 2], [3, 4]], [[1, 2], [2, 1]]], - [[[1, 3], [3, 1]], [[0, 1], [1, 3]]], - ] - ], - ids=[ - "[[[[1, 2], [3, 4]], [[1, 2], [2, 1]]], [[[1, 3], [3, 1]], [[0, 1], [1, 3]]]]" - ], -) -def test_putmask3(arr, mask, vals): - a = numpy.array(arr) - ia = dpnp.array(a) - m = numpy.array(mask) - im = dpnp.array(m) - v = numpy.array(vals) - iv = dpnp.array(v) - numpy.putmask(a, m, v) - dpnp.putmask(ia, im, iv) - assert_array_equal(a, ia) + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + def test_transpose(self, dt): + a = generate_random_numpy_array((3, 4), dtype=dt) + ia = dpnp.array(a) + a, ia = a.T, ia.T + mask = generate_random_numpy_array(a.shape, dtype=dpnp.bool) + vals = generate_random_numpy_array((3,), dtype=dt) + imask, ivals = dpnp.array(mask), dpnp.array(vals) + + numpy.putmask(a, mask, vals) + dpnp.putmask(ia, imask, ivals) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + def test_scalar_values(self, dt): + a = generate_random_numpy_array((2, 3), dtype=dt) + mask = generate_random_numpy_array((2, 3), dtype=dpnp.bool) + ia, imask = dpnp.array(a), dpnp.array(mask) + + numpy.putmask(a, mask, 5) + dpnp.putmask(ia, imask, 5) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("mask_dt", get_integer_dtypes()) + def test_integer_mask(self, mask_dt): + a = numpy.array([1, 2, 3, 3]) + mask = numpy.array([0, 1, 0, 2], dtype=mask_dt) + ia, imask = dpnp.array(a), dpnp.array(mask) + + numpy.putmask(a, mask, 0) + dpnp.putmask(ia, imask, 0) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("order", ["C", "F"]) + def test_empty_values(self, order): + dt = dpnp.default_float_type() + a = generate_random_numpy_array((2, 3), dtype=dt, order=order) + mask = generate_random_numpy_array((2, 3), dtype=dpnp.bool, order=order) + ia = dpnp.array(a, order=order) + imask = dpnp.array(mask, order=order) + vals = numpy.array([], dtype=dt) + ivals = dpnp.array(vals) + + numpy.putmask(a, mask, vals) + dpnp.putmask(ia, imask, ivals) + assert_array_equal(ia, a) + + def test_empty_array(self): + a = numpy.array([], dtype=dpnp.default_float_type()) + mask = numpy.array([], dtype=dpnp.bool) + ia, imask = dpnp.array(a), dpnp.array(mask) + + numpy.putmask(a, mask, numpy.array([1, 2], dtype=a.dtype)) + dpnp.putmask(ia, imask, dpnp.array([1, 2], dtype=a.dtype)) + assert_array_equal(ia, a) + + def test_0d(self): + a = numpy.array(5, dtype=dpnp.default_float_type()) + mask = numpy.array(True) + ia, imask = dpnp.array(a), dpnp.array(mask) + + numpy.putmask(a, mask, numpy.array([7], dtype=a.dtype)) + dpnp.putmask(ia, imask, dpnp.array([7], dtype=a.dtype)) + assert_array_equal(ia, a) + + def test_errors(self): + ia = dpnp.arange(6, dtype="i4") + + # unsupported types for the array and the mask + assert_raises(TypeError, dpnp.putmask, dpnp.asnumpy(ia), ia > 2, 0) + assert_raises(TypeError, dpnp.putmask, ia, dpnp.asnumpy(ia) > 2, 0) + + # array and mask must have the same shape + assert_raises( + ValueError, dpnp.putmask, ia, dpnp.array([True, False]), 0 + ) + + # values cannot be safely cast to the array data type + vals = dpnp.arange(2, dtype="i8") + assert_raises(TypeError, dpnp.putmask, ia, ia > 2, vals) @pytest.mark.parametrize("m", [None, 0, 1, 2, 3, 4]) diff --git a/dpnp/tests/third_party/cupy/indexing_tests/test_insert.py b/dpnp/tests/third_party/cupy/indexing_tests/test_insert.py index 3b23b32fe3b2..033196e3a105 100644 --- a/dpnp/tests/third_party/cupy/indexing_tests/test_insert.py +++ b/dpnp/tests/third_party/cupy/indexing_tests/test_insert.py @@ -231,7 +231,6 @@ def test_putmask_int_mask_scalar_values(self, xp): class TestPutmaskDifferentDtypes(unittest.TestCase): - @pytest.mark.skip("putmask() is not fully supported") @testing.for_all_dtypes_combination(names=["a_dtype", "val_dtype"]) def test_putmask_differnt_dtypes_raises(self, a_dtype, val_dtype): shape = (2, 3)