vtk-m: correct cuda_arch variant behavior (#38697)

Co-authored-by: eugeneswalker <eugenesunsetwalker@gmail.com>

2 files changed, 48 insertions, 0 deletions

diff --git a/lib/spack/spack/build_systems/cmake.py b/lib/spack/spack/build_systems/cmake.py
index 8e819f3958..7fff3771e9 100644
--- a/lib/spack/spack/build_systems/cmake.py
+++ b/lib/spack/spack/build_systems/cmake.py
@@ -296,9 +296,47 @@ class CMakeBuilder(BaseBuilder):
                 define("CMAKE_PREFIX_PATH", spack.build_environment.get_cmake_prefix_path(pkg)),
             ]
         )
+
         return args
 
     @staticmethod
+    def define_cuda_architectures(pkg):
+        """Returns the str ``-DCMAKE_CUDA_ARCHITECTURES:STRING=(expanded cuda_arch)``.
+
+        ``cuda_arch`` is variant composed of a list of target CUDA architectures and
+        it is declared in the cuda package.
+
+        This method is no-op for cmake<3.18 and when ``cuda_arch`` variant is not set.
+
+        """
+        cmake_flag = str()
+        if "cuda_arch" in pkg.spec.variants and pkg.spec.satisfies("^cmake@3.18:"):
+            cmake_flag = CMakeBuilder.define(
+                "CMAKE_CUDA_ARCHITECTURES", pkg.spec.variants["cuda_arch"].value
+            )
+
+        return cmake_flag
+
+    @staticmethod
+    def define_hip_architectures(pkg):
+        """Returns the str ``-DCMAKE_HIP_ARCHITECTURES:STRING=(expanded amdgpu_target)``.
+
+        ``amdgpu_target`` is variant composed of a list of the target HIP
+        architectures and it is declared in the rocm package.
+
+        This method is no-op for cmake<3.18 and when ``amdgpu_target`` variant is
+        not set.
+
+        """
+        cmake_flag = str()
+        if "amdgpu_target" in pkg.spec.variants and pkg.spec.satisfies("^cmake@3.21:"):
+            cmake_flag = CMakeBuilder.define(
+                "CMAKE_HIP_ARCHITECTURES", pkg.spec.variants["amdgpu_target"].value
+            )
+
+        return cmake_flag
+
+    @staticmethod
     def define(cmake_var, value):
         """Return a CMake command line argument that defines a variable.
 
diff --git a/lib/spack/spack/test/build_systems.py b/lib/spack/spack/test/build_systems.py
index afb9f51b77..54ed485f25 100644
--- a/lib/spack/spack/test/build_systems.py
+++ b/lib/spack/spack/test/build_systems.py
@@ -311,6 +311,16 @@ class TestCMakePackage:
         with pytest.raises(KeyError, match="not a variant"):
             s.package.define_from_variant("NONEXISTENT")
 
+    def test_cmake_std_args_cuda(self, default_mock_concretization):
+        s = default_mock_concretization("vtk-m +cuda cuda_arch=70 ^cmake@3.23")
+        option = spack.build_systems.cmake.CMakeBuilder.define_cuda_architectures(s.package)
+        assert "-DCMAKE_CUDA_ARCHITECTURES:STRING=70" == option
+
+    def test_cmake_std_args_hip(self, default_mock_concretization):
+        s = default_mock_concretization("vtk-m +rocm amdgpu_target=gfx900 ^cmake@3.23")
+        option = spack.build_systems.cmake.CMakeBuilder.define_hip_architectures(s.package)
+        assert "-DCMAKE_HIP_ARCHITECTURES:STRING=gfx900" == option
+
 
 @pytest.mark.usefixtures("config", "mock_packages")
 class TestDownloadMixins: