From 39e51d9aa9d2c79eff4a40337f3183c954b0cc19 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sat, 21 Oct 2023 19:45:18 -0400 Subject: [PATCH 1/3] WIP: Create installation functions for GPUs --- diffeqpy/__init__.py | 39 ++++++++++++++++++++++++++++++++++++++ diffeqpy/install_amdgpu.jl | 4 ++++ diffeqpy/install_cuda.jl | 4 ++++ diffeqpy/install_metal.jl | 4 ++++ diffeqpy/install_oneapi.jl | 4 ++++ 5 files changed, 55 insertions(+) create mode 100644 diffeqpy/install_amdgpu.jl create mode 100644 diffeqpy/install_cuda.jl create mode 100644 diffeqpy/install_metal.jl create mode 100644 diffeqpy/install_oneapi.jl diff --git a/diffeqpy/__init__.py b/diffeqpy/__init__.py index 24691c1..fa4efde 100644 --- a/diffeqpy/__init__.py +++ b/diffeqpy/__init__.py @@ -30,6 +30,45 @@ def install(*, confirm=False): env["PYTHON"] = sys.executable subprocess.check_call([julia, os.path.join(script_dir, "install.jl")], env=env) +def install_cuda(): + julia = _find_julia() + if not julia: + raise RuntimeError( + "Julia must be installed before adding CUDA. Please run `diffeqpy.install()` first" + ) + env = os.environ.copy() + env["PYTHON"] = sys.executable + subprocess.check_call([julia, os.path.join(script_dir, "install_cuda.jl")], env=env) + +def install_amdgpu(): + julia = _find_julia() + if not julia: + raise RuntimeError( + "Julia must be installed before adding AMDGPU. Please run `diffeqpy.install()` first" + ) + env = os.environ.copy() + env["PYTHON"] = sys.executable + subprocess.check_call([julia, os.path.join(script_dir, "install_amdgpu.jl")], env=env) + +def install_metal(): + julia = _find_julia() + if not julia: + raise RuntimeError( + "Julia must be installed before adding Metal. Please run `diffeqpy.install()` first" + ) + env = os.environ.copy() + env["PYTHON"] = sys.executable + subprocess.check_call([julia, os.path.join(script_dir, "install_metal.jl")], env=env) + +def install_oneapi(): + julia = _find_julia() + if not julia: + raise RuntimeError( + "Julia must be installed before adding oneAPI. Please run `diffeqpy.install()` first" + ) + env = os.environ.copy() + env["PYTHON"] = sys.executable + subprocess.check_call([julia, os.path.join(script_dir, "install_oneapi.jl")], env=env) def _ensure_installed(*kwargs): if not _find_julia(): diff --git a/diffeqpy/install_amdgpu.jl b/diffeqpy/install_amdgpu.jl new file mode 100644 index 0000000..450e2db --- /dev/null +++ b/diffeqpy/install_amdgpu.jl @@ -0,0 +1,4 @@ +using Pkg +Pkg.activate("diffeqpy", shared=true) +Pkg.add(["DiffEqGPU", "AMDGPU"]) +using DiffEqGPU, AMDGPU # Precompile diff --git a/diffeqpy/install_cuda.jl b/diffeqpy/install_cuda.jl new file mode 100644 index 0000000..b2e4869 --- /dev/null +++ b/diffeqpy/install_cuda.jl @@ -0,0 +1,4 @@ +using Pkg +Pkg.activate("diffeqpy", shared=true) +Pkg.add(["DiffEqGPU", "CUDA"]) +using DiffEqGPU, CUDA # Precompile diff --git a/diffeqpy/install_metal.jl b/diffeqpy/install_metal.jl new file mode 100644 index 0000000..64f50b5 --- /dev/null +++ b/diffeqpy/install_metal.jl @@ -0,0 +1,4 @@ +using Pkg +Pkg.activate("diffeqpy", shared=true) +Pkg.add(["DiffEqGPU", "Metal"]) +using DiffEqGPU, Metal # Precompile diff --git a/diffeqpy/install_oneapi.jl b/diffeqpy/install_oneapi.jl new file mode 100644 index 0000000..ec557c6 --- /dev/null +++ b/diffeqpy/install_oneapi.jl @@ -0,0 +1,4 @@ +using Pkg +Pkg.activate("diffeqpy", shared=true) +Pkg.add(["DiffEqGPU", "oneAPI"]) +using DiffEqGPU, oneAPI # Precompile From b31b7e78ac576312384753eb232651a50c0b3e95 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 23 Oct 2023 12:30:09 +0200 Subject: [PATCH 2/3] fix namespacing --- diffeqpy/amdgpu.py | 2 +- diffeqpy/cuda.py | 2 +- diffeqpy/metal.py | 2 +- diffeqpy/oneapi.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/diffeqpy/amdgpu.py b/diffeqpy/amdgpu.py index 8c97cb9..f6e0bb2 100644 --- a/diffeqpy/amdgpu.py +++ b/diffeqpy/amdgpu.py @@ -2,5 +2,5 @@ from . import load_julia_packages amdgpu, _ = load_julia_packages("DiffEqGPU, AMDGPU") from juliacall import Main -de.AMDGPUBackend = Main.seval("AMDGPU.AMDGPUBackend") # kinda hacky +amdgpu.AMDGPUBackend = Main.seval("AMDGPU.AMDGPUBackend") # kinda hacky sys.modules[__name__] = amdgpu # mutate myself diff --git a/diffeqpy/cuda.py b/diffeqpy/cuda.py index c10ac81..4cb0365 100644 --- a/diffeqpy/cuda.py +++ b/diffeqpy/cuda.py @@ -2,5 +2,5 @@ from . import load_julia_packages cuda, _ = load_julia_packages("DiffEqGPU, CUDA") from juliacall import Main -de.CUDABackend = Main.seval("CUDA.CUDABackend") # kinda hacky +cuda.CUDABackend = Main.seval("CUDA.CUDABackend") # kinda hacky sys.modules[__name__] = cuda # mutate myself diff --git a/diffeqpy/metal.py b/diffeqpy/metal.py index 9c61925..80913e6 100644 --- a/diffeqpy/metal.py +++ b/diffeqpy/metal.py @@ -2,5 +2,5 @@ from . import load_julia_packages metal, _ = load_julia_packages("DiffEqGPU, Metal") from juliacall import Main -de.MetalBackend = Main.seval("Metal.MetalBackend") # kinda hacky +metal.MetalBackend = Main.seval("Metal.MetalBackend") # kinda hacky sys.modules[__name__] = metal # mutate myself \ No newline at end of file diff --git a/diffeqpy/oneapi.py b/diffeqpy/oneapi.py index d79283e..5033b11 100644 --- a/diffeqpy/oneapi.py +++ b/diffeqpy/oneapi.py @@ -2,5 +2,5 @@ from . import load_julia_packages oneapi, _ = load_julia_packages("DiffEqGPU, oneAPI") from juliacall import Main -de.oneAPIBackend = Main.seval("oneAPI.oneAPIBackend") # kinda hacky +oneapi.oneAPIBackend = Main.seval("oneAPI.oneAPIBackend") # kinda hacky sys.modules[__name__] = oneapi # mutate myself From 867429e8c7b79d8137ee7a67a426e253549b4bd2 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Mon, 23 Oct 2023 12:32:12 +0200 Subject: [PATCH 3/3] Update README.md --- README.md | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8c7c443..c71e594 100644 --- a/README.md +++ b/README.md @@ -551,13 +551,19 @@ sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10000,sa ``` To add GPUs to the mix, we need to bring in [DiffEqGPU](https://github.com/SciML/DiffEqGPU.jl). -The `cuda` submodule will install CUDA for you and bring all of the bindings into the returned object: +The `diffeqpy.install_cuda()` will install CUDA for you and bring all of the bindings into the returned object: + +```py +diffeqpy.install_cuda() +``` + +then run the cuda import: ```py from diffeqpy import cuda ``` -#### Note: `from diffeqpy import cuda` can take awhile to run the first time as it installs the drivers! +#### Note: `diffeqpy.install_cuda()` and `from diffeqpy import cuda` can take awhile to run the first time as it installs the drivers! Now we simply use `EnsembleGPUKernel(cuda.CUDABackend())` with a GPU-specialized ODE solver `cuda.GPUTsit5()` to solve 10,000 ODEs on the GPU in