From 76f6938bafeb81a4ca41b8dac2b9c83e1286fa95 Mon Sep 17 00:00:00 2001 From: Guangda Lai Date: Thu, 25 Jan 2018 23:59:19 -0800 Subject: Set up TensorRT configurations for external use, and add a test. PiperOrigin-RevId: 183347199 --- tensorflow/tensorflow.bzl | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) (limited to 'tensorflow/tensorflow.bzl') diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 383c97344a..7fe9c98726 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -10,6 +10,10 @@ load( "tf_additional_xla_deps_py", "if_static", ) +load( + "@local_config_tensorrt//:build_defs.bzl", + "if_tensorrt", +) load( "@local_config_cuda//cuda:build_defs.bzl", "if_cuda", @@ -197,6 +201,7 @@ def tf_copts(android_optimization_level_override="-O2", is_external=False): "-fno-exceptions", "-ftemplate-depth=900"]) + if_cuda(["-DGOOGLE_CUDA=1"]) + + if_tensorrt(["-DGOOGLE_TENSORRT=1"]) + if_mkl(["-DINTEL_MKL=1", "-DEIGEN_USE_VML", "-fopenmp",]) + if_android_arm(["-mfpu=neon"]) + if_linux_x86_64(["-msse3"]) @@ -861,9 +866,11 @@ def tf_cuda_library(deps=None, cuda_deps=None, copts=tf_copts(), **kwargs): When the library is built with --config=cuda: - - both deps and cuda_deps are used as dependencies - - the cuda runtime is added as a dependency (if necessary) - - The library additionally passes -DGOOGLE_CUDA=1 to the list of copts + - Both deps and cuda_deps are used as dependencies. + - The cuda runtime is added as a dependency (if necessary). + - The library additionally passes -DGOOGLE_CUDA=1 to the list of copts. + - In addition, when the library is also built with TensorRT enabled, it + additionally passes -DGOOGLE_TENSORRT=1 to the list of copts. Args: - cuda_deps: BUILD dependencies which will be linked if and only if: @@ -882,7 +889,8 @@ def tf_cuda_library(deps=None, cuda_deps=None, copts=tf_copts(), **kwargs): clean_dep("//tensorflow/core:cuda"), "@local_config_cuda//cuda:cuda_headers" ]), - copts=copts + if_cuda(["-DGOOGLE_CUDA=1"]) + if_mkl(["-DINTEL_MKL=1"]), + copts=(copts + if_cuda(["-DGOOGLE_CUDA=1"]) + if_mkl(["-DINTEL_MKL=1"]) + + if_tensorrt(["-DGOOGLE_TENSORRT=1"])), **kwargs) register_extension_info( -- cgit v1.2.3