diff options
Diffstat (limited to 'tensorflow/core/kernels/shape_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/shape_ops.cc | 83 |
1 files changed, 82 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/shape_ops.cc b/tensorflow/core/kernels/shape_ops.cc index 6bc0b4560b..177a32464b 100644 --- a/tensorflow/core/kernels/shape_ops.cc +++ b/tensorflow/core/kernels/shape_ops.cc @@ -201,6 +201,7 @@ REGISTER_KERNEL_BUILDER(Name("Rank").Device(DEVICE_CPU).HostMemory("output"), .HostMemory("output"), \ RankOp); REGISTER_SYCL_KERNEL(float); +REGISTER_SYCL_KERNEL(double); #undef REGISTER_SYCL_KERNEL // A special GPU kernel for int32 and bool. @@ -297,6 +298,43 @@ REGISTER_KERNEL_BUILDER(Name("Size") SizeOp<int64>); #endif +#ifdef TENSORFLOW_USE_SYCL +#define REGISTER_SYCL_KERNEL(type) \ + REGISTER_KERNEL_BUILDER(Name("Size") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("out_type") \ + .HostMemory("output"), \ + SizeOp<int32>); \ + REGISTER_KERNEL_BUILDER(Name("Size") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int64>("out_type") \ + .HostMemory("output"), \ + SizeOp<int64>); +REGISTER_SYCL_KERNEL(float); +REGISTER_SYCL_KERNEL(double); +#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") + .TypeConstraint<int32>("out_type") + .HostMemory("input") + .HostMemory("output"), + SizeOp<int32>); +REGISTER_KERNEL_BUILDER(Name("Size") + .Device(DEVICE_SYCL) + .TypeConstraint<int32>("T") + .TypeConstraint<int64>("out_type") + .HostMemory("input") + .HostMemory("output"), + SizeOp<int64>); +#endif // TENSORFLOW_USE_SYCL + // ExpandDims ------------------------------------ REGISTER_KERNEL_BUILDER(Name("ExpandDims") .Device(DEVICE_CPU) @@ -323,7 +361,30 @@ REGISTER_KERNEL_BUILDER(Name("ExpandDims") .HostMemory("dim") .HostMemory("output"), ExpandDimsOp); -#endif +#endif // GOOGLE_CUDA + +#ifdef TENSORFLOW_USE_SYCL +#define REGISTER_SYCL_KERNEL(type) \ + REGISTER_KERNEL_BUILDER(Name("ExpandDims") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<type>("T") \ + .TypeConstraint<int32>("Tdim") \ + .HostMemory("dim"), \ + ExpandDimsOp); +REGISTER_SYCL_KERNEL(float) +REGISTER_SYCL_KERNEL(double) + +#undef REGISTER_SYCL_KERNEL + +REGISTER_KERNEL_BUILDER(Name("ExpandDims") + .Device(DEVICE_SYCL) + .TypeConstraint<int32>("T") + .TypeConstraint<int32>("Tdim") + .HostMemory("input") + .HostMemory("dim") + .HostMemory("output"), + ExpandDimsOp); +#endif // TENSORFLOW_USE_SYCL // Squeeze --------------------------------------- REGISTER_KERNEL_BUILDER(Name("Squeeze").Device(DEVICE_CPU), SqueezeOp); @@ -347,4 +408,24 @@ REGISTER_KERNEL_BUILDER(Name("Squeeze") SqueezeOp); #endif +#if TENSORFLOW_USE_SYCL +#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); +#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 + } // namespace tensorflow |