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/array_ops.cc. 17 18 #include "tensorflow/core/kernels/shape_ops.h" 19 #include "tensorflow/core/framework/register_types.h" 20 21 namespace tensorflow { 22 23 // Shape ---------------------------------------- 24 REGISTER_KERNEL_BUILDER(Name("Shape") 25 .Device(DEVICE_CPU) 26 .HostMemory("output") 27 .TypeConstraint<int32>("out_type"), 28 ShapeOp<int32>); 29 REGISTER_KERNEL_BUILDER(Name("Shape") 30 .Device(DEVICE_CPU) 31 .HostMemory("output") 32 .TypeConstraint<int64>("out_type"), 33 ShapeOp<int64>); 34 35 #ifdef TENSORFLOW_USE_SYCL 36 #define REGISTER_SYCL_KERNEL(type) \ 37 REGISTER_KERNEL_BUILDER(Name("Shape") \ 38 .Device(DEVICE_SYCL) \ 39 .HostMemory("output") \ 40 .TypeConstraint<int32>("out_type") \ 41 .TypeConstraint<type>("T"), \ 42 ShapeOp<int32>); \ 43 REGISTER_KERNEL_BUILDER(Name("Shape") \ 44 .Device(DEVICE_SYCL) \ 45 .HostMemory("output") \ 46 .TypeConstraint<int64>("out_type") \ 47 .TypeConstraint<type>("T"), \ 48 ShapeOp<int64>); 49 50 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 51 TF_CALL_bool(REGISTER_SYCL_KERNEL); 52 #undef REGISTER_SYCL_KERNEL 53 54 REGISTER_KERNEL_BUILDER(Name("Shape") 55 .Device(DEVICE_SYCL) 56 .HostMemory("input") 57 .HostMemory("output") 58 .TypeConstraint<int32>("T") 59 .TypeConstraint<int32>("out_type"), 60 ShapeOp<int32>); 61 REGISTER_KERNEL_BUILDER(Name("Shape") 62 .Device(DEVICE_SYCL) 63 .HostMemory("input") 64 .HostMemory("output") 65 .TypeConstraint<int32>("T") 66 .TypeConstraint<int64>("out_type"), 67 ShapeOp<int64>); 68 #endif // TENSORFLOW_USE_SYCL 69 70 #if GOOGLE_CUDA 71 #define REGISTER_GPU_KERNEL(type) \ 72 REGISTER_KERNEL_BUILDER(Name("Shape") \ 73 .Device(DEVICE_GPU) \ 74 .HostMemory("output") \ 75 .TypeConstraint<int32>("out_type") \ 76 .TypeConstraint<type>("T"), \ 77 ShapeOp<int32>); \ 78 REGISTER_KERNEL_BUILDER(Name("Shape") \ 79 .Device(DEVICE_GPU) \ 80 .HostMemory("output") \ 81 .TypeConstraint<int64>("out_type") \ 82 .TypeConstraint<type>("T"), \ 83 ShapeOp<int64>); 84 85 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 86 TF_CALL_bool(REGISTER_GPU_KERNEL); 87 TF_CALL_variant(REGISTER_GPU_KERNEL); 88 #undef REGISTER_GPU_KERNEL 89 90 // A special GPU kernel for int32. 91 // TODO(b/25387198): Also enable int32 in device memory. This kernel 92 // registration requires all int32 inputs and outputs to be in host memory. 93 REGISTER_KERNEL_BUILDER(Name("Shape") 94 .Device(DEVICE_GPU) 95 .HostMemory("input") 96 .HostMemory("output") 97 .TypeConstraint<int32>("T") 98 .TypeConstraint<int32>("out_type"), 99 ShapeOp<int32>); 100 REGISTER_KERNEL_BUILDER(Name("Shape") 101 .Device(DEVICE_GPU) 102 .HostMemory("input") 103 .HostMemory("output") 104 .TypeConstraint<int32>("T") 105 .TypeConstraint<int64>("out_type"), 106 ShapeOp<int64>); 107 108 #endif // GOOGLE_CUDA 109 110 // ShapeN --------------------------------------- 111 REGISTER_KERNEL_BUILDER(Name("ShapeN") 112 .Device(DEVICE_CPU) 113 .HostMemory("output") 114 .TypeConstraint<int32>("out_type"), 115 ShapeNOp<int32>); 116 REGISTER_KERNEL_BUILDER(Name("ShapeN") 117 .Device(DEVICE_CPU) 118 .HostMemory("output") 119 .TypeConstraint<int64>("out_type"), 120 ShapeNOp<int64>); 121 122 #if GOOGLE_CUDA 123 #define REGISTER_GPU_KERNEL(type) \ 124 REGISTER_KERNEL_BUILDER(Name("ShapeN") \ 125 .Device(DEVICE_GPU) \ 126 .HostMemory("output") \ 127 .TypeConstraint<int32>("out_type") \ 128 .TypeConstraint<type>("T"), \ 129 ShapeNOp<int32>); \ 130 REGISTER_KERNEL_BUILDER(Name("ShapeN") \ 131 .Device(DEVICE_GPU) \ 132 .HostMemory("output") \ 133 .TypeConstraint<int64>("out_type") \ 134 .TypeConstraint<type>("T"), \ 135 ShapeNOp<int64>) 136 137 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 138 TF_CALL_bool(REGISTER_GPU_KERNEL); 139 #undef REGISTER_GPU_KERNEL 140 141 // A special GPU kernel for int32. 142 // TODO(b/25387198): Also enable int32 in device memory. This kernel 143 // registration requires all int32 inputs and outputs to be in host memory. 144 REGISTER_KERNEL_BUILDER(Name("ShapeN") 145 .Device(DEVICE_GPU) 146 .HostMemory("input") 147 .HostMemory("output") 148 .TypeConstraint<int32>("T") 149 .TypeConstraint<int32>("out_type"), 150 ShapeNOp<int32>); 151 REGISTER_KERNEL_BUILDER(Name("ShapeN") 152 .Device(DEVICE_GPU) 153 .HostMemory("input") 154 .HostMemory("output") 155 .TypeConstraint<int32>("T") 156 .TypeConstraint<int64>("out_type"), 157 ShapeNOp<int64>); 158 #endif // GOOGLE_CUDA 159 160 #ifdef TENSORFLOW_USE_SYCL 161 #define REGISTER_SYCL_KERNEL(type) \ 162 REGISTER_KERNEL_BUILDER(Name("ShapeN") \ 163 .Device(DEVICE_SYCL) \ 164 .HostMemory("output") \ 165 .TypeConstraint<int32>("out_type") \ 166 .TypeConstraint<type>("T"), \ 167 ShapeNOp<int32>); \ 168 REGISTER_KERNEL_BUILDER(Name("ShapeN") \ 169 .Device(DEVICE_SYCL) \ 170 .HostMemory("output") \ 171 .TypeConstraint<int64>("out_type") \ 172 .TypeConstraint<type>("T"), \ 173 ShapeNOp<int64>) 174 175 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 176 TF_CALL_bool(REGISTER_SYCL_KERNEL); 177 #undef REGISTER_SYCL_KERNEL 178 179 REGISTER_KERNEL_BUILDER(Name("ShapeN") 180 .Device(DEVICE_SYCL) 181 .HostMemory("input") 182 .HostMemory("output") 183 .TypeConstraint<int32>("T") 184 .TypeConstraint<int32>("out_type"), 185 ShapeNOp<int32>); 186 REGISTER_KERNEL_BUILDER(Name("ShapeN") 187 .Device(DEVICE_SYCL) 188 .HostMemory("input") 189 .HostMemory("output") 190 .TypeConstraint<int32>("T") 191 .TypeConstraint<int64>("out_type"), 192 ShapeNOp<int64>); 193 #endif // TENSORFLOW_USE_SYCL 194 195 // Rank ------------------------------------------ 196 REGISTER_KERNEL_BUILDER(Name("Rank").Device(DEVICE_CPU).HostMemory("output"), 197 RankOp); 198 199 #ifdef TENSORFLOW_USE_SYCL 200 #define REGISTER_SYCL_KERNEL(type) \ 201 REGISTER_KERNEL_BUILDER(Name("Rank") \ 202 .Device(DEVICE_SYCL) \ 203 .TypeConstraint<type>("T") \ 204 .HostMemory("output"), \ 205 RankOp); 206 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 207 #undef REGISTER_SYCL_KERNEL 208 209 REGISTER_KERNEL_BUILDER(Name("Rank") 210 .Device(DEVICE_SYCL) 211 .TypeConstraint<int32>("T") 212 .HostMemory("input") 213 .HostMemory("output"), 214 RankOp); 215 216 REGISTER_KERNEL_BUILDER(Name("Rank") 217 .Device(DEVICE_SYCL) 218 .TypeConstraint<bool>("T") 219 .HostMemory("input") 220 .HostMemory("output"), 221 RankOp); 222 #endif // TENSORFLOW_USE_SYCL 223 224 #if GOOGLE_CUDA 225 #define REGISTER_GPU_KERNEL(type) \ 226 REGISTER_KERNEL_BUILDER(Name("Rank") \ 227 .Device(DEVICE_GPU) \ 228 .TypeConstraint<type>("T") \ 229 .HostMemory("output"), \ 230 RankOp); 231 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 232 TF_CALL_variant(REGISTER_GPU_KERNEL); 233 #undef REGISTER_GPU_KERNEL 234 235 // A special GPU kernel for int32 and bool. 236 // TODO(b/25387198): Also enable int32 in device memory. This kernel 237 // registration requires all int32 inputs and outputs to be in host memory. 238 REGISTER_KERNEL_BUILDER(Name("Rank") 239 .Device(DEVICE_GPU) 240 .TypeConstraint<int32>("T") 241 .HostMemory("input") 242 .HostMemory("output"), 243 RankOp); 244 245 REGISTER_KERNEL_BUILDER(Name("Rank") 246 .Device(DEVICE_GPU) 247 .TypeConstraint<bool>("T") 248 .HostMemory("input") 249 .HostMemory("output"), 250 RankOp); 251 252 #endif // GOOGLE_CUDA 253 254 // Size ------------------------------------------ 255 REGISTER_KERNEL_BUILDER(Name("Size") 256 .Device(DEVICE_CPU) 257 .HostMemory("output") 258 .TypeConstraint<int32>("out_type"), 259 SizeOp<int32>); 260 REGISTER_KERNEL_BUILDER(Name("Size") 261 .Device(DEVICE_CPU) 262 .HostMemory("output") 263 .TypeConstraint<int64>("out_type"), 264 SizeOp<int64>); 265 266 #if GOOGLE_CUDA 267 #define REGISTER_GPU_KERNEL(type) \ 268 REGISTER_KERNEL_BUILDER(Name("Size") \ 269 .Device(DEVICE_GPU) \ 270 .TypeConstraint<type>("T") \ 271 .TypeConstraint<int32>("out_type") \ 272 .HostMemory("output"), \ 273 SizeOp<int32>); \ 274 REGISTER_KERNEL_BUILDER(Name("Size") \ 275 .Device(DEVICE_GPU) \ 276 .TypeConstraint<type>("T") \ 277 .TypeConstraint<int64>("out_type") \ 278 .HostMemory("output"), \ 279 SizeOp<int64>); 280 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 281 TF_CALL_bool(REGISTER_GPU_KERNEL); 282 TF_CALL_variant(REGISTER_GPU_KERNEL); 283 #undef REGISTER_GPU_KERNEL 284 285 // A special GPU kernel for int32. 286 // TODO(b/25387198): Also enable int32 in device memory. This kernel 287 // registration requires all int32 inputs and outputs to be in host memory. 288 REGISTER_KERNEL_BUILDER(Name("Size") 289 .Device(DEVICE_GPU) 290 .TypeConstraint<int32>("T") 291 .TypeConstraint<int32>("out_type") 292 .HostMemory("input") 293 .HostMemory("output"), 294 SizeOp<int32>); 295 REGISTER_KERNEL_BUILDER(Name("Size") 296 .Device(DEVICE_GPU) 297 .TypeConstraint<int32>("T") 298 .TypeConstraint<int64>("out_type") 299 .HostMemory("input") 300 .HostMemory("output"), 301 SizeOp<int64>); 302 303 #endif // GOOGLE_CUDA 304 305 #ifdef TENSORFLOW_USE_SYCL 306 #define REGISTER_SYCL_KERNEL(type) \ 307 REGISTER_KERNEL_BUILDER(Name("Size") \ 308 .Device(DEVICE_SYCL) \ 309 .TypeConstraint<type>("T") \ 310 .TypeConstraint<int32>("out_type") \ 311 .HostMemory("output"), \ 312 SizeOp<int32>); \ 313 REGISTER_KERNEL_BUILDER(Name("Size") \ 314 .Device(DEVICE_SYCL) \ 315 .TypeConstraint<type>("T") \ 316 .TypeConstraint<int64>("out_type") \ 317 .HostMemory("output"), \ 318 SizeOp<int64>); 319 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 320 TF_CALL_bool(REGISTER_SYCL_KERNEL); 321 #undef REGISTER_SYCL_KERNEL 322 323 REGISTER_KERNEL_BUILDER(Name("Size") 324 .Device(DEVICE_SYCL) 325 .TypeConstraint<int32>("T") 326 .TypeConstraint<int32>("out_type") 327 .HostMemory("input") 328 .HostMemory("output"), 329 SizeOp<int32>); 330 REGISTER_KERNEL_BUILDER(Name("Size") 331 .Device(DEVICE_SYCL) 332 .TypeConstraint<int32>("T") 333 .TypeConstraint<int64>("out_type") 334 .HostMemory("input") 335 .HostMemory("output"), 336 SizeOp<int64>); 337 #endif // TENSORFLOW_USE_SYCL 338 339 // ExpandDims ------------------------------------ 340 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 341 .Device(DEVICE_CPU) 342 .HostMemory("dim") 343 .TypeConstraint<int32>("Tdim"), 344 ExpandDimsOp<int32>); 345 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 346 .Device(DEVICE_CPU) 347 .HostMemory("dim") 348 .TypeConstraint<int64>("Tdim"), 349 ExpandDimsOp<int64>); 350 351 #if GOOGLE_CUDA 352 #define REGISTER_GPU_KERNEL(type) \ 353 REGISTER_KERNEL_BUILDER(Name("ExpandDims") \ 354 .Device(DEVICE_GPU) \ 355 .TypeConstraint<type>("T") \ 356 .TypeConstraint<int32>("Tdim") \ 357 .HostMemory("dim"), \ 358 ExpandDimsOp<int32>); \ 359 REGISTER_KERNEL_BUILDER(Name("ExpandDims") \ 360 .Device(DEVICE_GPU) \ 361 .TypeConstraint<type>("T") \ 362 .TypeConstraint<int64>("Tdim") \ 363 .HostMemory("dim"), \ 364 ExpandDimsOp<int64>); 365 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 366 TF_CALL_bool(REGISTER_GPU_KERNEL); 367 #undef REGISTER_GPU_KERNEL 368 369 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 370 .Device(DEVICE_GPU) 371 .TypeConstraint<int32>("T") 372 .TypeConstraint<int32>("Tdim") 373 .HostMemory("input") 374 .HostMemory("dim") 375 .HostMemory("output"), 376 ExpandDimsOp<int32>); 377 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 378 .Device(DEVICE_GPU) 379 .TypeConstraint<int32>("T") 380 .TypeConstraint<int64>("Tdim") 381 .HostMemory("input") 382 .HostMemory("dim") 383 .HostMemory("output"), 384 ExpandDimsOp<int64>); 385 #endif // GOOGLE_CUDA 386 387 #ifdef TENSORFLOW_USE_SYCL 388 #define REGISTER_SYCL_KERNEL(type) \ 389 REGISTER_KERNEL_BUILDER(Name("ExpandDims") \ 390 .Device(DEVICE_SYCL) \ 391 .TypeConstraint<type>("T") \ 392 .TypeConstraint<int32>("Tdim") \ 393 .HostMemory("dim"), \ 394 ExpandDimsOp<int32>); \ 395 REGISTER_KERNEL_BUILDER(Name("ExpandDims") \ 396 .Device(DEVICE_SYCL) \ 397 .TypeConstraint<type>("T") \ 398 .TypeConstraint<int64>("Tdim") \ 399 .HostMemory("dim"), \ 400 ExpandDimsOp<int64>); 401 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 402 TF_CALL_bool(REGISTER_SYCL_KERNEL); 403 #undef REGISTER_SYCL_KERNEL 404 405 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 406 .Device(DEVICE_SYCL) 407 .TypeConstraint<int32>("T") 408 .TypeConstraint<int32>("Tdim") 409 .HostMemory("input") 410 .HostMemory("dim") 411 .HostMemory("output"), 412 ExpandDimsOp<int32>); 413 REGISTER_KERNEL_BUILDER(Name("ExpandDims") 414 .Device(DEVICE_SYCL) 415 .TypeConstraint<int32>("T") 416 .TypeConstraint<int64>("Tdim") 417 .HostMemory("input") 418 .HostMemory("dim") 419 .HostMemory("output"), 420 ExpandDimsOp<int64>); 421 #endif // TENSORFLOW_USE_SYCL 422 423 // Squeeze --------------------------------------- 424 REGISTER_KERNEL_BUILDER(Name("Squeeze").Device(DEVICE_CPU), SqueezeOp); 425 426 #if GOOGLE_CUDA 427 #define REGISTER_GPU_KERNEL(type) \ 428 REGISTER_KERNEL_BUILDER( \ 429 Name("Squeeze").Device(DEVICE_GPU).TypeConstraint<type>("T"), \ 430 SqueezeOp); 431 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL); 432 TF_CALL_bool(REGISTER_GPU_KERNEL); 433 #undef REGISTER_GPU_KERNEL 434 435 // A special GPU kernel for int32. 436 // TODO(b/25387198): Also enable int32 in device memory. This kernel 437 // registration requires all int32 inputs and outputs to be in host memory. 438 REGISTER_KERNEL_BUILDER(Name("Squeeze") 439 .Device(DEVICE_GPU) 440 .TypeConstraint<int32>("T") 441 .HostMemory("input") 442 .HostMemory("output"), 443 SqueezeOp); 444 #endif // GOOGLE_CUDA 445 446 #if TENSORFLOW_USE_SYCL 447 #define REGISTER_SYCL_KERNEL(type) \ 448 REGISTER_KERNEL_BUILDER( \ 449 Name("Squeeze").Device(DEVICE_SYCL).TypeConstraint<type>("T"), \ 450 SqueezeOp); 451 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_SYCL_KERNEL); 452 TF_CALL_bool(REGISTER_SYCL_KERNEL); 453 #undef REGISTER_SYCL_KERNEL 454 455 REGISTER_KERNEL_BUILDER(Name("Squeeze") 456 .Device(DEVICE_SYCL) 457 .TypeConstraint<int32>("T") 458 .HostMemory("input") 459 .HostMemory("output"), 460 SqueezeOp); 461 #endif // TENSORFLOW_USE_SYCL 462 463 } // namespace tensorflow 464
