In order to resolve gh-2156 move definition of the mask_positions

and _cumsum_1d functions to _tensor_accumulations_impl

Changed Python scripts accordingly, as well as CMake scripts to
add implementation cpp file to the list of source files for the
_tensor_accumulations_impl MODULE library.

Also moved find_package(Python) to find Module.Development component
before pybind11 is being activated to resolve CMake warning.

Incidentally, this change also results in reduced binary size and improved
compilation tiles, since accumulation kernels are not being generated
in duplicates (once for _tensor_ctor module, and once for _tensor_accumulation_impl
module).

Author: sycloid

CMakeLists.txt

@@ -114,6 +114,9 @@ install(DIRECTORY
     FILES_MATCHING REGEX "\\.h(pp)?$"
 )
 
+# find Python before enabling pybind11
+find_package(Python REQUIRED COMPONENTS Development.Module)
+
 # Define CMAKE_INSTALL_xxx: LIBDIR, INCLUDEDIR
 include(GNUInstallDirs)
 

dpctl/CMakeLists.txt

@@ -1,4 +1,4 @@
-find_package(Python REQUIRED COMPONENTS Development.Module NumPy)
+find_package(Python REQUIRED COMPONENTS NumPy)
 
 # -t is to only Cythonize sources with timestamps newer than existing CXX files (if present)
 # -w is to set working directory (and correctly set __pyx_f[] array of filenames)

dpctl/tensor/CMakeLists.txt

@@ -171,6 +171,7 @@ set(_accumulator_sources
 )
 set(_tensor_accumulation_impl_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_accumulation.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators.cpp
     ${_accumulator_sources}
 )
 

dpctl/tensor/_copy_utils.py

@@ -23,6 +23,7 @@
 import dpctl.memory as dpm
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
+from dpctl.tensor._tensor_accumulation_impl import mask_positions
 import dpctl.utils
 from dpctl.tensor._data_types import _get_dtype
 from dpctl.tensor._device import normalize_queue_device
