diff options
| author |  Guangda Lai <laigd@google.com> | 2018-07-19 15:53:13 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-19 16:04:58 -0700 |
| commit | 3647625e531e713ad9a7fb0f3c5b68863ae4e7b8 (patch) | |
| tree | 3172c6569b500597d3505d7a4d6acd7230825475 /tensorflow/contrib/tensorrt/BUILD | |
| parent | 716d15118f62c17c29bbec4d006fd3055bb56812 (diff) | |
Implement something similar to std::align since it's not supported.
PiperOrigin-RevId: 205316355
Diffstat (limited to 'tensorflow/contrib/tensorrt/BUILD')
| -rw-r--r-- | tensorflow/contrib/tensorrt/BUILD | 30 |
1 files changed, 28 insertions, 2 deletions
diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index 70ce4a499c..a9378e9ad6 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -80,6 +80,7 @@ cc_library( copts = tf_copts(), visibility = ["//visibility:public"], deps = [ + ":trt_allocator", ":trt_logging", ":trt_plugins", ":trt_resources", @@ -195,17 +196,16 @@ tf_py_wrap_cc( tf_cuda_library( name = "trt_resources", srcs = [ - "resources/trt_allocator.cc", "resources/trt_int8_calibrator.cc", "resources/trt_resource_manager.cc", ], hdrs = [ - "resources/trt_allocator.h", "resources/trt_int8_calibrator.h", "resources/trt_resource_manager.h", "resources/trt_resources.h", ], deps = [ + ":trt_allocator", ":trt_logging", ":utils", "//tensorflow/core:framework_headers_lib", @@ -216,6 +216,31 @@ tf_cuda_library( ]), ) +tf_cuda_library( + name = "trt_allocator", + srcs = ["resources/trt_allocator.cc"], + hdrs = ["resources/trt_allocator.h"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//tensorflow/core:framework_lite", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), +) + +tf_cc_test( + name = "trt_allocator_test", + size = "small", + srcs = ["resources/trt_allocator_test.cc"], + tags = ["no_windows"], + deps = [ + ":trt_allocator", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + # Library for the node-level conversion portion of TensorRT operation creation tf_cuda_library( name = "trt_conversion", @@ -231,6 +256,7 @@ tf_cuda_library( ], deps = [ ":segment", + ":trt_allocator", ":trt_plugins", ":trt_logging", ":trt_resources", |