CUDA array Interface in ITK¶

The CUDA array interface allows seamless interoperability between ITK’s CudaImage and GPU arrays from libraries like CuPy and PyTorch. This enables zero-copy data sharing, efficient GPU operations, and integration with deep learning pipelines.

Overview¶

ITK’s CudaImage implements the __cuda_array_interface__ protocol, allowing it to be used directly as a CuPy array or PyTorch tensor without copying data to/from the GPU.

Key Functions¶

  • itk.cuda_image_from_cuda_array(array)
    Creates an ITK CudaImage from a CUDA array (CuPy ndarray or PyTorch tensor).

    • Parameters: array - A CuPy or PyTorch array on GPU.

    • Returns: A CudaImage with the same data (shared memory).

  • cp.asarray(cuda_image) or torch.as_tensor(cuda_image, device="cuda")
    Creates a zero-copy view of the CudaImage as a CuPy array or PyTorch tensor.

    • Parameters: cuda_image - An ITK CudaImage.

    • Returns: A view of the GPU data (no copy).

Examples¶

import itk
import cupy as cp

GPUImageType = itk.CudaImage[itk.F, 2]
img = GPUImageType.New()

img.SetSpacing([0.7, 0.7])
img.SetOrigin([10.0, 20.0])
img.SetRegions((16, 20))
img.Allocate()
img.FillBuffer(50)

# convert to a CuPy array
a_view = cp.asarray(img)

# Do some operations with CuPy
a_view += 2
a_view *= 3

# Back to ITK CudaImage
img2 = itk.cuda_image_from_cuda_array(a_view)

import itk
import cupy as cp

GPUImageType = itk.CudaImage[itk.F, 2]
img = GPUImageType.New()

img.SetSpacing([0.7, 0.7])
img.SetOrigin([10.0, 20.0])
img.SetRegions((16, 20))
img.Allocate()
img.FillBuffer(50)

# convert to a CuPy array
a_view = cp.asarray(img)

# Do some operations with CuPy
a_view += 2
a_view *= 3

# Back to ITK CudaImage
img2 = itk.cuda_image_from_cuda_array(a_view)