Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions dpctl/tensor/_dlpack.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ cdef extern from "dlpack/dlpack.h" nogil:
int device_WebGPU "kDLWebGPU"
int device_Hexagon "kDLHexagon"
int device_MAIA "kDLMAIA"
int device_Trn "kDLTrn"

cpdef object to_dlpack_capsule(usm_ndarray array) except +
cpdef object to_dlpack_versioned_capsule(
Expand Down
12 changes: 12 additions & 0 deletions dpctl/tensor/_dlpack.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ cdef extern from "dlpack/dlpack.h" nogil:
kDLWebGPU
kDLHexagon
kDLMAIA
kDLTrn

ctypedef struct DLDevice:
DLDeviceType device_type
Expand All @@ -88,6 +89,17 @@ cdef extern from "dlpack/dlpack.h" nogil:
kDLBfloat
kDLComplex
kDLBool
kDLFloat8_e3m4
kDLFloat8_e4m3
kDLFloat8_e4m3b11fnuz
kDLFloat8_e4m3fn
kDLFloat8_e4m3fnuz
kDLFloat8_e5m2
kDLFloat8_e5m2fnuz
kDLFloat8_e8m0fnu
kDLFloat6_e2m3fn
kDLFloat6_e3m2fn
kDLFloat4_e2m1fn

ctypedef struct DLDataType:
uint8_t code
Expand Down
3 changes: 3 additions & 0 deletions dpctl/tensor/_usmarray.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ class DLDeviceType(IntEnum):
Qualcomm Hexagon DSP
``kDLMAIA``:
Microsoft MAIA device
``kDLTrn``:
AWS Trainium device
"""
kDLCPU = c_dlpack.device_CPU
kDLCUDA = c_dlpack.device_CUDA
Expand All @@ -101,6 +103,7 @@ class DLDeviceType(IntEnum):
kDLWebGPU = c_dlpack.device_WebGPU
kDLHexagon = c_dlpack.device_Hexagon
kDLMAIA = c_dlpack.device_MAIA
kDLTrn = c_dlpack.device_Trn


cdef class InternalUSMArrayError(Exception):
Expand Down
Loading
Loading