1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_KERNELS_STRIDED_SLICE_OP_H_ 17 #define TENSORFLOW_KERNELS_STRIDED_SLICE_OP_H_ 18 19 // Functor definition for StridedSliceOp, must be compilable by nvcc. 20 21 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 22 #include "tensorflow/core/framework/resource_handle.h" 23 #include "tensorflow/core/framework/tensor_types.h" 24 #include "tensorflow/core/framework/variant_encode_decode.h" 25 #include "tensorflow/core/platform/types.h" 26 27 namespace tensorflow { 28 namespace functor { 29 30 template <typename Device, typename T, int NDIMS> 31 struct StridedSlice { 32 void operator()(const Device& d, typename TTypes<T, NDIMS>::Tensor output, 33 typename TTypes<T, NDIMS>::ConstTensor input, 34 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices, 35 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& stop_indices, 36 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& strides) { 37 const bool use_64bit = input.size() > Eigen::NumTraits<int>::highest(); 38 if (!use_64bit && 39 Eigen::internal::is_same<Device, Eigen::GpuDevice>::value) { 40 Eigen::DSizes<int, NDIMS> start_i, stop_i, strides_i; 41 for (int i = 0; i < NDIMS; ++i) { 42 start_i[i] = start_indices[i]; 43 stop_i[i] = stop_indices[i]; 44 strides_i[i] = strides[i]; 45 } 46 To32Bit(output).device(d) = 47 To32Bit(input).stridedSlice(start_i, stop_i, strides_i); 48 } else { 49 output.device(d) = 50 input.stridedSlice(start_indices, stop_indices, strides); 51 } 52 } 53 }; 54 55 template <typename T, int NDIMS, typename Device> 56 struct InitOutput { 57 static void run(const Device& d, typename TTypes<T, NDIMS>::Tensor output) { 58 output.device(d) = output.constant(T(0)); 59 } 60 }; 61 62 template <int NDIMS, typename Device> 63 struct InitOutput<ResourceHandle, NDIMS, Device> { 64 static void run(const Device& d, 65 typename TTypes<ResourceHandle, NDIMS>::Tensor output) { 66 output.device(d) = output.constant(ResourceHandle()); 67 } 68 }; 69 70 template <int NDIMS, typename Device> 71 struct InitOutput<string, NDIMS, Device> { 72 static void run(const Device& d, 73 typename TTypes<string, NDIMS>::Tensor output) { 74 output.device(d) = output.constant(string()); 75 } 76 }; 77 78 template <typename Device, typename T, int NDIMS> 79 struct StridedSliceGrad { 80 void operator()(const Device& d, typename TTypes<T, NDIMS>::Tensor output, 81 typename TTypes<T, NDIMS>::ConstTensor input, 82 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices, 83 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& stop_indices, 84 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& strides) { 85 InitOutput<T, NDIMS, Device>::run(d, output); 86 const bool use_64bit = input.size() > Eigen::NumTraits<int>::highest(); 87 if (!use_64bit && 88 Eigen::internal::is_same<Device, Eigen::GpuDevice>::value) { 89 Eigen::DSizes<int, NDIMS> start_i, stop_i, strides_i; 90 for (int i = 0; i < NDIMS; ++i) { 91 start_i[i] = start_indices[i]; 92 stop_i[i] = stop_indices[i]; 93 strides_i[i] = strides[i]; 94 } 95 To32Bit(output).stridedSlice(start_i, stop_i, strides_i).device(d) = 96 input; 97 } else { 98 output.stridedSlice(start_indices, stop_indices, strides).device(d) = 99 input; 100 } 101 } 102 }; 103 104 template <typename Device, typename T, int NDIMS> 105 struct StridedSliceAssign { 106 void operator()(const Device& d, typename TTypes<T, NDIMS>::Tensor output, 107 typename TTypes<T, NDIMS>::ConstTensor input, 108 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices, 109 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& stop_indices, 110 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& strides) { 111 const bool use_64bit = input.size() > Eigen::NumTraits<int>::highest(); 112 if (!use_64bit && 113 Eigen::internal::is_same<Device, Eigen::GpuDevice>::value) { 114 Eigen::DSizes<int, NDIMS> start_i, stop_i, strides_i; 115 for (int i = 0; i < NDIMS; ++i) { 116 start_i[i] = start_indices[i]; 117 stop_i[i] = stop_indices[i]; 118 strides_i[i] = strides[i]; 119 } 120 To32Bit(output).stridedSlice(start_i, stop_i, strides_i).device(d) = 121 To32Bit(input); 122 } else { 123 output.stridedSlice(start_indices, stop_indices, strides).device(d) = 124 input; 125 } 126 } 127 }; 128 129 template <typename Device, typename T> 130 struct StridedSliceAssignScalar { 131 void operator()(const Device& d, typename TTypes<T, 1>::Tensor output, 132 typename TTypes<T, 1>::ConstTensor input) { 133 output.device(d) = input; 134 } 135 }; 136 137 } // namespace functor 138 } // namespace tensorflow 139 140 #endif // TENSORFLOW_KERNELS_SLICE_OP_H_ 141
