Kokkos complex_align variant, Trilinos+PETSc enforcement for Kokkos~complex_align (#47686)

3 files changed, 23 insertions, 18 deletions

diff --git a/var/spack/repos/builtin/packages/kokkos/package.py b/var/spack/repos/builtin/packages/kokkos/package.py
index 9b9b78e117..6492906ad7 100644
--- a/var/spack/repos/builtin/packages/kokkos/package.py
+++ b/var/spack/repos/builtin/packages/kokkos/package.py
@@ -179,6 +179,7 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
     options_variants = {
         "aggressive_vectorization": [False, "Aggressively vectorize loops"],
         "compiler_warnings": [False, "Print all compiler warnings"],
+        "complex_align": [True, "Align complex numbers"],
         "cuda_constexpr": [False, "Activate experimental constexpr features"],
         "cuda_lambda": [False, "Activate experimental lambda features"],
         "cuda_ldg_intrinsic": [False, "Use CUDA LDG intrinsics"],
diff --git a/var/spack/repos/builtin/packages/petsc/package.py b/var/spack/repos/builtin/packages/petsc/package.py
index 918b2e8367..9d983d254f 100644
--- a/var/spack/repos/builtin/packages/petsc/package.py
+++ b/var/spack/repos/builtin/packages/petsc/package.py
@@ -374,6 +374,7 @@ class Petsc(Package, CudaPackage, ROCmPackage):
     depends_on("libyaml", when="+libyaml")
     depends_on("hwloc", when="+hwloc")
     depends_on("kokkos", when="+kokkos")
+    depends_on("kokkos~complex_align", when="+kokkos+complex")
     depends_on("kokkos-kernels", when="+kokkos")
     for cuda_arch in CudaPackage.cuda_arch_values:
         depends_on(
diff --git a/var/spack/repos/builtin/packages/trilinos/package.py b/var/spack/repos/builtin/packages/trilinos/package.py
index bb12c6caa3..7966cb3e67 100644
--- a/var/spack/repos/builtin/packages/trilinos/package.py
+++ b/var/spack/repos/builtin/packages/trilinos/package.py
@@ -407,23 +407,25 @@ class Trilinos(CMakePackage, CudaPackage, ROCmPackage):
     # ###################### Dependencies ##########################
 
     # External Kokkos
-    depends_on("kokkos@4.4.01", when="@master: +kokkos")
-    depends_on("kokkos@4.3.01", when="@16.0.0 +kokkos")
-    depends_on("kokkos@4.2.01", when="@15.1.0:15.1.1 +kokkos")
-    depends_on("kokkos@4.1.00", when="@14.4.0:15.0.0 +kokkos")
-
-    depends_on("kokkos +wrapper", when="trilinos@14.4.0: +kokkos +wrapper")
-    depends_on("kokkos ~wrapper", when="trilinos@14.4.0: +kokkos ~wrapper")
-
-    for a in CudaPackage.cuda_arch_values:
-        arch_str = "+cuda cuda_arch={0}".format(a)
-        kokkos_spec = "kokkos {0}".format(arch_str)
-        depends_on(kokkos_spec, when="@14.4.0: +kokkos {0}".format(arch_str))
-
-    for a in ROCmPackage.amdgpu_targets:
-        arch_str = "+rocm amdgpu_target={0}".format(a)
-        kokkos_spec = "kokkos {0}".format(arch_str)
-        depends_on(kokkos_spec, when="@14.4.0: +kokkos {0}".format(arch_str))
+    with when("@14.4: +kokkos"):
+        depends_on("kokkos+wrapper", when="+wrapper")
+        depends_on("kokkos~wrapper", when="~wrapper")
+        depends_on("kokkos~complex_align")
+        depends_on("kokkos@4.4.01", when="@master:")
+        depends_on("kokkos@4.3.01", when="@16")
+        depends_on("kokkos@4.2.01", when="@15.1:15")
+        depends_on("kokkos@4.1.00", when="@14.4:15.0")
+        depends_on("kokkos-kernels@4.4.01", when="@master:")
+        depends_on("kokkos-kernels@4.3.01", when="@16")
+        depends_on("kokkos-kernels@4.2.01", when="@15.1:15")
+        depends_on("kokkos-kernels@4.1.00", when="@14.4:15.0")
+
+        for a in CudaPackage.cuda_arch_values:
+            arch_str = f"+cuda cuda_arch={a}"
+            depends_on(f"kokkos{arch_str}", when=arch_str)
+        for a in ROCmPackage.amdgpu_targets:
+            arch_str = f"+rocm amdgpu_target={a}"
+            depends_on(f"kokkos{arch_str}", when=arch_str)
 
     depends_on("adios2", when="+adios2")
     depends_on("binder@1.3:", when="@15: +python", type="build")
@@ -899,8 +901,9 @@ class Trilinos(CMakePackage, CudaPackage, ROCmPackage):
             define_tpl(tpl_name, dep_name, dep_name in spec)
 
         # External Kokkos
-        if spec.satisfies("@14.4.0 +kokkos"):
+        if spec.satisfies("@14.4.0: +kokkos"):
             options.append(define_tpl_enable("Kokkos"))
+            options.append(define_tpl_enable("KokkosKernels", True))
 
         # MPI settings
         options.append(define_tpl_enable("MPI"))