Add tc_config.h.in

Author: nicolasvasilache

CMakeLists.txt

@@ -139,9 +139,6 @@ if(WITH_CUDA)
   find_package(CUDA REQUIRED)
   include_directories(BEFORE ${CUDA_TOOLKIT_ROOT_DIR}/include)
 
-  # modified CUB
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DWITH_CUDA -DCUDA_HOME=\"\\\"${CUDA_INCLUDE_DIRS}\\\"\" -DCUB_HOME=\"\\\"${CUB_INSTALL_DIR}\\\"\" ")
-
   # Inherited from Torch, see
   # https://github.com/torch/cutorch/blob/master/lib/THC/cmake/select_compute_arch.cmake
   INCLUDE(cmake/select_compute_arch.cmake)
@@ -259,6 +256,18 @@ include(cmake/GetGitRevisionDescription.cmake)
 ################################################################################
 # Finally, build
 ################################################################################
+# Variables for tc_config.h.in
+set(TC_DIR ${TC_DIR})
+if (WITH_CUDA)
+   # CUDA-specific variables for tc_config.h.in
+   set(TC_WITH_CUDA 1)
+   set(TC_CUB_INCLUDE_DIR ${CUB_INSTALL_DIR})
+   set(TC_CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT_ROOT_DIR})
+   set(TC_CUDA_INCLUDE_DIR ${CUDA_INCLUDE_DIRS})
+else()
+   set(TC_WITH_CUDA 0)
+endif()
+configure_file("tc/tc_config.h.in" "${CMAKE_CURRENT_BINARY_DIR}/tc/tc_config.h")
 
 ################################################################################
 # Compile flags
@@ -281,13 +290,12 @@ elseif(${CMAKE_BUILD_TYPE} MATCHES "Release")
 endif()
 message(STATUS "CMAKE_INSTALL_PREFIX is ${CMAKE_INSTALL_PREFIX}")
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTC_DIR=\"\\\"${TC_DIR}\\\"\" ")
-
 include_directories(BEFORE ${PROJECT_SOURCE_DIR})
+include_directories(BEFORE ${CMAKE_CURRENT_BINARY_DIR})
 add_subdirectory(tc)
 
 # At the moment pybind is only supported in CUDA mode and compilation fails
-# for non-CUDA mode (CUDA_HOME and CUB_HOME undefined error).
+# for non-CUDA mode (CUDA_INCLUDE_DIR and CUB_INCLUDE_DIR undefined error).
 # Once the core CPU mapper is stabilized we can worry about pybind, deactivate
 # conditionally for now
 if (WITH_CUDA)

tc/core/cuda/cuda.h

@@ -15,13 +15,15 @@
  */
 #pragma once
 
-#ifndef CUDA_HOME
-#error "CUDA_HOME must be defined"
-#endif // CUDA_HOME
+#include "tc/tc_config.h"
 
-#ifndef CUB_HOME
-#error "CUB_HOME must be defined"
-#endif // CUB_HOME
+#ifndef TC_CUDA_INCLUDE_DIR
+#error "TC_CUDA_INCLUDE_DIR must be defined"
+#endif // TC_CUDA_INCLUDE_DIR
+
+#ifndef TC_CUB_INCLUDE_DIR
+#error "TC_CUB_INCLUDE_DIR must be defined"
+#endif // TC_CUB_INCLUDE_DIR
 
 #include <sstream>
 #include <stdexcept>

tc/core/cuda/cuda_rtc.cc

@@ -91,8 +91,10 @@ std::unique_ptr<CudaRTCFunction> CudaRTCFunction::Compile(
 
   // Compile the program.
   const char* nvrtc_debug_opts[] = {"-G", "-lineinfo"};
-  std::string cudaHome = std::string("-I ") + std::string(CUDA_HOME);
-  std::string cubHome = std::string("-I ") + std::string(CUB_HOME);
+  std::string cudaHome =
+      std::string("-I ") + std::string(TC_STRINGIFY(TC_CUDA_INCLUDE_DIR));
+  std::string cubHome =
+      std::string("-I ") + std::string(TC_STRINGIFY(TC_CUB_INCLUDE_DIR));
   std::vector<const char*> nvrtcts = {arch.c_str(),
                                       "--use_fast_math",
                                       "-std=c++11",

tc/core/gpu.h

@@ -15,9 +15,11 @@
  */
 #pragma once
 
+#include "tc/tc_config.h"
+
 // Conditionally include CUDA-specific headers.  This file should compile even
 // without them.
-#if defined(CUDA_HOME)
+#if TC_WITH_CUDA
 #include "tc/core/cuda/cuda.h"
 #endif
 
@@ -27,7 +29,7 @@ namespace tc {
 /// The call is forwarded to the appropriate GPU driver (CUDA in particular).
 /// If a thread has no associated GPU device, return 0.
 inline size_t querySharedMemorySize() {
-#if defined(CUDA_HOME) && !defined(NO_CUDA_SDK)
+#if TC_WITH_CUDA && !defined(NO_CUDA_SDK)
   return CudaGPUInfo::GPUInfo().SharedMemorySize();
 #else
   return 0;

tc/tc_config.h.in

@@ -0,0 +1,27 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#define TC_STRINGIFY_INNER(s) #s
+#define TC_STRINGIFY(s) TC_STRINGIFY_INNER(s)
+
+// clang-format off
+#define TC_DIR @TC_DIR@
+#define TC_WITH_CUDA @TC_WITH_CUDA@
+#define TC_CUDA_TOOLKIT_ROOT_DIR @TC_CUDA_TOOLKIT_ROOT_DIR@
+#define TC_CUDA_INCLUDE_DIR @TC_CUDA_INCLUDE_DIR@
+#define TC_CUB_INCLUDE_DIR @TC_CUB_INCLUDE_DIR@
+// clang-format on

test/cuda/test_basic_gpu.cc

@@ -51,7 +51,8 @@ std::vector<char> jitCompile(
   std::string arch = arch_param.str();
 
   // Compile the program.
-  std::string cudaHome = std::string("-I ") + std::string(CUDA_HOME);
+  std::string cudaHome =
+      std::string("-I ") + std::string(TC_STRINGIFY(TC_CUDA_INCLUDE_DIR));
   std::vector<const char*> nvrtcts = {arch.c_str(),
                                       "--use_fast_math",
                                       "-std=c++11",
@@ -116,7 +117,7 @@ __global__ void foo(int N)
 }
 
 TEST(BasicGpuTest, CubReduce) {
-  std::string path(CUB_HOME);
+  std::string path(TC_STRINGIFY(TC_CUB_INCLUDE_DIR));
   std::string include = std::string("-I ") + path;
   auto PTX = jitCompile(
       R"CUDA(

test/test_lang.cc

@@ -26,6 +26,7 @@
 #include "tc/lang/parser.h"
 #include "tc/lang/sema.h"
 #include "tc/lang/tc_format.h"
+#include "tc/tc_config.h"
 
 using namespace lang;
 
@@ -34,7 +35,7 @@ using namespace lang;
 #endif
 
 const std::string expected_file_path =
-    std::string(TC_DIR) + "/tc/lang/test_expected/";
+    std::string(TC_STRINGIFY(TC_DIR)) + "/tc/lang/test_expected/";
 
 static inline void barf(const char* fmt, ...) {
   char msg[2048];