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 // See docs in ../ops/data_flow_ops.cc. 17 18 #include "tensorflow/core/kernels/stack.h" 19 20 #include <limits.h> 21 #include <atomic> 22 #include <vector> 23 24 #include "tensorflow/core/common_runtime/device.h" 25 #include "tensorflow/core/framework/device_base.h" 26 #include "tensorflow/core/framework/op_kernel.h" 27 #include "tensorflow/core/framework/register_types.h" 28 #include "tensorflow/core/framework/resource_mgr.h" 29 #include "tensorflow/core/framework/tensor.h" 30 #include "tensorflow/core/framework/tensor_shape.h" 31 #include "tensorflow/core/framework/types.h" 32 #include "tensorflow/core/lib/core/errors.h" 33 #include "tensorflow/core/lib/core/refcount.h" 34 #include "tensorflow/core/lib/gtl/map_util.h" 35 #include "tensorflow/core/platform/logging.h" 36 #include "tensorflow/core/platform/macros.h" 37 #include "tensorflow/core/platform/mutex.h" 38 #include "tensorflow/core/platform/thread_annotations.h" 39 #include "tensorflow/core/platform/types.h" 40 41 namespace tensorflow { 42 43 REGISTER_KERNEL_BUILDER(Name("Stack").Device(DEVICE_CPU), StackOp); 44 REGISTER_KERNEL_BUILDER(Name("Stack").Device(DEVICE_GPU).HostMemory("handle"), 45 StackOp); 46 REGISTER_KERNEL_BUILDER(Name("StackV2").Device(DEVICE_CPU), StackOp); 47 REGISTER_KERNEL_BUILDER(Name("StackV2") 48 .Device(DEVICE_GPU) 49 .HostMemory("max_size") 50 .HostMemory("handle"), 51 StackOp); 52 #ifdef TENSORFLOW_USE_SYCL 53 REGISTER_KERNEL_BUILDER(Name("Stack").Device(DEVICE_SYCL).HostMemory("handle"), 54 StackOp); 55 REGISTER_KERNEL_BUILDER(Name("StackV2") 56 .Device(DEVICE_SYCL) 57 .HostMemory("max_size") 58 .HostMemory("handle"), 59 StackOp); 60 #endif // TENSORFLOW_USE_SYCL 61 62 REGISTER_KERNEL_BUILDER(Name("StackPush").Device(DEVICE_CPU), 63 TemplatedStackPushOp</*allow_swapping=*/false>); 64 REGISTER_KERNEL_BUILDER(Name("StackPushV2").Device(DEVICE_CPU), 65 TemplatedStackPushOp</*allow_swapping=*/false>); 66 67 #define REGISTER_GPU_KERNEL(type) \ 68 REGISTER_KERNEL_BUILDER(Name("StackPush") \ 69 .Device(DEVICE_GPU) \ 70 .HostMemory("handle") \ 71 .TypeConstraint<type>("T"), \ 72 TemplatedStackPushOp</*allow_swapping=*/true>); \ 73 REGISTER_KERNEL_BUILDER(Name("StackPushV2") \ 74 .Device(DEVICE_GPU) \ 75 .HostMemory("handle") \ 76 .TypeConstraint<type>("T"), \ 77 TemplatedStackPushOp</*allow_swapping=*/true>); 78 79 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 80 #undef REGISTER_GPU_KERNEL 81 82 // Special GPU kernels for int32 and bool. 83 // TODO(b/25387198): Also enable int32 in device memory. This kernel 84 // registration requires all int32 inputs and outputs to be in host memory. 85 #define REGISTER_GPU_HOST_KERNEL(type) \ 86 REGISTER_KERNEL_BUILDER(Name("StackPush") \ 87 .Device(DEVICE_GPU) \ 88 .HostMemory("handle") \ 89 .HostMemory("elem") \ 90 .HostMemory("output") \ 91 .TypeConstraint<type>("T"), \ 92 TemplatedStackPushOp</*allow_swapping=*/true>); \ 93 REGISTER_KERNEL_BUILDER(Name("StackPushV2") \ 94 .Device(DEVICE_GPU) \ 95 .HostMemory("handle") \ 96 .HostMemory("elem") \ 97 .HostMemory("output") \ 98 .TypeConstraint<type>("T"), \ 99 TemplatedStackPushOp</*allow_swapping=*/true>); 100 101 REGISTER_GPU_HOST_KERNEL(int32); 102 REGISTER_GPU_HOST_KERNEL(bool); 103 104 #undef REGISTER_GPU_HOST_KERNEL 105 106 #ifdef TENSORFLOW_USE_SYCL 107 #define REGISTER_SYCL_KERNEL(type) \ 108 REGISTER_KERNEL_BUILDER(Name("StackPush") \ 109 .Device(DEVICE_SYCL) \ 110 .HostMemory("handle") \ 111 .TypeConstraint<type>("T"), \ 112 TemplatedStackPushOp</*allow_swapping=*/true>); 113 114 TF_CALL_GPU_NUMBER_TYPES(REGISTER_SYCL_KERNEL); 115 116 #define REGISTER_SYCL_HOST_KERNEL(type) \ 117 REGISTER_KERNEL_BUILDER(Name("StackPush") \ 118 .Device(DEVICE_SYCL) \ 119 .HostMemory("handle") \ 120 .HostMemory("elem") \ 121 .HostMemory("output") \ 122 .TypeConstraint<type>("T"), \ 123 TemplatedStackPushOp</*allow_swapping=*/true>) 124 125 REGISTER_SYCL_HOST_KERNEL(int32); 126 REGISTER_SYCL_HOST_KERNEL(bool); 127 #undef REGISTER_SYCL_KERNEL 128 #undef REGISTER_SYCL_HOST_KERNEL 129 #endif // TENSORFLOW_USE_SYCL 130 131 REGISTER_KERNEL_BUILDER(Name("StackPop").Device(DEVICE_CPU), StackPopOp); 132 REGISTER_KERNEL_BUILDER(Name("StackPopV2").Device(DEVICE_CPU), StackPopOp); 133 134 #define REGISTER_GPU_KERNEL(type) \ 135 REGISTER_KERNEL_BUILDER(Name("StackPop") \ 136 .Device(DEVICE_GPU) \ 137 .HostMemory("handle") \ 138 .TypeConstraint<type>("elem_type"), \ 139 StackPopOp); \ 140 REGISTER_KERNEL_BUILDER(Name("StackPopV2") \ 141 .Device(DEVICE_GPU) \ 142 .HostMemory("handle") \ 143 .TypeConstraint<type>("elem_type"), \ 144 StackPopOp); 145 146 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 147 #undef REGISTER_GPU_KERNEL 148 149 // Special GPU kernels for int32 and bool. 150 // TODO(b/25387198): Also enable int32 in device memory. This kernel 151 // registration requires all int32 inputs and outputs to be in host memory. 152 #define REGISTER_GPU_HOST_KERNEL(type) \ 153 REGISTER_KERNEL_BUILDER(Name("StackPop") \ 154 .Device(DEVICE_GPU) \ 155 .HostMemory("handle") \ 156 .HostMemory("elem") \ 157 .TypeConstraint<type>("elem_type"), \ 158 StackPopOp); \ 159 REGISTER_KERNEL_BUILDER(Name("StackPopV2") \ 160 .Device(DEVICE_GPU) \ 161 .HostMemory("handle") \ 162 .HostMemory("elem") \ 163 .TypeConstraint<type>("elem_type"), \ 164 StackPopOp); 165 166 REGISTER_GPU_HOST_KERNEL(int32); 167 REGISTER_GPU_HOST_KERNEL(bool); 168 169 #undef REGISTER_GPU_HOST_KERNEL 170 171 #ifdef TENSORFLOW_USE_SYCL 172 #define REGISTER_SYCL_KERNEL(type) \ 173 REGISTER_KERNEL_BUILDER(Name("StackPop") \ 174 .Device(DEVICE_SYCL) \ 175 .HostMemory("handle") \ 176 .TypeConstraint<type>("elem_type"), \ 177 StackPopOp) 178 179 TF_CALL_GPU_NUMBER_TYPES(REGISTER_SYCL_KERNEL); 180 181 #define REGISTER_SYCL_HOST_KERNEL(type) \ 182 REGISTER_KERNEL_BUILDER(Name("StackPop") \ 183 .Device(DEVICE_SYCL) \ 184 .HostMemory("handle") \ 185 .HostMemory("elem") \ 186 .TypeConstraint<type>("elem_type"), \ 187 StackPopOp) 188 189 REGISTER_SYCL_HOST_KERNEL(int32); 190 REGISTER_SYCL_HOST_KERNEL(bool); 191 192 #undef REGISTER_SYCL_KERNEL 193 #undef REGISTER_SYCL_HOST_KERNEL 194 #endif // TENSORFLOW_USE_SYCL 195 196 REGISTER_KERNEL_BUILDER(Name("StackClose").Device(DEVICE_CPU), StackCloseOp); 197 REGISTER_KERNEL_BUILDER( 198 Name("StackClose").Device(DEVICE_GPU).HostMemory("handle"), StackCloseOp); 199 REGISTER_KERNEL_BUILDER(Name("StackCloseV2").Device(DEVICE_CPU), StackCloseOp); 200 REGISTER_KERNEL_BUILDER( 201 Name("StackCloseV2").Device(DEVICE_GPU).HostMemory("handle"), StackCloseOp); 202 #ifdef TENSORFLOW_USE_SYCL 203 REGISTER_KERNEL_BUILDER( 204 Name("StackClose").Device(DEVICE_SYCL).HostMemory("handle"), StackCloseOp); 205 REGISTER_KERNEL_BUILDER( 206 Name("StackCloseV2").Device(DEVICE_SYCL).HostMemory("handle"), 207 StackCloseOp); 208 #endif // TENSORFLOW_USE_SYCL 209 210 } // namespace tensorflow 211