From dc5889b717f1712a52b5f4ad34cdd9c5d57550c8 Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Wed, 24 Feb 2021 14:17:20 -0500
Subject: [PATCH] Simplify cmake cuda architectures handling (#709)

Based on #706 as they both modify `cmake/Modules/SetGPUArchs.cmake`

This brings rmm's handling of `CMAKE_CUDA_ARCHITECTURES` to match that is proposed for cudf in https://github.com/rapidsai/cudf/pull/7391

Authors:
  - Robert Maynard (@robertmaynard)

Approvers:
  - Mark Harris (@harrism)
  - Keith Kraus (@kkraus14)

URL: https://github.com/rapidsai/rmm/pull/709
---
 CMakeLists.txt                   | 16 ++++++++++++++--
 cmake/Modules/EvalGPUArchs.cmake |  2 --
 cmake/Modules/SetGPUArchs.cmake  | 31 +++++++++++--------------------
 3 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 11a855ff9..e15263086 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -13,6 +13,19 @@
 # =============================================================================
 cmake_minimum_required(VERSION 3.18...3.18 FATAL_ERROR)
 
+# If `CMAKE_CUDA_ARCHITECTURES` is not defined, build for all supported architectures. If
+# `CMAKE_CUDA_ARCHITECTURES` is set to an empty string (""), build for only the current
+# architecture. If `CMAKE_CUDA_ARCHITECTURES` is specified by the user, use user setting.
+
+# This needs to be run before enabling the CUDA language due to the default initialization behavior
+# of `CMAKE_CUDA_ARCHITECTURES`, https://gitlab.kitware.com/cmake/cmake/-/issues/21302
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+  set(RMM_BUILD_FOR_ALL_ARCHS TRUE)
+elseif(CMAKE_CUDA_ARCHITECTURES STREQUAL "")
+  unset(CMAKE_CUDA_ARCHITECTURES CACHE)
+  set(RMM_BUILD_FOR_DETECTED_ARCHS TRUE)
+endif()
+
 project(
   RMM
   VERSION 0.19.0
@@ -85,10 +98,9 @@ message(STATUS "RMM: RMM_LOGGING_LEVEL = '${RMM_LOGGING_LEVEL}'")
 
 if(BUILD_TESTS OR BUILD_BENCHMARKS)
   # Auto-detect available GPU compute architectures
+  enable_language(CUDA)
   include(${RMM_SOURCE_DIR}/cmake/Modules/SetGPUArchs.cmake)
   message(STATUS "RMM: Building benchmarks with GPU Architectures: ${CMAKE_CUDA_ARCHITECTURES}")
-  # Only enable the CUDA language after including SetGPUArchs.cmake
-  enable_language(CUDA)
 endif()
 
 # optionally build tests
diff --git a/cmake/Modules/EvalGPUArchs.cmake b/cmake/Modules/EvalGPUArchs.cmake
index c3afb8797..1d9f72f1f 100644
--- a/cmake/Modules/EvalGPUArchs.cmake
+++ b/cmake/Modules/EvalGPUArchs.cmake
@@ -10,8 +10,6 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 
-enable_language(CUDA)
-
 # Function uses the CUDA runtime API to query the compute capability of the device, so if a user
 # doesn't pass any architecture options to CMake we only build the current architecture
 function(evaluate_gpu_archs gpu_archs)
diff --git a/cmake/Modules/SetGPUArchs.cmake b/cmake/Modules/SetGPUArchs.cmake
index 86f92350e..f43799a75 100644
--- a/cmake/Modules/SetGPUArchs.cmake
+++ b/cmake/Modules/SetGPUArchs.cmake
@@ -48,27 +48,18 @@ if(CUDAToolkit_VERSION_MAJOR LESS 9)
   list(REMOVE_ITEM SUPPORTED_CUDA_ARCHITECTURES "70")
 endif()
 
-# If `CMAKE_CUDA_ARCHITECTURES` is not defined, build for all supported architectures. If
-# `CMAKE_CUDA_ARCHITECTURES` is set to an empty string (""), build for only the current
-# architecture. If `CMAKE_CUDA_ARCHITECTURES` is specified by the user, use user setting.
-
-# This needs to be run before enabling the CUDA language due to the default initialization behavior
-# of `CMAKE_CUDA_ARCHITECTURES`, https://gitlab.kitware.com/cmake/cmake/-/issues/21302
-
-if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+if(${PROJECT_NAME}_BUILD_FOR_ALL_ARCHS)
   set(CMAKE_CUDA_ARCHITECTURES ${SUPPORTED_CUDA_ARCHITECTURES})
-endif(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
 
-if(CMAKE_CUDA_ARCHITECTURES STREQUAL "")
-  unset(CMAKE_CUDA_ARCHITECTURES)
-  unset(CMAKE_CUDA_ARCHITECTURES CACHE)
-  include(${RMM_SOURCE_DIR}/cmake/Modules/EvalGPUArchs.cmake)
+  # CMake architecture list entry of "80" means to build compute and sm. What we want is for the
+  # newest arch only to build that way while the rest built only for sm.
+  list(POP_BACK CMAKE_CUDA_ARCHITECTURES latest_arch)
+  list(TRANSFORM CMAKE_CUDA_ARCHITECTURES APPEND "-real")
+  list(APPEND CMAKE_CUDA_ARCHITECTURES ${latest_arch})
+
+elseif(${PROJECT_NAME}_BUILD_FOR_DETECTED_ARCHS)
+  include(${PROJECT_SOURCE_DIR}/cmake/Modules/EvalGPUArchs.cmake)
   evaluate_gpu_archs(CMAKE_CUDA_ARCHITECTURES)
-endif(CMAKE_CUDA_ARCHITECTURES STREQUAL "")
 
-# CMake architecture list entry of "80" means to build compute and sm. What we want is for the
-# newest arch only to build that way while the rest built only for sm.
-list(SORT CMAKE_CUDA_ARCHITECTURES ORDER ASCENDING)
-list(POP_BACK CMAKE_CUDA_ARCHITECTURES latest_arch)
-list(TRANSFORM CMAKE_CUDA_ARCHITECTURES APPEND "-real")
-list(APPEND CMAKE_CUDA_ARCHITECTURES ${latest_arch})
+  list(TRANSFORM CMAKE_CUDA_ARCHITECTURES APPEND "-real")
+endif()