Skip to content

Commit

Permalink
ENH: Add support for __cuda_array_interface__
Browse files Browse the repository at this point in the history
Adapt ITK's PyBuffer approach to inject __cuda_array_interface__ in wrapped
code.
Closes RTKConsortium#10
  • Loading branch information
LucasGandel committed Aug 9, 2024
1 parent 2e0f434 commit b6ac52f
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 3 deletions.
17 changes: 17 additions & 0 deletions wrapping/CudaImage.i.in
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
%extend itkCudaImage@CudaImageTypes@{
%pythoncode %{
@property
def __cuda_array_interface__(self):
_pixelType = "@PixelType@"
_typestr = _get_type_string(_pixelType)
return {
'shape': (self.GetLargestPossibleRegion().GetSize()[0], self.GetLargestPossibleRegion().GetSize()[1], self.GetLargestPossibleRegion().GetSize()[2]),
'data': (int(self.GetCudaDataManager().GetGPUBufferPointer()), False),
'typestr': _typestr,
'descr': [('', _typestr)],
'version': 3,
'stream': None,
'strides': None
}
%}
};
37 changes: 37 additions & 0 deletions wrapping/CudaImage.i.init
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
%pythoncode %{
def _get_type_string(itk_Image_type):
"""Returns the type string of the ITK PixelType as defined in the NumPy array interface."""

# This is a Mapping from system byte order to typestr byte order.
_byteorder_typestr = {
"big":">",
"little":"<",
}
import sys
_byteorder = _byteorder_typestr[sys.byteorder]

# This is a Mapping from itk pixel types to typestr.
_itk_typestr = {
"UC":"|u1",
"US":_byteorder + "u2",
"UI":_byteorder + "u4",
"UL":_byteorder + "u8",
"ULL":_byteorder + "u8",
"SC":"|i1",
"SS":_byteorder + "i2",
"SI":_byteorder + "i4",
"SL":_byteorder + "i8",
"SLL":_byteorder + "i8",
"F":_byteorder + "f4",
"D":_byteorder + "f8",
}
import os
if os.name == 'nt':
_itk_typestr['UL'] = _byteorder + "u4"
_itk_typestr['SL'] = _byteorder + "i4"

try:
return _itk_typestr[itk_Image_type]
except KeyError as e:
raise e
%}
30 changes: 27 additions & 3 deletions wrapping/itkCudaImage.wrap
Original file line number Diff line number Diff line change
@@ -1,9 +1,15 @@
itk_wrap_class("itk::CudaImage" POINTER_WITH_CONST_POINTER)
configure_file("${CMAKE_CURRENT_SOURCE_DIR}/CudaImage.i.init" "${CMAKE_CURRENT_BINARY_DIR}/CudaImage.i" @ONLY)

itk_wrap_class("itk::CudaImage" POINTER_WITH_CONST_POINTER)
UNIQUE(types "UC;UL;${ITKM_IT};${WRAP_ITK_SCALAR}")
foreach(d ${ITK_WRAP_IMAGE_DIMS})
foreach(t ${types})
foreach(t ${types})
set(PixelType ${t})
foreach(d ${ITK_WRAP_IMAGE_DIMS})
itk_wrap_template("${t}${d}" "${ITKT_${t}}, ${d}")
set(CudaImageTypes ${t}${d})
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/CudaImage.i.in ${CMAKE_CURRENT_BINARY_DIR}/CudaImage.i.temp @ONLY)
file(READ ${CMAKE_CURRENT_BINARY_DIR}/CudaImage.i.temp CudaImageInterfaceTemp)
file(APPEND ${CMAKE_CURRENT_BINARY_DIR}/CudaImage.i ${CudaImageInterfaceTemp})
endforeach()
endforeach()

Expand All @@ -17,3 +23,21 @@ itk_wrap_class("itk::CudaImage" POINTER_WITH_CONST_POINTER)
endforeach()

itk_end_wrap_class()

# Add library files to be included at a submodule level and copy them into
# ITK's wrapping typedef directory.
# Another approach is to add CudaImage.i to the WRAPPER_SWIG_LIBRARY_FILES list
# but then the %pythoncode from CudaImage.i.init gets only included in
# itkCudaDataManagerPython.py even if the WRAPPER_SUBMODULE_ORDER is set.
# Prefer using ITK_WRAP_PYTHON_SWIG_EXT to make sure the block is included in
# the right file exclusively.
set(ITK_WRAP_PYTHON_SWIG_EXT
"%include CudaImage.i\n${ITK_WRAP_PYTHON_SWIG_EXT}")

file(COPY "${CMAKE_CURRENT_BINARY_DIR}/CudaImage.i"
DESTINATION "${WRAPPER_MASTER_INDEX_OUTPUT_DIR}")

# Make sure to rebuild the python file when CudaImage.i is modified.
# Touching CudaImage.i directly does not force a rebuild because it is just
# appended to the ITK_WRAP_PYTHON_SWIG_EXT variable
file(TOUCH ${WRAPPER_MASTER_INDEX_OUTPUT_DIR}/itkCudaImage.i)

0 comments on commit b6ac52f

Please sign in to comment.