1 /* Copyright 2021 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/compiler/tf2xla/mlir_xla_op_kernel.h"
17
18 #include "tensorflow/compiler/jit/xla_compilation_cache.h"
19 #include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
20 #include "tensorflow/compiler/mlir/utils/array_container_utils.h"
21
22 namespace tensorflow {
23
ContextToXlaArgs(XlaOpKernelContext * ctx,std::vector<XlaCompiler::Argument> & xla_args)24 Status MlirXlaOpKernel::ContextToXlaArgs(
25 XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>& xla_args) {
26 // Collect arguments that are registered as CompileTimeConstantInput.
27 std::vector<int> registered_consts_vec;
28 TF_RETURN_IF_ERROR(tensorflow::XlaOpRegistry::CompileTimeConstantInputs(
29 *this, ®istered_consts_vec));
30 llvm::SmallDenseSet<int, 4> registered_consts;
31 registered_consts.insert(registered_consts_vec.begin(),
32 registered_consts_vec.end());
33
34 int num_inputs = ctx->num_inputs();
35 xla_args.reserve(num_inputs);
36 for (int i = 0; i < num_inputs; ++i) {
37 // TODO(b/180448774): Handle kResource and kTensorList.
38 XlaExpression::Kind ctx_kind_i = ctx->InputExpression(i).kind();
39 if (ctx_kind_i != XlaExpression::Kind::kXlaOp &&
40 ctx_kind_i != XlaExpression::Kind::kConstant)
41 return tensorflow::errors::InvalidArgument(
42 absl::StrCat("Input ", i, " to an MlirXlaOpKernel is invalid: ",
43 ctx->InputExpression(i).HumanString()));
44 XlaCompiler::Argument arg;
45 arg.type = ctx->input_type(i);
46 arg.shape = ctx->InputXlaShape(i).ValueOrDie();
47 arg.name = absl::StrCat("_arg", i);
48 if (registered_consts.count(i)) {
49 arg.kind = XlaCompiler::Argument::kConstant;
50 TF_ASSIGN_OR_RETURN(arg.constant_value, ctx->ConstantInputTensor(i));
51 } else {
52 arg.kind = XlaCompiler::Argument::kParameter;
53 }
54 xla_args.push_back(arg);
55 }
56 return OkStatus();
57 }
58
MlirXlaOpKernel(OpKernelConstruction * ctx)59 MlirXlaOpKernel::MlirXlaOpKernel(OpKernelConstruction* ctx)
60 : XlaOpKernel(ctx),
61 // Since this kernel implements lowering for a single TF operation, we
62 // disable MLIR threading for efficiency purpose (avoid starting a large
63 // number of threads eagerly).
64 mlir_ctx_(mlir::MLIRContext::Threading::DISABLED) {}
65
ConstructXlaOp(XlaOpKernelContext * ctx)66 Status MlirXlaOpKernel::ConstructXlaOp(XlaOpKernelContext* ctx) {
67 // Create input XlaArguments.
68 std::vector<XlaCompiler::Argument> xla_args;
69 TF_RETURN_IF_ERROR(ContextToXlaArgs(ctx, xla_args));
70
71 // Create input XlaOps.
72 llvm::SmallVector<xla::XlaOp, 4> xla_params(ctx->num_inputs());
73 for (int i = 0, end = ctx->num_inputs(); i < end; ++i) {
74 xla_params[i] = ctx->Input(i);
75 }
76
77 // Create outputs.
78 std::vector<DataType> result_dtypes(ctx->num_outputs());
79 for (int i = 0, end = result_dtypes.size(); i < end; ++i) {
80 result_dtypes[i] = ctx->expected_output_dtype(i);
81 }
82
83 // When there are no data-flow outputs from the node, the node is used as a
84 // control output by the graph to TensorflowDialect importer.
85 std::vector<std::string> control_rets;
86 if (result_dtypes.empty()) {
87 control_rets.push_back(def().name());
88 }
89
90 // Get the context's device.
91 auto device = dynamic_cast<Device*>(ctx->op_kernel_context()->device());
92 if (!device) {
93 return tensorflow::errors::InvalidArgument(
94 "Expected the XlaOpKernelContext argument's device to have type "
95 "Device.");
96 }
97
98 // Create a graph that wraps the kernel.
99 TF_ASSIGN_OR_RETURN(auto graph, CreateGraph(def(), xla_args, result_dtypes));
100
101 // Compile the graph to HLO.
102 GraphDebugInfo debug_info;
103 std::vector<xla::XlaOp> returns(1);
104 TF_RETURN_IF_ERROR(BuildHloFromGraph(
105 *graph, *ctx->builder(), mlir_ctx_, xla_params, returns,
106 mlir::SpanToArrayRef<XlaCompiler::Argument>(xla_args), control_rets,
107 device->device_type(),
108 *ctx->function_library()->GetFunctionLibraryDefinition(), debug_info,
109 {}));
110
111 // Set context outputs.
112 for (int i = 0, end = returns.size(); i < end; ++i) {
113 ctx->SetOutput(i, returns[i]);
114 }
115
116 return OkStatus();
117 }
118
Compile(XlaOpKernelContext * ctx)119 void MlirXlaOpKernel::Compile(XlaOpKernelContext* ctx) {
120 auto status = ConstructXlaOp(ctx);
121 if (!status.ok()) {
122 errors::AppendToMessage(&status, "Failure to legalize ", def().name(),
123 " using MlirXlaOpKernel in the tf2xla bridge.");
124 }
125 OP_REQUIRES_OK(ctx, status);
126 }
127
128 } // namespace tensorflow
129