From e2694bd4bc2ae61075da98cad47f688b592d53bc Mon Sep 17 00:00:00 2001 From: Benoit Steiner Date: Thu, 3 Nov 2016 10:31:00 -0700 Subject: [PATCH] Don't create SYCL kernels unless TENSORFLOW_USE_SYCL is set --- tensorflow/core/kernels/identity_op.cc | 20 +++++++++++++++++++- tensorflow/core/kernels/no_op.cc | 3 +++ tensorflow/core/kernels/sendrecv_ops.cc | 6 ++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/identity_op.cc b/tensorflow/core/kernels/identity_op.cc index ced8280e3bc..45d27dd19eb 100644 --- a/tensorflow/core/kernels/identity_op.cc +++ b/tensorflow/core/kernels/identity_op.cc @@ -24,7 +24,6 @@ limitations under the License. namespace tensorflow { REGISTER_KERNEL_BUILDER(Name("Identity").Device(DEVICE_CPU), IdentityOp); -REGISTER_KERNEL_BUILDER(Name("Identity").Device(DEVICE_SYCL), IdentityOp); // StopGradient does the same thing as Identity, but has a different // gradient registered. REGISTER_KERNEL_BUILDER(Name("StopGradient").Device(DEVICE_CPU), IdentityOp); @@ -35,6 +34,24 @@ REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault").Device(DEVICE_CPU), REGISTER_KERNEL_BUILDER(Name("RefIdentity").Device(DEVICE_CPU), IdentityOp); +#if TENSORFLOW_USE_SYCL +#define REGISTER_SYCL_KERNEL(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("Identity").Device(DEVICE_SYCL).TypeConstraint("T"), \ + IdentityOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("RefIdentity").Device(DEVICE_SYCL).TypeConstraint("T"), \ + IdentityOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("StopGradient").Device(DEVICE_SYCL).TypeConstraint("T"),\ + IdentityOp) + +TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +REGISTER_SYCL_KERNEL(bfloat16); + +#undef REGISTER_SYCL_KERNEL +#endif + #define REGISTER_GPU_KERNEL(type) \ REGISTER_KERNEL_BUILDER( \ Name("Identity").Device(DEVICE_GPU).TypeConstraint("T"), \ @@ -51,6 +68,7 @@ REGISTER_GPU_KERNEL(bfloat16); #undef REGISTER_GPU_KERNEL + #if GOOGLE_CUDA // A special GPU kernel for int32 and bool. // TODO(b/25387198): Also enable int32 in device memory. This kernel diff --git a/tensorflow/core/kernels/no_op.cc b/tensorflow/core/kernels/no_op.cc index e4c7f48e773..0993e6e1fcf 100644 --- a/tensorflow/core/kernels/no_op.cc +++ b/tensorflow/core/kernels/no_op.cc @@ -19,6 +19,9 @@ namespace tensorflow { REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE_CPU), NoOp); REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE_GPU), NoOp); + +#if TENSORFLOW_USE_SYCL REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE_SYCL), NoOp); +#endif } // namespace tensorflow diff --git a/tensorflow/core/kernels/sendrecv_ops.cc b/tensorflow/core/kernels/sendrecv_ops.cc index 38b0bf5756e..9e9cdda3826 100644 --- a/tensorflow/core/kernels/sendrecv_ops.cc +++ b/tensorflow/core/kernels/sendrecv_ops.cc @@ -77,7 +77,10 @@ void SendOp::Compute(OpKernelContext* ctx) { REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_CPU), SendOp); REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_GPU), SendOp); + +#if TENSORFLOW_USE_SYCL REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE_SYCL), SendOp); +#endif REGISTER_KERNEL_BUILDER(Name("_HostSend").Device(DEVICE_CPU), SendOp); REGISTER_KERNEL_BUILDER( @@ -136,7 +139,10 @@ void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_CPU), RecvOp); REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_GPU), RecvOp); + +#if TENSORFLOW_USE_SYCL REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_SYCL), RecvOp); +#endif REGISTER_KERNEL_BUILDER(Name("_HostRecv").Device(DEVICE_CPU), RecvOp); REGISTER_KERNEL_BUILDER(