Write a GPU/CPU agnostic package with Flux

I am currently developing a Julia package based on Flux (namely an implementation of Deepmind’s AlphaZero algorithm) and I am looking for the best way to make it accessible to people with and without a GPU.

Ideally, enabling GPU usage should be as simple as providing a boolean parameter when calling the training algorithm (an error message would be displayed if no GPU is detected). I am having some trouble to make this happen though.

I looked at the solution used by model-zoo, which consists in providing two distinct project files for each model. Then, a model can be run on GPU by including its source file after running the following snippet:

using Pkg; Pkg.activate("cuda"); Pkg.instantiate()
using CuArrays

I do not see how to use this solution within a package though (as a package cannot have several project files).

I could imagine a solution looking something like this:

module AlphaZero
  using CuArrays
  using Flux
  const gpu = gpu_detected() ? Flux.gpu : identity
  # ...

Here, the package CuArrays would be imported but never actually used on a machine without a GPU.

There are two problems with this though. First, I am not aware of any function such as gpu_detected. There would be Flux.use_cuda but this variable is set based on whether or not CuArrays is imported. Second, importing CuArrays on a machine without a GPU yields an error:

[ Info: Precompiling CuArrays [3a865a2d-5b23-5a0f-bc46-62713ec82fae]
ERROR: LoadError: Could not find CUDA driver library

I would really appreciate your ideas on how to best write CPU/GPU agnostic packages!

1 Like

Right now if you understand you correctly you want some pointers on how to allow switching of resources after writing algorithms (the code for the algorithms doesn’t care about gpu or not anyway according to that: link). Then two packages come to my mind.

  1. Plots.jl which is somewhat like the general interface to many plotting backends.
  2. ComputationalResources.jl which is used in Images IIRC to decide whether computations shall be threaded, single core, multicore or GPU.
1 Like

Flux exports the function ‘gpu’ which wraps all arrays in CuArrays if cuda is available.

It has a (not exported iirc) function called has_cuda (or maybe has_gpu) which it uses to determine if cuda is available. It basically uses the method from this thread to do so Weak dependence on CUDA packages, 2019 version

The thread you linked to addresses my issue perfectly. Thanks!

Exciting! Repo? I want to contribute

Thanks for your interest!
The package is still under heavy development but you can already have a look at the repo. Hopefully, I’ll be able to make a public release with working examples and proper documentation soon. :slight_smile:

1 Like