diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a8197744359fde..77db3205344912 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -2015,6 +2015,19 @@ cc_library( hdrs = ["xla_batch_matcher.h"], deps = [ "//tensorflow/core/platform:logging", + "@com_google_absl//absl/types:span", "@local_xla//xla:debug_options_flags", + "@local_xla//xla:shape_util", + ], +) + +tf_cc_test( + name = "xla_batch_matcher_test", + srcs = ["xla_batch_matcher_test.cc"], + deps = [ + ":xla_batch_matcher", + "//tensorflow/core:test", + "@com_google_googletest//:gtest_main", + "@local_xla//xla:shape_util", ], -) \ No newline at end of file +) diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index 5a6f741a01e972..3c068333926614 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -167,6 +167,15 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { Flag("tf_xla_enable_dynamic_sizes", &mark_for_compilation_flags->tf_xla_enable_dynamic_sizes, "Enable dynamic sizes support."), + Flag("tf_xla_disable_dynamic_size_padding", + &mark_for_compilation_flags->tf_xla_disable_dynamic_size_padding, + "Disable padding dynamic runtime sizes up to an XLA compile batch " + "when dynamic sizes support is enabled."), + Flag("tf_xla_enable_dynamic_solve_majority_vote", + &mark_for_compilation_flags + ->tf_xla_enable_dynamic_solve_majority_vote, + "Enable solve-time majority vote when dynamic expression solving " + "observes conflicting candidates."), Flag("tf_xla_enable_symbolic_content", &mark_for_compilation_flags->tf_xla_enable_symbolic_content, "Enable symbolic content propagation."), @@ -265,6 +274,9 @@ void AllocateAndParseFlags() { ->tf_xla_disable_resource_variable_safety_checks_for_debugging = false; mark_for_compilation_flags->tf_xla_deterministic_cluster_names = false; mark_for_compilation_flags->tf_xla_enable_dynamic_sizes = false; + mark_for_compilation_flags->tf_xla_disable_dynamic_size_padding = false; + mark_for_compilation_flags->tf_xla_enable_dynamic_solve_majority_vote = + false; mark_for_compilation_flags->tf_xla_enable_symbolic_content = false; mark_for_compilation_flags->tf_xla_persistent_cache_directory = ""; mark_for_compilation_flags->tf_xla_persistent_cache_device_types = ""; diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index 480e20dc7ec296..cbef2eb629d6be 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -102,6 +102,13 @@ struct MarkForCompilationPassFlags { // If true enables support of dynamic sizes. bool tf_xla_enable_dynamic_sizes; + // If true disables padding runtime dynamic sizes up to an XLA compile batch. + bool tf_xla_disable_dynamic_size_padding; + + // If true enables solve-time majority vote when dynamic expression solving + // observes conflicting candidates. + bool tf_xla_enable_dynamic_solve_majority_vote; + // If true enables symbolic content propagation. bool tf_xla_enable_symbolic_content; diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 2fb97d66b3ac63..ac3af0e4e5c280 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -36,6 +36,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" @@ -205,6 +206,715 @@ using PjRtExecutableClosure = using PjRtExecutableClosureStore = ExecutableClosureStore; +struct DynamicBatchResolutionResult { + bool can_run = true; + bool has_batch_size = false; + int64_t batch_size = 0; + bool batch_size_resource_not_found = false; + std::string diagnostic; +}; + +struct DynamicSolveCandidate { + int64_t solved_value; + int64_t observed_dim_value; + bool may_be_broadcast_singleton; + std::string context; +}; + +struct IgnoredDynamicArgumentOccurrence { + enum class Source { + kShapeDimension, + kConstantValueElement, + }; + + Source source; + int arg_index; + int dim_or_index; + int64_t observed_value; + int64_t solved_value; + std::string expr; +}; + +struct DynamicSolveFilterDecision { + bool can_run = true; + std::string diagnostic; + std::vector ignored_occurrences; +}; + +std::optional GetConstantArgumentElementValue( + const XlaCompiler::Argument& arg, int index) { + if (index < 0 || index >= arg.constant_value.NumElements()) { + return std::nullopt; + } + switch (arg.constant_value.dtype()) { + case DT_INT32: + return static_cast(arg.constant_value.flat()(index)); + case DT_INT64: + return static_cast(arg.constant_value.flat()(index)); + default: + return std::nullopt; + } +} + +void SetConstantArgumentExpressionToLiteralValue(XlaCompiler::Argument* arg, + int index) { + if (arg == nullptr || index < 0 || + index >= arg->constant_value_expressions.size()) { + return; + } + std::optional value = GetConstantArgumentElementValue(*arg, index); + if (!value.has_value()) { + arg->constant_value_expressions[index].Clear(); + return; + } + arg->constant_value_expressions[index].Clear(); + arg->constant_value_expressions[index].set_constant_value(*value); +} + +DynamicSolveFilterDecision AnalyzeIgnoredDynamicArgumentOccurrences( + absl::Span args) { + struct Candidate { + int64_t solved_value; + int64_t observed_value; + bool may_be_broadcast_singleton; + IgnoredDynamicArgumentOccurrence occurrence; + }; + + DynamicSolveFilterDecision result; + std::map> expr_to_candidates; + + for (int arg_index = 0; arg_index < args.size(); ++arg_index) { + const XlaCompiler::Argument& arg = args[arg_index]; + if (absl::holds_alternative(arg.shape)) { + const TensorShape& shape = std::get(arg.shape); + for (int dim = 0; dim < shape.get_expressions().size(); ++dim) { + const xla::DExpr& expr = shape.get_expression(dim); + if (!(expr && expr->is_dynamic())) { + continue; + } + xla::DExpr simplified_expr = expr.simplify(); + std::optional solved_value = + simplified_expr->solve(shape.dim_size(dim)); + if (!solved_value.has_value()) { + continue; + } + const std::string expr_string = DExprToString(simplified_expr); + const bool may_be_broadcast_singleton = + dim < arg.may_be_broadcast_singleton_dimensions.size() && + arg.may_be_broadcast_singleton_dimensions[dim]; + expr_to_candidates[expr_string].push_back(Candidate{ + *solved_value, + shape.dim_size(dim), + may_be_broadcast_singleton, + IgnoredDynamicArgumentOccurrence{ + IgnoredDynamicArgumentOccurrence::Source::kShapeDimension, + arg_index, + dim, + shape.dim_size(dim), + *solved_value, + expr_string, + }, + }); + } + } + + if (arg.kind != XlaCompiler::Argument::kConstant || + arg.constant_value_expressions.empty()) { + continue; + } + for (int element_index = 0; + element_index < arg.constant_value_expressions.size(); + ++element_index) { + xla::DExpr expr = + xla::DExprFromProto(arg.constant_value_expressions[element_index]); + if (!(expr && expr->is_dynamic())) { + continue; + } + std::optional observed_value = + GetConstantArgumentElementValue(arg, element_index); + if (!observed_value.has_value()) { + continue; + } + xla::DExpr simplified_expr = expr.simplify(); + std::optional solved_value = + simplified_expr->solve(*observed_value); + if (!solved_value.has_value()) { + continue; + } + const std::string expr_string = DExprToString(simplified_expr); + expr_to_candidates[expr_string].push_back(Candidate{ + *solved_value, + *observed_value, + false, + IgnoredDynamicArgumentOccurrence{ + IgnoredDynamicArgumentOccurrence::Source::kConstantValueElement, + arg_index, + element_index, + *observed_value, + *solved_value, + expr_string, + }, + }); + } + } + + std::vector expr_summaries; + for (const auto& [expr, candidates] : expr_to_candidates) { + std::set singleton_values; + std::set nonsingleton_values; + for (const auto& candidate : candidates) { + if (candidate.observed_value == 1 && + candidate.may_be_broadcast_singleton) { + singleton_values.insert(candidate.solved_value); + } else { + nonsingleton_values.insert(candidate.solved_value); + } + } + + std::optional chosen_value; + if (!nonsingleton_values.empty()) { + if (nonsingleton_values.size() != 1) { + result.can_run = false; + expr_summaries.push_back(absl::StrCat( + "expr=", expr, " reason=conflicting strong candidates", + " values={", absl::StrJoin(nonsingleton_values, ", "), "}")); + continue; + } + chosen_value = *nonsingleton_values.begin(); + for (const auto& candidate : candidates) { + if (candidate.observed_value == 1 && + candidate.may_be_broadcast_singleton && + candidate.solved_value != *chosen_value) { + LOG(INFO) << "Ignoring marked broadcast-singleton dynamic " + << "occurrence during XLA signature filtering: expr=" + << expr + << " chosen_value=" << *chosen_value + << " ignored_observed_value=" << candidate.observed_value + << " ignored_solved_value=" << candidate.solved_value + << " arg_index=" << candidate.occurrence.arg_index + << (candidate.occurrence.source == + IgnoredDynamicArgumentOccurrence::Source:: + kShapeDimension + ? " dim=" + : " element=") + << candidate.occurrence.dim_or_index; + result.ignored_occurrences.push_back(candidate.occurrence); + } + } + } else if (!singleton_values.empty()) { + if (singleton_values.size() != 1) { + result.can_run = false; + expr_summaries.push_back(absl::StrCat( + "expr=", expr, + " reason=conflicting marked broadcast-singleton candidates", + " values={", absl::StrJoin(singleton_values, ", "), "}")); + continue; + } + chosen_value = *singleton_values.begin(); + } + + expr_summaries.push_back(absl::StrCat( + "expr=", expr, " chosen=", + chosen_value.has_value() ? std::to_string(*chosen_value) + : std::string(""), + " marked_broadcast_singleton_values={", + absl::StrJoin(singleton_values, ", "), "} strong_values={", + absl::StrJoin(nonsingleton_values, ", "), + "}")); + } + + if (!result.can_run) { + result.diagnostic = absl::StrCat( + "Failed to recover a unique XLA dynamic batch size after solve-time " + "candidate filtering. solved_expressions=[", + absl::StrJoin(expr_summaries, " | "), "]"); + return result; + } + + if (!result.ignored_occurrences.empty()) { + std::vector ignored_summaries; + ignored_summaries.reserve(result.ignored_occurrences.size()); + for (const auto& occurrence : result.ignored_occurrences) { + ignored_summaries.push_back(absl::StrCat( + occurrence.source == + IgnoredDynamicArgumentOccurrence::Source::kShapeDimension + ? "shape_arg=" + : "const_arg=", + occurrence.arg_index, + occurrence.source == + IgnoredDynamicArgumentOccurrence::Source::kShapeDimension + ? " dim=" + : " element=", + occurrence.dim_or_index, " expr=", occurrence.expr, + " observed=", occurrence.observed_value, + " solved=", occurrence.solved_value)); + } + result.diagnostic = absl::StrCat( + "Ignoring marked broadcast-singleton dynamic occurrences for XLA " + "signature/HLO: ", + absl::StrJoin(ignored_summaries, "; ")); + } + return result; +} + +void StripIgnoredDynamicArgumentOccurrences( + const std::vector& ignored_occurrences, + std::vector* args) { + if (args == nullptr) { + return; + } + for (const auto& occurrence : ignored_occurrences) { + if (occurrence.arg_index < 0 || occurrence.arg_index >= args->size()) { + continue; + } + XlaCompiler::Argument& arg = (*args)[occurrence.arg_index]; + if (occurrence.source == + IgnoredDynamicArgumentOccurrence::Source::kShapeDimension) { + if (!absl::holds_alternative(arg.shape)) { + continue; + } + LOG(INFO) << "Dropping ignored dynamic shape expression from XLA " + << "signature/HLO: arg_index=" << occurrence.arg_index + << " dim=" << occurrence.dim_or_index + << " expr=" << occurrence.expr + << " observed_value=" << occurrence.observed_value + << " solved_value=" << occurrence.solved_value; + TensorShape& shape = std::get(arg.shape); + shape.set_expression(occurrence.dim_or_index, xla::DExpr()); + continue; + } + + LOG(INFO) << "Dropping ignored dynamic constant-value expression from XLA " + << "signature/HLO: arg_index=" << occurrence.arg_index + << " element=" << occurrence.dim_or_index + << " expr=" << occurrence.expr + << " observed_value=" << occurrence.observed_value + << " solved_value=" << occurrence.solved_value; + SetConstantArgumentExpressionToLiteralValue(&arg, occurrence.dim_or_index); + } +} + +DynamicBatchResolutionResult ResolveDynamicBatchSizeFromRuntimeInputs( + OpKernelContext* ctx, const XlaCompiler::CompilationResult& comp_result, + int num_constant_args, bool log_solves, absl::string_view cluster_name, + const NodeDef& op_def) { + DynamicBatchResolutionResult result; + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + if (!flags->tf_xla_enable_dynamic_sizes) { + return result; + } + const bool enable_dynamic_solve_majority_vote = + flags->tf_xla_enable_dynamic_solve_majority_vote; + + auto resolve_runtime_input_index = [&](int xla_input_index) -> int { + const bool has_runtime_key = + ctx->num_inputs() > 0 && + ctx->input_dtype(ctx->num_inputs() - 1) == DT_STRING && + ctx->num_inputs() == comp_result.xla_input_shapes.size() + 1; + if (has_runtime_key) { + return xla_input_index; + } + return num_constant_args + xla_input_index; + }; + + std::set dyn_vals; + std::map> expr_to_candidates; + + for (int i = 0; i < comp_result.xla_input_shapes.size(); ++i) { + const auto& xla_shape = comp_result.xla_input_shapes[i]; + const bool has_input_mapping = i < comp_result.input_mapping.size(); + const int mapped_input = has_input_mapping ? comp_result.input_mapping[i] + : -1; + const int input_idx = resolve_runtime_input_index(i); + if (log_solves) { + const bool input_in_range = + has_input_mapping && input_idx >= 0 && input_idx < ctx->num_inputs(); + const std::string runtime_input_name = + input_in_range && input_idx < op_def.input_size() + ? op_def.input(input_idx) + : (input_in_range ? std::string("") + : std::string("")); + VLOG(1) << "[XLA INPUT MAP DEBUG] cluster=" << cluster_name + << " xla_input_index=" << i + << " mapped_input=" << mapped_input + << " runtime_input_index=" << (input_in_range ? input_idx : -1) + << " runtime_input_name=" << runtime_input_name + << " xla_shape=" << xla_shape.DebugString(); + } + if (!xla_shape.IsArray() || xla_shape.expressions().empty()) { + continue; + } + + for (int dim = 0; dim < xla_shape.expressions().size(); ++dim) { + const auto& expr = xla_shape.expressions(dim); + if (!(expr && expr->is_dynamic())) { + continue; + } + xla::DExpr simplified_expr = expr.simplify(); + if (!has_input_mapping) { + continue; + } + if (input_idx < 0 || input_idx >= ctx->num_inputs()) { + continue; + } + const bool is_runtime_key_input = + ctx->input_dtype(input_idx) == DT_STRING && + input_idx == op_def.input_size() - 1; + const std::string runtime_input_name = + input_idx < op_def.input_size() ? op_def.input(input_idx) + : std::string(""); + if (is_runtime_key_input) { + if (log_solves) { + VLOG(1) << "Skipping dynamic expression solve for runtime key input " + << "cluster=" << cluster_name + << " xla_input_index=" << i + << " runtime_input_index=" << input_idx + << " input_name=" << runtime_input_name; + } + continue; + } + int64_t size = ctx->input(input_idx).shape().dim_size(dim); + std::optional dyn_val = simplified_expr->solve(size); + if (log_solves) { + LOG(INFO) << "Dynamic solve: cluster=" << cluster_name + << " runtime_input_index=" << input_idx + << " xla_input_index=" << i << " dim=" << dim + << " input_name=" << runtime_input_name + << " expr=" << DExprToString(simplified_expr) + << " target_size=" << size << " result=" + << (dyn_val.has_value() ? std::to_string(*dyn_val) + : std::string("")); + } + if (!dyn_val.has_value()) { + if (log_solves) { + xla::StringPrinter printer; + simplified_expr->print(&printer); + LOG(INFO) << "Failed to solve dynamic expression: " + << "xla_input_index=" << i + << " runtime_input_index=" << input_idx + << " input_name=" << runtime_input_name + << " dim=" << dim + << " runtime_dim_size=" << size + << " expr=" << std::move(printer).ToString(); + } + continue; + } + const std::string expr_string = DExprToString(simplified_expr); + const bool may_be_broadcast_singleton = + i < comp_result + .xla_input_may_be_broadcast_singleton_dimensions.size() && + dim < comp_result + .xla_input_may_be_broadcast_singleton_dimensions[i] + .size() && + comp_result.xla_input_may_be_broadcast_singleton_dimensions[i][dim]; + const std::string context = absl::StrCat( + "xla_input_index=", i, " runtime_input_index=", input_idx, + " input_name=", runtime_input_name, " dim=", dim, + " runtime_dim_size=", size, + " solved_dynamic_value=", *dyn_val, + " may_be_broadcast_singleton=", + may_be_broadcast_singleton ? "true" : "false"); + expr_to_candidates[expr_string].push_back( + DynamicSolveCandidate{*dyn_val, size, may_be_broadcast_singleton, + context}); + } + } + + std::vector expr_summaries; + std::optional> expected_dyn_ids; + bool mismatched_dyn_ids = false; + bool candidate_filter_rejected = false; + bool majority_vote_used = false; + bool ambiguous_majority_vote = false; + std::vector voted_expr_summaries; + for (const auto& [expr, candidates] : expr_to_candidates) { + std::set singleton_values; + std::set nonsingleton_values; + std::vector singleton_contexts; + std::vector nonsingleton_contexts; + for (const auto& candidate : candidates) { + if (candidate.observed_dim_value == 1 && + candidate.may_be_broadcast_singleton) { + singleton_values.insert(candidate.solved_value); + singleton_contexts.push_back(candidate.context); + } else { + nonsingleton_values.insert(candidate.solved_value); + nonsingleton_contexts.push_back(candidate.context); + } + } + + std::optional chosen_value; + std::string filter_reason; + if (!nonsingleton_values.empty()) { + if (nonsingleton_values.size() == 1) { + chosen_value = *nonsingleton_values.begin(); + filter_reason = + singleton_values.empty() + ? "used strong evidence" + : "preferred strong evidence over marked " + "broadcast-singleton candidates"; + } else { + if (!enable_dynamic_solve_majority_vote) { + candidate_filter_rejected = true; + filter_reason = + "conflicting strong candidates; majority vote disabled for " + "debugging"; + LOG(INFO) << "Dynamic solve encountered conflicting strong " + "candidates for expr=" + << expr + << " but majority vote is disabled for debugging."; + } else { + std::map> vote_contexts; + for (const auto& candidate : candidates) { + vote_contexts[candidate.solved_value].push_back(candidate.context); + } + int64_t winner_value = 0; + size_t winner_count = 0; + size_t runner_up_count = 0; + bool have_winner = false; + bool tie_for_winner = false; + std::vector tally_parts; + for (const auto& [value, contexts] : vote_contexts) { + tally_parts.push_back(absl::StrCat( + value, " x", contexts.size(), " [", + absl::StrJoin(contexts, "; "), "]")); + if (!have_winner || contexts.size() > winner_count) { + runner_up_count = have_winner ? winner_count : 0; + winner_value = value; + winner_count = contexts.size(); + have_winner = true; + tie_for_winner = false; + } else if (contexts.size() == winner_count) { + tie_for_winner = true; + } else if (contexts.size() > runner_up_count) { + runner_up_count = contexts.size(); + } + } + const bool clear_majority = !tie_for_winner && winner_count >= 2 && + winner_count > runner_up_count; + if (clear_majority) { + chosen_value = winner_value; + majority_vote_used = true; + filter_reason = + "singleton filtering conflicted; majority vote kept value"; + std::vector dropped_parts; + for (const auto& [value, contexts] : vote_contexts) { + if (value == winner_value) continue; + dropped_parts.push_back( + absl::StrCat(value, " x", contexts.size())); + } + LOG(INFO) << "Dynamic solve majority vote kept value for expr=" + << expr << " kept=" << winner_value + << " kept_count=" << winner_count + << " runner_up_count=" << runner_up_count + << " dropped=[" + << (dropped_parts.empty() ? std::string() + : absl::StrJoin(dropped_parts, ", ")) + << "] tallies=[" << absl::StrJoin(tally_parts, " | ") + << "]"; + voted_expr_summaries.push_back(absl::StrCat( + "expr=", expr, " kept=", winner_value, " count=", winner_count, + " dropped=[", + dropped_parts.empty() ? std::string() + : absl::StrJoin(dropped_parts, ", "), + "]")); + } else { + candidate_filter_rejected = true; + ambiguous_majority_vote = true; + filter_reason = + "conflicting strong candidates and no clear majority vote"; + LOG(INFO) << "Dynamic solve majority vote could not choose a unique " + "value for expr=" + << expr << " tallies=[" << absl::StrJoin(tally_parts, " | ") + << "]"; + } + } + } + } else if (!singleton_values.empty()) { + if (singleton_values.size() == 1) { + chosen_value = *singleton_values.begin(); + filter_reason = "used marked broadcast-singleton evidence only"; + } else { + if (!enable_dynamic_solve_majority_vote) { + candidate_filter_rejected = true; + filter_reason = + "conflicting marked broadcast-singleton candidates; majority " + "vote disabled for debugging"; + LOG(INFO) << "Dynamic solve encountered conflicting marked " + "broadcast-singleton candidates for expr=" + << expr + << " but majority vote is disabled for debugging."; + } else { + std::map> vote_contexts; + for (const auto& candidate : candidates) { + vote_contexts[candidate.solved_value].push_back(candidate.context); + } + int64_t winner_value = 0; + size_t winner_count = 0; + size_t runner_up_count = 0; + bool have_winner = false; + bool tie_for_winner = false; + std::vector tally_parts; + for (const auto& [value, contexts] : vote_contexts) { + tally_parts.push_back(absl::StrCat( + value, " x", contexts.size(), " [", + absl::StrJoin(contexts, "; "), "]")); + if (!have_winner || contexts.size() > winner_count) { + runner_up_count = have_winner ? winner_count : 0; + winner_value = value; + winner_count = contexts.size(); + have_winner = true; + tie_for_winner = false; + } else if (contexts.size() == winner_count) { + tie_for_winner = true; + } else if (contexts.size() > runner_up_count) { + runner_up_count = contexts.size(); + } + } + const bool clear_majority = !tie_for_winner && winner_count >= 2 && + winner_count > runner_up_count; + if (clear_majority) { + chosen_value = winner_value; + majority_vote_used = true; + filter_reason = + "singleton-only evidence conflicted; majority vote kept value"; + std::vector dropped_parts; + for (const auto& [value, contexts] : vote_contexts) { + if (value == winner_value) continue; + dropped_parts.push_back( + absl::StrCat(value, " x", contexts.size())); + } + LOG(INFO) << "Dynamic solve majority vote kept value for expr=" + << expr << " kept=" << winner_value + << " kept_count=" << winner_count + << " runner_up_count=" << runner_up_count + << " dropped=[" + << (dropped_parts.empty() ? std::string() + : absl::StrJoin(dropped_parts, ", ")) + << "] tallies=[" << absl::StrJoin(tally_parts, " | ") + << "]"; + voted_expr_summaries.push_back(absl::StrCat( + "expr=", expr, " kept=", winner_value, " count=", winner_count, + " dropped=[", + dropped_parts.empty() ? std::string() + : absl::StrJoin(dropped_parts, ", "), + "]")); + } else { + candidate_filter_rejected = true; + ambiguous_majority_vote = true; + filter_reason = + "conflicting marked broadcast-singleton candidates and no clear " + "majority vote"; + LOG(INFO) << "Dynamic solve majority vote could not choose a unique " + "value for expr=" + << expr << " tallies=[" << absl::StrJoin(tally_parts, " | ") + << "]"; + } + } + } + } else { + filter_reason = "no dynamic values were solved"; + } + + if (chosen_value.has_value()) { + dyn_vals.insert(*chosen_value); + } + + expr_summaries.push_back(absl::StrCat( + "expr=", expr, " chosen=", + chosen_value.has_value() ? std::to_string(*chosen_value) + : std::string(""), + " marked_broadcast_singleton_values={", + absl::StrJoin(singleton_values, ", "), "} strong_values={", + absl::StrJoin(nonsingleton_values, ", "), "} reason=", filter_reason, + " marked_broadcast_singleton_contexts=[", + absl::StrJoin(singleton_contexts, "; "), "] strong_contexts=[", + absl::StrJoin(nonsingleton_contexts, "; "), "]")); + } + for (int i = 0; i < comp_result.xla_input_shapes.size(); ++i) { + const auto& xla_shape = comp_result.xla_input_shapes[i]; + if (!xla_shape.IsArray() || xla_shape.expressions().empty()) continue; + for (int dim = 0; dim < xla_shape.expressions().size(); ++dim) { + const auto& expr = xla_shape.expressions(dim); + if (!(expr && expr->is_dynamic())) { + continue; + } + std::set ids = expr->get_all_ids(); + if (!expected_dyn_ids.has_value()) { + expected_dyn_ids = ids; + } else if (*expected_dyn_ids != ids) { + mismatched_dyn_ids = true; + } + } + } + + if (dyn_vals.size() == 1 && expected_dyn_ids.has_value() && + expected_dyn_ids->size() == 1 && !mismatched_dyn_ids) { + result.has_batch_size = true; + result.batch_size = *dyn_vals.begin(); + result.diagnostic = absl::StrCat( + "shared variable ids={", absl::StrJoin(*expected_dyn_ids, ", "), + "} value=", result.batch_size); + return result; + } + + if (expr_to_candidates.empty()) { + result.diagnostic = "no dynamic values were solved"; + } else if (candidate_filter_rejected) { + result.can_run = false; + result.diagnostic = absl::StrCat( + "Failed to recover a unique XLA dynamic batch size after solve-time " + "candidate filtering. solved_expressions=[", + absl::StrJoin(expr_summaries, " | "), "] voted_expressions=[", + absl::StrJoin(voted_expr_summaries, " | "), "] all_values={", + absl::StrJoin(dyn_vals, ", "), "}", + ambiguous_majority_vote ? " reason=no clear majority vote" : ""); + return result; + } else if (dyn_vals.size() == 1 && mismatched_dyn_ids) { + result.diagnostic = absl::StrCat( + "solved dynamic expressions do not share the same variable ids: ", + absl::StrJoin(expr_summaries, " | "), " voted_expressions=[", + absl::StrJoin(voted_expr_summaries, " | "), "]"); + } else if (dyn_vals.size() == 1) { + result.diagnostic = absl::StrCat( + "solved dynamic expressions do not map to exactly one shared dynamic " + "variable id: ", + absl::StrJoin(expr_summaries, " | "), " voted_expressions=[", + absl::StrJoin(voted_expr_summaries, " | "), "]"); + } else { + result.can_run = false; + result.diagnostic = absl::StrCat( + "Failed to recover a unique XLA dynamic batch size from runtime " + "input expressions. solved_expressions=[", + absl::StrJoin(expr_summaries, " | "), "] voted_expressions=[", + absl::StrJoin(voted_expr_summaries, " | "), "] all_values={", + absl::StrJoin(dyn_vals, ", "), "}", + majority_vote_used ? " after majority-vote filtering" : ""); + return result; + } + + BatchSizeResource* bsr = nullptr; + ScopedStepContainer* step_container = ctx->step_container(); + absl::Status st = step_container->Lookup( + ctx->resource_manager(), BatchSizeResourceName, &bsr); + + if (st.ok()) { + CHECK(bsr != nullptr); + result.has_batch_size = true; + result.batch_size = bsr->GetBatchSize(); + result.diagnostic = absl::StrCat( + "BatchSizeResource fallback value=", result.batch_size); + bsr->Unref(); + } else if (IsNotFound(st)) { + result.batch_size_resource_not_found = true; + } else { + result.can_run = false; + result.diagnostic = st.ToString(); + } + + return result; +} + se::Stream* GetStream(OpKernelContext* ctx) { return ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; @@ -427,7 +1137,8 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) { return xla::DExpr::Const(ac->value()); } case DimExpr::Kind::kVariable: { - return xla::DExpr::Var(1); + auto* av = static_cast(e); + return xla::DExpr::Var(av->id()); } case DimExpr::Kind::kAdd: { auto* ee = static_cast(e); @@ -449,6 +1160,168 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) { return xla::DExpr::Unknown(); } +int ExprProtoNodeCount(const xla::ExpressionProto& proto) { + switch (proto.node_type_case()) { + case xla::ExpressionProto::kConstantValue: + case xla::ExpressionProto::kVariableId: + case xla::ExpressionProto::NODE_TYPE_NOT_SET: + return 1; + case xla::ExpressionProto::kAddNode: + return 1 + ExprProtoNodeCount(proto.add_node().lhs()) + + ExprProtoNodeCount(proto.add_node().rhs()); + case xla::ExpressionProto::kSubNode: + return 1 + ExprProtoNodeCount(proto.sub_node().lhs()) + + ExprProtoNodeCount(proto.sub_node().rhs()); + case xla::ExpressionProto::kMulNode: + return 1 + ExprProtoNodeCount(proto.mul_node().lhs()) + + ExprProtoNodeCount(proto.mul_node().rhs()); + case xla::ExpressionProto::kDivNode: + return 1 + ExprProtoNodeCount(proto.div_node().lhs()) + + ExprProtoNodeCount(proto.div_node().rhs()); + } + return 1; +} + +void CollectCoveringSubexpressions(const xla::DExpr& expr, + const std::set& ids, + std::vector* matches) { + CHECK(matches != nullptr); + if (!expr || !expr->is_dynamic()) { + return; + } + if (expr->get_all_ids() == ids) { + matches->push_back(expr); + } + + xla::ExpressionProto proto; + expr.to_proto(&proto); + auto recurse = [&](const xla::ExpressionProto& child) { + xla::DExpr child_expr = xla::DExprFromProto(child); + CollectCoveringSubexpressions(child_expr, ids, matches); + }; + + switch (proto.node_type_case()) { + case xla::ExpressionProto::kAddNode: + recurse(proto.add_node().lhs()); + recurse(proto.add_node().rhs()); + return; + case xla::ExpressionProto::kSubNode: + recurse(proto.sub_node().lhs()); + recurse(proto.sub_node().rhs()); + return; + case xla::ExpressionProto::kMulNode: + recurse(proto.mul_node().lhs()); + recurse(proto.mul_node().rhs()); + return; + case xla::ExpressionProto::kDivNode: + recurse(proto.div_node().lhs()); + recurse(proto.div_node().rhs()); + return; + case xla::ExpressionProto::kConstantValue: + case xla::ExpressionProto::kVariableId: + case xla::ExpressionProto::NODE_TYPE_NOT_SET: + return; + } +} + +xla::DExpr FindSmallestCoveringSubexpression(const xla::DExpr& expr) { + CHECK(expr); + std::set ids = expr->get_all_ids(); + CHECK(!ids.empty()); + std::vector matches; + CollectCoveringSubexpressions(expr, ids, &matches); + CHECK(!matches.empty()); + auto best_it = std::min_element( + matches.begin(), matches.end(), + [](const xla::DExpr& lhs, const xla::DExpr& rhs) { + xla::ExpressionProto lhs_proto; + xla::ExpressionProto rhs_proto; + lhs.to_proto(&lhs_proto); + rhs.to_proto(&rhs_proto); + return ExprProtoNodeCount(lhs_proto) < ExprProtoNodeCount(rhs_proto); + }); + return *best_it; +} + +void ReplaceDynamicSubexpressionProto(xla::ExpressionProto* expr_proto, + const xla::DExpr& target, + int replacement_var_id) { + CHECK(expr_proto != nullptr); + xla::DExpr current = xla::DExprFromProto(*expr_proto); + if (current && current == target) { + expr_proto->Clear(); + expr_proto->set_variable_id(replacement_var_id); + return; + } + + switch (expr_proto->node_type_case()) { + case xla::ExpressionProto::kAddNode: + ReplaceDynamicSubexpressionProto(expr_proto->mutable_add_node()->mutable_lhs(), + target, replacement_var_id); + ReplaceDynamicSubexpressionProto(expr_proto->mutable_add_node()->mutable_rhs(), + target, replacement_var_id); + return; + case xla::ExpressionProto::kSubNode: + ReplaceDynamicSubexpressionProto(expr_proto->mutable_sub_node()->mutable_lhs(), + target, replacement_var_id); + ReplaceDynamicSubexpressionProto(expr_proto->mutable_sub_node()->mutable_rhs(), + target, replacement_var_id); + return; + case xla::ExpressionProto::kMulNode: + ReplaceDynamicSubexpressionProto(expr_proto->mutable_mul_node()->mutable_lhs(), + target, replacement_var_id); + ReplaceDynamicSubexpressionProto(expr_proto->mutable_mul_node()->mutable_rhs(), + target, replacement_var_id); + return; + case xla::ExpressionProto::kDivNode: + ReplaceDynamicSubexpressionProto(expr_proto->mutable_div_node()->mutable_lhs(), + target, replacement_var_id); + ReplaceDynamicSubexpressionProto(expr_proto->mutable_div_node()->mutable_rhs(), + target, replacement_var_id); + return; + case xla::ExpressionProto::kConstantValue: + case xla::ExpressionProto::kVariableId: + case xla::ExpressionProto::NODE_TYPE_NOT_SET: + return; + } +} + +xla::DExpr ReplaceDynamicSubexpression(const xla::DExpr& expr, + const xla::DExpr& target, + int replacement_var_id) { + if (!expr) { + return expr; + } + xla::ExpressionProto proto; + expr.to_proto(&proto); + ReplaceDynamicSubexpressionProto(&proto, target, replacement_var_id); + return xla::DExprFromProto(proto).simplify(); +} + +std::vector CollectReshapeOutputExpressions( + const FunctionDef& function) { + std::vector expressions; + for (const NodeDef& node : function.node_def()) { + if (node.op() != "Reshape") { + continue; + } + const auto shapes_it = + node.attr().find(kXlaInferredOutputTensorShapesAttrName); + if (shapes_it == node.attr().end()) { + continue; + } + for (const TensorShapeProto& shape : shapes_it->second.list().shape()) { + for (const ExpressionProto& expression : shape.expressions()) { + xla::DExpr dynamic_expression = + DimExprToDExpr(ExprFromProto(expression).get()).simplify(); + if (dynamic_expression && dynamic_expression->is_dynamic()) { + expressions.push_back(std::move(dynamic_expression)); + } + } + } + } + return expressions; +} absl::Status CompileToLocalExecutable( OpKernelContext* ctx, const NameAttrList& function, bool has_ref_vars, @@ -500,6 +1373,8 @@ absl::Status CompileToLocalExecutable( MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); if (flags->tf_xla_enable_dynamic_sizes) { + const bool disable_dynamic_size_padding = + flags->tf_xla_disable_dynamic_size_padding; // Rewriting the argument with expressions if they have dynamic // dimension, detecting dynamic dimension via either _dynamic_dim or the // inferred-output-shapes attr attached during encapsulation. @@ -511,7 +1386,44 @@ absl::Status CompileToLocalExecutable( int64_t dynamic_dim_value = 0; XlaBatchMatcher* xla_batch_matcher = xla_device_compiler->xla_batch_matcher(); + const FunctionDef* cluster_fdef = + options.flib_def == nullptr + ? nullptr + : options.flib_def->Find(function.name()); + const std::vector reshape_output_expressions = + cluster_fdef == nullptr + ? std::vector() + : CollectReshapeOutputExpressions(*cluster_fdef); + const int64_t dynamic_padding_alignment = + GetDynamicPaddingAlignment(reshape_output_expressions); + if (dynamic_padding_alignment > 1) { + VLOG(1) << "Using dynamic compile-batch alignment " + << dynamic_padding_alignment << " for cluster " + << function.name() << " based on internal Reshape expressions"; + } std::optional dynamic_dim_expr; + std::optional shared_dynamic_subexpr; + auto normalize_dynamic_expr = + [&](xla::DExpr expr, absl::string_view context) { + if (!expr || !expr->is_dynamic()) { + return expr; + } + if (!shared_dynamic_subexpr.has_value()) { + shared_dynamic_subexpr = FindSmallestCoveringSubexpression(expr); + LOG(INFO) << "Using shared dynamic subexpression " + << DExprToString(*shared_dynamic_subexpr) + << " for XLA dynamic input normalization. context=" + << context; + } + xla::DExpr normalized_expr = ReplaceDynamicSubexpression( + expr, *shared_dynamic_subexpr, /*replacement_var_id=*/1); + LOG(INFO) << "Rewriting dynamic input expression " + << DExprToString(expr) + << " to shared-core form " + << DExprToString(normalized_expr) + << " before XLA compilation. context=" << context; + return normalized_expr; + }; auto maybe_attach_shape_contents_from_attrs = [&](int arg_index, const auto& attr_map, const std::string& node_name) { @@ -558,7 +1470,11 @@ absl::Status CompileToLocalExecutable( xla::ExpressionProto expr; const xla::DExpr& dim_expr = inferred_shape.get_expression(i); if (dim_expr && dim_expr->is_dynamic()) { - dim_expr->to_proto(&expr); + xla::DExpr normalized_expr = normalize_dynamic_expr( + dim_expr, + absl::StrCat("const_arg=", arg_index, " node=", node_name, + " dim=", i)); + normalized_expr->to_proto(&expr); } else if (arg.constant_value.dtype() == DT_INT32) { expr.set_constant_value(arg.constant_value.flat()(i)); } else if (arg.constant_value.dtype() == DT_INT64) { @@ -576,6 +1492,17 @@ absl::Status CompileToLocalExecutable( VLOG(1) << "XlaCompileOp recovered " << arg.constant_value_expressions.size() << " constant_value_expressions for const arg " << arg_index << " node=" << node_name << " from user_inferred_shape"; + std::vector expr_parts; + expr_parts.reserve(arg.constant_value_expressions.size()); + for (const auto& expr_proto : arg.constant_value_expressions) { + expr_parts.push_back( + DExprToString(xla::DExprFromProto(expr_proto).simplify())); + } + VLOG(1) << "[XLA ARG CONTENT DEBUG] const arg_index=" << arg_index + << " node=" << node_name << " tensor_shape=" + << arg.constant_value.shape().DebugString() + << " content_exprs=[" << absl::StrJoin(expr_parts, ", ") + << "]"; }; auto record_dynamic_dim_value = [&](int64_t dim_size, xla::DExpr expr) { if (!saw_dynamic_dim_value) { @@ -601,6 +1528,9 @@ absl::Status CompileToLocalExecutable( if (shape_derived_attr != attr_map.end()) { VLOG(1) << "XlaCompileOp retrieved shape-derived marker for arg " << arg_index << " node=" << node_name; + VLOG(1) << "[XLA ARG INJECT DEBUG] arg_index=" << arg_index + << " node=" << node_name + << " source=shape_derived_attr"; } maybe_attach_shape_contents_from_attrs(arg_index, attr_map, node_name); @@ -611,18 +1541,34 @@ absl::Status CompileToLocalExecutable( std::get(norm_args[arg_index].shape); const AttrValue& v = dyn_dim_attr->second; int64_t idx = v.i(); - record_dynamic_dim_value(shp.dim_size(idx), xla::DExpr::Var(1)); - if (!filled_batch && xla_batch_matcher) { + xla::DExpr dyn_dim_expr = normalize_dynamic_expr( + xla::DExpr::Var(1), + absl::StrCat("arg=", arg_index, " dim=", idx, + " dynamic_dim_attr")); + record_dynamic_dim_value(shp.dim_size(idx), dyn_dim_expr); + if (!disable_dynamic_size_padding && !filled_batch && + xla_batch_matcher) { filled_batch = - xla_batch_matcher->get_xla_compile_batch(shp.dim_size(idx)); + xla_batch_matcher->get_xla_compile_batch( + shp.dim_size(idx), dynamic_padding_alignment); } std::vector dyn_exprs; for (int d : shp.dim_sizes()) { dyn_exprs.push_back(xla::DExpr::Const(d)); } - dyn_exprs[idx] = xla::DExpr::Var(1); + dyn_exprs[idx] = dyn_dim_expr; + VLOG(1) << "[XLA ARG INJECT DEBUG] arg_index=" << arg_index + << " node=" << node_name + << " source=_dynamic_dim" + << " target_dim=" << idx + << " before_shape=" << shp.DebugString() + << " injected_expr=" << DExprToString(dyn_dim_expr); shp.set_expressions(std::move(dyn_exprs)); + VLOG(1) << "[XLA ARG INJECT DEBUG] arg_index=" << arg_index + << " node=" << node_name + << " source=_dynamic_dim" + << " after_shape=" << shp.DebugString(); continue; } auto it = attr_map.find(kXlaInferredOutputShapesAttrName); @@ -631,15 +1577,51 @@ absl::Status CompileToLocalExecutable( const TensorShapeProto& proto = it->second.list().shape(0); const auto& exp = proto.expressions(); TensorShape& shp = std::get(norm_args[arg_index].shape); - - if (!filled_batch && xla_batch_matcher) { + norm_args[arg_index] + .may_be_broadcast_singleton_dimensions.assign(shp.dims(), + false); + for (int dim = 0; dim < proto.dim_size() && dim < shp.dims(); + ++dim) { + const bool may_be_broadcast_singleton = + proto.dim(dim).may_be_broadcast_singleton(); + norm_args[arg_index] + .may_be_broadcast_singleton_dimensions[dim] = + may_be_broadcast_singleton; + if (may_be_broadcast_singleton) { + LOG(INFO) << "XLA argument dimension may be a broadcast " + << "singleton: arg_index=" << arg_index + << " node=" << node_name << " dim=" << dim + << " observed_size=" << shp.dim_size(dim); + } + } + VLOG(1) << "[XLA ARG INJECT DEBUG] arg_index=" << arg_index + << " node=" << node_name + << " source=inferred_output_shape" + << " proto=" << proto.DebugString() + << " before_shape=" << shp.DebugString(); + + if (!disable_dynamic_size_padding && !filled_batch && + xla_batch_matcher) { for (int idx = 0; idx < exp.size(); ++idx) { // Look for dynamic expression. If found then compute padding // value and exit loop. auto e = DimExprToDExpr(ExprFromProto(exp[idx]).get()).simplify(); + e = normalize_dynamic_expr( + e, absl::StrCat("arg=", arg_index, " dim=", idx, + " before_fill_batch")) + .simplify(); if (e->is_dynamic()) { + LOG(INFO) << "Calling dynamic expression solve for compile " + << "argument " << arg_index << " dimension " << idx + << " expr=" << DExprToString(e) + << " target_size=" << shp.dim_size(idx); std::optional solved_value = e->solve(shp.dim_size(idx)); + LOG(INFO) << "Dynamic expression solve for compile argument " + << arg_index << " dimension " << idx << " returned " + << (solved_value.has_value() + ? std::to_string(*solved_value) + : std::string("")); int64_t var_value; if (!solved_value.has_value()) { LOG(WARNING) @@ -655,7 +1637,8 @@ absl::Status CompileToLocalExecutable( } record_dynamic_dim_value(var_value, e); filled_batch = - xla_batch_matcher->get_xla_compile_batch(var_value); + xla_batch_matcher->get_xla_compile_batch( + var_value, dynamic_padding_alignment); break; } } @@ -668,10 +1651,65 @@ absl::Status CompileToLocalExecutable( for (int j = 0; j < exp.size(); ++j) { auto e = DimExprToDExpr(ExprFromProto(exp[j]).get()); if (e->is_dynamic()) { + e = normalize_dynamic_expr( + e, absl::StrCat("arg=", arg_index, " dim=", j, + " input_shape")); + VLOG(1) << "[XLA ARG INJECT DEBUG] arg_index=" << arg_index + << " node=" << node_name + << " source=inferred_output_shape" + << " dim=" << j + << " injected_expr=" << DExprToString(e); dyn_exprs[j] = e; } } shp.set_expressions(std::move(dyn_exprs)); + VLOG(1) << "[XLA ARG CONTENT DEBUG] tensor arg_index=" << arg_index + << " node=" << node_name + << " normalized_shape=" << shp.DebugString(); + } + } + } + + DynamicSolveFilterDecision solve_filter_decision = + AnalyzeIgnoredDynamicArgumentOccurrences(norm_args); + if (!solve_filter_decision.can_run) { + if (compile_mode == DeviceCompileMode::kLazy) { + return errors::Unimplemented(solve_filter_decision.diagnostic); + } + return errors::InvalidArgument(solve_filter_decision.diagnostic); + } + if (!solve_filter_decision.ignored_occurrences.empty()) { + LOG(INFO) << solve_filter_decision.diagnostic; + StripIgnoredDynamicArgumentOccurrences( + solve_filter_decision.ignored_occurrences, &norm_args); + saw_dynamic_dim_value = false; + has_multiple_dynamic_dim_values = false; + dynamic_dim_value = 0; + dynamic_dim_expr.reset(); + filled_batch = 0; + for (int arg_index = 0; arg_index < norm_args.size(); ++arg_index) { + if (!absl::holds_alternative(norm_args[arg_index].shape)) { + continue; + } + TensorShape& shape = std::get(norm_args[arg_index].shape); + for (int dim = 0; dim < shape.get_expressions().size(); ++dim) { + xla::DExpr expr = shape.get_expression(dim); + if (!(expr && expr->is_dynamic())) { + continue; + } + xla::DExpr simplified_expr = expr.simplify(); + std::optional solved_value = + simplified_expr->solve(shape.dim_size(dim)); + if (!solved_value.has_value()) { + continue; + } + record_dynamic_dim_value(*solved_value, simplified_expr); + if (!disable_dynamic_size_padding && !filled_batch && + xla_batch_matcher) { + filled_batch = + xla_batch_matcher->get_xla_compile_batch( + *solved_value, dynamic_padding_alignment); + } } } } @@ -771,7 +1809,25 @@ absl::Status CompileToLocalExecutable( int64_t old = shp.dim_size(j); old_vars.push_back({i, j, old}); xla::DExpr padded_expr = xla::DExpr::Const(filled_batch); - xla::DExpr subst_expr = e.substitute(1, padded_expr).simplify(); + const std::set ids = e->get_all_ids(); + if (ids.size() != 1) { + return errors::InvalidArgument( + "Dynamic shape padding expected exactly one dynamic " + "variable for argument ", + i, ", dimension ", j, ", but found ", ids.size(), + " variables in expression ", DExprToString(e)); + } + const int substitute_var_id = *ids.begin(); + LOG(INFO) << "Calling dynamic expression substitute for compile " + << "argument " << i << " dimension " << j + << " expr=" << DExprToString(e) + << " substitute Var(" << substitute_var_id + << ")=" << filled_batch; + xla::DExpr subst_expr = + e.substitute(substitute_var_id, padded_expr).simplify(); + LOG(INFO) << "Dynamic expression substitute for compile argument " + << i << " dimension " << j + << " returned " << DExprToString(subst_expr); if (!subst_expr->is_constant()) { return errors::InvalidArgument( "Dynamic shape padding substitution did not produce an " @@ -1209,6 +2265,27 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { } } + if ((executable || pjrt_executable) && kernel != nullptr && + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) { + DynamicBatchResolutionResult resolution = + ResolveDynamicBatchSizeFromRuntimeInputs(ctx, *kernel, constants_.size(), + /*log_solves=*/true, + function_.name(), def()); + if (!resolution.can_run) { + const std::string error_message = absl::StrCat( + "Rejecting XLA cluster at compile time because dynamic expressions " + "cannot be solved consistently for this request. ", + resolution.diagnostic); + if (must_compile_) { + OP_REQUIRES(ctx, false, errors::InvalidArgument(error_message)); + } + LOG(WARNING) << error_message; + executable = nullptr; + pjrt_executable = nullptr; + kernel = nullptr; + } + } + AllocatorAttributes host_alloc_attrs; host_alloc_attrs.set_gpu_compatible(true); host_alloc_attrs.set_on_host(true); @@ -1268,6 +2345,14 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat()(0); PjRtExecutableClosure closure = PjRtExecutableClosureStore::Global()->Consume(key); + const std::string cluster_name = + closure.compilation_result() != nullptr && + closure.compilation_result()->computation != nullptr + ? closure.compilation_result()->computation->name() + : std::string(""); + LOG(INFO) << "Entering XLA cluster: cluster=" << cluster_name + << " op=" << def().name() << " key=" << key + << " mode=pjrt step_id=" << ctx->step_id(); // Fetch inputs from the OpKernelContext. Inputs are the same as the ones // for XlaCompile, except that the must-be-constant inputs that appear in @@ -1304,6 +2389,14 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { XlaExecutableClosure closure = XlaExecutableClosureStore::Global()->Consume(key); + const std::string cluster_name = + closure.compilation_result() != nullptr && + closure.compilation_result()->computation != nullptr + ? closure.compilation_result()->computation->name() + : std::string(""); + LOG(INFO) << "Entering XLA cluster: cluster=" << cluster_name + << " op=" << def().name() << " key=" << key + << " mode=local step_id=" << ctx->step_id(); std::shared_ptr allocator = GetAllocator(ctx->device(), GetStream(ctx), platform_info_); XlaComputationLaunchContext launch_context = @@ -1342,71 +2435,33 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); if (flags->tf_xla_enable_dynamic_sizes) { + DynamicBatchResolutionResult batch_resolution = + ResolveDynamicBatchSizeFromRuntimeInputs( + ctx, *closure.compilation_result(), closure.num_constant_args(), + /*log_solves=*/true, cluster_name, def()); bool is_set = false; - std::set dyn_vals; - const auto* comp_result = closure.compilation_result(); - const int num_constant_args = closure.num_constant_args(); - for (int i = 0; i < comp_result->xla_input_shapes.size(); i++) { - const auto& xla_shape = closure.compilation_result()->xla_input_shapes[i]; - if (!xla_shape.IsArray() || xla_shape.expressions().empty()) continue; - - for (int dim = 0; dim < xla_shape.expressions().size(); dim++) { - const auto& expr = xla_shape.expressions(dim); - if (expr && expr->is_dynamic()) { - xla::DExpr simplified_expr = expr.simplify(); - int input_idx = comp_result->input_mapping[i] - num_constant_args; - if (input_idx < 0 || input_idx >= ctx->num_inputs()) { - VLOG(1) << "Warning: Input index is out of range"; - continue; - } - VLOG(1) << "input shape is " << ctx->input(input_idx).shape() - << ", corresponding xla input shape is " << xla_shape; - int64_t size = ctx->input(input_idx).shape().dim_size(dim); - std::optional dyn_val = - simplified_expr->solve( - size); // TODO: check if the result is correct later. - if (dyn_val.has_value()) { - VLOG(1) << "Found dynamic input. Real size is: " << size - << ", solved dynamic value is " << *dyn_val; - } else { - xla::StringPrinter printer; - simplified_expr->print(&printer); - VLOG(1) << "Warning: Failed to solve the expression " - << std::move(printer).ToString(); - continue; - } - dyn_vals.insert(*dyn_val); - } - } + if (!batch_resolution.can_run) { + LOG(ERROR) << batch_resolution.diagnostic; + ctx->CtxFailure(errors::InvalidArgument(batch_resolution.diagnostic)); + return; } - - if (dyn_vals.size() == 1) { - run_options.set_batch_size(*(dyn_vals.begin())); + if (batch_resolution.has_batch_size) { + LOG(INFO) << "Setting run_options.batch_size " + << batch_resolution.diagnostic; + run_options.set_batch_size(batch_resolution.batch_size); is_set = true; } else { - // Found multiple variables - VLOG(1) << "Warning: Found multiple variables"; + LOG(INFO) << "Not setting run_options.batch_size because " + << batch_resolution.diagnostic; } - if (!is_set) { - // TODO: Fallback to BatchSizeResource for now. Remove it later. - BatchSizeResource* bsr = nullptr; - ScopedStepContainer* step_container = ctx->step_container(); - - absl::Status st = step_container->Lookup( - ctx->resource_manager(), BatchSizeResourceName, &bsr); - - if (st.ok()) { - run_options.set_batch_size(bsr->GetBatchSize()); - VLOG(1) << "run_options.batch_size is set to: " - << run_options.batch_size() << ". step_id: " << ctx->step_id(); - bsr->Unref(); - - } else if (IsNotFound(st)) { - VLOG(1) << "Warning: Not found BatchSizeResource in step_container."; - } else { - OP_REQUIRES_OK(ctx, st); - } + LOG(WARNING) << "Entering XLA cluster without run_options.batch_size " + << "being set. op=" << def().name() << " closure_key=" + << key << " step_id=" << ctx->step_id() + << " current_run_options_batch_size=" + << run_options.batch_size() + << " batch_size_resource_not_found=" + << batch_resolution.batch_size_resource_not_found; } } diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 566bab23a11867..1611f6dfd7298c 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -25,6 +25,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -267,10 +268,18 @@ class MarkForCompilationPassImpl { dim_vars_.insert(dim_vars.begin(), dim_vars.end()); } const std::set& dim_vars() const { return dim_vars_; } + void add_dim_expr(const xla::DExpr& dim_expr) { + dim_exprs_.push_back(dim_expr); + } + void merge_dim_exprs(const std::vector& dim_exprs) { + dim_exprs_.insert(dim_exprs_.end(), dim_exprs.begin(), dim_exprs.end()); + } + const std::vector& dim_exprs() const { return dim_exprs_; } private: int annotated_id_ = -1; std::set dim_vars_; + std::vector dim_exprs_; int chain_id_ = -1; int cluster_size_ = 1; int cycles_graph_node_id_; @@ -345,6 +354,10 @@ class MarkForCompilationPassImpl { absl::Status AssignAnnotatedClusterIDs(); absl::Status AssignDimVars(); + std::optional CheckDynamicExpressionCompatibility( + absl::Span exprs); + bool DynamicNodeExpressionsAreCompatible(const Node& node, + std::string* reason); void collectInputNodes(std::set &path_nodes); void collectMergeNodes(const std::vector& nodeSet, std::set &merger_nodes); @@ -645,6 +658,8 @@ void MarkForCompilationPassImpl::Cluster::Merge(Cluster* other) { merge_dim_vars(other->dim_vars_); other->dim_vars_.clear(); + merge_dim_exprs(other->dim_exprs_); + other->dim_exprs_.clear(); resource_var_operation_node_ids_.reserve( resource_var_operation_node_ids_.size() + @@ -783,6 +798,126 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) { return xla::DExpr(); } +std::string DExprToString(const xla::DExpr& expr) { + std::ostringstream oss; + oss << expr.get(); + return oss.str(); +} + +std::optional CheckDynamicExpressionCompatibilityImpl( + absl::Span exprs) { + const xla::DExpr* anchor_source = nullptr; + for (const xla::DExpr& expr : exprs) { + if (expr && expr->is_dynamic() && !expr->get_all_ids().empty()) { + anchor_source = &expr; + break; + } + } + if (anchor_source == nullptr) { + return std::nullopt; + } + + const std::set expected_ids = (*anchor_source)->get_all_ids(); + std::optional expected_core; + if (expected_ids.size() > 1) { + expected_core = + anchor_source->find_smallest_subexpression_covering_all_variables(); + } + + for (const xla::DExpr& expr : exprs) { + if (!expr || !expr->is_dynamic() || expr->get_all_ids().empty()) { + continue; + } + const std::set expr_ids = expr->get_all_ids(); + if (expr_ids != expected_ids) { + return absl::StrCat( + "dynamic expressions do not share the same variable ids: " + "expected={", + absl::StrJoin(expected_ids, ", "), "}, expr=", + DExprToString(expr), ", ids={", + absl::StrJoin(expr_ids, ", "), "}"); + } + if (!expected_core.has_value()) { + continue; + } + const xla::DExpr core = + expr.find_smallest_subexpression_covering_all_variables(); + if (!(core == *expected_core)) { + return absl::StrCat( + "dynamic expressions do not share the same clusterable core: " + "expected=", + DExprToString(*expected_core), ", expr=", DExprToString(expr), + ", core=", DExprToString(core)); + } + } + + return std::nullopt; +} + +std::optional +MarkForCompilationPassImpl::CheckDynamicExpressionCompatibility( + absl::Span exprs) { + return CheckDynamicExpressionCompatibilityImpl(exprs); +} + +bool MarkForCompilationPassImpl::DynamicNodeExpressionsAreCompatible( + const Node& node, std::string* reason) { + std::vector node_exprs; + + for (const Edge* edge : node.in_edges()) { + if (edge->IsControlEdge()) { + continue; + } + + const Node* src = edge->src(); + auto it = expr_map.find(src->name()); + if (it == expr_map.end()) { + continue; + } + + const int output_index = edge->src_output(); + if (output_index < 0 || + output_index >= static_cast(it->second.size())) { + continue; + } + + for (const auto& expr_ptr : it->second[output_index]) { + if (expr_ptr == nullptr) { + continue; + } + xla::DExpr dyn = DimExprToDExpr(expr_ptr.get()); + if (!dyn || !dyn->is_dynamic() || dyn->get_all_ids().empty()) { + continue; + } + node_exprs.push_back(std::move(dyn)); + } + } + + auto output_it = expr_map.find(node.name()); + if (output_it != expr_map.end()) { + for (const auto& output_exprs : output_it->second) { + for (const auto& expr_ptr : output_exprs) { + if (expr_ptr == nullptr) { + continue; + } + xla::DExpr dyn = DimExprToDExpr(expr_ptr.get()); + if (!dyn || !dyn->is_dynamic() || dyn->get_all_ids().empty()) { + continue; + } + node_exprs.push_back(std::move(dyn)); + } + } + } + + const std::optional incompatibility = + CheckDynamicExpressionCompatibility(node_exprs); + if (!incompatibility.has_value()) { + return true; + } + *reason = *incompatibility; + return false; +} + // Runs Grappler static inference and logs any ExpressionProto found in output // tensor shapes (from GraphProperties, not from _output_shapes attrs). void LogExpressionsViaGraphProperties(tensorflow::Graph& graph) { @@ -793,6 +928,8 @@ void LogExpressionsViaGraphProperties(tensorflow::Graph& graph) { using tensorflow::grappler::GraphProperties; using tensorflow::grappler::GrapplerItem; + expr_map.clear(); + GraphDef graph_def; graph.ToGraphDef(&graph_def); auto node_name_index = graph.BuildNodeNameIndex(); @@ -884,6 +1021,10 @@ absl::StatusOr MarkForCompilationPassImpl::Initialize() { TF_RET_CHECK(!initialized_ && !edges_contracted_ && !clusters_created_); initialized_ = true; + if (debug_options_.enable_dynamic_sizes) { + LogExpressionsViaGraphProperties(*graph_); + } + TF_RETURN_IF_ERROR(FindCompilationCandidates()); if (compilation_candidates_.empty()) { @@ -921,7 +1062,6 @@ absl::StatusOr MarkForCompilationPassImpl::Initialize() { TF_RETURN_IF_ERROR(AssignAnnotatedClusterIDs()); } if (debug_options_.enable_dynamic_sizes) { - LogExpressionsViaGraphProperties(*graph_); TF_RETURN_IF_ERROR(AssignDimVars()); auto has_dynamic_input_expression = [&](const Node* n) { for (const Edge* edge : n->in_edges()) { @@ -1726,6 +1866,15 @@ absl::Status MarkForCompilationPassImpl::FindCompilationCandidates() { continue; } + if (debug_options_.enable_dynamic_sizes) { + std::string reason; + if (!DynamicNodeExpressionsAreCompatible(*node, &reason)) { + VLOG(1) << "Rejecting " << node->name() + << " from XLA clustering: " << reason; + continue; + } + } + if (compile_time_const_nodes[node->id()]) { const OpDef* op_def; TF_RETURN_IF_ERROR( @@ -1910,6 +2059,10 @@ absl::Status MarkForCompilationPassImpl::AssignDimVars(void) { for (auto& pDim: (it->second)[output_index]) { DimExpr * d= pDim.get(); xla::DExpr dyn = DimExprToDExpr(d); + if (!dyn || !dyn->is_dynamic()) { + continue; + } + cluster->add_dim_expr(dyn); auto new_ids = dyn->get_all_ids(); for (auto id : new_ids) { cluster->add_dim_var(id); @@ -1917,6 +2070,21 @@ absl::Status MarkForCompilationPassImpl::AssignDimVars(void) { } } } + auto output_it = expr_map.find(node_name); + if (output_it != expr_map.end()) { + for (const auto& output_exprs : output_it->second) { + for (const auto& expr_ptr : output_exprs) { + if (expr_ptr == nullptr) { + continue; + } + xla::DExpr dyn = DimExprToDExpr(expr_ptr.get()); + if (!dyn || !dyn->is_dynamic()) { + continue; + } + cluster->add_dim_expr(dyn); + } + } + } // create a for loop for each dim vars in cluster and print each dim var if (VLOG_IS_ON(2)) { if (cluster->dim_vars().empty()) { @@ -2161,29 +2329,13 @@ absl::StatusOr MarkForCompilationPassImpl::TryToContractEdge( } if (debug_options_.enable_dynamic_sizes) { - if (from->dim_vars().size() > 1 || to->dim_vars().size() > 1) { - std::string from_str = "from_vars: "; - for (auto id : from->dim_vars()) { - from_str += std::to_string(id) + ", "; - } - std::string to_str = "to_vars: "; - for (auto id : to->dim_vars()) { - to_str += std::to_string(id) + ", "; - } - return LogNotContractableAndReturnFalse( - from, to, absl::StrCat("the two nodes have multiple dynamic dimensions: ", - from_str, " and ", to_str)); - } - if (from->dim_vars().size() == 1 && to->dim_vars().size() == 1 && - from->dim_vars() != to->dim_vars()) { - return LogNotContractableAndReturnFalse( - from, to, - absl::StrCat("the two nodes have different dynamic dimensions: ", - from->dim_vars().size() == 1 - ? std::to_string(*from->dim_vars().begin()) : "none", - " and ", - to->dim_vars().size() == 1 - ? std::to_string(*to->dim_vars().begin()) : "none")); + std::vector combined_exprs = from->dim_exprs(); + combined_exprs.insert(combined_exprs.end(), to->dim_exprs().begin(), + to->dim_exprs().end()); + std::optional incompatibility = + CheckDynamicExpressionCompatibility(combined_exprs); + if (incompatibility.has_value()) { + return LogNotContractableAndReturnFalse(from, to, *incompatibility); } } @@ -2679,6 +2831,11 @@ void ResetClusterSequenceNumber() { ClusterSequenceNumberGenerator::Global().Reset(); } +std::optional CheckDynamicExpressionCompatibilityForTest( + absl::Span exprs) { + return CheckDynamicExpressionCompatibilityImpl(exprs); +} + absl::flat_hash_set GetKnownXLAAllowlistOp() { absl::flat_hash_set result{ "AdjustContrastv2", diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.h b/tensorflow/compiler/jit/mark_for_compilation_pass.h index 558912f2eee2e0..79e3c88074fcdf 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.h +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.h @@ -20,10 +20,17 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_ #define TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_ +#include + #include "absl/container/flat_hash_set.h" +#include "absl/types/span.h" #include "tensorflow/compiler/jit/compilability_check_util.h" #include "tensorflow/core/common_runtime/optimization_registry.h" +namespace xla { +class DExpr; +} + namespace tensorflow { // The attribute that marks nodes to be grouped into functions by the @@ -57,6 +64,11 @@ void ResetClusterSequenceNumber(); // Return a list of operation that we choose not to put into the allowlist. absl::flat_hash_set GetKnownXLAAllowlistOp(); + +// Returns an explanation when dynamic expressions cannot use the same +// cluster-level dynamic value. +std::optional CheckDynamicExpressionCompatibilityForTest( + absl::Span exprs); } // namespace testing } // namespace tensorflow diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index 1a120791206369..f0a0058dbd16b4 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -53,6 +53,7 @@ limitations under the License. #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/test.h" +#include "xla/shape_expr.h" using ::tensorflow::testing::FindNodeByName; @@ -117,6 +118,52 @@ absl::flat_hash_map> GetClusterSets( return cluster_sets; } +// Expressions may differ outside the smallest subtree containing all variables. +TEST(XlaCompilationTest, + DynamicExpressionCompatibilityUsesSmallestCoveringSubexpression) { + const xla::DExpr shared_core = + xla::DExpr::Var(1) + xla::DExpr::Var(2); + const std::vector compatible_exprs = { + shared_core + 2, shared_core + 3}; + + EXPECT_FALSE( + testing::CheckDynamicExpressionCompatibilityForTest(compatible_exprs) + .has_value()); + + const std::vector incompatible_exprs = { + shared_core + 2, + (xla::DExpr::Var(1) - xla::DExpr::Var(2)) + 3}; + const std::optional incompatibility = + testing::CheckDynamicExpressionCompatibilityForTest(incompatible_exprs); + + ASSERT_TRUE(incompatibility.has_value()); + EXPECT_NE(incompatibility->find("same clusterable core"), + std::string::npos); +} + +TEST(XlaCompilationTest, + DynamicExpressionCompatibilityRejectsDifferentVariableSets) { + const std::vector exprs = { + xla::DExpr::Var(1) + xla::DExpr::Var(2), + xla::DExpr::Var(1) + xla::DExpr::Var(3)}; + + const std::optional incompatibility = + testing::CheckDynamicExpressionCompatibilityForTest(exprs); + + ASSERT_TRUE(incompatibility.has_value()); + EXPECT_NE(incompatibility->find("same variable ids"), std::string::npos); +} + +TEST(XlaCompilationTest, + DynamicExpressionCompatibilityAcceptsSingleVariableExpressions) { + const std::vector exprs = { + xla::DExpr::Var(1) + 2, + xla::DExpr::Var(1) * xla::DExpr::Const(3)}; + + EXPECT_FALSE( + testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value()); +} + TEST(XlaCompilationTest, Chains) { std::unique_ptr graph(new Graph(OpRegistry::Global())); { diff --git a/tensorflow/compiler/jit/xla_batch_matcher.cc b/tensorflow/compiler/jit/xla_batch_matcher.cc index c6f2f50fecf3e1..854f4d45c4214c 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.cc +++ b/tensorflow/compiler/jit/xla_batch_matcher.cc @@ -1,8 +1,105 @@ #include "tensorflow/compiler/jit/xla_batch_matcher.h" + +#include +#include + #include "xla/debug_options_flags.h" #include "tensorflow/core/platform/logging.h" namespace tensorflow { +namespace { + +bool MergeAlignment(int64_t factor, int64_t* alignment) { + factor = std::abs(factor); + if (factor <= 1) { + return true; + } + const int64_t common = std::gcd(*alignment, factor); + const int64_t multiplier = factor / common; + if (*alignment > kMaxBatch / multiplier) { + return false; + } + *alignment *= multiplier; + return true; +} + +bool CollectDivisorAlignment(xla::DynExpr* expr, int64_t* alignment) { + if (expr == nullptr) { + return true; + } + + xla::DynExpr* lhs = nullptr; + xla::DynExpr* rhs = nullptr; + switch (expr->kind()) { + case xla::DExpr::Kind::kUnknown: + case xla::DExpr::Kind::kConstant: + case xla::DExpr::Kind::kVariable: + return true; + case xla::DExpr::Kind::kAdd: { + auto* add = static_cast(expr); + lhs = add->get_lhs(); + rhs = add->get_rhs(); + break; + } + case xla::DExpr::Kind::kSub: { + auto* sub = static_cast(expr); + lhs = sub->get_lhs(); + rhs = sub->get_rhs(); + break; + } + case xla::DExpr::Kind::kMul: { + auto* mul = static_cast(expr); + lhs = mul->get_lhs(); + rhs = mul->get_rhs(); + break; + } + case xla::DExpr::Kind::kDiv: { + auto* div = static_cast(expr); + lhs = div->get_lhs(); + rhs = div->get_rhs(); + if (lhs->is_dynamic() && rhs->is_constant() && + !MergeAlignment(rhs->get_val(), alignment)) { + return false; + } + break; + } + } + + return CollectDivisorAlignment(lhs, alignment) && + CollectDivisorAlignment(rhs, alignment); +} + +} // namespace + +int64_t GetDynamicPaddingAlignment( + absl::Span exact_shape_expressions) { + int64_t alignment = 1; + for (const xla::DExpr& expression : exact_shape_expressions) { + if (expression && + !CollectDivisorAlignment(expression.get(), &alignment)) { + return 0; + } + } + return alignment; +} + +int64_t GetAlignedPowerOfTwoBatch(int64_t real_batch, int64_t alignment) { + if (real_batch <= 0 || real_batch > kMaxBatch || alignment <= 0 || + alignment > kMaxBatch) { + return real_batch; + } + + const int64_t units = (real_batch + alignment - 1) / alignment; + int64_t power = 1; + while (power < units && power <= kMaxBatch / 2) { + power <<= 1; + } + if (power < units || power > kMaxBatch / alignment) { + const int64_t rounded = units * alignment; + return rounded <= kMaxBatch ? rounded : real_batch; + } + return power * alignment; +} XlaBatchMatcher::XlaBatchMatcher() { env_str_ = xla::GetDebugOptionsFromFlags().xla_compile_batch_sizes(); @@ -121,55 +218,48 @@ void XlaBatchMatcher::parse_env_config() { return; } -// Calculate the smallest power of two greater than the real batch -static int64_t GetNextPowerOfTwo(int64_t real_batch) { - // If real_batch is already a power of two, return real_batch directly - if ((real_batch & (real_batch - 1)) == 0) { - return real_batch; - } - - int64_t power = 1; - while (power < real_batch) { - power <<= 1; - } - return power; -} - -int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) { +int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch, + int64_t alignment) { if (real_batch <= 0 || real_batch > kMaxBatch) { LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch; return real_batch; } + if (alignment <= 0) { + LOG(WARNING) << "[XLA_BATCH_WARN] Dynamic padding alignment cannot be " + "represented; using the exact runtime batch " + << real_batch; + return real_batch; + } if (all_batches_.empty()) { - // Return the next power of two directly without modifying all_batches_ - return GetNextPowerOfTwo(real_batch); + return GetAlignedPowerOfTwoBatch(real_batch, alignment); } - // Edge case 1: Real value < the smallest batch, use smallest - if (real_batch < all_batches_.front()) { - return all_batches_.front(); + auto it = std::lower_bound(all_batches_.begin(), all_batches_.end(), + real_batch); + for (; it != all_batches_.end(); ++it) { + if (alignment <= 1 || *it % alignment == 0) { + return *it; + } } - // Edge case 2: Real value ≥ the largest batch, use the nearest power of two - if (real_batch > all_batches_.back()) { - int64_t val = GetNextPowerOfTwo(real_batch); + + int64_t val = GetAlignedPowerOfTwoBatch(real_batch, alignment); + if (alignment <= 1) { all_batches_.emplace_back(val); print_all_batches(); - return val; } - - // Find first batch larger than real value (binary search via lower_bound) - auto it = std::lower_bound(all_batches_.begin(), all_batches_.end(), real_batch); - return (it != all_batches_.end()) ? *it : all_batches_.back(); + return val; } -int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) { +int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch, + int64_t alignment) { // Match target batch size - int64_t selected = find_min_larger_batch(real_batch); + int64_t selected = find_min_larger_batch(real_batch, alignment); if (real_batch != last_batch_ || all_batches_.empty()) { last_batch_ = real_batch; VLOG(2) << "[XLA_BATCH_INFO] Real batch: " << real_batch - << " -> Selected compile batch: " << selected; + << " -> Selected compile batch: " << selected + << " with alignment: " << alignment; } return selected; } diff --git a/tensorflow/compiler/jit/xla_batch_matcher.h b/tensorflow/compiler/jit/xla_batch_matcher.h index 7a294f5430221a..f9968e09cfd099 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.h +++ b/tensorflow/compiler/jit/xla_batch_matcher.h @@ -11,24 +11,36 @@ #include #include +#include "absl/types/span.h" +#include "xla/shape_expr.h" + namespace tensorflow { // Define the maximum allowed batch size: 2147483648 >> 1 = 1073741824 // Exceeding this value will make the next power of two exceed the safe range constexpr int kMaxBatch = 2147483648ULL >> 1; +// Returns the alignment required to preserve the integer divisions in exact +// dynamic shape expressions. Returns zero if no valid aligned bucket can be +// represented. +int64_t GetDynamicPaddingAlignment( + absl::Span exact_shape_expressions); + +// Selects a power-of-two bucket in units of `alignment`. +int64_t GetAlignedPowerOfTwoBatch(int64_t real_batch, int64_t alignment); + class XlaBatchMatcher { public: XlaBatchMatcher(); virtual ~XlaBatchMatcher() = default; - int64_t get_xla_compile_batch(int64_t real_batch); + int64_t get_xla_compile_batch(int64_t real_batch, int64_t alignment = 1); std::vector get_all_batches() { return all_batches_; } private: void parse_env_config(); void print_all_batches(); std::vector parse_single_item(const std::string& item); - int64_t find_min_larger_batch(int64_t real_batch); + int64_t find_min_larger_batch(int64_t real_batch, int64_t alignment); std::vector all_batches_; std::string env_str_; diff --git a/tensorflow/compiler/jit/xla_batch_matcher_test.cc b/tensorflow/compiler/jit/xla_batch_matcher_test.cc new file mode 100644 index 00000000000000..eed95858ee35df --- /dev/null +++ b/tensorflow/compiler/jit/xla_batch_matcher_test.cc @@ -0,0 +1,53 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/jit/xla_batch_matcher.h" + +#include +#include + +#include "tensorflow/core/platform/test.h" +#include "xla/shape_expr.h" + +namespace tensorflow { +namespace { + +TEST(XlaBatchMatcherTest, AlignsBucketForDynamicReshapeDivision) { + std::vector reshape_expressions = { + xla::DExpr::Var(1) / 12}; + + const int64_t alignment = + GetDynamicPaddingAlignment(reshape_expressions); + + EXPECT_EQ(GetAlignedPowerOfTwoBatch(300, /*alignment=*/1), 512); + EXPECT_EQ(alignment, 12); + EXPECT_EQ(GetAlignedPowerOfTwoBatch(300, alignment), 384); + EXPECT_EQ(GetAlignedPowerOfTwoBatch(384, alignment), 384); + EXPECT_EQ(GetAlignedPowerOfTwoBatch(385, alignment), 768); +} + +TEST(XlaBatchMatcherTest, CombinesIndependentReshapeDivisors) { + std::vector reshape_expressions = { + xla::DExpr::Var(1) / 12, xla::DExpr::Var(1) / 18}; + + const int64_t alignment = + GetDynamicPaddingAlignment(reshape_expressions); + + EXPECT_EQ(alignment, 36); + EXPECT_EQ(GetAlignedPowerOfTwoBatch(300, alignment), 576); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index ea2a50e241ebe2..e883ace189c30c 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -460,8 +460,30 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs( has_dynamic = true; VLOG(1) << "Current expression is " << expr; if (run_options) { - xla::DExpr batch_size = xla::DExpr::Const(run_options->batch_size()); - xla::DExpr subst_expr = expr.substitute(1, batch_size).simplify(); + const int64_t run_options_batch_size = run_options->batch_size(); + LOG(INFO) << "PopulateOutputs read run_options->batch_size()=" + << run_options_batch_size << " for output " << i + << " dimension " << dim; + xla::DExpr batch_size = xla::DExpr::Const(run_options_batch_size); + const std::set ids = expr->get_all_ids(); + if (ids.size() != 1) { + return absl::InvalidArgumentError(absl::StrCat( + "Runtime shape substitution expected exactly one dynamic " + "variable for output ", + i, ", dimension ", dim, ", but found ", ids.size(), + " variables in expression ", DExprToString(expr))); + } + const int substitute_var_id = *ids.begin(); + LOG(INFO) << "Calling output shape substitute with run_options for " + << "output " << i << " dimension " << dim + << " expr=" << DExprToString(expr) + << " substitute Var(" << substitute_var_id << ")=" + << DExprToString(batch_size); + xla::DExpr subst_expr = + expr.substitute(substitute_var_id, batch_size).simplify(); + LOG(INFO) << "Output shape substitute with run_options for output " + << i << " dimension " << dim << " returned " + << DExprToString(subst_expr); if (!subst_expr->is_constant()) { return absl::InvalidArgumentError(absl::StrCat( "Runtime shape substitution did not produce an integer " @@ -471,14 +493,37 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs( shape.set_dim(dim, subst_expr->get_val()); } else { // TODO: Fallback to BatchSizeResource for now. Remove it later. - VLOG(1) << "Warning: Didn't find run_options"; + LOG(INFO) << "PopulateOutputs did not receive run_options for " + << "output " << i << " dimension " << dim; BatchSizeResource* bsr = nullptr; ScopedStepContainer* step_container = ctx->step_container(); TF_RETURN_IF_ERROR(step_container->Lookup( ctx->resource_manager(), BatchSizeResourceName, &bsr)); + if (bsr == nullptr) { + return errors::Internal( + "BatchSizeResource lookup succeeded but returned null"); + } xla::DExpr batch_size = xla::DExpr::Const(bsr->GetBatchSize()); - // Just substitute Var(1) for now. - xla::DExpr subst_expr = expr.substitute(1, batch_size).simplify(); + const std::set ids = expr->get_all_ids(); + if (ids.size() != 1) { + return absl::InvalidArgumentError(absl::StrCat( + "Runtime shape substitution expected exactly one dynamic " + "variable for output ", + i, ", dimension ", dim, ", but found ", ids.size(), + " variables in expression ", DExprToString(expr))); + } + const int substitute_var_id = *ids.begin(); + LOG(INFO) << "Calling output shape substitute with " + << "BatchSizeResource for output " << i + << " dimension " << dim + << " expr=" << DExprToString(expr) + << " substitute Var(" << substitute_var_id + << ")=" << DExprToString(batch_size); + xla::DExpr subst_expr = + expr.substitute(substitute_var_id, batch_size).simplify(); + LOG(INFO) << "Output shape substitute with BatchSizeResource for " + << "output " << i << " dimension " << dim + << " returned " << DExprToString(subst_expr); if (!subst_expr->is_constant()) { return absl::InvalidArgumentError(absl::StrCat( "Runtime shape substitution did not produce an integer " diff --git a/tensorflow/compiler/tf2xla/xla_argument.h b/tensorflow/compiler/tf2xla/xla_argument.h index f037cb8610619e..0bdb7923edd4ec 100644 --- a/tensorflow/compiler/tf2xla/xla_argument.h +++ b/tensorflow/compiler/tf2xla/xla_argument.h @@ -83,6 +83,11 @@ struct XlaArgument { // This is only used for shape-like integer tensors crossing cluster // boundaries. std::vector constant_value_expressions; + + // Per-dimension provenance from TF shape inference. A marked dimension may + // legally be a runtime singleton broadcast to the shared dynamic dimension. + std::vector may_be_broadcast_singleton_dimensions; + // The upper bounds of the value. std::optional value_bound; diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 219c51175b5743..d21f4fadd9fe07 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -1548,6 +1548,18 @@ absl::Status XlaCompiler::CompileGraph( *graph, real_args, options.use_tuple_arg, builder.get(), context, arg_shardings, &arg_expressions, &result->input_mapping, &result->xla_input_shapes, options.is_entry_computation)); + result->xla_input_may_be_broadcast_singleton_dimensions.clear(); + result->xla_input_may_be_broadcast_singleton_dimensions.reserve( + result->input_mapping.size()); + for (int arg_index : result->input_mapping) { + if (arg_index >= 0 && + arg_index < static_cast(real_args.size())) { + result->xla_input_may_be_broadcast_singleton_dimensions.push_back( + real_args[arg_index].may_be_broadcast_singleton_dimensions); + } else { + result->xla_input_may_be_broadcast_singleton_dimensions.emplace_back(); + } + } context->set_args(std::move(arg_expressions)); PushNodeTokenMapping(); diff --git a/tensorflow/compiler/tf2xla/xla_helpers.h b/tensorflow/compiler/tf2xla/xla_helpers.h index 38f01c83db8251..78b8653639d76e 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.h +++ b/tensorflow/compiler/tf2xla/xla_helpers.h @@ -150,6 +150,11 @@ struct XlaCompilationResult { // the flattened shapes. std::vector xla_input_shapes; + // Per-dimension broadcast-singleton provenance aligned with + // xla_input_shapes. + std::vector> + xla_input_may_be_broadcast_singleton_dimensions; + // Output shape in XLA format. The output shape is always a tuple. If we // are flattening outputs, these are the flattened shapes. xla::Shape xla_output_shape; diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc index 42cf77245b8e1b..1c16ea849d0797 100644 --- a/tensorflow/core/framework/common_shape_fns.cc +++ b/tensorflow/core/framework/common_shape_fns.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" #include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/logging.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/lib/core/errors.h" @@ -2208,6 +2209,11 @@ absl::Status ReductionShape(InferenceContext* c) { absl::Status ConcatShapeHelper(InferenceContext* c, int start_value_index, int end_value_index, int dim_index) { + auto dim_expr_string = [&](DimensionHandle dim) -> string { + DimExpr* expr = c->GetDimExpr(dim); + return expr == nullptr ? "" : expr->DebugString(); + }; + ShapeHandle unused; TF_RETURN_IF_ERROR(c->WithRank(c->input(dim_index), 0, &unused)); const Tensor* concat_dim_t = c->input_tensor(dim_index); @@ -2355,12 +2361,14 @@ absl::Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c, *out = c->UnknownShape(); return absl::OkStatus(); } + c->MarkMayBeBroadcastSingleton(dim_y); dims.push_back(dim_x); } else if (c->Value(dim_y) > 1) { if (!incompatible_shape_error) { *out = c->UnknownShape(); return absl::OkStatus(); } + c->MarkMayBeBroadcastSingleton(dim_x); dims.push_back(dim_y); } else if (c->Value(dim_x) == 1) { dims.push_back(dim_y); @@ -2370,7 +2378,7 @@ absl::Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c, dims.push_back(dim_x); } else if (!c->ValueKnown(dim_x) && !c->ValueKnown(dim_y)) { DimensionHandle merged; - absl::Status s = c->Merge(dim_x, dim_y, &merged); + absl::Status s = c->MergeForBroadcast(dim_x, dim_y, &merged); if (s.ok()) { dims.push_back(merged); } else { diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc index b1336bf4398844..781b76ba051b4c 100644 --- a/tensorflow/core/framework/shape_inference.cc +++ b/tensorflow/core/framework/shape_inference.cc @@ -327,7 +327,7 @@ string InferenceContext::DebugString(ShapeHandle s) { } string InferenceContext::DebugString(DimensionHandle d) { - return ValueKnown(d) ? strings::StrCat(Value(d), strings::StrCat("~",DynamicRatio(d))) : "?"; + return ValueKnown(d) ? strings::StrCat(Value(d)) : "?"; } string InferenceContext::DebugString() const { @@ -468,6 +468,34 @@ absl::Status InferenceContext::Merge(DimensionHandle d0, DimensionHandle d1, } } +absl::Status InferenceContext::MergeForBroadcast(DimensionHandle d0, + DimensionHandle d1, + DimensionHandle* out) { + if (d0.SameHandle(d1) || ValueKnown(d0) || ValueKnown(d1)) { + return Merge(d0, d1, out); + } + + // Neither unknown operand is necessarily equal to the output: either may be + // the runtime singleton while the other determines the output dimension. + DimExpr* output_expr = GetDimExpr(d0); + if (output_expr == nullptr) { + output_expr = GetDimExpr(d1); + } + *out = shape_manager_.MakeDim(kUnknownDim, /*dynamic_ratio=*/0, output_expr); + broadcast_merged_dims_.emplace_back(d0, *out); + broadcast_merged_dims_.emplace_back(d1, *out); + MarkMayBeBroadcastSingleton(d0); + MarkMayBeBroadcastSingleton(d1); + return absl::OkStatus(); +} + +void InferenceContext::MarkMayBeBroadcastSingleton(DimensionHandle dim) { + if (!dim.IsSet() || ValueKnown(dim)) { + return; + } + may_be_broadcast_singleton_dims_.push_back(dim); +} + absl::Status InferenceContext::MergePrefix(ShapeHandle s, ShapeHandle prefix, ShapeHandle* s_out, ShapeHandle* prefix_out) { diff --git a/tensorflow/core/framework/shape_inference.h b/tensorflow/core/framework/shape_inference.h index 70d09d8fde5e93..e43a1a47f10df2 100644 --- a/tensorflow/core/framework/shape_inference.h +++ b/tensorflow/core/framework/shape_inference.h @@ -505,6 +505,14 @@ class InferenceContext { absl::Status Merge(DimensionHandle d0, DimensionHandle d1, DimensionHandle* out); + // Records broadcast compatibility without treating the input dimensions as + // equal. For two distinct unknown dimensions, the result is a fresh unknown + // dimension and either input may legally be a singleton at runtime. + absl::Status MergeForBroadcast(DimensionHandle d0, DimensionHandle d1, + DimensionHandle* out); + + void MarkMayBeBroadcastSingleton(DimensionHandle dim); + // Returns in <*out> a sub-shape of with dimensions [start:]. // can be negative to index from the end of the shape. If > // rank of , then an empty subshape is returned. @@ -761,6 +769,13 @@ class InferenceContext { const { return merged_dims_; } + const std::vector>& + BroadcastMergedDims() const { + return broadcast_merged_dims_; + } + const std::vector& MayBeBroadcastSingletonDims() const { + return may_be_broadcast_singleton_dims_; + } // Adds new outputs; useful when mutating the graph. absl::Status ExpandOutputs(int new_output_size); @@ -856,6 +871,8 @@ class InferenceContext { void ForgetMerges() { merged_shapes_.clear(); merged_dims_.clear(); + broadcast_merged_dims_.clear(); + may_be_broadcast_singleton_dims_.clear(); } // Helper method for MakeShapeFromTensor and MakeShapeFromShapeTensor. @@ -909,6 +926,9 @@ class InferenceContext { // known shapes or dims here. std::vector> merged_shapes_; std::vector> merged_dims_; + std::vector> + broadcast_merged_dims_; + std::vector may_be_broadcast_singleton_dims_; InferenceContext(const InferenceContext&) = delete; void operator=(const InferenceContext&) = delete; diff --git a/tensorflow/core/framework/tensor_shape.proto b/tensorflow/core/framework/tensor_shape.proto index f69b4228a7fb31..bac1e97c27a5f3 100644 --- a/tensorflow/core/framework/tensor_shape.proto +++ b/tensorflow/core/framework/tensor_shape.proto @@ -26,6 +26,10 @@ message TensorShapeProto { //Symbolic expression for this dimension when size == -1. //Allows tracking relationships between unknown dimensions. ExpressionProto expr = 3; + + // This occurrence may be a singleton that is broadcast to another + // dimension. It is not an equality constraint on the symbolic expression. + bool may_be_broadcast_singleton = 4; }; // Dimensions of the tensor, such as {"input", 30}, {"output", 40} @@ -81,4 +85,4 @@ message MulNode { message DivNode { ExpressionProto lhs = 1; ExpressionProto rhs = 2; -} \ No newline at end of file +} diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc index 774c628ce07c34..04144ef2441d53 100644 --- a/tensorflow/core/grappler/costs/graph_properties.cc +++ b/tensorflow/core/grappler/costs/graph_properties.cc @@ -2396,9 +2396,8 @@ class SymbolicShapeManager { if (InferenceContext::Rank(s1) > 0 && InferenceContext::Rank(s2) > 0) { CHECK_EQ(InferenceContext::Rank(s1), InferenceContext::Rank(s2)); for (int i = 0; i < InferenceContext::Rank(s1); ++i) { - TF_RETURN_IF_ERROR( - MergeDimsWithExpr(InferenceContext::DimKnownRank(s1, i), - InferenceContext::DimKnownRank(s2, i))); + TF_RETURN_IF_ERROR(Merge(InferenceContext::DimKnownRank(s1, i), + InferenceContext::DimKnownRank(s2, i))); } } return absl::OkStatus(); @@ -2407,9 +2406,30 @@ class SymbolicShapeManager { if (!d1.IsSet() || !d2.IsSet()) { return absl::OkStatus(); } + TF_RETURN_IF_ERROR(strong_dims_.Merge(d1, d2)); return MergeDimsWithExpr(d1, d2); } + absl::Status MergeBroadcast(DimensionHandle d1, DimensionHandle d2) { + if (!d1.IsSet() || !d2.IsSet()) { + return absl::OkStatus(); + } + return MergeDimsWithExpr(d1, d2); + } + + void MarkMayBeBroadcastSingleton(DimensionHandle dim) { + if (dim.IsSet()) { + may_be_broadcast_singleton_dims_.push_back(dim); + } + } + + void FinalizeBroadcastSingletonComponents() { + may_be_broadcast_singleton_roots_.clear(); + for (DimensionHandle dim : may_be_broadcast_singleton_dims_) { + may_be_broadcast_singleton_roots_.insert(strong_dims_.RootId(dim)); + } + } + void AsTensorProperties(const ShapeHandle& shape, const DataType& type, OpInfo::TensorProperties* properties) { properties->set_dtype(type); @@ -2434,6 +2454,10 @@ class SymbolicShapeManager { expr->ToProto(out_dim->mutable_expr()); // TODO: Apply simplification? } + out_dim->set_may_be_broadcast_singleton( + may_be_broadcast_singleton_roots_.find( + strong_dims_.RootId(dim)) != + may_be_broadcast_singleton_roots_.end()); } } } @@ -2595,6 +2619,10 @@ class SymbolicShapeManager { // Map from union-find root pointer to the best expression for that set. absl::flat_hash_map dim_root_expr_; DisjointSet dims_; + DisjointSet strong_dims_; + std::vector + may_be_broadcast_singleton_dims_; + absl::flat_hash_set may_be_broadcast_singleton_roots_; }; // Checks whether there is any conflict in merged shapes and dims in @@ -3110,6 +3138,17 @@ absl::Status GraphProperties::InferStatically( break; } } + for (const auto& merged_dims : node_ctx->BroadcastMergedDims()) { + if (!shape_manager + ->MergeBroadcast(merged_dims.first, merged_dims.second) + .ok()) { + found_error = true; + break; + } + } + for (const auto& dim : node_ctx->MayBeBroadcastSingletonDims()) { + shape_manager->MarkMayBeBroadcastSingleton(dim); + } if (found_error) { // The shapes aren't consistent, we can't infer safely: discard all the // information discovered so far. @@ -3118,6 +3157,8 @@ absl::Status GraphProperties::InferStatically( } } + shape_manager->FinalizeBroadcastSingletonComponents(); + TF_RETURN_IF_ERROR(ValidateSymbolicShapeManager(item_.graph, refiner.get(), shape_manager.get())); diff --git a/third_party/xla/xla/shape_expr.cc b/third_party/xla/xla/shape_expr.cc index 21dad324e8d44f..35a3b5f5eda976 100644 --- a/third_party/xla/xla/shape_expr.cc +++ b/third_party/xla/xla/shape_expr.cc @@ -23,6 +23,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -39,6 +40,125 @@ Constant* AsConstant(DynExpr* expr) { : nullptr; } +struct CoveringSubexpression { + std::set ids; + int node_count = 1; + DynExpr* smallest = nullptr; + int smallest_node_count = 0; +}; + +CoveringSubexpression FindSmallestCoveringSubexpression( + DynExpr* expr, const std::set& target_ids) { + CHECK(expr != nullptr); + CoveringSubexpression result; + DynExpr* lhs = nullptr; + DynExpr* rhs = nullptr; + + switch (expr->kind()) { + case DExpr::Kind::kUnknown: + case DExpr::Kind::kConstant: + break; + case DExpr::Kind::kVariable: + result.ids.insert(static_cast(expr)->get_id()); + break; + case DExpr::Kind::kAdd: { + auto* add = static_cast(expr); + lhs = add->get_lhs(); + rhs = add->get_rhs(); + break; + } + case DExpr::Kind::kSub: { + auto* sub = static_cast(expr); + lhs = sub->get_lhs(); + rhs = sub->get_rhs(); + break; + } + case DExpr::Kind::kMul: { + auto* mul = static_cast(expr); + lhs = mul->get_lhs(); + rhs = mul->get_rhs(); + break; + } + case DExpr::Kind::kDiv: { + auto* div = static_cast(expr); + lhs = div->get_lhs(); + rhs = div->get_rhs(); + break; + } + } + + if (lhs != nullptr) { + CoveringSubexpression lhs_result = + FindSmallestCoveringSubexpression(lhs, target_ids); + CoveringSubexpression rhs_result = + FindSmallestCoveringSubexpression(rhs, target_ids); + result.node_count += lhs_result.node_count + rhs_result.node_count; + result.ids.insert(lhs_result.ids.begin(), lhs_result.ids.end()); + result.ids.insert(rhs_result.ids.begin(), rhs_result.ids.end()); + + if (lhs_result.smallest != nullptr) { + result.smallest = lhs_result.smallest; + result.smallest_node_count = lhs_result.smallest_node_count; + } + if (rhs_result.smallest != nullptr && + (result.smallest == nullptr || + rhs_result.smallest_node_count < result.smallest_node_count)) { + result.smallest = rhs_result.smallest; + result.smallest_node_count = rhs_result.smallest_node_count; + } + } + + if (result.ids == target_ids && + (result.smallest == nullptr || + result.node_count < result.smallest_node_count)) { + result.smallest = expr; + result.smallest_node_count = result.node_count; + } + return result; +} + +std::unique_ptr ReplaceSubexpression(DynExpr* expr, DynExpr* target, + DynExpr* replacement) { + CHECK(expr != nullptr); + CHECK(target != nullptr); + CHECK(replacement != nullptr); + if (DynExpr::equal(expr, target)) { + return replacement->clone(); + } + + switch (expr->kind()) { + case DExpr::Kind::kUnknown: + case DExpr::Kind::kConstant: + case DExpr::Kind::kVariable: + return expr->clone(); + case DExpr::Kind::kAdd: { + auto* add = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(add->get_lhs(), target, replacement).release(), + ReplaceSubexpression(add->get_rhs(), target, replacement).release()); + } + case DExpr::Kind::kSub: { + auto* sub = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(sub->get_lhs(), target, replacement).release(), + ReplaceSubexpression(sub->get_rhs(), target, replacement).release()); + } + case DExpr::Kind::kMul: { + auto* mul = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(mul->get_lhs(), target, replacement).release(), + ReplaceSubexpression(mul->get_rhs(), target, replacement).release()); + } + case DExpr::Kind::kDiv: { + auto* div = static_cast(expr); + return std::make_unique
( + ReplaceSubexpression(div->get_lhs(), target, replacement).release(), + ReplaceSubexpression(div->get_rhs(), target, replacement).release()); + } + } + return expr->clone(); +} + void NormalizeFraction(int64_t* numerator, int64_t* denominator) { CHECK(denominator != nullptr); CHECK(*denominator != 0); @@ -347,6 +467,25 @@ std::unique_ptr SimplifyCanonical(const DynExpr* expr) { } // namespace +DExpr DExpr::find_smallest_subexpression_covering_all_variables() const { + CHECK(expr_ != nullptr); + const std::set ids = expr_->get_all_ids(); + CHECK(!ids.empty()); + CoveringSubexpression result = + FindSmallestCoveringSubexpression(expr_.get(), ids); + CHECK(result.smallest != nullptr); + return DExpr::Adopt(result.smallest->clone().release()); +} + +DExpr DExpr::replace_subexpression(const DExpr& target, + const DExpr& replacement) const { + if (expr_ == nullptr) { + return DExpr(); + } + return DExpr(ReplaceSubexpression(expr_.get(), target.get(), + replacement.get())); +} + const DExpr& Shape::MissingExpression() { static const DExpr missing = DExpr::Unknown(kMissingExpressionSentinel); return missing; diff --git a/third_party/xla/xla/shape_expr.h b/third_party/xla/xla/shape_expr.h index feecb054bfbc92..2120cbea78a6b3 100644 --- a/third_party/xla/xla/shape_expr.h +++ b/third_party/xla/xla/shape_expr.h @@ -141,6 +141,12 @@ class DExpr { DExpr substitute(int id, const DExpr& value) const { return expr_ == nullptr ? DExpr() : Adopt(expr_->substitute(id, value.get())); } + // Returns the smallest subtree containing every variable in this expression. + // The expression must contain at least one variable. + DExpr find_smallest_subexpression_covering_all_variables() const; + // Replaces every subtree equivalent to `target` with `replacement`. + DExpr replace_subexpression(const DExpr& target, + const DExpr& replacement) const; template friend H AbslHashValue(H h, const DExpr& expr) { diff --git a/third_party/xla/xla/shape_test.cc b/third_party/xla/xla/shape_test.cc index b6e4bcd79c81bb..65d8a6a5025f9b 100644 --- a/third_party/xla/xla/shape_test.cc +++ b/third_party/xla/xla/shape_test.cc @@ -53,9 +53,10 @@ class ShapeTest : public ::testing::Test { const Shape nested_tuple_ = ShapeUtil::MakeTupleShape({tuple_, matrix_, token_}); const Shape dynamic_matrix_ = - ShapeUtil::MakeShape(S32, {5, 2}, {true, false}); + ShapeUtil::MakeShape(S32, {5, 2}, std::vector{true, false}, {}); const Shape unbounded_ = - ShapeUtil::MakeShape(F32, {Shape::kUnboundedSize, 784}, {true, false}); + ShapeUtil::MakeShape(F32, {Shape::kUnboundedSize, 784}, + std::vector{true, false}, {}); }; // Tests that if the dynamic_dimensions parameter empty in the Shape @@ -105,8 +106,8 @@ TEST_F(ShapeTest, ShapeToString) { } TEST_F(ShapeTest, DynamicShapeToString) { - Shape array_shape = - ShapeUtil::MakeShape(F32, {23, 44, 55}, {true, false, true}); + Shape array_shape = ShapeUtil::MakeShape( + F32, {23, 44, 55}, std::vector{true, false, true}, {}); EXPECT_EQ("f32[<=23,44,<=55]", array_shape.ToString()); array_shape.set_dynamic_dimension(2, false); @@ -125,6 +126,24 @@ TEST_F(ShapeTest, DExprSimplifyCombinesEqualFractions) { EXPECT_EQ("A", DExprToString(expr.simplify())); } +TEST_F(ShapeTest, DExprFindsSmallestSubexpressionCoveringAllVariables) { + const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2); + const DExpr expr = (shared_core + 2) * (shared_core + 3); + + EXPECT_EQ( + shared_core, + expr.find_smallest_subexpression_covering_all_variables()); +} + +TEST_F(ShapeTest, DExprReplacesEveryMatchingSubexpression) { + const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2); + const DExpr expr = (shared_core + 2) * (shared_core + 3); + const DExpr replacement = DExpr::Var(7); + + EXPECT_EQ((replacement + 2) * (replacement + 3), + expr.replace_subexpression(shared_core, replacement)); +} + TEST_F(ShapeTest, DeleteDimensions) { Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 3, 2, 7, 9}, {2, 0, 1, 4, 3});