Cuda array interface

CMakeLists.txt

@@ -14,6 +14,19 @@ endforeach()
 
 project(CudaCommon)
 
+## CudaCommon Version
+set(CudaCommon_VERSION_MAJOR "2")
+set(CudaCommon_VERSION_MINOR "0")
+set(CudaCommon_VERSION_PATCH "0")
+set(CudaCommon_VERSION_STRING "${CudaCommon_VERSION_MAJOR}.${CudaCommon_VERSION_MINOR}")
+
+#=========================================================
+# Installation variables
+#=========================================================
+if(NOT CudaCommon_INSTALL_INCLUDE_DIR)
+  set(CudaCommon_INSTALL_INCLUDE_DIR include/CudaCommon)
+endif()
+
 if(NOT ITK_SOURCE_DIR)
   include(itk-module-init.cmake)
 endif()
@@ -27,6 +40,23 @@ if(NOT ITK_SOURCE_DIR)
   find_package(ITK 5.2 REQUIRED)
 endif()
 
+# Propagate cmake options in a header file
+configure_file(${CudaCommon_SOURCE_DIR}/cudaCommonConfiguration.h.in
+  ${CudaCommon_BINARY_DIR}/cudaCommonConfiguration.h)
+
+list(APPEND CudaCommon_INCLUDE_DIRS ${CudaCommon_BINARY_DIR})
+
+set(CudaCommon_EXPORT_CODE_BUILD "
+# Pass source dir to allow other modules, e.g. RTK, to use CudaImage.i.init and CudaImage.i.in
+set(CudaCommon_SOURCE_DIR ${CudaCommon_SOURCE_DIR})
+
+# Also pass CudaCommon version
+set(CudaCommon_VERSION_MAJOR ${CudaCommon_VERSION_MAJOR})
+set(CudaCommon_VERSION_MINOR ${CudaCommon_VERSION_MINOR})
+set(CudaCommon_VERSION_PATCH ${CudaCommon_VERSION_PATCH})
+set(CudaCommon_VERSION_STRING ${CudaCommon_VERSION_STRING})
+")
+
 if(NOT ITK_SOURCE_DIR)
   if(NOT EXISTS ${ITK_CMAKE_DIR}/ITKModuleMacros.cmake)
     message(FATAL_ERROR "Modules can only be built against an ITK build tree; they cannot be built against an ITK install tree.")
@@ -38,3 +68,6 @@ else()
   set(ITK_DIR ${CMAKE_BINARY_DIR})
   itk_module_impl()
 endif()
+
+# Install configuration file
+install(FILES ${CudaCommon_BINARY_DIR}/cudaCommonConfiguration.h DESTINATION ${CudaCommon_INSTALL_INCLUDE_DIR})

cudaCommonConfiguration.h.in

@@ -0,0 +1,26 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *=========================================================================*/
+#ifndef __cudaCommonConfiguration_h
+#define __cudaCommonConfiguration_h
+
+#define CudaCommon_VERSION_MAJOR @CudaCommon_VERSION_MAJOR@
+#define CudaCommon_VERSION_MINOR @CudaCommon_VERSION_MINOR@
+#define CudaCommon_VERSION_PATCH @CudaCommon_VERSION_PATCH@
+#define CudaCommon_VERSION_STRING "@CudaCommon_VERSION_STRING@"
+
+#endif

include/itkCudaDataManager.h

@@ -162,6 +162,12 @@ class CudaCommon_EXPORT CudaDataManager : public Object
   void
   SetGPUDirtyFlag(bool isDirty);
 
+  /** Controls whether GPU memory should be released when dirty. On by default.
+   * When turning it off, one must call Free() to release the GPU memory. */
+  itkGetConstMacro(ReleaseDirtyGPUBuffer, bool);
+  itkSetMacro(ReleaseDirtyGPUBuffer, bool);
+  itkBooleanMacro(ReleaseDirtyGPUBuffer);
+
   /** Make GPU up-to-date and mark CPU as dirty.
    * Call this function when you want to modify CPU data */
   void
@@ -213,6 +219,10 @@ class CudaCommon_EXPORT CudaDataManager : public Object
   void *
   GetGPUBufferPointer();
 
+  /** Get pointer to the Cuda buffer pointer */
+  void *
+  GetGPUBufferPointerPtr();
+
   /** Get CPU buffer pointer */
   void *
   GetCPUBufferPointer();

include/itkCudaSquareImageFilter.hxx

@@ -31,8 +31,8 @@ namespace itk
       size[i] = this->GetInput()->GetBufferedRegion().GetSize()[i];
     }
 
