Revert "Stop depending on libtensorflow_framework.so just for InitMain()."

README.md

@@ -147,7 +147,7 @@ Below is more information about TensorFlow-related build arguments.
     * Default: None.
 * `tensorflow-host-include-dir`: A directory containing custom TensorFlow headers.
     * Default value: None.
-* `tensorflow-host-lib-dir`: A directory containing custom TensorFlow shared libraries (`libtensorflow.so`).
+* `tensorflow-host-lib-dir`: A directory containing custom TensorFlow shared libraries (`libtensorflow.so` and `libtensorflow_framework.so`).
     * Default value: None.
 * `tensorflow-swift-bindings`: A generated TensorFlow Swift bindings file (`RawOpsGenerated.swift`) obtained from [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings).
     * Default value: `tensorflow-swift-bindings/RawOpsGenerated.swift` if the [tensorflow/swift-bindings](https://github.com/tensorflow/swift-bindings) repository is cloned. Otherwise, none.

cmake/modules/FindTensorFlow.cmake

@@ -10,7 +10,10 @@ find_path(TF_INCLUDE_DIR
 find_library(TF_LIBRARY
   NAMES tensorflow
   HINTS ${SWIFT_TENSORFLOW_TARGET_LIB_DIR} /usr/lib /usr/local/lib)
-set(TF_LIBRARIES ${TF_LIBRARY})
+find_library(TF_FRAMEWORK_LIBRARY
+  NAMES tensorflow_framework
+  HINTS ${SWIFT_TENSORFLOW_TARGET_LIB_DIR} /usr/lib /usr/local/lib)
+set(TF_LIBRARIES ${TF_LIBRARY} ${TF_FRAMEWORK_LIBRARY})
 
 find_package_handle_standard_args(TensorFlow DEFAULT_MSG TF_INCLUDE_DIR TF_LIBRARIES)
 mark_as_advanced(${TF_INCLUDE_DIR} ${TF_LIBRARIES})

stdlib/public/CTensorFlow/ctensorflow_init.cpp

@@ -3,6 +3,7 @@
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/c_api_experimental.h"
 #include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/platform/init_main.h"
 
 #include <assert.h>
 #include <vector>
@@ -30,7 +31,7 @@ void InitTensorFlowRuntime(unsigned char enable_debug_logging,
   int my_argc = my_argv.size();
   char** tmpArgv = my_argv.data();
   // Initialize GPU devices.
-  TF_InitMain(/*usage=*/nullptr, &my_argc, &tmpArgv);
+  tensorflow::port::InitMain(/*usage=*/nullptr, &my_argc, &tmpArgv);
 }
 
 static bool setValue(TF_DataType tfDtype, int64_t val, void *ptr) {

test/lit.cfg

@@ -331,7 +331,7 @@ config.swift_test_options = '-swift-version ' + swift_version
 swift_tensorflow_extra_options = ""
 swift_tensorflow_path = lit_config.params.get('swift_tensorflow_path', None)
 if swift_tensorflow_path:
-    swift_tensorflow_extra_options += ('-L%s -ltensorflow' % swift_tensorflow_path)
+    swift_tensorflow_extra_options += ('-L%s -ltensorflow -ltensorflow_framework' % swift_tensorflow_path)
 swift_tensorflow_gpu = lit_config.params.get('swift_tensorflow_gpu', None)
 # Can use #if on CUDA to selectively disable tests for GPU.
 if swift_tensorflow_gpu:

utils/build-script

@@ -266,7 +266,7 @@ class BuildScriptInvocation(object):
 
         # SWIFT_ENABLE_TENSORFLOW
         if args.enable_tensorflow:
-            tensorflow_libs = ["libtensorflow.so"]
+            tensorflow_libs = ["libtensorflow.so", "libtensorflow_framework.so"]
             tensorflow_dirs = [
                 ("tensorflow_host_lib_dir", True),
                 ("tensorflow_host_include_dir", False),

utils/build-script-impl

@@ -1318,7 +1318,7 @@ function verify_tensorflow_directories() {
 
                 # Check that "LIB" arguments contain the necessary libraries.
                 if [[ "dir_type" == "LIB" ]]; then
-                    for lib_name in tensorflow; do
+                    for lib_name in tensorflow tensorflow_framework; do
                         lib="lib${lib_name}.so"
                         if [[ ! -e "${!dir_name}/${lib}" ]]; then
                             echo "Error: '${lib}' does not exist in TensorFlow ${lower_platform} ${lower_dir_type} directory (was '${!dir_name}')"
@@ -1954,7 +1954,7 @@ function set_lldb_xcodebuild_options() {
         lldb_xcodebuild_options=(
             "${lldb_xcodebuild_options[@]}"
             ENABLE_TENSORFLOW=1
-            EXTRA_OTHER_LDFLAGS="-L\$(LLDB_PATH_TO_SWIFT_BUILD)/lib/swift/\$(PLATFORM_NAME) ${RPATHS} -ltensorflow"
+            EXTRA_OTHER_LDFLAGS="-L\$(LLDB_PATH_TO_SWIFT_BUILD)/lib/swift/\$(PLATFORM_NAME) ${RPATHS} -ltensorflow -ltensorflow_framework"
         )
     fi
 }
@@ -2386,7 +2386,7 @@ for host in "${ALL_HOSTS[@]}"; do
                     # Bazel-built libraries do not have write permission, which is
                     # problematic for overwriting/stripping symbols. Thus, write
                     # permission is added here.
-                    for lib_name in tensorflow; do
+                    for lib_name in tensorflow tensorflow_framework; do
                         lib="lib${lib_name}.so"
                         rm -f "${TF_LIB_DIR}/${lib}"
                         chmod +w "${TENSORFLOW_HOST_LIB_DIR}/${lib}"
@@ -2887,7 +2887,7 @@ for host in "${ALL_HOSTS[@]}"; do
                 # Configure TensorFlow and build using bazel.
                 echo yes "" | "${source_dir}/configure"
                 with_pushd "${source_dir}" \
-                    call bazel build -c opt "${TENSORFLOW_BAZEL_OPTIONS[@]}" //tensorflow:libtensorflow.so
+                    call bazel build -c opt "${TENSORFLOW_BAZEL_OPTIONS[@]}" //tensorflow:libtensorflow.so //tensorflow:libtensorflow_framework.so
 
                 # Set TensorFlow include/lib directories.
                 TENSORFLOW_HOST_INCLUDE_DIR="${source_dir}"

utils/build_swift/driver_arguments.py

@@ -954,7 +954,7 @@ def create_argument_parser():
     option('--tensorflow-host-lib-dir', store_path,
            default=None,
            help='Path to a directory containing TensorFlow libraries '
-                '(libtensorflow.so). '
+                '(libtensorflow.so and libtensorflow_framework.so). '
                 'Used for linking swiftc.')
     option('--tensorflow-host-include-dir', store_path,
            default=None,
@@ -963,7 +963,7 @@ def create_argument_parser():
     option('--tensorflow-target-lib-dir', store_path,
            default=None,
            help='Path to a directory containing TensorFlow libraries '
-                '(libtensorflow.so). '
+                '(libtensorflow.so and libtensorflow_framework.so). '
                 'Used for linking Swift programs.')
     option('--tensorflow-target-include-dir', store_path,
            default=None,