[SE] Let users specify CUDA path

Summary: Add logic to allow users to specify the CUDA path at configuration time.

Reviewers: jlebar

Subscribers: beanz, mgorny, jlebar, jprice, parallel_libs-commits

Differential Revision: https://reviews.llvm.org/D24580

llvm-svn: 281626

4 files changed, 63 insertions, 51 deletions

diff --git a/parallel-libs/streamexecutor/CMakeLists.txt b/parallel-libs/streamexecutor/CMakeLists.txt
index e5a2950438b0..b1862c5e5245 100644
--- a/parallel-libs/streamexecutor/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/CMakeLists.txt
@@ -2,10 +2,18 @@ cmake_minimum_required(VERSION 3.1)
 
 option(STREAM_EXECUTOR_UNIT_TESTS "enable unit tests" ON)
 option(STREAM_EXECUTOR_ENABLE_DOXYGEN "enable StreamExecutor doxygen" ON)
-option(STREAM_EXECUTOR_ENABLE_CONFIG_TOOL "enable building streamexecutor-config tool" ON)
-option(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM "enable building the CUDA StreamExecutor platform" OFF)
-
-configure_file("include/streamexecutor/PlatformOptions.h.in" "include/streamexecutor/PlatformOptions.h")
+option(
+    STREAM_EXECUTOR_ENABLE_CONFIG_TOOL
+    "enable building streamexecutor-config tool"
+    ON)
+option(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM
+    "enable building the CUDA StreamExecutor platform \
+(see CMake's 'FindCUDA' documentation for info on specifying the CUDA path)"
+    OFF)
+
+configure_file(
+    "include/streamexecutor/PlatformOptions.h.in"
+    "include/streamexecutor/PlatformOptions.h")
 
 # First find includes relative to the streamexecutor top-level source path.
 include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include)
@@ -64,6 +72,24 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 # Add warning flags.
 set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wno-unused-parameter")
 
+# Check for CUDA if it is enabled.
+if(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM)
+    find_package(CUDA REQUIRED)
+    include_directories(${CUDA_INCLUDE_DIRS})
+    find_library(CUDA_DRIVER_LIBRARY cuda)
+    if(NOT CUDA_DRIVER_LIBRARY)
+        message(FATAL_ERROR
+            "could not find libcuda, \
+is the CUDA driver is installed on your system?")
+    endif()
+    set(
+        STREAM_EXECUTOR_CUDA_PLATFORM_TARGET_OBJECT
+        $<TARGET_OBJECTS:streamexecutor_cuda_platform>)
+    set(
+        STREAM_EXECUTOR_LIBCUDA_LIBRARIES
+        ${CUDA_DRIVER_LIBRARY})
+endif(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM)
+
 add_subdirectory(lib)
 add_subdirectory(examples)
 
diff --git a/parallel-libs/streamexecutor/lib/CMakeLists.txt b/parallel-libs/streamexecutor/lib/CMakeLists.txt
index 4209c78ab6ff..6157654a97a2 100644
--- a/parallel-libs/streamexecutor/lib/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/lib/CMakeLists.txt
@@ -3,24 +3,6 @@ macro(add_se_library name)
   set_target_properties(${name} PROPERTIES FOLDER "streamexecutor libraries")
 endmacro(add_se_library)
 
