DataContainer backend and cupy experiments

[paskino]

Before it gets lost, I am copying here a discord discussion.

I worked at extending the CIL DataContainer to handle different backends than numpy, currently only cupy.

I made an experiment comparing CIL's TotalVariation with backend C, numpy and cupy with the regularisation toolkit FGP_TV with the CUDA implementation. The C backend refers to the calculation of the GradientOperator inside TotalVariation, for the rest it uses numpy.

Experiment

Image size 128x128, number of TV iterations 1000. To reproduce these results I used the following script.

Algoritm	run time
FGP_TV CUDA	0.37s
TotalVariation + cupy	5s
TotalVariation + C or numpy backend	22s

notice that the image is really small

Results

1. FGP_TV is 59 times faster than standard CIL TotalVariation
2. FGP_TV is 13.5 times faster than cupy CIL TotalVariation
3. The cupy backend for CIL's TotalVariation is 4.4 times faster than the standard implementation with numpy
4. The C or numpy implementation of the GradientOperator does not seem to have any impact on the calculation of TotalVariation (proximal)
5. FGP_TV does not have the warm start functionality

Changes to CIL

I made a few changes to CIL to achieve these results.

23e59b6...paskino:CIL:cupy_array_copy

The biggest differences are:

• modified ImageGeometry/AcquisitionGeometry.allocate and DataContainer.fill methods as numpy and cupy have slightly different ways of copying/creating data.

• added the concept of backend and relative member to DataContainer (probably should be inferred by the member array, as currently it risks to get out of sync?)

• modified __init__ of ImageData . This should be done for AcquisitionData too. Notice that the backend is not passed on to the base class.

• added backend parameter to ImageGeometry.allocate to be able to allocate either numpy or cupy array. In turn this uses the __init__ of the DataContainer. Notice that the same work should be done for AcquisitionGeometry.

• some other modifications to make sure that TotalVariation knows it has to run on either backend.

[paskino]

Basic concept

At the Hackathon on Hackathon on Efficient integration of SIRF/STIR/CIL with Pytorch we planned to work at making CIL's data containers agnostic on the actual array they wrap, by using the Python Array API

The work in progress can be checked at

• DataContainer agnostic on backend #2132

I focused on making use of array_api_compat which in theory should avoid the problems/workarounds I had to do previously (at creation/copy). See above.

• I made a huge improvement in fill which was using some ancient/horrible code.
• I had to make many little changes to the DataContainer algebraic methods similar to below, where I had to define xp using array_namespace
  

  CIL/Wrappers/Python/cil/framework/data_container.py

  Lines 447 to 452 in 50c2ffa

  	def add(self, other, *args, **kwargs):
  	if hasattr(other, '__container_priority__') and \
  	self.__class__.__container_priority__ < other.__class__.__container_priority__:
  	return other.add(self, *args, **kwargs)
  	xp = array_namespace(self.array)
  	return self.pixel_wise_binary(xp.add, other, *args, **kwargs)

• some methods do not exist in the Python Array API standard, such as power, and others behave differently from the ones in NumPy. I was therefore forced to create a CIL array_api_compat file to make up for the differences. Currently I had to add:
  • squeeze
  • expand_dims
  • allclose
  • dtype_namespace to try to infer the namespace from the object's dtype. This when I just have a type, not an object. This fails with cupy and numpy objects as their type is the same: cupy.float32 = numpy.float32
• I created extra unittest to test with various backends.

Thanks to @purepani I found out that other people started to patch the Python Array API to add the missing functionality, see https://data-apis.org/array-api-extra/

## Problems

Once the original unittests passed, I added specific tests for torch and cupy and made sure they passed.

Then, I tried to run the same TotalVariation tests as earlier, however there are problems especially with allocate when we cannot/do not infer the right type and we default to a NumPy array.

It seems we would need to be able to specify the backend as earlier.

[purepani]

This may be a better place to ping @lucascolley due to the discussion