experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Don't create SYCL kernels unless TENSORFLOW_USE_SYCL is set

Browse Source
This commit is contained in:
Benoit Steiner 2016-11-03 10:31:00 -07:00
parent 5ea3f1830c
commit e2694bd4bc
3 changed files with 28 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tensorflow/core/kernels/identity_op.cc
View File

@ -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<type>("T"), \
IdentityOp); \
REGISTER_KERNEL_BUILDER( \
Name("RefIdentity").Device(DEVICE_SYCL).TypeConstraint<type>("T"), \
IdentityOp); \
REGISTER_KERNEL_BUILDER( \
Name("StopGradient").Device(DEVICE_SYCL).TypeConstraint<type>("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<type>("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

3
tensorflow/core/kernels/no_op.cc
View File

@ -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

6
tensorflow/core/kernels/sendrecv_ops.cc
View File

@ -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(