diff options
Diffstat (limited to 'tensorflow/core/kernels/transpose_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/transpose_op.cc | 29 |
1 files changed, 29 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/transpose_op.cc b/tensorflow/core/kernels/transpose_op.cc index 67300c1e96..4d303f0173 100644 --- a/tensorflow/core/kernels/transpose_op.cc +++ b/tensorflow/core/kernels/transpose_op.cc @@ -82,6 +82,15 @@ REGISTER_KERNEL_BUILDER(Name("InvertPermutation") .HostMemory("y"), InvertPermutationOp); +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER(Name("InvertPermutation") + .Device(DEVICE_SYCL) + .TypeConstraint<int32>("T") + .HostMemory("x") + .HostMemory("y"), + InvertPermutationOp); +#endif // TENSORFLOW_USE_SYCL + // output = TransposeOp(T<any> input, T<int32> perm) takes a tensor // of type T and rank N, and a permutation of 0, 1, ..., N-1. It // shuffles the dimensions of the input tensor according to permutation. @@ -201,4 +210,24 @@ TF_CALL_POD_TYPES(REGISTER); #undef REGISTER #endif +#ifdef TENSORFLOW_USE_SYCL +Status TransposeSyclOp::DoTranspose(OpKernelContext* ctx, const Tensor& in, + gtl::ArraySlice<int32> perm, Tensor* out) { + typedef Eigen::SyclDevice SYCLDevice; + return ::tensorflow::DoTranspose(ctx->eigen_device<SYCLDevice>(), in, perm, + out); +} +#define REGISTER(T) \ + REGISTER_KERNEL_BUILDER(Name("Transpose") \ + .Device(DEVICE_SYCL) \ + .TypeConstraint<T>("T") \ + .TypeConstraint<int32>("Tperm") \ + .HostMemory("perm"), \ + TransposeSyclOp); +REGISTER(float); +REGISTER(bool); +REGISTER(int32); +#undef REGISTER +#endif + } // namespace tensorflow |