diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 566bab23a11867..b71b719bd40c76 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -1711,6 +1711,49 @@ absl::Status MarkForCompilationPassImpl::FindCompilationCandidates() { continue; } + if (debug_options_.enable_dynamic_sizes) { + // Under dynamic sizes, keep ops in TF if XLA expects a compile-time + // constant input that is actually computed at runtime. + const OpDef* op_def; + TF_RETURN_IF_ERROR( + graph_->op_registry()->LookUpOpDef(node->type_string(), &op_def)); + TF_ASSIGN_OR_RETURN(std::vector compile_time_constant_input_indices, + XlaOpRegistry::CompileTimeConstantInputs(node->def(), + *op_def)); + bool rejected_for_nonconst_compile_time_input = false; + for (int input_index : compile_time_constant_input_indices) { + bool found_input_edge = false; + for (const Edge* edge : node->in_edges()) { + if (edge->IsControlEdge() || edge->dst_input() != input_index) { + continue; + } + found_input_edge = true; + if (!edge->src()->IsConstant()) { + VLOG(1) << "Rejecting " << node->name() + << " from XLA clustering: compile-time constant input " + << input_index << " is not fed by a Const. producer=" + << edge->src()->name() << " producer_op=" + << edge->src()->type_string(); + rejected_for_nonconst_compile_time_input = true; + } + break; + } + if (!found_input_edge) { + VLOG(1) << "Rejecting " << node->name() + << " from XLA clustering: missing data edge for " + "compile-time constant input " + << input_index; + rejected_for_nonconst_compile_time_input = true; + } + if (rejected_for_nonconst_compile_time_input) { + break; + } + } + if (rejected_for_nonconst_compile_time_input) { + continue; + } + } + if (node->type_string() == "Const") { // Skip Const op with type DT_STRING, since XLA autoclustering doesn't // support it.