diff options
author | 2017-08-24 19:16:03 -0700 | |
---|---|---|
committer | 2017-08-24 19:20:28 -0700 | |
commit | 80341a3ac5e14d36aad4c5cd9e9a6d4479235017 (patch) | |
tree | 340f5e38a90394953d848d54298ead28e82016b6 /tensorflow/core/kernels/strided_slice_op.cc | |
parent | b486f75b4cf6fc937f4ac1a408be6fb6470150d2 (diff) |
StridedSliceOp: GPU kernel for 64-bit indices
The TypeConstraint was unnecessary. The kernel implementation
was correctly converting both DT_INT32 and DT_INT64 tensors to
an InlinedVector<int64, 4>.
PiperOrigin-RevId: 166427269
Diffstat (limited to 'tensorflow/core/kernels/strided_slice_op.cc')
-rw-r--r-- | tensorflow/core/kernels/strided_slice_op.cc | 32 |
1 files changed, 8 insertions, 24 deletions
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc index 4655503e26..8fc40db3cc 100644 --- a/tensorflow/core/kernels/strided_slice_op.cc +++ b/tensorflow/core/kernels/strided_slice_op.cc @@ -398,8 +398,7 @@ REGISTER_STRIDED_SLICE(bfloat16); .TypeConstraint<type>("T") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceOp<GPUDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") \ .Device(DEVICE_GPU) \ @@ -407,16 +406,14 @@ REGISTER_STRIDED_SLICE(bfloat16); .HostMemory("shape") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceGradOp<GPUDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") \ .Device(DEVICE_GPU) \ .TypeConstraint<type>("T") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceAssignOp<GPUDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") \ .Device(DEVICE_GPU) \ @@ -424,8 +421,7 @@ REGISTER_STRIDED_SLICE(bfloat16); .HostMemory("ref") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceAssignOp<GPUDevice, type>) TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU); @@ -438,7 +434,6 @@ TF_CALL_complex128(REGISTER_GPU); REGISTER_KERNEL_BUILDER(Name("StridedSlice") .Device(DEVICE_GPU) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("input") .HostMemory("begin") .HostMemory("end") @@ -448,7 +443,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSlice") REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") .Device(DEVICE_GPU) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("shape") .HostMemory("begin") .HostMemory("end") @@ -459,7 +453,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") .Device(DEVICE_GPU) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("ref") .HostMemory("begin") .HostMemory("end") @@ -468,7 +461,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") .Device(DEVICE_GPU) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("ref") .HostMemory("begin") .HostMemory("end") @@ -485,8 +477,7 @@ REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") .TypeConstraint<type>("T") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceOp<SYCLDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") \ .Device(DEVICE_SYCL) \ @@ -494,16 +485,14 @@ REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") .HostMemory("shape") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceGradOp<SYCLDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") \ .Device(DEVICE_SYCL) \ .TypeConstraint<type>("T") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceAssignOp<SYCLDevice, type>) \ REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") \ .Device(DEVICE_SYCL) \ @@ -511,8 +500,7 @@ REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") .HostMemory("ref") \ .HostMemory("begin") \ .HostMemory("end") \ - .HostMemory("strides") \ - .TypeConstraint<int32>("Index"), \ + .HostMemory("strides"), \ StridedSliceAssignOp<SYCLDevice, type>) TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SYCL); @@ -520,7 +508,6 @@ TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SYCL); REGISTER_KERNEL_BUILDER(Name("StridedSlice") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("input") .HostMemory("begin") .HostMemory("end") @@ -530,7 +517,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSlice") REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("shape") .HostMemory("begin") .HostMemory("end") @@ -541,7 +527,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSliceGrad") REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("ref") .HostMemory("begin") .HostMemory("end") @@ -550,7 +535,6 @@ REGISTER_KERNEL_BUILDER(Name("StridedSliceAssign") REGISTER_KERNEL_BUILDER(Name("ResourceStridedSliceAssign") .Device(DEVICE_SYCL) .TypeConstraint<int32>("T") - .TypeConstraint<int32>("Index") .HostMemory("ref") .HostMemory("begin") .HostMemory("end") |