Merge pull request #127 from SciML/gpu_install

Create installation functions for GPUs

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 

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():

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

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

diffeqpy/install_amdgpu.jl

@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("diffeqpy", shared=true)
+Pkg.add(["DiffEqGPU", "AMDGPU"])
+using DiffEqGPU, AMDGPU # Precompile

diffeqpy/install_cuda.jl

@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("diffeqpy", shared=true)
+Pkg.add(["DiffEqGPU", "CUDA"])
+using DiffEqGPU, CUDA # Precompile

diffeqpy/install_metal.jl

@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("diffeqpy", shared=true)
+Pkg.add(["DiffEqGPU", "Metal"])
+using DiffEqGPU, Metal # Precompile

diffeqpy/install_oneapi.jl

@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("diffeqpy", shared=true)
+Pkg.add(["DiffEqGPU", "oneAPI"])
+using DiffEqGPU, oneAPI # Precompile

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

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