Skip to content

Commit 330746b

Browse files
Adding test that DLPack can support sharing on sub-device
1 parent e3d0838 commit 330746b

File tree

1 file changed

+32
-0
lines changed

1 file changed

+32
-0
lines changed

dpctl/tests/test_usm_ndarray_dlpack.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,3 +197,35 @@ def test_from_dlpack_fortran_contig_array_roundtripping():
197197

198198
assert dpt.all(dpt.equal(ar2d_f, ar2d_r))
199199
assert dpt.all(dpt.equal(ar2d_c, ar2d_r))
200+
201+
202+
def test_dlpack_from_subdevice():
203+
"""
204+
This test checks that array allocated on a sub-device,
205+
with memory bound to platform-default SyclContext can be
206+
exported and imported via DLPack.
207+
"""
208+
n = 64
209+
try:
210+
dev = dpctl.SyclDevice()
211+
except dpctl.SyclDeviceCreationError:
212+
pytest.skip("No default device available")
213+
try:
214+
sdevs = dev.create_sub_devices(partition="next_partitionable")
215+
except dpctl.SyclSubDeviceCreationError:
216+
sdevs = None
217+
try:
218+
sdevs = (
219+
dev.create_sub_devices(partition=[1, 1]) if sdevs is None else sdevs
220+
)
221+
except dpctl.SyclSubDeviceCreationError:
222+
pytest.skip("Default device can not be partitioned")
223+
assert isinstance(sdevs, list) and len(sdevs) > 0
224+
try:
225+
q = dpctl.SyclQueue(sdevs[0].sycl_platform.default_context, sdevs[0])
226+
except dpctl.SyclQueueCreationError:
227+
pytest.skip("Default device can not be partitioned")
228+
229+
ar = dpt.arange(n, dtype=dpt.int32, sycl_queue=q)
230+
ar2 = dpt.from_dlpack(ar)
231+
assert ar2.sycl_device == sdevs[0]

0 commit comments

Comments
 (0)