Implement interface to control CPU and GPU setup
- Dominant language
- C++
- Stars
- 607
- Forks
- 67
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 1
Description
As discussed by @goldingn and @cboettig in https://github.com/greta-dev/greta/issues/306
I thought it might be good to have an issue open on this discussion so I can knock it off as a task.
What we require from this are:
- CPU use by default
- Coarse level control of control over either CPU or GPU
- Finer level control of using CPU or GPU on specific operations.
The default interface would be something like:
(@goldingn - I'm not sure if we need to specify this for `extra_samples` as well? )
```r
draws <- mcmc(
...,
compute_options = cpu_only()
)
```
Which would have the output:
```
running 4 chains simultaneously on up to 8 CPU cores
```
```r
samples <- calculate(
...,
compute_options = cpu_only()
)
extra_samples <- extra_samples(
...,
compute_options = cpu_only()
)
opt_res <- opt(
...,
compute_options = cpu_only()
)
```
With the idea being that you could turn on the GPU like so:
```r
draws <- mcmc(
...,
compute_options = gpu_only()
)
```
And allow for finer grained use like so:
```r
beta <- normal(0, 1, dim = 10)
eta <- X %*% beta
m <- model(eta)
draws <- mcmc(
m,
...,
compute_options = compute_setup(gpu_operations = list(
with_gpu(0, eta),
with_gpu(1, some_other_op)
))
)
```
Which would have output like:
```
running 4 chains simultaneously on up to 8 CPU cores and 2 GPU devices
GPU execution enabled for the operation greta array 'eta' (matrix multiply)
CPU execution enabled for all other greta operations
GPU devices in use: 0 and 1
```
To approach this API, I think the first release should probably focus on just getting the `cpu_only()` and `gpu_only()` approaches, but that we should design this so we can extend it easily - as Carl and Nick G both said.
In terms of order of operations, where's what I think we need:
- [x] `compute_options` argument for `mcmc`, `calculate`, `opt`, and `extra_samples`
- [x] `device_used()` (or similar) function which returns the name and type of the current device used - e.g., CPU or GPU. We can use this to capture the device used and then set it up at the end with `on.exit()`.
- [x] `cpu_only()` function - uses CPU and returns to default state after operation
- [x] `gpu_only()` function - uses GPU and returns to default state after operation
Reading https://www.tensorflow.org/guide/gpu
they say:
> If a TensorFlow operation has both CPU and GPU implementations, by default, the GPU device is prioritized when the operation is assigned.
They also say:
> Note: Use `tf.config.list_physical_devices('GPU')` to confirm that TensorFlow is using the GPU.
So if we can confirm that CPU and GPU are available, then we can "safely" (I think?) return to using GPU
Contributor guide
Assessment
This issue has not been assessed yet.