XPU and MPS take 3

[jwallwork23]

Closes #127.
Builds upon #125 and #209.
(Contains changes from #268 so that will need to be merged first.)

This PR adds support XPU and MPS. Unfortunately, it ended up requiring an overhaul of the pt2ts scripts, too.

Notable changes:

• Switching from ENABLE_CUDA to the more general and extensible GPU_DEVICE=<NONE/CUDA/XPU/MPS>.
• Pre-processor directives for handling different GPU types.
• Support for XPU under GPU device code 12.
• Support for MPS under GPU device code 13.
• Use of argparse for reading command line arguments into Python scripts, rather than sys.argv.
• Updates to docs.

Checklist

• Test on 2 Nvidia GPUs
• Test on 2 XPUs
• Test on 1 MPS device

[jwallwork23]

To get the Fortran linting check to pass, I had to ignore errors of the form reported in https://github.com/Cambridge-ICCS/FTorch/actions/runs/13240946802/job/36956002877, as well as ones of the form

src/ftorch.F90:153:57: T041 'tensor_shape' has assumed size
    |
151 |       type(c_ptr), value, intent(in)    :: data
152 |       integer(c_int), value, intent(in) :: ndims
153 |       integer(c_int64_t), intent(in)    :: tensor_shape(*)
    |                                                         ^ T041
154 |       integer(c_int64_t), intent(in)    :: strides(*)
155 |       integer(c_int), value, intent(in) :: dtype

because we established that the cases where we use assumed size in ftorch.F90 are required.

[jwallwork23]

Offline testing for CUDA version of MultiGPU example with 2 devices passed on Ampere. In the queue for XPU testing on PVC.

[jatkinson1000]

Thanks @jwallwork23 I have only done a quick pass of this so far and will need to schedule time for a closer look, but I suspect you know my first comment - can you update the docs at utils/README etc. to reflect the new args and usage of pt2ts?

I would also like to see it documented somewhere how the device enums are managed from CMake. Perhaps under the developer docs. As, whilst a very nifty solution, it's slightly abstract if you are not the one who came up with it 😉

[jatkinson1000]

Also looks like you may want a rebase after #268

[jwallwork23]

Thanks @jwallwork23 I have only done a quick pass of this so far and will need to schedule time for a closer look, but I suspect you know my first comment - can you update the docs at utils/README etc. to reflect the new args and usage of pt2ts?

Oh, I'd forgotten about that README actually. Addressed here 50bd953.

I would also like to see it documented somewhere how the device enums are managed from CMake. Perhaps under the developer docs. As, whilst a very nifty solution, it's slightly abstract if you are not the one who came up with it 😉

Ah yes... done in 838f3ae.

Also looks like you may want a rebase after #268

Done!

[jwallwork23]

I don't understand why the clang-tidy errors didn't crop up previously...?

For one of them, I can fix with

+#include <stdbool.h>

For the device related one, I could fix with

+// NOTE: These need to be defined here for clang-tidy to pass
+#ifndef GPU_DEVICE_NONE
+#define GPU_DEVICE_NONE 0
+#endif
+#ifndef GPU_DEVICE_CUDA
+#define GPU_DEVICE_CUDA 1
+#endif
+#ifndef GPU_DEVICE_XPU
+#define GPU_DEVICE_XPU 12
+#endif
+#ifndef GPU_DEVICE_MPS
+#define GPU_DEVICE_MPS 13
+#endif

(although would prefer to avoid duplicating these codes)

But the other torch/script one - is clang-tidy meant to be compiler-aware or something?

[jatkinson1000]

I have just run this on a Mac using Mps and made the relevant tweaks to get it to work.
@jwallwork23 if you can check the diff in c863cc7 to see if you are happy that would be great. Hopefully everything should still work OK on XPU/CUDA though I had to adjust how num_devices was set.

I'll look to schedule time to review the rest of this.

[jwallwork23]

A few very minor suggestions but otherwise those additions looks good - thanks for testing on MPS @jatkinson1000. Still waiting for my XPU job to run.

[jwallwork23]

Good news - test passed on 2 XPU devices on Dawn! 🥳

[jwallwork23]

Re-tested on Dawn with latest version of branch - all good

[ElliottKasoar]

It's probably worth at least mentioning Intel/Mac GPUs in the GPU Support section of the README too?

[jwallwork23]

Ah, good point! Addressed in 0897b20.