@@ -792,7 +793,7 @@ def _extract_impl(ary, ary_mask, axis=0):
     exec_q = cumsum.sycl_queue
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
-    mask_count = ti.mask_positions(
+    mask_count = mask_positions(
         ary_mask, cumsum, sycl_queue=exec_q, depends=dep_evs
     )
     dst_shape = ary.shape[:pp] + (mask_count,) + ary.shape[pp + mask_nd :]
@@ -828,7 +829,7 @@ def _nonzero_impl(ary):
     )
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
-    mask_count = ti.mask_positions(
+    mask_count = mask_positions(
         ary, cumsum, sycl_queue=exec_q, depends=dep_evs
     )
     indexes_dt = ti.default_device_index_type(exec_q.sycl_device)
@@ -1050,7 +1051,7 @@ def _place_impl(ary, ary_mask, vals, axis=0):
     exec_q = cumsum.sycl_queue
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_ev = _manager.submitted_events
-    mask_count = ti.mask_positions(
+    mask_count = mask_positions(
         ary_mask, cumsum, sycl_queue=exec_q, depends=dep_ev
     )
     expected_vals_shape = (

dpctl/tensor/_indexing_functions.py

@@ -19,6 +19,7 @@
 import dpctl
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
+from dpctl.tensor._tensor_accumulation_impl import mask_positions
 import dpctl.utils
 
 from ._copy_utils import (
@@ -413,7 +414,7 @@ def place(arr, mask, vals):
     cumsum = dpt.empty(mask.size, dtype="i8", sycl_queue=exec_q)
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
-    nz_count = ti.mask_positions(
+    nz_count = mask_positions(
         mask, cumsum, sycl_queue=exec_q, depends=deps_ev
     )
     if nz_count == 0:

dpctl/tensor/_manipulation_functions.py

@@ -23,6 +23,7 @@
 import dpctl
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
+from dpctl.tensor._tensor_accumulation_impl import _cumsum_1d
 import dpctl.utils as dputils
 
 from ._copy_utils import _broadcast_strides
@@ -908,7 +909,7 @@ def repeat(x, repeats, /, *, axis=None):
                 sycl_queue=exec_q,
             )
             # _cumsum_1d synchronizes so `depends` ends here safely
-            res_axis_size = ti._cumsum_1d(
+            res_axis_size = _cumsum_1d(
                 rep_buf, cumsum, sycl_queue=exec_q, depends=[copy_ev]
             )
             if axis is not None:
@@ -940,7 +941,7 @@ def repeat(x, repeats, /, *, axis=None):
                 usm_type=usm_type,
                 sycl_queue=exec_q,
             )
-            res_axis_size = ti._cumsum_1d(
+            res_axis_size = _cumsum_1d(
                 repeats, cumsum, sycl_queue=exec_q, depends=dep_evs
             )
             if axis is not None:

dpctl/tensor/_set_functions.py

@@ -34,8 +34,8 @@
     _linspace_step,
     _take,
     default_device_index_type,
-    mask_positions,
 )
+from ._tensor_accumulation_impl import mask_positions
 from ._tensor_sorting_impl import (
     _argsort_ascending,
     _isin,

dpctl/tensor/libtensor/source/tensor_accumulation.cpp

@@ -24,12 +24,30 @@
 //===----------------------------------------------------------------------===//
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 
 #include "accumulators/accumulators_common.hpp"
+#include "accumulators.hpp"
 
 namespace py = pybind11;
 
+namespace py_int = dpctl::tensor::py_internal;
+
+using py_int::py_mask_positions;
+using py_int::py_cumsum_1d;
+
 PYBIND11_MODULE(_tensor_accumulation_impl, m)
 {
+    py_int::populate_mask_positions_dispatch_vectors();
+    py_int::populate_cumsum_1d_dispatch_vectors();
+
     dpctl::tensor::py_internal::init_accumulator_functions(m);
+
+    m.def("mask_positions", &py_mask_positions, "", py::arg("mask"),
+          py::arg("cumsum"), py::arg("sycl_queue"),
+          py::arg("depends") = py::list());
+
+    m.def("_cumsum_1d", &py_cumsum_1d, "", py::arg("src"), py::arg("cumsum"),
+          py::arg("sycl_queue"), py::arg("depends") = py::list());
+
 }

dpctl/tensor/libtensor/source/tensor_ctors.cpp

@@ -105,12 +105,12 @@ using dpctl::tensor::py_internal::usm_ndarray_put;
 using dpctl::tensor::py_internal::usm_ndarray_take;
 
 using dpctl::tensor::py_internal::py_extract;
-using dpctl::tensor::py_internal::py_mask_positions;
+// using dpctl::tensor::py_internal::py_mask_positions;
 using dpctl::tensor::py_internal::py_nonzero;
 using dpctl::tensor::py_internal::py_place;
 
 /* ================= Repeat ====================*/
-using dpctl::tensor::py_internal::py_cumsum_1d;
+// using dpctl::tensor::py_internal::py_cumsum_1d;
 using dpctl::tensor::py_internal::py_repeat_by_scalar;
 using dpctl::tensor::py_internal::py_repeat_by_sequence;
 
@@ -158,9 +158,6 @@ void init_dispatch_vectors(void)
     populate_masked_extract_dispatch_vectors();
     populate_masked_place_dispatch_vectors();
 
-    populate_mask_positions_dispatch_vectors();
-
-    populate_cumsum_1d_dispatch_vectors();
     init_repeat_dispatch_vectors();
 
     init_clip_dispatch_vectors();
@@ -402,13 +399,6 @@ PYBIND11_MODULE(_tensor_impl, m)
           py::arg("dst"), py::arg("k") = 0, py::arg("sycl_queue"),
           py::arg("depends") = py::list());
 
-    m.def("mask_positions", &py_mask_positions, "", py::arg("mask"),
-          py::arg("cumsum"), py::arg("sycl_queue"),
-          py::arg("depends") = py::list());
-
-    m.def("_cumsum_1d", &py_cumsum_1d, "", py::arg("src"), py::arg("cumsum"),
-          py::arg("sycl_queue"), py::arg("depends") = py::list());
-
     m.def("_extract", &py_extract, "", py::arg("src"), py::arg("cumsum"),
           py::arg("axis_start"), py::arg("axis_end"), py::arg("dst"),
           py::arg("sycl_queue"), py::arg("depends") = py::list());