Allow specification for GPU device index

[jwallwork23]

Closes #85.

The main change associated with this PR is allowing the GPU device index to be specified for the following functions and subroutines:

• torch_zeros (C++) / torch_tensor_zeros (Fortran)
• torch_ones (C++) / torch_tensor_zeros (Fortran)
• torch_empty (C++)
• torch_from_blob (C++) / torch_tensor_from_blob (Fotran)
• torch_jit_load (C++) / torch_module_load (Fortran)
• torch_tensor_from_array_${PREC}$_${RANK}$d (Fortran)

To avoid confusion/ambiguity, device is replaced by device_type in several places in the code, as device_type and device_index are consistent with the naming used in CUDA.

The GPU device index is specified using an additional argument, although this is made optional both in C++ and Fortran to ensure that the examples can be run without modification. In the case of torch_jit_load / torch_module_load, the device_type also needed to be added as an optional argument to support the new functionality.

If unset:

• device_type defaults to torch_kCPU
• device_index defaults to -1 if device_type is torch_kCPU and 0 if device_type is torch_kGPU.

New functions called torch_tensor_get_device_index are introduced so that we can test the new functionality.

[jwallwork23]

Here is the test that I used:

! Import precision info from iso
use, intrinsic :: iso_fortran_env, only : sp => real32

! Import our library for interfacing with PyTorch
use ftorch

! Import MPI
use mpi

implicit none

! Set precision for reals
integer, parameter :: wp = sp

! Set up Fortran data structures
real(wp), dimension(5), target :: in_data
integer :: tensor_layout(1) = [1]

! Set up Torch data structures
type(torch_tensor) :: in_tensor
integer :: device_type
integer :: device_index

! MPI configuration
integer rank, ierr

call mpi_init(ierr)
call mpi_comm_rank(mpi_comm_world, rank, ierr)

! Initialise data
in_data = [0.0, 1.0, 2.0, 3.0, 4.0]

! Loop over device type torch_kCPU and torch_kGPU
do device_type = 0, 1
  if (device_type == torch_kCPU) then
    device_index = - 1
  else
    device_index = rank
  end if

  ! Create Torch input tensor from the above arrays
  in_tensor = torch_tensor_from_array(in_data, tensor_layout, device_type, device_index)

  ! Print some information
  if (torch_tensor_get_device_index(in_tensor) == device_index) then
    write(*, *) rank, "PASS"
  else
    write(*, *) rank, "expected index ", device_index, "got ", torch_tensor_get_device_index(in_tensor)
  end if

  ! Cleanup
  call torch_tensor_delete(in_tensor)
end do
call mpi_finalize(ierr)

end program test_device_index

If run on my laptop (CPU-only), I get the output

           0 PASS
           1 PASS
           2 PASS
           3 PASS
[ERROR]: invalid device index 0 for device count [ERROR]: invalid device index 1 for device count 0, using zero instead
[ERROR]: invalid device index 2 for device count 0, using zero instead
[ERROR]: invalid device index 3 for device count 0, using zero instead
0, using zero instead
[ERROR]: PyTorch is not linked with support for cuda devices
[ERROR]: PyTorch is not linked with support for cuda devices
[ERROR]: PyTorch is not linked with support for cuda devices
[ERROR]: PyTorch is not linked with support for cuda devices

which confirms that the CPU case works, but obviously the GPU case isn't going to work.

If I run on Wilkes3 with four GPUs and four MPI processes, I get the output

           3 PASS
           2 PASS
           0 PASS
           1 PASS
           0 PASS
           1 PASS
           3 PASS
           2 PASS

which confirms that the GPU case works, too.

[jatkinson1000]

Amazing!
At a quick glance this looks good - I'll do a detailed review later when I have some time.
One quick comment before then - can you provide some simple instructions on how I can check/verify this is working on CSD3/elsewhere?

We will probably want an example adding to the examples/ and some info adding to the docs once the code is settled before it goes in.

[jwallwork23]

One quick comment before then - can you provide some simple instructions on how I can check/verify this is working on CSD3/elsewhere?

Sure. I created a new branch to demonstrate the testing: 85_gpu_device_number_test. Would you like me to include the Slurm scripts, too?

[jwallwork23]

Okay, this is ready for re-review! I added some docs and managed to get example 3 working on Wilkes3, giving the following output for 2 GPUs:

input on rank0: [  0.0,  1.0,  2.0,  3.0,  4.0]
input on rank1: [  1.0,  2.0,  3.0,  4.0,  5.0]
output on rank1: [  2.0,  4.0,  6.0,  8.0, 10.0]
output on rank0: [  0.0,  2.0,  4.0,  6.0,  8.0]

Will test it for 4 GPUs, too, but don't anticipate any issues.

[jatkinson1000]

This is great addition @jwallwork23
The new docs and example read really well.

Added a couple of points that I feel would make things clearer for me as an external reader, feel free to incorporate or not.

Once we've resolved these I think we're good to go!

[jwallwork23]

Thanks @jatkinson1000, this is now ready for re-review.

Will test it for 4 GPUs, too, but don't anticipate any issues.

I can confirm that this worked (with the updated 85_gpu_device_number_test branch) on Wilkes3, giving output

input on rank1: [  1.0,  2.0,  3.0,  4.0,  5.0]
input on rank2: [  2.0,  3.0,  4.0,  5.0,  6.0]
input on rank3: [  3.0,  4.0,  5.0,  6.0,  7.0]
input on rank0: [  0.0,  1.0,  2.0,  3.0,  4.0]
output on rank1: [  2.0,  4.0,  6.0,  8.0, 10.0]
output on rank2: [  4.0,  6.0,  8.0, 10.0, 12.0]
output on rank3: [  6.0,  8.0, 10.0, 12.0, 14.0]
output on rank0: [  0.0,  2.0,  4.0,  6.0,  8.0]

[jatkinson1000]

Thanks @jwallwork23 This is a great addition!

All looking good to me now so I'll squash and merge shortly.