A simple SWE solver on GPU using several different backends, including CuPy, PyTorch, and Legate NumPy. It can also run on CPU through PyTorch, vanilla NumPy, or Legate NumPy.
Everything is WIP, including documentation. But basically, to install:
$ pip install .
It installs a binary executable called TorchSWE.py
to your bin
path. Which
bin
path it installs to depends on your pip
.
After installing through pip
, only NumPy backend is available. To use other
backends, you may need to install them manually. Both PyTorch and CuPy can be
found from PyPI and Anaconda. Legate NumPy has to be installed manually
currently.
To see help
$ TorchSWE.py --help
To run a case with vanilla NumPy, go to a case folder and execute:
$ TorchSWE.py ./
To run with CuPy:
$ USE_CUPY=1 TorchSWE.py ./
To run with PyTorch on a GPU:
$ USE_TORCH=1 TorchSWE.py ./
To run with PyTorch on CPU:
$ USE_TORCH=1 TORCH_USE_CPU=1 TorchSWE.py ./
To run with Legate:
$ legate <flags> $(which TorchSWE.py) ./
For Legate, use its flags to control the target hardware. See Legate's
documentation. Legate does not know where to find the main Python script of the
solver, so we use $(which ...)
to provide the full path of TorchSWE.py
.
I have a very specific application in mind for this solver, so it's capability is somehow limited. And I intend to keep it simple, instead of generalizing it to general-purpose SWE solver. However, because it's implemented in Python, I believe it's not difficult to apply modifications for other applications.