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 | #include "tensorflow/core/framework/op_kernel.h" |
17 | #include "tensorflow/core/framework/tensor.h" |
18 | #include "tensorflow/core/framework/tensor_shape.h" |
19 | #include "tensorflow/core/framework/types.h" |
20 | #include "tensorflow/core/kernels/queue_op.h" |
21 | #include "tensorflow/core/lib/core/errors.h" |
22 | #include "tensorflow/core/platform/macros.h" |
23 | #include "tensorflow/core/platform/types.h" |
24 | |
25 | namespace tensorflow { |
26 | |
27 | REGISTER_KERNEL_BUILDER(Name("QueueEnqueue" ).Device(DEVICE_CPU), EnqueueOp); |
28 | REGISTER_KERNEL_BUILDER(Name("QueueEnqueueV2" ).Device(DEVICE_CPU), EnqueueOp); |
29 | |
30 | REGISTER_KERNEL_BUILDER(Name("QueueEnqueueMany" ).Device(DEVICE_CPU), |
31 | EnqueueManyOp); |
32 | REGISTER_KERNEL_BUILDER(Name("QueueEnqueueManyV2" ).Device(DEVICE_CPU), |
33 | EnqueueManyOp); |
34 | |
35 | REGISTER_KERNEL_BUILDER(Name("QueueDequeue" ).Device(DEVICE_CPU), DequeueOp); |
36 | REGISTER_KERNEL_BUILDER(Name("QueueDequeueV2" ).Device(DEVICE_CPU), DequeueOp); |
37 | |
38 | REGISTER_KERNEL_BUILDER(Name("QueueDequeueMany" ).Device(DEVICE_CPU), |
39 | DequeueManyOp); |
40 | REGISTER_KERNEL_BUILDER(Name("QueueDequeueManyV2" ).Device(DEVICE_CPU), |
41 | DequeueManyOp); |
42 | |
43 | REGISTER_KERNEL_BUILDER(Name("QueueDequeueUpTo" ).Device(DEVICE_CPU), |
44 | DequeueUpToOp); |
45 | REGISTER_KERNEL_BUILDER(Name("QueueDequeueUpToV2" ).Device(DEVICE_CPU), |
46 | DequeueUpToOp); |
47 | |
48 | REGISTER_KERNEL_BUILDER(Name("QueueClose" ).Device(DEVICE_CPU), QueueCloseOp); |
49 | REGISTER_KERNEL_BUILDER(Name("QueueCloseV2" ).Device(DEVICE_CPU), QueueCloseOp); |
50 | |
51 | REGISTER_KERNEL_BUILDER(Name("QueueSize" ).Device(DEVICE_CPU), QueueSizeOp); |
52 | REGISTER_KERNEL_BUILDER(Name("QueueSizeV2" ).Device(DEVICE_CPU), QueueSizeOp); |
53 | |
54 | REGISTER_KERNEL_BUILDER(Name("QueueIsClosed" ).Device(DEVICE_CPU), |
55 | QueueIsClosedOp); |
56 | REGISTER_KERNEL_BUILDER(Name("QueueIsClosedV2" ).Device(DEVICE_CPU), |
57 | QueueIsClosedOp); |
58 | |
59 | REGISTER_KERNEL_BUILDER( |
60 | Name("QueueEnqueueV2" ).Device(DEVICE_DEFAULT).HostMemory("handle" ), |
61 | EnqueueOp); |
62 | REGISTER_KERNEL_BUILDER( |
63 | Name("QueueDequeueV2" ).Device(DEVICE_DEFAULT).HostMemory("handle" ), |
64 | DequeueOp); |
65 | REGISTER_KERNEL_BUILDER( |
66 | Name("QueueCloseV2" ).Device(DEVICE_DEFAULT).HostMemory("handle" ), |
67 | QueueCloseOp); |
68 | REGISTER_KERNEL_BUILDER(Name("QueueSizeV2" ) |
69 | .Device(DEVICE_DEFAULT) |
70 | .HostMemory("size" ) |
71 | .HostMemory("handle" ), |
72 | QueueSizeOp); |
73 | REGISTER_KERNEL_BUILDER( |
74 | Name("QueueIsClosedV2" ).Device(DEVICE_DEFAULT).HostMemory("handle" ), |
75 | QueueIsClosedOp); |
76 | |
77 | class FakeQueueOp : public OpKernel { |
78 | public: |
79 | explicit FakeQueueOp(OpKernelConstruction* context) : OpKernel(context) { |
80 | OP_REQUIRES_OK( |
81 | context, context->allocate_temp(DT_STRING, TensorShape({2}), &tensor_)); |
82 | } |
83 | |
84 | void Compute(OpKernelContext* context) override { |
85 | const ResourceHandle& ref = context->input(0).flat<ResourceHandle>()(0); |
86 | tensor_.flat<tstring>()(0) = ref.container(); |
87 | tensor_.flat<tstring>()(1) = ref.name(); |
88 | context->set_output_ref(0, &mu_, &tensor_); |
89 | } |
90 | |
91 | private: |
92 | mutex mu_; |
93 | Tensor tensor_; |
94 | }; |
95 | |
96 | REGISTER_KERNEL_BUILDER(Name("FakeQueue" ).Device(DEVICE_CPU), FakeQueueOp); |
97 | |
98 | } // namespace tensorflow |
99 | |