Skip to content

Commit 4339cf3

Browse files
oleksandr-pavlykoleksandr-pavlyk
authored
Merge pull request #1300 from IntelPython/boolean-indexing-performance
Boolean indexing performance
2 parents dcca136 + 7ade829 commit 4339cf3

4 files changed

Lines changed: 287 additions & 93 deletions

File tree

‎dpctl/tensor/_copy_utils.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
":class:`dpctl.tensor.usm_ndarray`."
3232
)
3333

34+
int32_t_max = 2147483648
35+
3436

3537
def _copy_to_numpy(ary):
3638
if not isinstance(ary, dpt.usm_ndarray):
@@ -482,7 +484,8 @@ def _extract_impl(ary, ary_mask, axis=0):
482484
"Parameter p is inconsistent with input array dimensions"
483485
)
484486
mask_nelems = ary_mask.size
485-
cumsum = dpt.empty(mask_nelems, dtype=dpt.int64, device=ary_mask.device)
487+
cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
488+
cumsum = dpt.empty(mask_nelems, dtype=cumsum_dt, device=ary_mask.device)
486489
exec_q = cumsum.sycl_queue
487490
mask_count = ti.mask_positions(ary_mask, cumsum, sycl_queue=exec_q)
488491
dst_shape = ary.shape[:pp] + (mask_count,) + ary.shape[pp + mask_nd :]
@@ -509,8 +512,9 @@ def _nonzero_impl(ary):
509512
exec_q = ary.sycl_queue
510513
usm_type = ary.usm_type
511514
mask_nelems = ary.size
515+
cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
512516
cumsum = dpt.empty(
513-
mask_nelems, dtype=dpt.int64, sycl_queue=exec_q, order="C"
517+
mask_nelems, dtype=cumsum_dt, sycl_queue=exec_q, order="C"
514518
)
515519
mask_count = ti.mask_positions(ary, cumsum, sycl_queue=exec_q)
516520
indexes = dpt.empty(
@@ -604,7 +608,8 @@ def _place_impl(ary, ary_mask, vals, axis=0):
604608
"Parameter p is inconsistent with input array dimensions"
605609
)
606610
mask_nelems = ary_mask.size
607-
cumsum = dpt.empty(mask_nelems, dtype=dpt.int64, device=ary_mask.device)
611+
cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
612+
cumsum = dpt.empty(mask_nelems, dtype=cumsum_dt, device=ary_mask.device)
608613
exec_q = cumsum.sycl_queue
609614
mask_count = ti.mask_positions(ary_mask, cumsum, sycl_queue=exec_q)
610615
expected_vals_shape = (

‎dpctl/tensor/libtensor/include/kernels/boolean_advanced_indexing.hpp‎

Lines changed: 102 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -393,15 +393,24 @@ size_t mask_positions_contig_impl(sycl::queue q,
393393
throw std::bad_alloc();
394394
}
395395
sycl::event copy_e =
396-
q.copy<std::int64_t>(last_elem, last_elem_host_usm, 1, {comp_ev});
396+
q.copy<cumsumT>(last_elem, last_elem_host_usm, 1, {comp_ev});
397397
copy_e.wait();
398398
size_t return_val = static_cast<size_t>(*last_elem_host_usm);
399399
sycl::free(last_elem_host_usm, q);
400400

401401
return return_val;
402402
}
403403