-    typename ImageType::PixelType* pin0 = *(typename ImageType::PixelType**)(this->GetInput(0)->GetCudaDataManager()->GetGPUBufferPointer());
-    typename ImageType::PixelType* pout = *(typename ImageType::PixelType**)(this->GetOutput()->GetCudaDataManager()->GetGPUBufferPointer());
+    typename ImageType::PixelType* pin0 = (typename ImageType::PixelType*)(this->GetInput(0)->GetCudaDataManager()->GetGPUBufferPointer());
+    typename ImageType::PixelType* pout = (typename ImageType::PixelType*)(this->GetOutput()->GetCudaDataManager()->GetGPUBufferPointer());
 
     CudaSquareImage3D<typename ImageType::PixelType>(size, pin0, pout);
   }

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "itk-cudacommon"
-version = "1.1.0"
+version = "2.0.0"
 description = "Framework for processing images with Cuda"
 readme = "README.md"
 license = {file = "LICENSE"}

src/itkCudaDataManager.cxx

@@ -207,6 +207,13 @@ CudaDataManager::UpdateGPUBuffer()
 
 void *
 CudaDataManager::GetGPUBufferPointer()
+{
+  SetCPUBufferDirty();
+  return m_GPUBuffer->GetPointer();
+}
+
+void *
+CudaDataManager::GetGPUBufferPointerPtr()
 {
   SetCPUBufferDirty();
   return m_GPUBuffer->GetPointerPtr();

wrapping/CudaImage.i.in

@@ -0,0 +1,31 @@
+%extend itkCudaImage@CudaImageTypes@{
+    %pythoncode %{
+    @property
+    def __cuda_array_interface__(self):
+        _pixelType = "@PixelType@"
+        _typestr = _get_type_string(_pixelType)
+
+        _itksize = self.GetBufferedRegion().GetSize()
+        _dim     = len(_itksize)
+        _shape   = [int(_itksize[idx]) for idx in range(_dim)]
+
+        if self.GetNumberOfComponentsPerPixel() > 1:
+            _shape = [self.GetNumberOfComponentsPerPixel(), ] + _shape
+
+        # Reverse array to force C-order indexing. This is the reverse of how
+        # indices are specified in ITK, i.e. k,j,i versus i,j,k. However
+        # C-order indexing is expected by most algorithms in NumPy / SciPy.
+        _shape.reverse()
+        _shape = tuple(_shape)
+
+        return {
+            'shape': _shape,
+            'data': (int(self.GetCudaDataManager().GetGPUBufferPointer()), False),
+            'typestr': _typestr,
+            'descr': [('', _typestr)],
+            'version': 3,
+            'stream': None,
+            'strides': None
+        }
+    %}
+};

wrapping/CudaImage.i.init

@@ -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
+%}

wrapping/itkCudaImage.wrap

@@ -1,19 +1,46 @@
-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})
+  UNIQUE(vector_universe "${WRAP_ITK_VECTOR_REAL};${WRAP_ITK_COV_VECTOR_REAL}")
+  foreach(t ${types})
+    set(PixelType ${t})
+    string(REGEX MATCHALL "(V${t}|CV${t})" VectorTypes ${vector_universe})
+    foreach(d ${ITK_WRAP_IMAGE_DIMS})
       itk_wrap_template("${t}${d}" "${ITKT_${t}}, ${d}")
-    endforeach()
-  endforeach()
+      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})
 
-  UNIQUE(vector_types "${WRAP_ITK_VECTOR_REAL};${WRAP_ITK_COV_VECTOR_REAL}")
-  foreach(c ${ITK_WRAP_VECTOR_COMPONENTS})
-    foreach(d ${ITK_WRAP_IMAGE_DIMS})
-      foreach(vt ${vector_types})
+      foreach(c ${ITK_WRAP_VECTOR_COMPONENTS})
+        foreach(vt ${VectorTypes})
           itk_wrap_template("${ITKM_${vt}${c}}${d}" "${ITKT_${vt}${c}}, ${d}")
+          set(CudaImageTypes ${ITKM_${vt}${c}}${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()
-    endforeach()
-  endforeach()
+    endforeach(d)
+  endforeach(t)
 
 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)