Vtkm kokkos variant (#28363)

* VTK-m: Make vtk-m consistent with ROCmPackage

* VTKm: Add kokkos variant

Specifying +kokkos will enable kokkos backend.
Specifying +kokkos with +rocm will require a kokkos with a ROCm backend.
Specifying +cuda enables VTK-m native CUDA backend. VTK-m native cuda backend
  is not compatible with the kokkos +cuda backend.

* VTK-m: Add cuda_native variant

Required to allow specifying a vtk-m spec the selects a
cuda_arch and predictably propagate that to the underlying kokkos
dependency.

This also makes explicit selecting kokkos with a cuda backend or using
the VTK-m cuda backend.

1 files changed, 42 insertions, 35 deletions

diff --git a/var/spack/repos/builtin/packages/vtk-m/package.py b/var/spack/repos/builtin/packages/vtk-m/package.py
index 4308748d29..c6535b0140 100644
--- a/var/spack/repos/builtin/packages/vtk-m/package.py
+++ b/var/spack/repos/builtin/packages/vtk-m/package.py
@@ -11,7 +11,7 @@ import sys
 from spack import *
 
 
-class VtkM(CMakePackage, CudaPackage):
+class VtkM(CMakePackage, CudaPackage, ROCmPackage):
     """VTK-m is a toolkit of scientific visualization algorithms for emerging
     processor architectures. VTK-m supports the fine-grained concurrency for
     data analysis and visualization algorithms required to drive extreme scale
@@ -58,37 +58,45 @@ class VtkM(CMakePackage, CudaPackage):
     variant("testlib", default=False, description="build test library")
 
     # Device variants
-    variant("cuda", default=False, description="build cuda support")
+    # CudaPackage provides cuda variant
+    # ROCmPackage provides rocm variant
+    variant("kokkos", default=False, description="build using Kokkos backend")
+    variant("cuda_native", default=True, description="build using native cuda backend", when="+cuda")
     variant("openmp", default=(sys.platform != 'darwin'), description="build openmp support")
     variant("tbb", default=(sys.platform == 'darwin'), description="build TBB support")
-    variant("hip", default=False, description="build hip support")
-
-    # it doesn't look like spack has an amd gpu abstraction
-    # Going to have to restrict our set to ones that Kokkos supports
-    amdgpu_targets = (
-        'gfx900', 'gfx906', 'gfx908'
-    )
-
-    variant('amdgpu_target', default='none', multi=True, values=('none',) + amdgpu_targets)
-    conflicts("+hip", when="amdgpu_target=none")
 
     depends_on("cmake@3.12:", type="build")               # CMake >= 3.12
-    depends_on("cmake@3.18:", when="+hip", type="build")  # CMake >= 3.18
+    depends_on("cmake@3.18:", when="+rocm", type="build")  # CMake >= 3.18
 
     conflicts('%gcc@:4.10',
               msg='vtk-m requires gcc >= 5. Please install a newer version')
 
-    depends_on('cuda@10.1.0:', when='+cuda')
+    depends_on('cuda@10.1.0:', when='+cuda_native')
     depends_on("tbb", when="+tbb")
     depends_on("mpi", when="+mpi")
 
-    for amdgpu_value in amdgpu_targets:
-        depends_on("kokkos@develop +rocm amdgpu_target=%s" % amdgpu_value, when="amdgpu_target=%s" % amdgpu_value)
-
-    depends_on("rocm-cmake@3.7:", when="+hip")
-    depends_on("hip@3.7:", when="+hip")
-
-    conflicts("+hip", when="+cuda")
+    # VTK-m uses the default Kokkos backend
+    depends_on('kokkos', when='+kokkos')
+    # VTK-m native CUDA and Kokkos CUDA backends are not compatible
+    depends_on('kokkos ~cuda', when='+kokkos +cuda +cuda_native')
+    depends_on('kokkos +cuda', when='+kokkos +cuda ~cuda_native')
+    for cuda_arch in CudaPackage.cuda_arch_values:
+        depends_on("kokkos cuda_arch=%s" % cuda_arch, when="+kokkos +cuda ~cuda_native cuda_arch=%s" % cuda_arch)
+    # VTK-m uses the Kokkos HIP backend.
+    # If Kokkos provides multiple backends, the HIP backend may or
+    # may not be used for VTK-m depending on the default selected by Kokkos
+    depends_on('kokkos +rocm', when='+kokkos +rocm')
+    # Propagate AMD GPU target to kokkos for +rocm
+    for amdgpu_value in ROCmPackage.amdgpu_targets:
+        depends_on("kokkos amdgpu_target=%s" % amdgpu_value, when="+kokkos +rocm amdgpu_target=%s" % amdgpu_value)
+
+    depends_on("rocm-cmake@3.7:", when="+rocm")
+    depends_on("hip@3.7:", when="+rocm")
+
+    conflicts("+rocm", when="+cuda")
+    conflicts("+rocm", when="~kokkos", msg="VTK-m does not support HIP without Kokkos")
+
+    conflicts("+shared", when="+cuda_native")
 
     conflicts("+cuda", when="cuda_arch=none",
               msg="vtk-m +cuda requires that cuda_arch be set")
@@ -105,13 +113,10 @@ class VtkM(CMakePackage, CudaPackage):
             options = ["-DVTKm_ENABLE_TESTING:BOOL=OFF"]
             # shared vs static libs logic
             # force building statically with cuda
-            if "+cuda" in spec:
-                options.append('-DBUILD_SHARED_LIBS=OFF')
+            if "+shared" in spec:
+                options.append('-DBUILD_SHARED_LIBS=ON')
             else:
-                if "+shared" in spec:
-                    options.append('-DBUILD_SHARED_LIBS=ON')
-                else:
-                    options.append('-DBUILD_SHARED_LIBS=OFF')
+                options.append('-DBUILD_SHARED_LIBS=OFF')
 
             # double precision
             if "+doubleprecision" in spec:
@@ -170,7 +175,7 @@ class VtkM(CMakePackage, CudaPackage):
                 options.append("-DVTKm_NO_ASSERT:BOOL=ON")
 
             # cuda support
-            if "+cuda" in spec:
+            if "+cuda_native" in spec:
                 options.append("-DVTKm_ENABLE_CUDA:BOOL=ON")
                 options.append("-DCMAKE_CUDA_HOST_COMPILER={0}".format(
                                env["SPACK_CXX"]))
@@ -190,15 +195,12 @@ class VtkM(CMakePackage, CudaPackage):
                 options.append("-DVTKm_ENABLE_CUDA:BOOL=OFF")
 
             # hip support
-            if "+hip" in spec:
+            if "+rocm" in spec:
                 options.append("-DVTKm_NO_DEPRECATED_VIRTUAL:BOOL=ON")
-                options.append("-DVTKm_ENABLE_HIP:BOOL=ON")
 
                 archs = ",".join(self.spec.variants['amdgpu_target'].value)
                 options.append(
                     "-DCMAKE_HIP_ARCHITECTURES:STRING={0}".format(archs))
-            else:
-                options.append("-DVTKm_ENABLE_HIP:BOOL=OFF")
 
             # openmp support
             if "+openmp" in spec:
@@ -210,6 +212,11 @@ class VtkM(CMakePackage, CudaPackage):
             else:
                 options.append("-DVTKm_ENABLE_OPENMP:BOOL=OFF")
 
+            if "+kokkos" in spec:
+                options.append("-DVTKm_ENABLE_KOKKOS:BOOL=ON")
+            else:
+                options.append("-DVTKm_ENABLE_KOKKOS:BOOL=OFF")
+
             # tbb support
             if "+tbb" in spec:
                 # vtk-m detectes tbb via TBB_ROOT env var
@@ -348,10 +355,10 @@ vtkm_add_target_information(VTKmSmokeTest
                 except ProcessError:
                     output = ""
                 print(output)
-                if "+hip" in spec:
-                    expected_device = 'Kokkos'
-                elif "+cuda" in spec:
+                if "+cuda_native" in spec:
                     expected_device = 'Cuda'
+                elif "+kokkos" in spec:
+                    expected_device = 'Kokkos'
                 elif "+tbb" in spec:
                     expected_device = 'TBB'
                 elif "+openmp" in spec: