1 /* Copyright 2019 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/kernels/if_while_utils.h"
17
18 #include "tensorflow/compiler/tf2xla/const_analysis.h"
19 #include "tensorflow/compiler/tf2xla/literal_util.h"
20 #include "tensorflow/compiler/xla/literal.h"
21
22 namespace tensorflow {
23
24 const char kPropagateCompileTimeConsts[] = "_xla_propagate_compile_time_consts";
25
ConvertCompileTimeConstArgumentsToConst(XlaOpKernelContext * ctx,std::vector<XlaCompiler::Argument> * args,int xla_expression_offset,std::function<bool (int arg_idx)> should_resolve_constant)26 absl::InlinedVector<int, 5> ConvertCompileTimeConstArgumentsToConst(
27 XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>* args,
28 int xla_expression_offset,
29 std::function<bool(int arg_idx)> should_resolve_constant) {
30 absl::InlinedVector<int, 5> resolved_constant_idxs;
31 for (int i = 0; i < args->size(); i++) {
32 XlaCompiler::Argument* arg = &(*args)[i];
33 const XlaExpression& expression =
34 ctx->InputExpression(i + xla_expression_offset);
35 // If the input tensor is a compile time constant build a kConstant type
36 // argument.
37 if (should_resolve_constant(i)) {
38 // NOTE: We can not simply check that this is Kind::kConstant because
39 // this could be the output of a MetadataOnly op e.g. Size.
40 xla::StatusOr<absl::optional<Tensor>> maybe_constant =
41 expression.ResolveConstant(ctx->compiler()->client());
42 if (maybe_constant.ok() && maybe_constant.ValueOrDie().has_value()) {
43 xla::StatusOr<Tensor> values_are_dynamic =
44 expression.ResolveDynamism(ctx->compiler()->client());
45 bool all_values_are_static = false;
46 if (!values_are_dynamic.ok()) {
47 // Conservatiely assume all values are dynamic.
48 all_values_are_static = true;
49 } else {
50 xla::Literal literal =
51 HostTensorToLiteral(values_are_dynamic.ValueOrDie()).ValueOrDie();
52 all_values_are_static = literal.IsAll(0);
53 }
54
55 if (all_values_are_static) {
56 arg->kind = XlaCompiler::Argument::kConstant;
57 arg->type = expression.dtype();
58 arg->constant_value = std::move(maybe_constant.ValueOrDie().value());
59 arg->shape = expression.GetShape().ValueOrDie();
60 resolved_constant_idxs.push_back(i);
61 } else {
62 arg->value_bound.emplace(
63 std::move(maybe_constant.ValueOrDie().value()));
64 }
65 }
66 }
67 }
68 return resolved_constant_idxs;
69 }
70
FindMustBeConstNodes(XlaOpKernelContext * ctx,const NameAttrList & func_name,std::vector<bool> * must_be_const_nodes,const FunctionBody ** body)71 Status FindMustBeConstNodes(XlaOpKernelContext* ctx,
72 const NameAttrList& func_name,
73 std::vector<bool>* must_be_const_nodes,
74 const FunctionBody** body) {
75 TF_RETURN_IF_ERROR(ctx->compiler()->FindFunctionBody(func_name, body));
76 must_be_const_nodes->resize((*body)->graph->num_node_ids(), false);
77 return BackwardsConstAnalysis(*((*body)->graph),
78 /*compile_time_const_arg_indices=*/nullptr,
79 must_be_const_nodes, ctx->function_library());
80 }
81
82 } // namespace tensorflow
83