GPU support

[gully]

Tensorflow offers wide support for hardware acceleration with Graphics Processing Units (GPUs). In principle, wobble could leverage this native GPU support to make it even faster.

I explored the prospect of wobble with GPUs on Google colaboratory, with no effort to modify the existing wobble code to anticipate GPUs:

with tf.device('/gpu:0'):
  data = wobble.Data('/content/gdrive/My Drive/51peg_e2ds.hdf5', orders=[30, 31, 32, 33, 34, 35])
  results = wobble.Results(data=data)
  for r in range(len(data.orders)):
          model = wobble.Model(data, results, r)
          model.add_star('star')
          model.add_telluric('tellurics')
          wobble.optimize_order(model)

I observed the following error message:

InvalidArgumentError: Cannot assign a device for operation Interp: Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Registered kernels:
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_FLOAT]

I investigated this error message as arising from the wobble/interp custom C code that uses the eigen library. Within this path I spotted some suggestive code that could be extended to support GPUs:

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

[...]

#define REGISTER_CPU(type)                                                 \
  REGISTER_KERNEL_BUILDER(                                                 \
      Name("Interp").Device(DEVICE_CPU).TypeConstraint<type>("T"),         \
      InterpOp<CPUDevice, type>)

REGISTER_CPU(float);
REGISTER_CPU(double);

#undef REGISTER_CPU

To be clear, the speed of wobble is not its primary limitation--- it runs reasonably fast on modern CPUs with typical high bandwidth echelle spectra. This exploration was merely out of curiosity to see if effortless improvement could be obtained by flipping a GPU switch. It seems that GPU support for wobble requires some careful examination of the interp extensions, and some understanding of GPUs on tensorflow.

[dfm]

In principle GPUs do natively support linear interpolation (that's pretty much exactly what GPUs were designed to do!) and perhaps we could work out the reverse mode op for this using the same features, but this would probably take some serious CUDA know how that I don't have. It would definitely be a welcome addition, but don't hold your breath for an implementation coming from us :-)