build tensorflow_wrapper with tensorflow

.gitignore

@@ -1,5 +1,6 @@
 *~
 *.o
+*.so
 *.pyc
 *.d
 .DS_Store*
@@ -23,3 +24,8 @@ compile_commands.json
 .vscode/
 .cache/
 .mypy_cache/
+
+bazel-bin
+bazel-out
+bazel-testlogs
+bazel-tensorflow

CMakeLists.txt

@@ -190,25 +190,14 @@ target_link_libraries(backend PRIVATE backend_obj)
 
 ## backend: TensorFlow support ##
 if(USE_TENSORFLOW)
-    add_library(tensorflow_cc INTERFACE)
-    target_include_directories(tensorflow_cc INTERFACE ${LIBTF_DIR}/tensorflow/include)
-    target_include_directories(tensorflow_cc INTERFACE ${LIBTF_DIR}/protobuf/include)
-    target_link_libraries(tensorflow_cc INTERFACE ${LIBTF_DIR}/tensorflow/lib/libtensorflow_cc.so)
-    target_link_libraries(tensorflow_cc INTERFACE ${LIBTF_DIR}/tensorflow/lib/libtensorflow_framework.so)
-
-    add_library(tensorflow_wrapper src/nexus/backend/tensorflow_wrapper.cpp)
-    target_include_directories(tensorflow_wrapper PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
-    target_link_libraries(tensorflow_wrapper PUBLIC tensorflow_cc)
-
-    # TensorFlow is currently built with C++14. Some TensorFlow APIs expose
-    # pre-adopted std types with absl. To avoid weird ABI compatibility issues,
-    # compile the wrapper with the same C++ standard as TensorFlow.
-    target_compile_options(tensorflow_wrapper PRIVATE -std=c++14)
-
-    target_compile_definitions(backend_obj PUBLIC USE_TENSORFLOW)
+    add_library(tfwrapper INTERFACE)
+    target_include_directories(tfwrapper INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/tensorflow/include)
+    target_link_libraries(tfwrapper INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/tensorflow/lib/libtensorflow_wrapper.so)
+    target_compile_definitions(tfwrapper INTERFACE USE_TENSORFLOW)
+
     target_sources(backend_obj PRIVATE
             src/nexus/backend/tensorflow_model.cpp)
-    target_link_libraries(backend_obj PUBLIC tensorflow_wrapper)
+    target_link_libraries(backend_obj PUBLIC tfwrapper)
 endif()
 
 ## backend: Caffe2 support ##
@@ -279,7 +268,6 @@ endif()
 ## backend: CUDA support ##
 if(USE_GPU)
     target_compile_definitions(common PUBLIC USE_GPU)
-    target_compile_definitions(tensorflow_wrapper PUBLIC USE_GPU)
     target_include_directories(common PUBLIC ${CUDA_PATH}/include)
     target_link_libraries(common PUBLIC
             ${CUDA_PATH}/lib64/stubs/libcuda.so
@@ -339,12 +327,11 @@ target_link_libraries(stress_rankmt PUBLIC bench_dispatcher_obj)
 
 
 ###### tools/test_gpu_mem_sharing ######
-if(USE_GPU AND USE_TENSORFLOW)
-    add_executable(test_gpu_mem_sharing tools/test_gpu_mem_sharing.cpp)
-    target_link_libraries(test_gpu_mem_sharing PUBLIC common)
-    target_link_libraries(test_gpu_mem_sharing PUBLIC tensorflow_wrapper)
-    target_link_libraries(test_gpu_mem_sharing PUBLIC tensorflow_cc)
-endif()
+# if(USE_GPU AND USE_TENSORFLOW)
+#     add_executable(test_gpu_mem_sharing tools/test_gpu_mem_sharing.cpp)
+#     target_link_libraries(test_gpu_mem_sharing PUBLIC common)
+#     target_link_libraries(test_gpu_mem_sharing PUBLIC tfwrapper)
+# endif()
 
 
 

tensorflow/.bazelrc

@@ -0,0 +1,43 @@
+# Cherry-picked from https://github.com/tensorflow/tensorflow/blob/master/.bazelrc
+
+# TensorFlow common settings
+build --spawn_strategy=standalone
+build --announce_rc
+build --define=grpc_no_ares=true
+build --noincompatible_remove_legacy_whole_archive
+build --noincompatible_prohibit_aapt1
+common --experimental_repo_remote_exec
+
+# TensorFlow uses C++14
+build --cxxopt=-std=c++14
+build --host_cxxopt=-std=c++14
+
+# Compiler optimization
+build -c opt
+build --copt=-mavx2
+build --copt=-mfma
+
+# Suppress C++ compiler warnings
+build --copt=-w
+build --host_copt=-w
+
+# TensorFlow 2.x API
+build --define=tf_api_version=2 --action_env=TF2_BEHAVIOR=1
+
+# Monolithic
+build --define=framework_shared_object=false
+
+# Disable features
+build --define=with_xla_support=false
+build --define=no_aws_support=true
+build --define=no_gcp_support=true
+build --define=no_hdfs_support=true
+build --define=no_nccl_support=true
+build --define=build_with_mkl=false --define=enable_mkl=false --define=build_with_openmp=false
+build --repo_env TF_NEED_TENSORRT=0
+build --repo_env TF_NEED_ROCM=0 --define=using_rocm=false --define=using_rocm_hipcc=false
+
+# Options used to build with CUDA.
+build:cuda --repo_env TF_NEED_CUDA=1
+build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
+build:cuda --@local_config_cuda//:enable_cuda

tensorflow/BUILD

@@ -0,0 +1,26 @@
+load(
+    "@org_tensorflow//tensorflow:tensorflow.bzl",
+    "tf_cc_shared_object",
+    "tf_copts",
+)
+
+cc_library(
+    name = "wrapper",
+    srcs = ["tensorflow_wrapper.cpp"],
+    hdrs = ["include/nexus/backend/tensorflow_wrapper.h"],
+    copts = ["-Iinclude"] + tf_copts(),
+    deps = [
+        "@org_tensorflow//tensorflow/core:tensorflow",
+    ],
+)
+
+tf_cc_shared_object(
+    name = "libtensorflow_wrapper.so",
+    linkopts = [
+        "-z defs",
+        "-Wl,--version-script=version_script.lds",
+    ],
+    deps = [
+        ":wrapper",
+    ],
+)

tensorflow/Makefile

@@ -0,0 +1,58 @@
+# Default value from: https://www.tensorflow.org/install/source#gpu
+TF_CUDA_VERSION ?= 11.2
+TF_CUDNN_VERSION ?= 8
+TF_CUDA_COMPUTE_CAPABILITIES ?= sm_35,sm_50,sm_60,sm_70,sm_75,compute_80
+
+
+
+TF_VERSION != grep -Po '(?<="tensorflow-).*(?=")' WORKSPACE
+CPU_SO := libtensorflow_wrapper.so.$(TF_VERSION)-cpu
+GPU_SO := libtensorflow_wrapper.so.$(TF_VERSION)-cuda$(TF_CUDA_VERSION)-cudnn$(TF_CUDNN_VERSION)
+
+DOCKER_IMAGE := gcr.io/tensorflow-testing/nosla-cuda11.2-cudnn8.1-ubuntu18.04-manylinux2010-multipython
+DOCKER_NAME  := tensorflow_wrapper_builder
+DOCKER_START := docker run --rm -d --name $(DOCKER_NAME) -v $(shell pwd):/build --init $(DOCKER_IMAGE) sleep inf
+DOCKER_EXEC  := docker exec -t -w /build $(DOCKER_NAME)
+DOCKER_CLEAN := $(DOCKER_EXEC) rm -rf bazel-bin bazel-out bazel-testlogs bazel-tensorflow bazel-build
+DOCKER_STOP  := docker stop $(DOCKER_NAME)
+BAZEL_BUILD   = bazel build --color=yes --curses=yes \
+		--action_env=PYTHON_BIN_PATH=/usr/local/bin/python3.8 \
+		--crosstool_top=@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11.2:toolchain
+BAZEL_TARGET = //:libtensorflow_wrapper.so
+
+.NOTPARALLEL:
+.DEFAULT_GOAL: prompt
+.PHONY: prompt cpu gpu clean
+prompt:
+	@echo Please \`make gpu\` or \`make cpu\`
+
+gpu: lib/$(GPU_SO)
+	ln -sf $(GPU_SO) lib/libtensorflow_wrapper.so
+lib/$(GPU_SO):
+	$(DOCKER_START)
+	$(DOCKER_EXEC) $(BAZEL_BUILD) \
+		--config=cuda \
+		--action_env TF_CUDA_VERSION=$(TF_CUDA_VERSION) \
+		--action_env TF_CUDNN_VERSION=$(TF_CUDNN_VERSION) \
+		--action_env TF_CUDA_COMPUTE_CAPABILITIES=$(TF_CUDA_COMPUTE_CAPABILITIES) \
+		$(BAZEL_TARGET)
+	docker cp $(DOCKER_NAME):/build/bazel-bin/libtensorflow_wrapper.so lib/$(GPU_SO)
+	$(DOCKER_CLEAN)
+	$(DOCKER_STOP)
+
+
+cpu: lib/$(CPU_SO)
+	ln -sf $(CPU_SO) lib/libtensorflow_wrapper.so
+lib/$(CPU_SO):
+	$(DOCKER_START)
+	$(DOCKER_EXEC) $(BAZEL_BUILD) $(BAZEL_TARGET)
+	docker cp $(DOCKER_NAME):/build/bazel-bin/libtensorflow_wrapper.so lib/$(CPU_SO)
+	$(DOCKER_CLEAN)
+	$(DOCKER_STOP)
+
+clean:
+	rm -f bazel-bin/libtensorflow_wrapper.so
+	rm -f lib/libtensorflow_wrapper.so
+	rm -f lib/$(GPU_SO)
+	rm -f lib/$(CPU_SO)
+	rm -f bazel-out bazel-out bazel-testlogs bazel-tensorflow

tensorflow/WORKSPACE

@@ -0,0 +1,24 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "org_tensorflow",
+    build_file = "@//:BUILD",
+    sha256 = "e3d0ee227cc19bd0fa34a4539c8a540b40f937e561b4580d4bbb7f0e31c6a713",
+    strip_prefix = "tensorflow-2.5.0",
+    urls = ["https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.5.0.zip"],
+)
+
+load(
+    "@org_tensorflow//tensorflow:version_check.bzl",
+    "check_bazel_version_at_least"
+)
+check_bazel_version_at_least("3.7.2")
+
+load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
+tf_workspace3()
+load("@org_tensorflow//tensorflow:workspace2.bzl", "tf_workspace2")
+tf_workspace2()
+load("@org_tensorflow//tensorflow:workspace1.bzl", "tf_workspace1")
+tf_workspace1()
+load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
+tf_workspace0()

src/nexus/backend/tensorflow_wrapper.cpp → tensorflow/tensorflow_wrapper.cpp

@@ -3,7 +3,7 @@
 #include <memory>
 #include <string>
 
-#ifdef USE_GPU
+#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA)
 #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h"
 #else
 #include "tensorflow/core/common_runtime/process_state.h"
@@ -106,7 +106,7 @@ Session::Session(const std::string& visible_device_list,
                  const std::string& pb_path)
     : impl_(std::make_unique<Impl>()) {
   // Init session options
-#ifdef USE_GPU
+#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA)
   auto& tf_option = impl_->gpu_option;
   auto gpu_opt = impl_->gpu_option.config.mutable_gpu_options();
   gpu_opt->set_visible_device_list(visible_device_list);
@@ -140,7 +140,7 @@ Session::Session(const std::string& visible_device_list,
   }
 
   // Get the GPU allocator for creating input buffer
-#ifdef USE_GPU
+#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA)
   impl_->tf_allocator =
       tensorflow::GPUProcessState::singleton()->GetGPUAllocator(
           impl_->gpu_option.config.gpu_options(), tensorflow::TfDeviceId(0), 0,
@@ -175,6 +175,7 @@ Tensor Session::NewTensor(DataType dtype, const std::vector<size_t>& shape) {
     tf_shape.AddDim(dim);
   }
   tensorflow::Tensor tf_tensor(impl_->tf_allocator, tf_dtype, tf_shape);
+  CHECK(tf_tensor.data() != nullptr);
   return TensorProxy::CopyFromTensorFlow(tf_tensor);
 }
 

tensorflow/version_script.lds

@@ -0,0 +1,11 @@
+NEXUS_1.0 {
+  # Export nexus namespace
+  global:
+    extern "C++" {
+      nexus::*;
+    };
+
+  # Hide everything else.
+  local:
+    *;
+};