diff options
-rw-r--r-- | tensorflow/contrib/nccl/kernels/nccl_manager.h | 7 | ||||
-rw-r--r-- | tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc | 9 |
2 files changed, 12 insertions, 4 deletions
diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager.h b/tensorflow/contrib/nccl/kernels/nccl_manager.h index 57a96c5d33..09fad35d23 100644 --- a/tensorflow/contrib/nccl/kernels/nccl_manager.h +++ b/tensorflow/contrib/nccl/kernels/nccl_manager.h @@ -20,6 +20,13 @@ limitations under the License. #include <unordered_map> #include <vector> +// TODO(rmlarsen): Get rid of this workaround. "gpu_assert" is defined when +// setting EIGEN_USE_THREADS. But when defining EIGEN_USE_THREADS here, +// incAtomic and other CUDA specific symbols are no longer recognized. +#ifndef gpu_assert +#define gpu_assert(x) +#endif + #include "third_party/nccl/nccl.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" #include "tensorflow/core/framework/tensor.h" diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc index 2de7973750..11335d7da6 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc @@ -13,14 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + #include "tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.h" #include <vector> +#define EIGEN_USE_GPU #include "tensorflow/core/framework/op_kernel.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT #include "cuda/include/cuda_runtime_api.h" #include "tensorflow/core/platform/stream_executor.h" @@ -80,5 +81,5 @@ REGISTER_KERNEL_BUILDER(Name("IncPluginTRT").Device(DEVICE_GPU), IncPluginTRT); } // namespace tensorrt } // namespace tensorflow -#endif // GOOGLE_CUDA #endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA |