diff options
Diffstat (limited to 'tensorflow/core/kernels/shape_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/shape_ops.cc | 56 |
1 files changed, 22 insertions, 34 deletions
diff --git a/tensorflow/core/kernels/shape_ops.cc b/tensorflow/core/kernels/shape_ops.cc index d78c6d2639..c5e3164145 100644 --- a/tensorflow/core/kernels/shape_ops.cc +++ b/tensorflow/core/kernels/shape_ops.cc @@ -48,6 +48,7 @@ REGISTER_KERNEL_BUILDER(Name("Shape") ShapeOp<int64>); TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +TF_CALL_bool(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL REGISTER_KERNEL_BUILDER(Name("Shape") @@ -102,7 +103,7 @@ REGISTER_KERNEL_BUILDER(Name("Shape") .TypeConstraint<int32>("T") .TypeConstraint<int64>("out_type"), ShapeOp<int64>); -#endif +#endif // GOOGLE_CUDA // ShapeN --------------------------------------- REGISTER_KERNEL_BUILDER(Name("ShapeN") @@ -152,9 +153,9 @@ REGISTER_KERNEL_BUILDER(Name("ShapeN") .TypeConstraint<int32>("T") .TypeConstraint<int64>("out_type"), ShapeNOp<int64>); -#endif +#endif // GOOGLE_CUDA -#if TENSORFLOW_USE_SYCL +#ifdef TENSORFLOW_USE_SYCL #define REGISTER_SYCL_KERNEL(type) \ REGISTER_KERNEL_BUILDER(Name("ShapeN") \ .Device(DEVICE_SYCL) \ @@ -170,11 +171,9 @@ REGISTER_KERNEL_BUILDER(Name("ShapeN") ShapeNOp<int64>) TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +TF_CALL_bool(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL -// A special GPU kernel for int32. -// TODO(b/25387198): Also enable int32 in device memory. This kernel -// registration requires all int32 inputs and outputs to be in host memory. REGISTER_KERNEL_BUILDER(Name("ShapeN") .Device(DEVICE_SYCL) .HostMemory("input") @@ -202,13 +201,9 @@ REGISTER_KERNEL_BUILDER(Name("Rank").Device(DEVICE_CPU).HostMemory("output"), .TypeConstraint<type>("T") \ .HostMemory("output"), \ RankOp); -REGISTER_SYCL_KERNEL(float); -REGISTER_SYCL_KERNEL(double); +TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL -// A special GPU kernel for int32 and bool. -// TODO(b/25387198): Also enable int32 in device memory. This kernel -// registration requires all int32 inputs and outputs to be in host memory. REGISTER_KERNEL_BUILDER(Name("Rank") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") @@ -250,7 +245,7 @@ REGISTER_KERNEL_BUILDER(Name("Rank") .HostMemory("input") .HostMemory("output"), RankOp); -#endif +#endif // GOOGLE_CUDA // Size ------------------------------------------ REGISTER_KERNEL_BUILDER(Name("Size") @@ -299,7 +294,7 @@ REGISTER_KERNEL_BUILDER(Name("Size") .HostMemory("input") .HostMemory("output"), SizeOp<int64>); -#endif +#endif // GOOGLE_CUDA #ifdef TENSORFLOW_USE_SYCL #define REGISTER_SYCL_KERNEL(type) \ @@ -315,13 +310,10 @@ REGISTER_KERNEL_BUILDER(Name("Size") .TypeConstraint<int64>("out_type") \ .HostMemory("output"), \ SizeOp<int64>); -REGISTER_SYCL_KERNEL(float); -REGISTER_SYCL_KERNEL(double); +TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +TF_CALL_bool(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL -// A special GPU kernel for int32. -// TODO(b/25387198): Also enable int32 in device memory. This kernel -// registration requires all int32 inputs and outputs to be in host memory. REGISTER_KERNEL_BUILDER(Name("Size") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") @@ -336,7 +328,7 @@ REGISTER_KERNEL_BUILDER(Name("Size") .HostMemory("input") .HostMemory("output"), SizeOp<int64>); -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL // ExpandDims ------------------------------------ REGISTER_KERNEL_BUILDER(Name("ExpandDims") @@ -365,7 +357,7 @@ REGISTER_KERNEL_BUILDER(Name("ExpandDims") .HostMemory("dim") .HostMemory("output"), ExpandDimsOp); -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA #ifdef TENSORFLOW_USE_SYCL #define REGISTER_SYCL_KERNEL(type) \ @@ -375,9 +367,8 @@ REGISTER_KERNEL_BUILDER(Name("ExpandDims") .TypeConstraint<int32>("Tdim") \ .HostMemory("dim"), \ ExpandDimsOp); -REGISTER_SYCL_KERNEL(float) -REGISTER_SYCL_KERNEL(double) - +TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +TF_CALL_bool(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL REGISTER_KERNEL_BUILDER(Name("ExpandDims") @@ -388,7 +379,7 @@ REGISTER_KERNEL_BUILDER(Name("ExpandDims") .HostMemory("dim") .HostMemory("output"), ExpandDimsOp); -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL // Squeeze --------------------------------------- REGISTER_KERNEL_BUILDER(Name("Squeeze").Device(DEVICE_CPU), SqueezeOp); @@ -411,26 +402,23 @@ REGISTER_KERNEL_BUILDER(Name("Squeeze") .HostMemory("input") .HostMemory("output"), SqueezeOp); -#endif +#endif // GOOGLE_CUDA #if TENSORFLOW_USE_SYCL -#define REGISTER_SYCL_KERNEL(type) \ - REGISTER_KERNEL_BUILDER( \ - Name("Squeeze").Device(DEVICE_SYCL).TypeConstraint<type>("T"),\ +#define REGISTER_SYCL_KERNEL(type) \ + REGISTER_KERNEL_BUILDER( \ + Name("Squeeze").Device(DEVICE_SYCL).TypeConstraint<type>("T"), \ SqueezeOp); -REGISTER_SYCL_KERNEL(float); -REGISTER_SYCL_KERNEL(double); +TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); +TF_CALL_bool(REGISTER_SYCL_KERNEL); #undef REGISTER_SYCL_KERNEL -// A special GPU kernel for int32. -// TODO(b/25387198): Also enable int32 in device memory. This kernel -// registration requires all int32 inputs and outputs to be in host memory. REGISTER_KERNEL_BUILDER(Name("Squeeze") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") .HostMemory("input") .HostMemory("output"), SqueezeOp); -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow |