search

tum-pbs / PhiFlow

A differentiable PDE solving framework for machine learning
MIT License
1.39k stars 189 forks source link

MPS device/MAC GPU #119

Closed timmyL17 closed 1 year ago

timmyL17 commented 1 year ago

How hard would it be to enable the Mac M1/M2 GPU devices in phiflow? It seems that phiflow currently only searches the cuda backend for devices. Are the phiflow cuda kernels only used for tensorflow? Maybe it could allow the MPS backend in pytorch, if pytorch already includes kernels for grid interpolation?

holl- commented 1 year ago

You can use the GPU even without the custom CUDA kernels. If you get either TensorFlow, PyTorch or Jax to recognize the GPU, you can simply select that backend using the corresponding phi.<backend>.flow import and you're good to go. You can check whether the GPU is recognized using

import phi; phi.verify()