Skip to content

Commit 03c4822

Browse files
committed
_take and _put returned event changes
- Host_tasks now collected and used as dependencies for dec_ref of py arguments - Return temporaries deallocation event to further prevent dangling host_tasks
1 parent db84c42 commit 03c4822

File tree

1 file changed

+10
-7
lines changed

1 file changed

+10
-7
lines changed

dpctl/tensor/libtensor/source/integer_advanced_indexing.cpp

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -548,10 +548,12 @@ usm_ndarray_take(dpctl::tensor::usm_ndarray src,
548548
});
549549
});
550550

551-
sycl::event host_task_ev = keep_args_alive(
552-
exec_q, {src, py_ind, dst}, {take_generic_ev, temporaries_cleanup_ev});
551+
host_task_events.push_back(temporaries_cleanup_ev);
553552

554-
return std::make_pair(host_task_ev, take_generic_ev);
553+
sycl::event arg_cleanup_ev =
554+
keep_args_alive(exec_q, {src, py_ind, dst}, host_task_events);
555+
556+
return std::make_pair(arg_cleanup_ev, temporaries_cleanup_ev);
555557
}
556558

557559
std::pair<sycl::event, sycl::event>
@@ -857,7 +859,6 @@ usm_ndarray_put(dpctl::tensor::usm_ndarray dst,
857859
dst_offset, val_offset, packed_ind_offsets, all_deps);
858860

859861
// free packed temporaries
860-
861862
sycl::event temporaries_cleanup_ev = exec_q.submit([&](sycl::handler &cgh) {
862863
cgh.depends_on(put_generic_ev);
863864
auto ctx = exec_q.get_context();
@@ -872,10 +873,12 @@ usm_ndarray_put(dpctl::tensor::usm_ndarray dst,
872873
});
873874
});
874875

875-
sycl::event py_obj_cleanup_ev = keep_args_alive(
876-
exec_q, {dst, py_ind, val}, {put_generic_ev, temporaries_cleanup_ev});
876+
host_task_events.push_back(temporaries_cleanup_ev);
877+
878+
sycl::event arg_cleanup_ev =
879+
keep_args_alive(exec_q, {dst, py_ind, val}, host_task_events);
877880

878-
return std::make_pair(temporaries_cleanup_ev, put_generic_ev);
881+
return std::make_pair(arg_cleanup_ev, temporaries_cleanup_ev);
879882
}
880883

881884
void init_advanced_indexing_dispatch_tables(void)

0 commit comments

Comments
 (0)