404-
template <typename fnT, typename T> struct MaskPositionsContigFactory
404+
template <typename fnT, typename T> struct MaskPositionsContigFactoryForInt32
405+
{
406+
fnT get()
407+
{
408+
fnT fn = mask_positions_contig_impl<T, std::int32_t>;
409+
return fn;
410+
}
411+
};
412+
413+
template <typename fnT, typename T> struct MaskPositionsContigFactoryForInt64
405414
{
406415
fnT get()
407416
{
@@ -452,15 +461,24 @@ size_t mask_positions_strided_impl(sycl::queue q,
452461
throw std::bad_alloc();
453462
}
454463
sycl::event copy_e =
455-
q.copy<std::int64_t>(last_elem, last_elem_host_usm, 1, {comp_ev});
464+
q.copy<cumsumT>(last_elem, last_elem_host_usm, 1, {comp_ev});
456465
copy_e.wait();
457466
size_t return_val = static_cast<size_t>(*last_elem_host_usm);
458467
sycl::free(last_elem_host_usm, q);
459468

460469
return return_val;
461470
}
462471

463-
template <typename fnT, typename T> struct MaskPositionsStridedFactory
472+
template <typename fnT, typename T> struct MaskPositionsStridedFactoryForInt32
473+
{
474+
fnT get()
475+
{
476+
fnT fn = mask_positions_strided_impl<T, std::int32_t>;
477+
return fn;
478+
}
479+
};
480+
481+
template <typename fnT, typename T> struct MaskPositionsStridedFactoryForInt64
464482
{
465483
fnT get()
466484
{
@@ -611,7 +629,18 @@ sycl::event masked_extract_some_slices_strided_impl(
611629
return comp_ev;
612630
}
613631

614-
template <typename fnT, typename T> struct MaskExtractAllSlicesStridedFactory
632+
template <typename fnT, typename T>
633+
struct MaskExtractAllSlicesStridedFactoryForInt32
634+
{
635+
fnT get()
636+
{
637+
fnT fn = masked_extract_all_slices_strided_impl<T, std::int32_t>;
638+
return fn;
639+
}
640+
};
641+
642+
template <typename fnT, typename T>
643+
struct MaskExtractAllSlicesStridedFactoryForInt64
615644
{
616645
fnT get()
617646
{
@@ -620,7 +649,18 @@ template <typename fnT, typename T> struct MaskExtractAllSlicesStridedFactory
620649
}
621650
};
622651

623-
template <typename fnT, typename T> struct MaskExtractSomeSlicesStridedFactory
652+
template <typename fnT, typename T>
653+
struct MaskExtractSomeSlicesStridedFactoryForInt32
654+
{
655+
fnT get()
656+
{
657+
fnT fn = masked_extract_some_slices_strided_impl<T, std::int32_t>;
658+
return fn;
659+
}
660+
};
661+
662+
template <typename fnT, typename T>
663+
struct MaskExtractSomeSlicesStridedFactoryForInt64
624664
{
625665
fnT get()
626666
{
@@ -763,7 +803,18 @@ sycl::event masked_place_some_slices_strided_impl(
763803
return comp_ev;
764804
}
765805

766-
template <typename fnT, typename T> struct MaskPlaceAllSlicesStridedFactory
806+
template <typename fnT, typename T>
807+
struct MaskPlaceAllSlicesStridedFactoryForInt32
808+
{
809+
fnT get()
810+
{
811+
fnT fn = masked_place_all_slices_strided_impl<T, std::int32_t>;
812+
return fn;
813+
}
814+
};
815+
816+
template <typename fnT, typename T>
817+
struct MaskPlaceAllSlicesStridedFactoryForInt64
767818
{
768819
fnT get()
769820
{
@@ -772,7 +823,18 @@ template <typename fnT, typename T> struct MaskPlaceAllSlicesStridedFactory
772823
}
773824
};
774825

775-
template <typename fnT, typename T> struct MaskPlaceSomeSlicesStridedFactory
826+
template <typename fnT, typename T>
827+
struct MaskPlaceSomeSlicesStridedFactoryForInt32
828+
{
829+
fnT get()
830+
{
831+
fnT fn = masked_place_some_slices_strided_impl<T, std::int32_t>;
832+
return fn;
833+
}
834+
};
835+
836+
template <typename fnT, typename T>
837+
struct MaskPlaceSomeSlicesStridedFactoryForInt64
776838
{
777839
fnT get()
778840
{
@@ -783,7 +845,17 @@ template <typename fnT, typename T> struct MaskPlaceSomeSlicesStridedFactory
783845

784846
// Non-zero
785847

786-
class non_zero_indexes_krn;
848+
template <typename T1, typename T2> class non_zero_indexes_krn;
849+
850+
typedef sycl::event (*non_zero_indexes_fn_ptr_t)(
851+
sycl::queue,
852+
py::ssize_t,
853+
py::ssize_t,
854+
int,
855+
const char *,
856+
char *,
857+
const py::ssize_t *,
858+
std::vector<sycl::event> const &);
787859

788860
template <typename indT1, typename indT2>
789861
sycl::event non_zero_indexes_impl(sycl::queue exec_q,
@@ -800,28 +872,29 @@ sycl::event non_zero_indexes_impl(sycl::queue exec_q,
800872

801873
sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) {
802874
cgh.depends_on(depends);
803-
cgh.parallel_for<class non_zero_indexes_krn>(
804-
sycl::range<1>(iter_size), [=](sycl::id<1> idx) {
805-
auto i = idx[0];
806-
807-
auto cs_curr_val = cumsum_data[i] - 1;
808-
auto cs_prev_val = (i > 0) ? cumsum_data[i - 1] : indT1(0);
809-
bool cond = (cs_curr_val == cs_prev_val);
810-
811-
py::ssize_t i_ = static_cast<py::ssize_t>(i);
812-
for (int dim = nd; --dim > 0;) {
813-
auto sd = mask_shape[dim];
814-
py::ssize_t q = i_ / sd;
815-
py::ssize_t r = (i_ - q * sd);
816-
if (cond) {
817-
indexes_data[cs_curr_val + dim * nz_elems] =
818-
static_cast<indT2>(r);
819-
}
820-
i_ = q;
821-
}
875+
cgh.parallel_for<class non_zero_indexes_krn<indT1, indT2>>(
876+
sycl::range<1>(iter_size), [=](sycl::id<1> idx)
877+
{
878+
auto i = idx[0];
879+
880+
auto cs_curr_val = cumsum_data[i] - 1;
881+
auto cs_prev_val = (i > 0) ? cumsum_data[i - 1] : indT1(0);
882+
bool cond = (cs_curr_val == cs_prev_val);
883+
884+
py::ssize_t i_ = static_cast<py::ssize_t>(i);
885+
for (int dim = nd; --dim > 0;) {
886+
auto sd = mask_shape[dim];
887+
py::ssize_t q = i_ / sd;
888+
py::ssize_t r = (i_ - q * sd);
822889
if (cond) {
823-
indexes_data[cs_curr_val] = static_cast<indT2>(i_);
890+
indexes_data[cs_curr_val + dim * nz_elems] =
891+
static_cast<indT2>(r);
824892
}
893+
i_ = q;
894+
}
895+
if (cond) {
896+
indexes_data[cs_curr_val] = static_cast<indT2>(i_);
897+
}
825898
});
826899
});
827900

0 commit comments

Comments
 (0)