From f57b76dc53e0517f8e2ae405daedad3ee66590e4 Mon Sep 17 00:00:00 2001
From: Thomas Roeblitz <thomas.roblitz@uib.no>
Date: Wed, 15 May 2024 19:57:18 +0200
Subject: [PATCH] drop dependency on cuDNN to builddependency

---
 eb_hooks.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/eb_hooks.py b/eb_hooks.py
index a778f8b7c1..cf1c911b23 100644
--- a/eb_hooks.py
+++ b/eb_hooks.py
@@ -768,6 +768,25 @@ def inject_gpu_property(ec):
                 ec[key] = '\n'.join([ec_dict[key], value])
         else:
             ec[key] = value
+    # Check if cuDNN is in the dependencies, if so add the 'gpu' Lmod property
+    if ('cuDNN' in [dep[0] for dep in iter(ec_dict['dependencies'])]):
+        ec.log.info("Injecting gpu as Lmod arch property and envvar with cuDNN version")
+        key = 'modluafooter'
+        value = 'add_property("arch","gpu")'
+        cudnn_version = 0
+        for dep in iter(ec_dict['dependencies']):
+            # Make cuDNN a build dependency only (rpathing saves us from link errors)
+            if 'cuDNN' in dep[0]:
+                cudnn_version = dep[1]
+                ec_dict['dependencies'].remove(dep)
+                if dep not in ec_dict['builddependencies']:
+                    ec_dict['builddependencies'].append(dep)
+        value = '\n'.join([value, 'setenv("EESSICUDNNVERSION","%s")' % cudnn_version])
+        if key in ec_dict:
+            if not value in ec_dict[key]:
+                ec[key] = '\n'.join([ec_dict[key], value])
+        else:
+            ec[key] = value
     return ec