Allow compiling against slightly different tensorflow directory layouts. (#21199)

Author: pschuh

cmake/modules/FindTensorFlow.cmake

@@ -4,8 +4,15 @@
 include(FindPackageHandleStandardArgs)
 
 find_path(TF_INCLUDE_DIR
-  NAMES tensorflow/c
+  NAMES third_party/tensorflow/c tensorflow/c
   HINTS ${SWIFT_TENSORFLOW_TARGET_INCLUDE_DIR} /usr/include /usr/local/include)
+if (EXISTS ${TF_INCLUDE_DIR}/third_party/tensorflow/c/c_api.h)
+  # This is experimental and not covered by CI.
+  set(TF_PATH_ADJUSTMENT "third_party")
+else()
+  # Note: This is the normal workflow.
+  set(TF_PATH_ADJUSTMENT "")
+endif()
 
 find_library(TF_LIBRARY
   NAMES tensorflow

lib/AST/CMakeLists.txt

@@ -5,6 +5,9 @@ endif()
 
 if(SWIFT_ENABLE_TENSORFLOW)
   find_package(TensorFlow REQUIRED)
+  if (TF_PATH_ADJUSTMENT)
+    include_directories(BEFORE "${TF_INCLUDE_DIR}/${TF_PATH_ADJUSTMENT}")
+  endif()
   include_directories(BEFORE "${TF_INCLUDE_DIR}")
 endif()
 

lib/SIL/CMakeLists.txt

@@ -1,5 +1,8 @@
 if(SWIFT_ENABLE_TENSORFLOW)
   find_package(TensorFlow REQUIRED)
+  if (TF_PATH_ADJUSTMENT)
+    include_directories(BEFORE "${TF_INCLUDE_DIR}/${TF_PATH_ADJUSTMENT}" )
+  endif()
   include_directories(BEFORE "${TF_INCLUDE_DIR}" )
 endif()
 

lib/SILOptimizer/CMakeLists.txt

@@ -1,5 +1,8 @@
 if(SWIFT_ENABLE_TENSORFLOW)
   find_package(TensorFlow REQUIRED)
+  if (TF_PATH_ADJUSTMENT)
+    include_directories(BEFORE "${TF_INCLUDE_DIR}/${TF_PATH_ADJUSTMENT}")
+  endif()
   include_directories(BEFORE "${TF_INCLUDE_DIR}")
 endif()
 

stdlib/public/CTensorFlow/CMakeLists.txt

@@ -23,6 +23,9 @@ include("../../../cmake/modules/SwiftList.cmake")
 find_package(TensorFlow REQUIRED)
 message(STATUS "Building CTensorFlow.")
 
+if (TF_PATH_ADJUSTMENT)
+  include_directories(BEFORE "${TF_INCLUDE_DIR}/${TF_PATH_ADJUSTMENT}")
+endif()
 include_directories(BEFORE "${TF_INCLUDE_DIR}")
 
 # Get target SDKs.
@@ -53,20 +56,31 @@ foreach(sdk ${TARGET_SDKS})
 
     # Copy headers.
     macro(copy_tf_headers src dst)
+      if (TF_PATH_ADJUSTMENT)
+        set(copy_tf_headers_path_adjustment "${TF_PATH_ADJUSTMENT}/")
+      else()
+        set(copy_tf_headers_path_adjustment "")
+      endif()
       add_custom_command(
         OUTPUT "${output_dir}/${dst}"
         DEPENDS "${src}"
         COMMAND
           "${SWIFT_SOURCE_DIR}/utils/fix-tf-headers.sh"
           "${src}"
-          "${output_dir}/${dst}")
+          "${output_dir}/${dst}"
+          "${copy_tf_headers_path_adjustment}")
 
       list(APPEND outputs "${output_dir}/${dst}")
     endmacro(copy_tf_headers)
 
-    copy_tf_headers("${TF_INCLUDE_DIR}/tensorflow/c/c_api.h" c_api.h)
-    copy_tf_headers("${TF_INCLUDE_DIR}/tensorflow/c/c_api_experimental.h" c_api_experimental.h)
-    copy_tf_headers("${TF_INCLUDE_DIR}/tensorflow/c/eager/c_api.h" c_api_eager.h)
+    if (TF_PATH_ADJUSTMENT)
+      set(adjusted_tf_include_dir "${TF_INCLUDE_DIR}/${TF_PATH_ADJUSTMENT}")
+    else()
+      set(adjusted_tf_include_dir "${TF_INCLUDE_DIR}")
+    endif()
+    copy_tf_headers("${adjusted_tf_include_dir}/tensorflow/c/c_api.h" c_api.h)
+    copy_tf_headers("${adjusted_tf_include_dir}/tensorflow/c/c_api_experimental.h" c_api_experimental.h)
+    copy_tf_headers("${adjusted_tf_include_dir}/tensorflow/c/eager/c_api.h" c_api_eager.h)
     copy_tf_headers("${CMAKE_CURRENT_SOURCE_DIR}/ctensorflow_init.h" ctensorflow_init.h)
 
     # Put the output dir itself last so that it isn't considered the primary output.

stdlib/public/TensorFlow/CMakeLists.txt

@@ -21,16 +21,13 @@ endif()
 find_package(TensorFlow REQUIRED)
 message(STATUS "Building TensorFlow.")
 
-include_directories(BEFORE "${TF_INCLUDE_DIR}")
-
 set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
 set(swift_stdlib_compile_flags "${SWIFT_RUNTIME_SWIFT_COMPILE_FLAGS}")
 list(APPEND swift_stdlib_compile_flags "-Xllvm" "-sil-inline-generics")
 list(APPEND swift_stdlib_compile_flags "-Xllvm" "-sil-partial-specialization")
 list(APPEND swift_stdlib_compile_flags "-Xfrontend" "-enable-sil-ownership")
 list(APPEND swift_stdlib_compile_flags "-swift-version" "4")
 list(APPEND swift_stdlib_compile_flags "-force-single-frontend-invocation")
-list(APPEND swift_stdlib_compile_flags "-Xcc" "-I${TF_INCLUDE_DIR}")
 # FIXME(SR-7972): Some tests fail when TensorFlow is optimized.
 list(APPEND swift_stdlib_compile_flags "-Onone")
 

utils/fix-tf-headers.sh

@@ -7,9 +7,10 @@ set -x
 
 src="$1"
 dst="$2"
+path_adjustment="$3"
 
 mkdir -p "$(dirname $2)"
 cp "$1" "$2"
-sed -i -e 's#include "tensorflow/c/c_api.h"#include "c_api.h"#g' "$2"
-sed -i -e 's#include "tensorflow/c/c_api_experimental.h"#include "c_api_experimental.h"#g' "$2"
-sed -i -e 's#include "tensorflow/c/eager/c_api.h"#include "c_api_eager.h"#g' "$2"
+sed -i -e 's#include "'"$3"'tensorflow/c/c_api.h"#include "c_api.h"#g' "$2"
+sed -i -e 's#include "'"$3"'tensorflow/c/c_api_experimental.h"#include "c_api_experimental.h"#g' "$2"
+sed -i -e 's#include "'"$3"'tensorflow/c/eager/c_api.h"#include "c_api_eager.h"#g' "$2"