WIP: Create installation functions for GPUs

SciML · Oct 21, 2023 · 39e51d9 · 39e51d9
1 parent 42ce4d3
commit 39e51d9
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("diffeqpy", shared=true)
+Pkg.add(["DiffEqGPU", "oneAPI"])
+using DiffEqGPU, oneAPI # Precompile