testing #87

easystacks/software.eessi.io/2023.06/accel/nvidia/rebuilds/20250923-eb-5.1.1-rebuild-LAMMPS-for-cuda-sanity-check.yml

@@ -0,0 +1,6 @@
+# DO NOT MERGE
+# this is solely for testing the new hook on ARM + CUDA targets
+easyconfigs:
+  - LAMMPS-2Aug2023_update2-foss-2023a-kokkos-CUDA-12.1.1.eb:
+      options:
+        include-easyblocks-from-commit: 338c31724aa3dc58223ec35120d4c18d1bdf8d94

eb_hooks.py

@@ -1003,19 +1003,49 @@ def pre_configure_hook_wrf_aarch64(self, *args, **kwargs):
         raise EasyBuildError("WRF-specific hook triggered for non-WRF easyconfig?!")
 
 
-def pre_configure_hook_LAMMPS_zen4(self, *args, **kwargs):
+def pre_configure_hook_LAMMPS_zen4_and_aarch64_cuda(self, *args, **kwargs):
     """
     pre-configure hook for LAMMPS:
-    - set kokkos_arch on x86_64/amd/zen4
+    - set kokkos_arch on x86_64/amd/zen4 and aarch64/nvidia/grace
+    - Disable SIMD for Aarch64 + cuda builds
     """
 
+    # Get cpu_target for zen4 hook
     cpu_target = get_eessi_envvar('EESSI_SOFTWARE_SUBDIR')
+
+    # Get the dependencies to check CUDA
+    deps = self.cfg.dependencies()
+    cuda_versions = ('12.1.1')
+
     if self.name == 'LAMMPS':
+
+        # Set kokkos_arch for LAMMPS version which do not have support for versions that do not support ZEN4
         if self.version in ('2Aug2023_update2', '2Aug2023_update4', '29Aug2024'):
             if get_cpu_architecture() == X86_64:
                 if cpu_target == CPU_TARGET_ZEN4:
                     # There is no support for ZEN4 in LAMMPS yet so falling back to ZEN3
                     self.cfg['kokkos_arch'] = 'ZEN3'
+            elif get_cpu_architecture() == AARCH64:
+                if cpu_target == CPU_TARGET_NVIDIA_GRACE:
+                    # There is no support for NVIDA grace in LAMMPS yet so falling back to ARMV81
+                    self.cfg['kokkos_arch'] = 'ARMV81'
+                # To disable simd kokkos_arch need to be set to ARMV7
+                if self.cuda:
+                    self.cfg['kokkos_arch'] = 'ARMV70'
+
+        # Disable SIMD for specific CUDA versions
+        if self.version == '2Aug2023_update2':
+            if get_cpu_architecture() == AARCH64:
+                if ('CUDA' in [dep['name'] for dep in deps]):
+                    for dep in deps:
+                        if 'CUDA' == dep['name']:
+                            if dep['version'] in cuda_versions:
+                                cxxflags = os.getenv('CXXFLAGS', '')
+                                cxxflags = cxxflags.replace('-mcpu=native', '')
+                                cxxflags += ' -march=armv8-a+nosimd'
+                                self.log.info("Setting CXXFLAGS to disable NEON: %s", cxxflags)
+                                env.setvar('CXXFLAGS', cxxflags)
+
     else:
         raise EasyBuildError("LAMMPS-specific hook triggered for non-LAMMPS easyconfig?!")
 
@@ -1579,7 +1609,7 @@ def post_easyblock_hook(self, *args, **kwargs):
     'PMIx': pre_configure_hook_pmix_ipv6,
     'PRRTE': pre_configure_hook_prrte_ipv6,
     'WRF': pre_configure_hook_wrf_aarch64,
-    'LAMMPS': pre_configure_hook_LAMMPS_zen4,
+    'LAMMPS': pre_configure_hook_LAMMPS_zen4_and_aarch64_cuda,
     'Score-P': pre_configure_hook_score_p,
     'VSEARCH': pre_configure_hook_vsearch,
     'CMake': pre_configure_hook_cmake_system,