Adding GPU support documentation (#43)

* Adding GPU support documentation * add warning about gpu memory for large simulations --------- Co-authored-by: d-lucena <[email protected]>
agencyenterprise · May 25, 2023 · 7d944a5 · 7d944a5
1 parent eefb88a
commit 7d944a5
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 0 deletions.
diff --git a/docs/usage/gpu.md b/docs/usage/gpu.md
@@ -0,0 +1,38 @@
+In order to use your GPU to run NDK simulations you will have to install the [NVIDIA HPC SDK](https://developer.nvidia.com/hpc-sdk-downloads).
+
+!!! warning
+    Make sure the [HPC SDK environment variables](https://docs.nvidia.com/hpc-sdk/hpc-sdk-install-guide/index.html#install-linux-end-usr-env-settings) are exported.
+
+You will only have to export one environment variable to enable the GPU support for NDK:
+
+```bash
+export PLATFORM="nvidia-acc"
+```
+
+Now when running NDK simulations you should be able to see `platform=nvidiaX` in the execution output:
+
+```py
+import neurotechdevkit as ndk
+
+scenario = ndk.make('scenario-2-2d-v0')
+result = scenario.simulate_steady_state()
+result.render_steady_state_amplitudes()
+```
+
+
+Output:
+```bash
+...
+
+Operator `acoustic_iso_state` instance configuration:
+	 * subs={h_x: 0.0005, h_y: 0.0005}
+	 * opt=advanced
+	 * compiler=pgcc
+	 * language=openacc
+	 * platform=nvidiaX
+
+...
+```
+
+!!! warning
+	Simulations with high memory requirement (e.g. 3D simulations) may not fit in the GPU and running them with GPU acceleration might crash the simulation.
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -34,6 +34,7 @@ nav:
       - Simulation: usage/running_simulation.md
       - 3D Visualization: usage/3d.md
       - Troubleshooting: usage/troubleshooting.md
+      - GPU support: usage/gpu.md
   - API:
       - Make: api/make.md
       - Scenarios: api/scenarios.md

diff --git a/src/neurotechdevkit/scenarios/_base.py b/src/neurotechdevkit/scenarios/_base.py
@@ -2,6 +2,7 @@
 
 import abc
 import asyncio
+import os
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Mapping
@@ -852,6 +853,7 @@ def _execute_pde(
         devito_args = {}
         if n_jobs is not None:
             devito_args = dict(nthreads=n_jobs)
+
         assert sub_problem.shot is not None
         loop = asyncio.get_event_loop()
         return loop.run_until_complete(
@@ -864,6 +866,7 @@ def _execute_pde(
                 boundary_type="complex_frequency_shift_PML_2",
                 diff_source=True,
                 save_wavefield=True,
+                platform=os.environ.get("PLATFORM"),
                 save_bounds=save_bounds,
                 save_undersampling=save_undersampling,
                 wavefield_slice=wavefield_slice,