1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16// Common kernel registrations for XLA devices.
17
18#ifndef TENSORFLOW_COMPILER_JIT_XLA_DEVICE_OPS_H_
19#define TENSORFLOW_COMPILER_JIT_XLA_DEVICE_OPS_H_
20
21#include "tensorflow/core/framework/op_kernel.h"
22#include "tensorflow/core/framework/resource_mgr.h"
23#include "tensorflow/core/kernels/cast_op.h"
24#include "tensorflow/core/kernels/constant_op.h"
25#include "tensorflow/core/kernels/control_flow_ops.h"
26#include "tensorflow/core/kernels/identity_op.h"
27#include "tensorflow/core/kernels/no_op.h"
28#include "tensorflow/core/kernels/sendrecv_ops.h"
29#include "tensorflow/core/kernels/variable_ops.h"
30
31namespace tensorflow {
32
33// Dummy OpKernel, used for kernels assigned to an XLA device that should be
34// compiled. Should never be called at runtime since such ops should be
35// rewritten to a _XlaLaunch op. If it is called, it means the placer placed an
36// operator on an XLA device but the compiler did not compile it.
37class XlaDeviceDummyOp : public OpKernel {
38 public:
39  explicit XlaDeviceDummyOp(OpKernelConstruction* ctx);
40  void Compute(OpKernelContext* ctx) override;
41};
42
43#define REGISTER_XLA_LAUNCH_KERNEL(DEVICE, KERNEL, TYPES) \
44  REGISTER_KERNEL_BUILDER(Name("_XlaLaunch")              \
45                              .Device(DEVICE)             \
46                              .HostMemory("constants")    \
47                              .HostMemory("resources"),   \
48                          KERNEL);
49
50#define REGISTER_XLA_DEVICE_KERNELS(DEVICE, TYPES)                             \
51  REGISTER_KERNEL_BUILDER(Name("_Send").Device(DEVICE), SendOp);               \
52  REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE), RecvOp);               \
53  REGISTER_KERNEL_BUILDER(                                                     \
54      Name("_HostSend").Device(DEVICE).HostMemory("tensor"), SendOp);          \
55  REGISTER_KERNEL_BUILDER(                                                     \
56      Name("_HostRecv").Device(DEVICE).HostMemory("tensor"), RecvOp);          \
57  REGISTER_KERNEL_BUILDER(                                                     \
58      Name("_HostCast").Device(DEVICE).HostMemory("x").HostMemory("y"),        \
59      CpuCastOp);                                                              \
60  REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE), NoOp);                  \
61  REGISTER_KERNEL_BUILDER(                                                     \
62      Name("Const").Device(DEVICE).TypeConstraint("dtype", TYPES),             \
63      ConstantOp);                                                             \
64  REGISTER_KERNEL_BUILDER(                                                     \
65      Name("Identity").Device(DEVICE).TypeConstraint("T", TYPES), IdentityOp); \
66  REGISTER_KERNEL_BUILDER(Name("Placeholder").Device(DEVICE), PlaceholderOp);  \
67  REGISTER_KERNEL_BUILDER(Name("PlaceholderV2").Device(DEVICE),                \
68                          PlaceholderOp);                                      \
69                                                                               \
70  REGISTER_KERNEL_BUILDER(                                                     \
71      Name("VarHandleOp").Device(DEVICE).HostMemory("resource"),               \
72      ResourceHandleOp<Var>);
73
74}  // namespace tensorflow
75
76#endif  // TENSORFLOW_COMPILER_JIT_XLA_DEVICE_OPS_H_
77