Skip to content

Commit 5989277

Browse files
Added support for stream other than None
If `stream` keyword argument has type `dpctl.SyclQueue`, then submit a barrier into `self.sycl_queue`, and than submit another barrier into stream with dependency on the event returned by submission of the barrier into `self.sycl_queue`: ```python ev = self.sycl_queue.submit_barrier() stream.submit_barrier(dependent_events=[ev]) ``` All other types of stream are ignored, in particular integer values of stream.
1 parent 3d79ef9 commit 5989277

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

dpctl/tensor/_usmarray.pyx

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -779,13 +779,14 @@ cdef class usm_ndarray:
779779
NotImplementedError: when non-default value of `stream` keyword
780780
is used.
781781
"""
782-
if stream is None:
783-
return c_dlpack.to_dlpack_capsule(self)
782+
_caps = c_dlpack.to_dlpack_capsule(self)
783+
if (stream is None or type(stream) is not dpctl.SyclQueue or
784+
stream == self.sycl_queue):
785+
pass
784786
else:
785-
raise NotImplementedError(
786-
"Only stream=None is supported. "
787-
"Use `dpctl.SyclQueue.submit_barrier` to synchronize queues."
788-
)
787+
ev = self.sycl_queue.submit_barrier()
788+
stream.submit_barrier(dependent_events=[ev])
789+
return _caps
789790

790791
def __dlpack_device__(self):
791792
"""

0 commit comments

Comments
 (0)