-if(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM)
-    set(
-        CMAKE_MODULE_PATH
-        ${CMAKE_MODULE_PATH}
-        "${CMAKE_CURRENT_SOURCE_DIR}/platforms/cuda/cmake/modules/")
-
-    find_package(Libcuda REQUIRED)
-    include_directories(${LIBCUDA_INCLUDE_DIRS})
-
-    set(
-        STREAM_EXECUTOR_CUDA_PLATFORM_TARGET_OBJECT
-        $<TARGET_OBJECTS:streamexecutor_cuda_platform>)
-
-    set(
-        STREAM_EXECUTOR_LIBCUDA_LIBRARIES
-        ${LIBCUDA_LIBRARIES})
-endif(STREAM_EXECUTOR_ENABLE_CUDA_PLATFORM)
-
 add_subdirectory(platforms)
 
 add_se_library(
diff --git a/parallel-libs/streamexecutor/lib/platforms/cuda/cmake/modules/FindLibcuda.cmake b/parallel-libs/streamexecutor/lib/platforms/cuda/cmake/modules/FindLibcuda.cmake
deleted file mode 100644
index 6572ac09caa2..000000000000
--- a/parallel-libs/streamexecutor/lib/platforms/cuda/cmake/modules/FindLibcuda.cmake
+++ /dev/null
@@ -1,21 +0,0 @@
-# - Try to find the libcuda library
-# Once done this will define
-#  LIBCUDA_FOUND - System has libcuda
-#  LIBCUDA_INCLUDE_DIRS - The libcuda include directories
-#  LIBCUDA_LIBRARIES - The libraries needed to use libcuda
-
-# TODO(jhen): Allow users to specify a search path.
-find_path(LIBCUDA_INCLUDE_DIR cuda.h /usr/local/cuda/include)
-# TODO(jhen): Use the library that goes with the headers.
-find_library(LIBCUDA_LIBRARY cuda)
-
-include(FindPackageHandleStandardArgs)
-# handle the QUIETLY and REQUIRED arguments and set LIBCUDA_FOUND to TRUE if
-# all listed variables are TRUE
-find_package_handle_standard_args(
-    LIBCUDA DEFAULT_MSG LIBCUDA_INCLUDE_DIR LIBCUDA_LIBRARY)
-
-mark_as_advanced(LIBCUDA_INCLUDE_DIR LIBCUDA_LIBRARY)
-
-set(LIBCUDA_LIBRARIES ${LIBCUDA_LIBRARY})
-set(LIBCUDA_INCLUDE_DIRS ${LIBCUDA_INCLUDE_DIR})
diff --git a/parallel-libs/streamexecutor/tools/streamexecutor-config/streamexecutor-config.in b/parallel-libs/streamexecutor/tools/streamexecutor-config/streamexecutor-config.in
index 6fb4bec21e54..f3e15305d27c 100755
--- a/parallel-libs/streamexecutor/tools/streamexecutor-config/streamexecutor-config.in
+++ b/parallel-libs/streamexecutor/tools/streamexecutor-config/streamexecutor-config.in
@@ -32,20 +32,42 @@ import shlex
 import subprocess
 import sys
 
-def get_llvm_config_dir():
-  """Gets the path to the llvm-config executable.
+# The following functions are configured by cmake. They use raw triple-quoted
+# strings to surround values that are substituted by cmake at configure time.
+# This kind of quoting should allow for paths that contain spaces.
 
-  Surrounding the cmake replacement with triple quotes should allow for paths
-  that contain quotes."""
+def get_llvm_config_dir():
+  """Gets the path to the llvm-config executable."""
   return r"""@LLVM_BINARY_DIR@/bin"""
 
 def get_cmake_install_prefix():
-  """Gets the value of the cmake variable CMAKE_INSTALL_PREFIX.
-
-  Surrounding the cmake replacement with triple quotes should allow for paths
-  that contain quotes."""
+  """Gets the value of the cmake variable CMAKE_INSTALL_PREFIX."""
   return r"""@CMAKE_INSTALL_PREFIX@"""
 
+def convert_library_name(library_name):
+  """Converts a library name ending in '.framework' into a '-framework' flag.
+
+  This is used to support OS X.
+
+  >>> convert_library_name('')
+  ''
+
+  >>> convert_library_name('/usr/local/lib64/libcuda.so')
+  '/usr/local/lib64/libcuda.so'
+
+  >>> convert_library_name('/Library/Frameworks/cuda.framework')
+  '-framework cuda'
+  """
+  framework_suffix = '.framework'
+  if library_name.endswith(framework_suffix):
+    framework_name = os.path.basename(library_name)[:-len(framework_suffix)]
+    library_name = '-framework ' + framework_name
+  return library_name
+
+def get_cuda_driver_library():
+  """Gets the value of the cmake variable CUDA_DRIVER_LIBRARY."""
+  return convert_library_name(r"""@CUDA_DRIVER_LIBRARY@""")
+
 def cuddle_flag(flag, tokens):
   """If flag appears by itself in tokens, combines it with the next token.
 
@@ -195,6 +217,9 @@ def main():
     se_flag = '-lstreamexecutor'
     if not has_token(token=se_flag, string=llvm_flags):
       all_flags.append(se_flag)
+    cuda_driver_library = get_cuda_driver_library()
+    if cuda_driver_library:
+      all_flags.append(cuda_driver_library)
     all_flags.append(llvm_flags)
 
   if args.system_libs: