From 81f3190c900783edba4136c2e5f5a6827096fd18 Mon Sep 17 00:00:00 2001 From: Guillermo-Callaghan Date: Wed, 29 Jul 2026 17:17:40 +0100 Subject: [PATCH 01/21] Use shared reciprocal for broadcast scalar division Rewrite eligible scalar broadcast divisions to multiply by one shared reciprocal when fast math permits it. Preserve the original divide shape and cover both the simplifier and CPU compiler pipeline. --- .../simplifiers/algebraic_simplifier.cc | 23 ++++++++ .../simplifiers/algebraic_simplifier.h | 10 ++++ .../simplifiers/algebraic_simplifier_test.cc | 55 +++++++++++++++++++ .../xla/xla/service/cpu/cpu_compiler.cc | 10 ++++ .../cpu/cpu_compiler_internals_test.cc | 42 ++++++++++++++ 5 files changed, 140 insertions(+) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 9fcdf7ce1eb7b7..6493c56fb93dee 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -2317,6 +2317,29 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { return absl::OkStatus(); } + // A / broadcast(B) => A * broadcast(1 / B). + // + // This changes floating-point rounding and the handling of exceptional + // values, so backends must explicitly opt in when reciprocal substitution + // is permitted. + if (options_.enable_divide_by_broadcast_reciprocal() && + ShapeUtil::ElementIsFloating(divide->shape()) && + Match(divide, + m::Divide( + m::Op(&a), + m::Broadcast(m::Op(&b).WithShape(m::Shape().IsScalar()))))) { + TF_ASSIGN_OR_RETURN( + HloInstruction * reciprocal, + MakeBinaryHlo(HloOpcode::kDivide, MakeScalarLike(b, 1), b)); + HloInstruction* reciprocal_broadcast = + divide->mutable_operand(1)->AddInstruction( + HloInstruction::CreateBroadcast(divide->shape(), reciprocal, {})); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary(divide->shape(), + HloOpcode::kMultiply, a, + reciprocal_broadcast)); + } + // A / Const => A * (1 / Const) // // (Backends can do this transformation, but generally only if the constant is diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h index 1a2cab861649e3..a0f44014e5d712 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h @@ -314,6 +314,15 @@ class AlgebraicSimplifierOptions { } bool enable_fast_math() const { return enable_fast_math_; } + // If true, replace elementwise division by a broadcast scalar with + // multiplication by a broadcast scalar reciprocal. + void set_enable_divide_by_broadcast_reciprocal(bool enable) { + enable_divide_by_broadcast_reciprocal_ = enable; + } + bool enable_divide_by_broadcast_reciprocal() const { + return enable_divide_by_broadcast_reciprocal_; + } + void set_enable_broadcast_degenerate_dimension( bool enable_broadcast_degenerate_dimension) { enable_broadcast_degenerate_dimension_ = @@ -388,6 +397,7 @@ class AlgebraicSimplifierOptions { bool use_convert_constant_folding_{false}; bool disable_dynamic_slice_to_slice_conversion_{false}; bool enable_fast_math_{false}; + bool enable_divide_by_broadcast_reciprocal_{false}; bool enable_broadcast_degenerate_dimension_{true}; bool enable_remove_no_op_reduce_precision_{false}; bool enable_onednn_support_{ diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index 7cff2113e4aceb..14991fb8dca9cf 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -2016,6 +2016,61 @@ TEST_F(AlgebraicSimplifierTest, DivideByBroadcastedConstant) { m::Broadcast(m::Op().IsConstantScalar(1.0f / 256.0f))))); } +TEST_F(AlgebraicSimplifierTest, + DivideByBroadcastedRuntimeScalarUsingReciprocal) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[16] parameter(0) + p1 = f32[] parameter(1) + b = f32[16] broadcast(p1), dimensions={} + ROOT d = f32[16] divide(p0, b) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + + ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).value()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), + m::Broadcast(m::Parameter(1))))); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_divide_by_broadcast_reciprocal(true); + ASSERT_TRUE(AlgebraicSimplifier(options).Run(m.get()).value()); + EXPECT_THAT( + m->entry_computation()->root_instruction(), + GmockMatch(m::Multiply( + m::Parameter(0), + m::Broadcast(m::Divide(m::ConstantScalar(1), m::Parameter(1)))))); +} + +TEST_F(AlgebraicSimplifierTest, + DivideByBroadcastedRuntimeScalarPreservesResultLayout) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[2,3]{1,0} parameter(0) + p1 = f32[] parameter(1) + b = f32[2,3]{1,0} broadcast(p1), dimensions={} + ROOT d = f32[2,3]{0,1} divide(p0, b) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_divide_by_broadcast_reciprocal(true); + ASSERT_TRUE(AlgebraicSimplifier(options).Run(m.get()).value()); + + HloInstruction* root = m->entry_computation()->root_instruction(); + EXPECT_THAT( + root, + GmockMatch(m::Multiply( + m::Parameter(0), + m::Broadcast(m::Divide(m::ConstantScalar(1), m::Parameter(1)))))); + EXPECT_TRUE(LayoutUtil::Equal(root->shape().layout(), + LayoutUtil::MakeLayout({0, 1}))); +} + // pow(pow(A, X), Y) => pow(A, X*Y) TEST_F(AlgebraicSimplifierTest, PowerOfPower) { auto m = CreateNewVerifiedModule(); diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index fcffdafcf5e5ee..bdf2d19d9492e8 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -465,6 +465,11 @@ std::unique_ptr> CreateSimplificationPipeline( !module->config().debug_options().xla_cpu_enable_fast_min_max()); options.set_supports_non_canonical_dots(false); options.set_executing_on_cpu(true); + options.set_enable_divide_by_broadcast_reciprocal( + module->config().debug_options().xla_cpu_enable_fast_math() && + !module->config() + .debug_options() + .xla_cpu_fast_math_honor_division()); pipeline->AddPass(options); pipeline->AddPass(); pipeline->AddPass(); @@ -885,6 +890,11 @@ absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn( options.set_minmax_propagate_nan( !module->config().debug_options().xla_cpu_enable_fast_min_max()); options.set_executing_on_cpu(true); + options.set_enable_divide_by_broadcast_reciprocal( + module->config().debug_options().xla_cpu_enable_fast_math() && + !module->config() + .debug_options() + .xla_cpu_fast_math_honor_division()); pipeline.AddPass(options); pipeline.AddPass(); pipeline.AddPass(/*is_layout_sensitive=*/true); diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc index cefc0274106501..11e9e35a511cd7 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc @@ -27,7 +27,10 @@ limitations under the License. #include "llvm/IR/Module.h" #include "llvm/Support/Casting.h" #include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_config.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/service/llvm_compiler.h" #include "xla/tests/hlo_test_base.h" @@ -164,6 +167,45 @@ TEST_F(CpuCompilerInternalsTest, JustOneDylibWithThunks) { << "one dylib is allowed."; } +TEST_F(CpuCompilerInternalsTest, SharesBroadcastScalarReciprocal) { + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(R"( + HloModule test + ENTRY main { + p0 = f32[16] parameter(0) + p1 = f32[16] parameter(1) + scalar = f32[] parameter(2) + broadcast = f32[16] broadcast(scalar), dimensions={} + divide0 = f32[16] divide(p0, broadcast) + divide1 = f32[16] divide(p1, broadcast) + ROOT result = f32[16] add(divide0, divide1) + } + )")); + DebugOptions& debug_options = + hlo_module->mutable_config().mutable_debug_options(); + debug_options.set_xla_cpu_enable_fast_math(true); + debug_options.set_xla_cpu_fast_math_honor_division(false); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr optimized_module, + GetOptimizedModule(std::move(hlo_module))); + + int64_t divide_count = 0; + int64_t scalar_divide_count = 0; + for (const HloComputation* computation : optimized_module->computations()) { + for (const HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() != HloOpcode::kDivide) { + continue; + } + ++divide_count; + if (instruction->shape().dimensions().empty()) { + ++scalar_divide_count; + } + } + } + EXPECT_EQ(divide_count, 1); + EXPECT_EQ(scalar_divide_count, 1); +} + } // namespace } // namespace cpu } // namespace xla From 4e5c9bad6e4cf34db8cefd4220aeab5f4e7bf46f Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:03:09 +0100 Subject: [PATCH 02/21] Add shared dynamic shape expression representation Use the same owned expression representation for TensorFlow TensorShape dimensions and XLA Shape dimensions. Add serialization, simplification, substitution, solving, and focused expression tests. --- tensorflow/compiler/jit/flags.cc | 7 +- tensorflow/compiler/jit/flags.h | 3 + .../compiler/tf2xla/xla_compiler_test.cc | 3 +- tensorflow/core/framework/BUILD | 37 + tensorflow/core/framework/tensor_shape.cc | 136 ++- tensorflow/core/framework/tensor_shape.h | 78 ++ tensorflow/core/framework/tensor_shape.proto | 57 ++ .../core/framework/tensor_shape_expr.cc | 214 +++++ tensorflow/core/framework/tensor_shape_expr.h | 57 ++ .../core/framework/tensor_shape_expr_test.cc | 105 +++ third_party/xla/xla/BUILD | 2 + ...optimize_input_output_buffer_alias_test.cc | 6 +- .../translate/mhlo_to_hlo/type_to_shape.cc | 3 +- third_party/xla/xla/shape.cc | 73 +- third_party/xla/xla/shape.h | 67 +- third_party/xla/xla/shape_expr.cc | 773 +++++++++++++++++ third_party/xla/xla/shape_expr.h | 802 ++++++++++++++++++ third_party/xla/xla/shape_util.cc | 148 +++- third_party/xla/xla/shape_util.h | 32 +- third_party/xla/xla/xla_data.proto | 52 ++ 20 files changed, 2614 insertions(+), 41 deletions(-) create mode 100644 tensorflow/core/framework/tensor_shape_expr.cc create mode 100644 tensorflow/core/framework/tensor_shape_expr.h create mode 100644 tensorflow/core/framework/tensor_shape_expr_test.cc create mode 100644 third_party/xla/xla/shape_expr.cc create mode 100644 third_party/xla/xla/shape_expr.h diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index 10756ddf9de7b5..d8e0b416b2d638 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -155,6 +155,9 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { &mark_for_compilation_flags->tf_xla_deterministic_cluster_names, "Causes the function names assigned by auto clustering to be " "deterministic from run to run."), + Flag("tf_xla_enable_dynamic_sizes", + &mark_for_compilation_flags->tf_xla_enable_dynamic_sizes, + "Enable dynamic sizes support."), Flag("tf_xla_persistent_cache_directory", &mark_for_compilation_flags->tf_xla_persistent_cache_directory, "If non-empty, JIT-compiled executables are saved to and loaded " @@ -241,6 +244,7 @@ void AllocateAndParseFlags() { mark_for_compilation_flags ->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_persistent_cache_directory = ""; mark_for_compilation_flags->tf_xla_persistent_cache_device_types = ""; mark_for_compilation_flags->tf_xla_persistent_cache_read_only = false; @@ -404,7 +408,8 @@ void AllocateAndParseFlags() { "and creates TPUReshardVariables ops.")}); AppendMarkForCompilationPassFlagsInternal(flag_list); - xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list); + xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list, + /*reset_envvar=*/true); mlir_flags = new MlirCommonFlags; if (!enable_mlir_bridge_is_explicit) { diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index 0d0c5082cf9a82..b4aba197559475 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -93,6 +93,9 @@ struct MarkForCompilationPassFlags { // so that they remain stable from run to run of auto clusteing. bool tf_xla_deterministic_cluster_names; + // If true enables support of dynamic sizes. + bool tf_xla_enable_dynamic_sizes; + // If non-empty, JIT-compiled executables are saved to and loaded from the // specified file system directory path. std::string tf_xla_persistent_cache_directory; diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 5aef5601af61ca..1e598502ad65d0 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -274,7 +274,8 @@ TEST_F(XlaCompilerTest, SimpleDynamicShapeParameter) { args[0].type = DT_INT32; args[0].shape = xla::ShapeUtil::MakeShape(/*element_type=*/xla::S32, /*dimensions=*/{2}, - /*dynamic_dimensions=*/{true}); + /*dynamic_dimensions=*/{true}, + /*expressions=*/{}); args[1].kind = XlaCompiler::Argument::kParameter; args[1].type = DT_INT32; args[1].shape = TensorShape(/*dimensions=*/{2}); diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD index 09142e303e3e13..8a2c2032f3ab25 100644 --- a/tensorflow/core/framework/BUILD +++ b/tensorflow/core/framework/BUILD @@ -340,6 +340,8 @@ filegroup( "tensor_key.h", "tensor_shape.cc", "tensor_shape.h", + "tensor_shape_expr.cc", + "tensor_shape_expr.h", "tensor_types.h", "tensor_util.h", "tracking_allocator.h", @@ -702,6 +704,7 @@ cc_library( ], deps = [ ":bounds_check", + ":tensor_shape_expr", ":tensor_shape_proto_cc", ":types_proto_cc", "//tensorflow/core/lib/core:errors", @@ -718,10 +721,43 @@ cc_library( "//tensorflow/core/platform:statusor", "//tensorflow/core/util:overflow", "@eigen_archive//:eigen3", + "@local_xla//xla:shape_util", ], alwayslink = 1, ) +cc_library( + name = "tensor_shape_expr", + srcs = ["tensor_shape_expr.cc"], + hdrs = ["tensor_shape_expr.h"], + visibility = [ + "//tensorflow/core:__pkg__", + "//tensorflow/core/grappler:__subpackages__", + "//tensorflow/core/runtime_fallback:__subpackages__", + "//tensorflow/core/tfrt/utils:__subpackages__", + ], + deps = [ + ":tensor_shape_proto_cc", + "//tensorflow/core:lib", + "@local_xla//xla:parse_flags_from_env", + "@local_xla//xla:shape_util", + "@local_xla//xla/tsl/util:command_line_flags", + ], + alwayslink = 1, +) + +tf_cc_test( + name = "tensor_shape_expr_test", + size = "small", + srcs = ["tensor_shape_expr_test.cc"], + deps = [ + ":tensor_shape", + ":tensor_shape_expr", + "//tensorflow/core/platform:test", + "//tensorflow/core:test_main", + ], +) + cc_library( name = "resource_base", hdrs = ["resource_base.h"], @@ -921,6 +957,7 @@ cc_library( ":node_def_util", ":op_def_proto_cc", ":tensor_shape", + ":tensor_shape_expr", ":tensor_shape_proto_cc", "//tensorflow/core/lib/core:errors", "//tensorflow/core/lib/core:status", diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc index 35c628216ed3c6..4b06e0861d86cf 100644 --- a/tensorflow/core/framework/tensor_shape.cc +++ b/tensorflow/core/framework/tensor_shape.cc @@ -17,15 +17,24 @@ limitations under the License. #include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/util/overflow.h" +#include "xla/printer.h" namespace tensorflow { +std::string ExprToString(const xla::DExpr& e) { + if (!e && !e.is_unknown()) return ""; + xla::StringPrinter printer; + e->print(&printer); + return std::move(printer).ToString(); +} + // TensorShape and PartialTensorShape should have no fields beyond // TensorShapeRep. In particular, their sizes should be the same. static_assert(sizeof(TensorShapeRep) == sizeof(TensorShape), @@ -152,6 +161,11 @@ TensorShapeBase::TensorShapeBase(const TensorShapeProto& proto) { for (const auto& d : proto.dim()) { AddDim(d.size()); } + if (TensorShapeExpressionsEnabled()) { + for (const auto& e : proto.expressions()) { + AddExpression(DimExprFromProto(e)); + } + } } } @@ -191,6 +205,11 @@ absl::Status TensorShapeBase::BuildTensorShapeBase( } } } + if (TensorShapeExpressionsEnabled()) { + for (const auto& e : proto.expressions()) { + out->AddExpression(DimExprFromProto(e)); + } + } } return absl::OkStatus(); } @@ -375,10 +394,46 @@ void TensorShapeRep::Clear() { set_data_type(DT_INVALID); } +void TensorShapeRep::set_expression(int d, xla::DExpr expr) { + if (!TensorShapeExpressionsEnabled()) { + expressions_.clear(); + return; + } + if (expressions_.size() <= static_cast(d)) { + expressions_.resize(d + 1, + xla::DExpr::Unknown(xla::kMissingExpressionSentinel)); + } + expressions_[d] = expr ? std::move(expr) + : xla::DExpr::Unknown(xla::kMissingExpressionSentinel); +} + +void TensorShapeRep::AddExpression(xla::DExpr expr) { + if (!TensorShapeExpressionsEnabled()) { + return; + } + CHECK_LT(expressions_.size(), ndims_byte()); + expressions_.push_back(expr ? std::move(expr) + : xla::DExpr::Unknown( + xla::kMissingExpressionSentinel)); +} + +void TensorShapeRep::set_expressions(std::vector exprs) { + if (!TensorShapeExpressionsEnabled()) { + expressions_.clear(); + return; + } + CHECK_LE(exprs.size(), ndims_byte()); + for (auto& expr : exprs) { + if (!expr) expr = xla::DExpr::Unknown(xla::kMissingExpressionSentinel); + } + expressions_ = std::move(exprs); +} + void TensorShapeRep::ClearAllButDataType() { if (tag() == REP_OUT_OF_LINE) { delete as64()->dims_; } + expressions_.clear(); set_tag(REP16); set_ndims_byte(0); // Leaves data_type alone @@ -505,6 +560,9 @@ void TensorShapeBase::UnsafeAddDim(int64_t size, template void TensorShapeBase::AppendShape(const TensorShapeBase& shape) { for (auto d : shape) AddDim(d.size); + for (auto e : shape.get_expressions()){ + AddExpression(e); + } } template @@ -585,6 +643,13 @@ template void TensorShapeBase::set_dim(int d, int64_t size) { CHECK_GE(d, 0); CHECK_LT(d, dims()); + // After DExpr migration, missing slots may be normalized to Unknown(). + // Preserve those placeholders here instead of materializing them into a + // concrete constant just because the dimension size changed. + if (get_expressions().size() > d && + get_expression(d).kind() != xla::DExpr::Kind::kUnknown) { + set_expression(d, xla::DExpr::Const(size)); + } if (!kIsPartial) { CHECK_GE(size, 0); } @@ -646,6 +711,13 @@ absl::Status TensorShapeBase::SetDimWithStatus(int d, int64_t size) { } } + // After DExpr migration, missing slots may be normalized to Unknown(). + // Preserve those placeholders here instead of materializing them into a + // concrete constant just because the dimension size changed. + if (get_expressions().size() > d && + get_expression(d).kind() != xla::DExpr::Kind::kUnknown) { + set_expression(d, xla::DExpr::Const(size)); + } return RecomputeNumElements(); } @@ -661,11 +733,31 @@ void TensorShapeBase::RemoveDimRange(int begin, int end) { if (begin >= end) return; absl::InlinedVector vals; AppendTo(*this, &vals); + std::vector new_exprs(get_expressions().begin(), + get_expressions().end()); + if (begin < static_cast(new_exprs.size())) { + int64_t expr_end = end; + if (expr_end > static_cast(new_exprs.size())) { + expr_end = new_exprs.size(); + } + if (expr_end > begin) { + new_exprs.erase(new_exprs.begin() + begin, new_exprs.begin() + expr_end); + } + } + vals.erase(vals.begin() + begin, vals.begin() + end); + + // Truncate if the removed dims reduce rank below expression vector size. + const int64_t new_rank = vals.size(); + if (new_exprs.size() > static_cast(new_rank)) { + new_exprs.resize(new_rank); + } + ClearAllButDataType(); for (auto dval : vals) { AddDim(dval); } + set_expressions(new_exprs); TF_CHECK_OK(RecomputeNumElements()); } @@ -700,7 +792,26 @@ absl::Status TensorShapeBase::RemoveDimRangeWithStatus(int begin, absl::InlinedVector vals; AppendTo(*this, &vals); + + std::vector new_exprs(get_expressions().begin(), + get_expressions().end()); + if (begin < static_cast(new_exprs.size())) { + int64_t expr_end = end; + if (expr_end > static_cast(new_exprs.size())) { + expr_end = new_exprs.size(); + } + if (expr_end > begin) { + new_exprs.erase(new_exprs.begin() + begin, new_exprs.begin() + expr_end); + } + } + vals.erase(vals.begin() + begin, vals.begin() + end); + + const int64_t new_rank = vals.size(); + if (new_exprs.size() > static_cast(new_rank)) { + new_exprs.resize(new_rank); + } + ClearAllButDataType(); absl::Status s = absl::OkStatus(); @@ -710,7 +821,7 @@ absl::Status TensorShapeBase::RemoveDimRangeWithStatus(int begin, return s; } } - + set_expressions(new_exprs); return RecomputeNumElements(); } @@ -731,6 +842,12 @@ void TensorShapeBase::AsProto(TensorShapeProto* proto) const { for (int i = 0; i < dims(); i++) { proto->add_dim()->set_size(dim_size(i)); } + if (TensorShapeExpressionsEnabled()) { + for (int i = 0; i < get_expressions().size(); i++) { + ExpressionProto* eproto = proto->add_expressions(); + DimExprToProto(get_expression(i), eproto); + } + } } } @@ -764,6 +881,11 @@ string TensorShapeRep::DebugString() const { } else { strings::StrAppend(&s, dim); } + if (shape.get_expression(i)) { + strings::StrAppend(&s, "<"); + strings::StrAppend(&s, ExprToString(shape.get_expression(i))); + strings::StrAppend(&s, ">"); + } } strings::StrAppend(&s, "]"); return s; @@ -787,6 +909,16 @@ string TensorShapeRep::DebugString(const TensorShapeProto& proto) { first = false; } strings::StrAppend(&s, "]"); + if (TensorShapeExpressionsEnabled()) { + strings::StrAppend(&s, "<"); + first = true; + for (const auto& e : proto.expressions()) { + if (!first) strings::StrAppend(&s, ","); + strings::StrAppend(&s, ExprToString(DimExprFromProto(e))); + first = false; + } + strings::StrAppend(&s, ">"); + } return s; } @@ -950,6 +1082,8 @@ absl::Status PartialTensorShape::MergeWith(const PartialTensorShape& shape, return s; } } + result->set_expressions(std::vector(shape.get_expressions().begin(), + shape.get_expressions().end())); return absl::OkStatus(); } diff --git a/tensorflow/core/framework/tensor_shape.h b/tensorflow/core/framework/tensor_shape.h index 0bcf1fc54af844..9a9e5159a6556f 100644 --- a/tensorflow/core/framework/tensor_shape.h +++ b/tensorflow/core/framework/tensor_shape.h @@ -17,6 +17,7 @@ limitations under the License. #define TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_H_ #include +#include #include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive #include "tensorflow/core/framework/types.pb.h" @@ -28,6 +29,7 @@ limitations under the License. #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/platform/statusor.h" +#include "xla/shape_expr.h" namespace tensorflow { @@ -73,7 +75,63 @@ class TensorShapeRep { std::string DebugString() const; static std::string DebugString(const TensorShapeProto& proto); + void set_expression(int d, xla::DExpr expr); + + void AddExpression(xla::DExpr expr); + + // Set the array of dynamic multipliers. + void set_expressions(std::vector exprs); + + // Get the array of dynamic multipliers. + absl::Span get_expressions() const { + return expressions_; + } + + // Get the array of dynamic multipliers, filling missing entries with + // constant expressions derived from the concrete dimensions. + std::vector get_filled_expressions() const { + if (ndims_byte() == kUnknownRank) { + return {}; + } + std::vector exprs(ndims_byte()); + for (int i = 0; i < ndims_byte(); ++i) { + exprs[i] = get_filled_expression(i); + } + return exprs; + } + + // Return the multiplier for a specific dynamic dimension. + // -1 if the dimension is not dynamic. + const xla::DExpr& get_expression(int64_t dimension) const { + static const xla::DExpr kMissingExpression = + xla::DExpr::Unknown(xla::kMissingExpressionSentinel); + if (dimension < 0) return kMissingExpression; + const size_t dim = static_cast(dimension); + if (dim >= expressions_.size()) { + return kMissingExpression; + } + return expressions_[dim]; + } + + xla::DExpr get_filled_expression(int64_t dimension) const { + if (dimension < 0) { + return xla::DExpr::Unknown(xla::kMissingExpressionSentinel); + } + const size_t dim = static_cast(dimension); + if (dim < expressions_.size() && expressions_[dim]) { + return expressions_[dim]; + } + + if (ndims_byte() == kUnknownRank || dim >= ndims_byte()) { + return xla::DExpr::Unknown(xla::kMissingExpressionSentinel); + } + + return constant_expression_for_dim(dim); + } + protected: + std::vector expressions_; + // Constructable only via TensorShapeBase TensorShapeRep() = default; @@ -143,6 +201,20 @@ class TensorShapeRep { void set_num_elements(int64_t n) { num_elements_ = n; } private: + xla::DExpr constant_expression_for_dim(size_t dim) const { + int64_t dim_value = -1; + if (tag() == REP16) { + uint16 raw_dim = as16()->dims_[dim]; + dim_value = raw_dim == kUnknownRep16 ? -1 : raw_dim; + } else if (tag() == REP32) { + uint32 raw_dim = as32()->dims_[dim]; + dim_value = raw_dim == kUnknownRep32 ? -1 : raw_dim; + } else { + dim_value = (*as64()->dims_)[dim]; + } + return xla::DExpr::Const(dim_value); + } + void DestructorOutOfLine(); void SlowCopyFrom(const TensorShapeRep& b); @@ -710,6 +782,7 @@ absl::Status TensorShape::AsEigenDSizesWithPaddingWithStatus( inline TensorShapeRep::TensorShapeRep(const TensorShapeRep& b) { num_elements_ = b.num_elements_; + expressions_ = b.expressions_; if (b.tag() != REP_OUT_OF_LINE) { memcpy(buf(), b.buf(), sizeof(u_.buf)); // memcpy above Implicitly does: @@ -723,6 +796,7 @@ inline TensorShapeRep::TensorShapeRep(const TensorShapeRep& b) { inline TensorShapeRep::TensorShapeRep(TensorShapeRep&& b) { num_elements_ = b.num_elements_; + expressions_ = b.expressions_; memcpy(buf(), b.buf(), sizeof(u_.buf)); // memcpy above Implicitly does: // set_ndims_byte(b.ndims_byte()); @@ -738,6 +812,8 @@ inline TensorShapeRep::~TensorShapeRep() { inline void TensorShapeRep::operator=(const TensorShapeRep& b) { num_elements_ = b.num_elements_; + expressions_ = b.expressions_; + if (tag() != REP_OUT_OF_LINE && b.tag() != REP_OUT_OF_LINE) { memcpy(buf(), b.buf(), sizeof(u_.buf)); // memcpy above implicitly also does: @@ -753,6 +829,8 @@ inline void TensorShapeRep::operator=(TensorShapeRep&& b) { DestructorOutOfLine(); } num_elements_ = b.num_elements_; + expressions_ = b.expressions_; + memcpy(buf(), b.buf(), sizeof(u_.buf)); // memcpy above Implicitly does: // set_ndims_byte(b.ndims_byte()); diff --git a/tensorflow/core/framework/tensor_shape.proto b/tensorflow/core/framework/tensor_shape.proto index 45d5b78ecbbc4c..efddff41dc64ba 100644 --- a/tensorflow/core/framework/tensor_shape.proto +++ b/tensorflow/core/framework/tensor_shape.proto @@ -22,6 +22,10 @@ message TensorShapeProto { // Optional name of the tensor dimension. string name = 2; + //Only keep one symbolic expr. + //Symbolic expression for this dimension when size == -1. + //Allows tracking relationships between unknown dimensions. + ExpressionProto expr = 3; }; // Dimensions of the tensor, such as {"input", 30}, {"output", 40} @@ -43,4 +47,57 @@ message TensorShapeProto { // // If true, "dim.size()" must be 0. bool unknown_rank = 3; + + repeated ExpressionProto expressions = 4; + }; + +message ExpressionProto { + oneof node_type { + int32 constant_value = 1; // cons + int32 variable_id = 2; // var + AddNode add_node = 3; // exp + exp + SubNode sub_node = 4; // exp - exp + MulNode mul_node = 5; // exp * exp + DivNode div_node = 6; // exp / exp + MaxNode max_node = 7; // max(exp, exp) + GtNode gt_node = 8; // exp > exp + SelectNode select_node = 9; // select(pred, on_true, on_false) + } +} + +message AddNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message SubNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message MulNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message DivNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message MaxNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message GtNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message SelectNode { + ExpressionProto pred = 1; + ExpressionProto on_true = 2; + ExpressionProto on_false = 3; +} diff --git a/tensorflow/core/framework/tensor_shape_expr.cc b/tensorflow/core/framework/tensor_shape_expr.cc new file mode 100644 index 00000000000000..8406c550477c1b --- /dev/null +++ b/tensorflow/core/framework/tensor_shape_expr.cc @@ -0,0 +1,214 @@ +/* Copyright 2026 The TensorFlow Authors. + +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/core/framework/tensor_shape_expr.h" + +#include +#include +#include + +#include "xla/parse_flags_from_env.h" +#include "xla/printer.h" +#include "xla/tsl/util/command_line_flags.h" + +namespace tensorflow { + +namespace { + +bool ParseTensorShapeExpressionsEnabled() { + bool tf_xla_enable_dynamic_sizes = false; + std::vector flag_list = { + tsl::Flag("tf_xla_enable_dynamic_sizes", &tf_xla_enable_dynamic_sizes, + "XLA flag for enabling XLA dynamic sizes."), + }; + xla::ParseFlagsFromEnvAndIgnoreUnknown("TF_XLA_FLAGS", flag_list); + return tf_xla_enable_dynamic_sizes; +} + +std::optional& TensorShapeExpressionsEnabledOverride() { + static auto* enabled_override = new std::optional(); + return *enabled_override; +} + +void DynExprToTensorFlowProto(const xla::DynExpr& expr, + ExpressionProto* proto) { + proto->Clear(); + switch (expr.kind()) { + case xla::DExpr::Kind::kUnknown: + return; + case xla::DExpr::Kind::kConstant: + proto->set_constant_value( + static_cast(expr).get_val()); + return; + case xla::DExpr::Kind::kVariable: + proto->set_variable_id( + static_cast(expr).get_id()); + return; + case xla::DExpr::Kind::kAdd: { + const auto& add = static_cast(expr); + DynExprToTensorFlowProto(*add.get_lhs(), + proto->mutable_add_node()->mutable_lhs()); + DynExprToTensorFlowProto(*add.get_rhs(), + proto->mutable_add_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kSub: { + const auto& sub = static_cast(expr); + DynExprToTensorFlowProto(*sub.get_lhs(), + proto->mutable_sub_node()->mutable_lhs()); + DynExprToTensorFlowProto(*sub.get_rhs(), + proto->mutable_sub_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kMul: { + const auto& mul = static_cast(expr); + DynExprToTensorFlowProto(*mul.get_lhs(), + proto->mutable_mul_node()->mutable_lhs()); + DynExprToTensorFlowProto(*mul.get_rhs(), + proto->mutable_mul_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kDiv: { + const auto& div = static_cast(expr); + DynExprToTensorFlowProto(*div.get_lhs(), + proto->mutable_div_node()->mutable_lhs()); + DynExprToTensorFlowProto(*div.get_rhs(), + proto->mutable_div_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kMax: { + const auto& max = static_cast(expr); + DynExprToTensorFlowProto(*max.get_lhs(), + proto->mutable_max_node()->mutable_lhs()); + DynExprToTensorFlowProto(*max.get_rhs(), + proto->mutable_max_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kGt: { + const auto& gt = static_cast(expr); + DynExprToTensorFlowProto(*gt.get_lhs(), + proto->mutable_gt_node()->mutable_lhs()); + DynExprToTensorFlowProto(*gt.get_rhs(), + proto->mutable_gt_node()->mutable_rhs()); + return; + } + case xla::DExpr::Kind::kSelect: { + const auto& select = static_cast(expr); + DynExprToTensorFlowProto( + *select.get_pred(), proto->mutable_select_node()->mutable_pred()); + DynExprToTensorFlowProto( + *select.get_on_true(), + proto->mutable_select_node()->mutable_on_true()); + DynExprToTensorFlowProto( + *select.get_on_false(), + proto->mutable_select_node()->mutable_on_false()); + return; + } + } +} + +} // namespace + +bool TensorShapeExpressionsEnabled() { + if (TensorShapeExpressionsEnabledOverride().has_value()) { + return *TensorShapeExpressionsEnabledOverride(); + } + static const bool enabled = ParseTensorShapeExpressionsEnabled(); + return enabled; +} + +void SetTensorShapeExpressionsEnabledForTesting(std::optional enabled) { + TensorShapeExpressionsEnabledOverride() = enabled; +} + +DimExpr DimExprFromProto(const ExpressionProto& proto) { + switch (proto.node_type_case()) { + case ExpressionProto::kConstantValue: + return DimExpr::Const(proto.constant_value()); + case ExpressionProto::kVariableId: + return DimExpr::Var(proto.variable_id()); + case ExpressionProto::kAddNode: { + return DimExprFromProto(proto.add_node().lhs()) + + DimExprFromProto(proto.add_node().rhs()); + } + case ExpressionProto::kSubNode: { + return DimExprFromProto(proto.sub_node().lhs()) - + DimExprFromProto(proto.sub_node().rhs()); + } + case ExpressionProto::kMulNode: { + return DimExprFromProto(proto.mul_node().lhs()) * + DimExprFromProto(proto.mul_node().rhs()); + } + case ExpressionProto::kDivNode: { + return DimExprFromProto(proto.div_node().lhs()) / + DimExprFromProto(proto.div_node().rhs()); + } + case ExpressionProto::kMaxNode: { + return DimExpr::Max(DimExprFromProto(proto.max_node().lhs()), + DimExprFromProto(proto.max_node().rhs())); + } + case ExpressionProto::kGtNode: { + return DimExpr::Gt(DimExprFromProto(proto.gt_node().lhs()), + DimExprFromProto(proto.gt_node().rhs())); + } + case ExpressionProto::kSelectNode: { + return DimExpr::Select( + DimExprFromProto(proto.select_node().pred()), + DimExprFromProto(proto.select_node().on_true()), + DimExprFromProto(proto.select_node().on_false())); + } + case ExpressionProto::NODE_TYPE_NOT_SET: + default: + return DimExpr::Unknown(xla::kMissingExpressionSentinel); + } +} + +void DimExprToProto(const DimExpr& expr, ExpressionProto* proto) { + if (!expr) { + proto->Clear(); + return; + } + DynExprToTensorFlowProto(*expr, proto); +} + +std::string DimExprDebugString(const DimExpr& expr) { + if (!expr) return "_"; + xla::StringPrinter printer; + expr->print(&printer); + return std::move(printer).ToString(); +} + +DimExpr* SimplifyExpr(DimExpr* expr, + std::vector>* arena) { + if (expr == nullptr) return nullptr; + auto owned = std::make_unique(expr->simplify()); + DimExpr* result = owned.get(); + arena->push_back(std::move(owned)); + return result; +} + +bool IsDynamicDimExpr(const ExpressionProto& proto) { + DimExpr expr = DimExprFromProto(proto); + return expr && expr->is_dynamic(); +} + +bool HasDynamicDimExprs(const TensorShapeProto& proto) { + for (const auto& expr : proto.expressions()) { + if (IsDynamicDimExpr(expr)) return true; + } + return false; +} + +} // namespace tensorflow diff --git a/tensorflow/core/framework/tensor_shape_expr.h b/tensorflow/core/framework/tensor_shape_expr.h new file mode 100644 index 00000000000000..7c61d0d079eb25 --- /dev/null +++ b/tensorflow/core/framework/tensor_shape_expr.h @@ -0,0 +1,57 @@ +/* Copyright 2026 The TensorFlow Authors. + +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. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_ +#define TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_ + +#include +#include +#include +#include + +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "xla/shape_expr.h" + +namespace tensorflow { + +// TensorFlow shape inference and XLA use the same owning expression value. +// TensorFlow-specific helpers below only bridge TensorShapeProto's protobuf. +using DimExpr = xla::DExpr; + +DimExpr DimExprFromProto(const ExpressionProto& proto); +void DimExprToProto(const DimExpr& expr, ExpressionProto* proto); +std::string DimExprDebugString(const DimExpr& expr); + +// Simplifies through xla::DExpr and stores the returned value in `arena`. +DimExpr* SimplifyExpr(DimExpr* expr, + std::vector>* arena); + +// Returns whether TensorShape should preserve symbolic expressions. The +// shape-expression support follows the `tf_xla_enable_dynamic_sizes` flag. +bool TensorShapeExpressionsEnabled(); + +// Overrides TensorShapeExpressionsEnabled for tests. Passing std::nullopt +// restores the default environment-derived behavior. +void SetTensorShapeExpressionsEnabledForTesting(std::optional enabled); + +// Returns true if the expression proto depends on a symbolic variable. +bool IsDynamicDimExpr(const ExpressionProto& proto); + +// Returns true if any expression attached to the TensorShapeProto is dynamic. +bool HasDynamicDimExprs(const TensorShapeProto& proto); + +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_ diff --git a/tensorflow/core/framework/tensor_shape_expr_test.cc b/tensorflow/core/framework/tensor_shape_expr_test.cc new file mode 100644 index 00000000000000..55e0f1d138f768 --- /dev/null +++ b/tensorflow/core/framework/tensor_shape_expr_test.cc @@ -0,0 +1,105 @@ +/* Copyright 2026 The TensorFlow Authors. + +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/core/framework/tensor_shape_expr.h" + +#include +#include +#include + +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { + +class TensorShapeExpressionsEnabledTest : public ::testing::Test { + protected: + void SetUp() override { + SetTensorShapeExpressionsEnabledForTesting(true); + } + + void TearDown() override { + SetTensorShapeExpressionsEnabledForTesting(std::nullopt); + } +}; + +TEST(TensorShapeExprTest, UsesXlaCanonicalization) { + ExpressionProto proto; + auto* add = proto.mutable_add_node(); + add->mutable_lhs()->mutable_div_node()->mutable_lhs()->set_variable_id(1); + add->mutable_lhs()->mutable_div_node()->mutable_rhs()->set_constant_value(2); + add->mutable_rhs()->mutable_div_node()->mutable_lhs()->set_variable_id(1); + add->mutable_rhs()->mutable_div_node()->mutable_rhs()->set_constant_value(2); + + auto expr = DimExprFromProto(proto); + EXPECT_EQ(DimExprDebugString(expr.simplify()), "A"); +} + +TEST(TensorShapeExprTest, TensorFlowProtoRoundTripsSharedExpression) { + DimExpr original = (DimExpr::Var(7) + DimExpr::Const(3)) / 2; + ExpressionProto proto; + DimExprToProto(original, &proto); + + auto round_tripped = DimExprFromProto(proto); + EXPECT_TRUE(xla::DynExpr::equal(original.get(), round_tripped.get())); +} + +TEST(TensorShapeExprTest, TensorFlowProtoRoundTripsConditionalExpression) { + DimExpr variable = DimExpr::Var(7); + DimExpr original = DimExpr::Select( + DimExpr::Gt(variable, DimExpr::Const(3)), + DimExpr::Max(variable, DimExpr::Const(8)), + (variable + DimExpr::Const(3)) / 2); + ExpressionProto proto; + DimExprToProto(original, &proto); + + auto round_tripped = DimExprFromProto(proto); + EXPECT_TRUE(xla::DynExpr::equal(original.get(), round_tripped.get())); +} + +TEST(TensorShapeExprTest, SimplifyExprUsesSharedImplementation) { + DimExpr original = + (DimExpr::Var(1) / 2) + (DimExpr::Var(1) / 2); + std::vector> arena; + + DimExpr* simplified = SimplifyExpr(&original, &arena); + + ASSERT_NE(simplified, nullptr); + EXPECT_EQ(DimExprDebugString(*simplified), "A"); +} + +TEST_F(TensorShapeExpressionsEnabledTest, + SetExpressionsRejectsEntriesBeyondRank) { + TensorShape shape({2, 3}); + EXPECT_DEATH( + shape.set_expressions( + {xla::DExpr::Var(1), xla::DExpr::Var(2), xla::DExpr::Var(3)}), + ""); +} + +TEST_F(TensorShapeExpressionsEnabledTest, + RemoveDimRangePreservesRemainingExpressions) { + TensorShape shape({2, 3}); + shape.set_expressions({xla::DExpr::Var(1), xla::DExpr::Var(2)}); + + shape.RemoveDim(0); + + ASSERT_EQ(shape.get_expressions().size(), 1); + EXPECT_TRUE(shape.get_expression(0) == xla::DExpr::Var(2)); +} + +} // namespace +} // namespace tensorflow diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index e94208f90f4537..3988745801be74 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -460,6 +460,7 @@ cc_library( "layout_util.cc", "primitive_util.cc", "shape.cc", + "shape_expr.cc", "shape_partition.cc", "shape_util.cc", ], @@ -469,6 +470,7 @@ cc_library( "layout_util.h", "primitive_util.h", "shape.h", + "shape_expr.h", "shape_partition.h", "shape_util.h", ], diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc index ea348d312eaaa6..d8e03d1f98de09 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc @@ -40,9 +40,11 @@ class OptimizeInputOutputBufferAliasTest r2f32_ = ShapeUtil::MakeShape(F32, {4, 5}); r3f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6}); r4f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6, 7}); - d1f32_ = ShapeUtil::MakeShape(F32, {256}, /*dynamic_dimensions=*/{true}); + d1f32_ = ShapeUtil::MakeShape(F32, {256}, /*dynamic_dimensions=*/{true}, + /*expressions=*/{}); d2f32_ = ShapeUtil::MakeShape(F32, {128, 128}, - /*dynamic_dimensions=*/{false, true}); + /*dynamic_dimensions=*/{false, true}, + /*expressions=*/{}); // Static shape with same size as dynamic shape `d1f32_`. d3f32_ = ShapeUtil::MakeShape(F32, {512}); } diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc index 7bfbbff39956ce..b057d644df8a9f 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc @@ -191,7 +191,8 @@ Shape TypeToShape(mlir::Type type) { return sparse_shape; } - return ShapeUtil::MakeShape(primitive_type, shape, is_dynamic); + return ShapeUtil::MakeShape(primitive_type, shape, is_dynamic, + /*expressions=*/{}); } else if (auto tuple_type = mlir::dyn_cast(type)) { llvm::SmallVector shapes; shapes.reserve(tuple_type.size()); diff --git a/third_party/xla/xla/shape.cc b/third_party/xla/xla/shape.cc index 1cee38146fb07d..e12d07a25ecfe5 100644 --- a/third_party/xla/xla/shape.cc +++ b/third_party/xla/xla/shape.cc @@ -17,6 +17,9 @@ limitations under the License. #include #include +#include +#include +#include #include #include #include @@ -40,7 +43,6 @@ limitations under the License. #include "xla/xla_data.pb.h" namespace xla { - // Defined in .cc file to avoid inlining these large routines Shape::Shape() = default; Shape::~Shape() = default; @@ -97,13 +99,18 @@ Shape::Shape(const ShapeProto& shape_proto) { } absl::StatusOr Shape::FromProto(const ShapeProto& shape_proto) { + + // LOG(INFO) << "FROM PROTO:\n" << shape_proto.DebugString() << std::endl; + Shape shape; shape.set_element_type(shape_proto.element_type()); if (auto* const state = shape.if_array_state()) { const int num_dims = shape_proto.dimensions_size(); const int num_is_dynamic_dims = shape_proto.is_dynamic_dimension_size(); + const int num_expressions = shape_proto.expressions_size(); state->dimensions.reserve(num_dims); state->dynamic_dimensions.reserve(num_dims); + state->expressions.reserve(num_dims); if (num_is_dynamic_dims != 0) { TF_RET_CHECK(num_dims == num_is_dynamic_dims) << "Malformed shape proto: number of is_dynamic_dimension " @@ -111,6 +118,13 @@ absl::StatusOr Shape::FromProto(const ShapeProto& shape_proto) { << num_is_dynamic_dims << ") does not match number of dimension " << "fields (" << num_dims << ")."; } + if (num_expressions != 0) { + TF_RET_CHECK(num_dims == num_expressions) + << "Malformed shape proto: number of expressions " + "fields (" + << num_expressions << ") does not match number of dimension " + << "fields (" << num_dims << ")."; + } for (int i = 0; i < num_dims; ++i) { const bool is_dynamic = (i < num_is_dynamic_dims) && shape_proto.is_dynamic_dimension(i); @@ -118,7 +132,11 @@ absl::StatusOr Shape::FromProto(const ShapeProto& shape_proto) { // UnsafeAddDimension. We expect that the caller will eventually call a // validation routine that will detect the error in case the dimension // value is invalid. - shape.UnsafeAddDimension(shape_proto.dimensions(i), is_dynamic); + DExpr expression = + (i < num_expressions) ? DExprFromProto(shape_proto.expressions(i)) + : DExpr::Const(shape_proto.dimensions(i)); + shape.UnsafeAddDimension(shape_proto.dimensions(i), is_dynamic, + expression); } } else if (auto* const state = shape.if_tuple_state()) { state->tuple_shapes.reserve(shape_proto.tuple_shapes_size()); @@ -136,6 +154,7 @@ absl::StatusOr Shape::FromProto(const ShapeProto& shape_proto) { TF_ASSIGN_OR_RETURN(*shape.mutable_layout(), Layout::FromProto(shape_proto.layout())); } + // LOG(INFO) << "FROM PROTO " << shape << "\n"; return shape; } @@ -143,6 +162,8 @@ ShapeProto Shape::ToProto() const { ShapeProto proto; proto.set_element_type(element_type_); + // LOG(INFO) << "TO PROTO " << ToString() << "\n"; + if (const auto* const state = if_array_state()) { proto.mutable_dimensions()->Reserve(state->dimensions.size()); for (const int64_t dimension : state->dimensions) { @@ -151,6 +172,11 @@ ShapeProto Shape::ToProto() const { for (const bool dynamic : state->dynamic_dimensions) { proto.add_is_dynamic_dimension(dynamic); } + for (const DExpr& e : state->expressions) { + ExpressionProto* eproto = proto.add_expressions(); + CHECK(e.get() != nullptr) << "Missing expression in expression list."; + e->to_proto(eproto); + } if (state->layout.has_value()) { *proto.mutable_layout() = state->layout->ToProto(); } @@ -163,6 +189,7 @@ ShapeProto Shape::ToProto() const { proto.mutable_tuple_shapes()->Reserve(1); *proto.add_tuple_shapes() = state->buffer_shape[0].ToProto(); } + // LOG(INFO) << "DEBUG VIEW:\n" << proto.DebugString() << std::endl; return proto; } @@ -245,14 +272,16 @@ bool Shape::AreAllLeavesIntegers() const { return primitive_util::IsIntegralType(element_type()); } -void Shape::add_dimensions(int64_t value, bool is_dynamic) { +void Shape::add_dimensions(int64_t value, bool is_dynamic, DExpr expr) { if (value < 0) { CHECK(is_dynamic) << "static dimension must have size >= 0 instead of " << value << "."; CHECK_EQ(value, kUnboundedSize) << "dynamic dimension must have size == kUnboundedSize or >= 0."; } - UnsafeAddDimension(value, is_dynamic); + UnsafeAddDimension( + value, is_dynamic, + expr ? std::move(expr) : DExpr::Const(value)); } void Shape::set_dynamic_dimension(int dimension, bool is_dynamic) { @@ -262,6 +291,26 @@ void Shape::set_dynamic_dimension(int dimension, bool is_dynamic) { state.dynamic_dimensions[dimension] = is_dynamic; } +void Shape::set_expression(int dimension, DExpr e) { + auto& state = array_state(); + state.expressions[dimension] = + e ? std::move(e) : DExpr::Const(state.dimensions[dimension]); +} + +void Shape::set_expressions(std::vector exps) { + auto& state = array_state(); + CHECK_LE(exps.size(), state.dimensions.size()); + state.expressions.resize(state.dimensions.size()); + for (size_t i = 0; i < state.dimensions.size(); ++i) { + state.expressions[i] = i < exps.size() + ? std::move(exps[i]) + : DExpr::Const(state.dimensions[i]); + if (!state.expressions[i]) { + state.expressions[i] = DExpr::Const(state.dimensions[i]); + } + } +} + void Shape::set_dimensions(int index, int64_t size, std::optional is_dynamic) { auto& state = array_state(); @@ -270,6 +319,7 @@ void Shape::set_dimensions(int index, int64_t size, CheckDimensionSize(index, size, dynamic); state.dimensions[index] = size; state.dynamic_dimensions[index] = dynamic; + state.expressions[index] = DExpr::Const(size); } void Shape::set_dimensions_minor(int index, int64_t size, @@ -291,12 +341,15 @@ void Shape::CheckDimensionSize(int dim_index, int64_t size, bool is_dynamic) { } } -void Shape::UnsafeAddDimension(int64_t value, bool is_dynamic) { +void Shape::UnsafeAddDimension(int64_t value, bool is_dynamic, DExpr exp) { auto& state = array_state(); CHECK_EQ(state.dimensions.size(), state.dynamic_dimensions.size()) << "where the shape is " << ToString(); + CHECK_EQ(state.dimensions.size(), state.expressions.size()) + << "where the shape is " << ToString(); state.dimensions.push_back(value); state.dynamic_dimensions.push_back(is_dynamic); + state.expressions.push_back(exp ? std::move(exp) : DExpr::Const(value)); } bool Shape::is_static() const { @@ -345,6 +398,8 @@ void Shape::DeleteDimension(int64_t dim_to_delete) { state.dimensions.erase(state.dimensions.begin() + dim_to_delete); state.dynamic_dimensions.erase(state.dynamic_dimensions.begin() + dim_to_delete); + state.expressions.erase(state.expressions.begin() + + dim_to_delete); if (LayoutUtil::HasLayout(*this)) { state.layout->DeleteDimension(dim_to_delete); // NOLINT: optional-access } @@ -358,6 +413,8 @@ void Shape::DeleteDimensions(absl::Span dims_to_delete) { state.dimensions = RemoveElements(sorted_dims_to_delete, state.dimensions); state.dynamic_dimensions = RemoveElements(sorted_dims_to_delete, state.dynamic_dimensions); + state.expressions = + RemoveElements(sorted_dims_to_delete, state.expressions); if (LayoutUtil::HasLayout(*this)) { for (auto it = sorted_dims_to_delete.rbegin(); it != sorted_dims_to_delete.rend(); ++it) { @@ -370,6 +427,7 @@ void Shape::CheckStateIsEmpty() const { if (const auto* const state = if_array_state()) { CHECK(state->dimensions.empty()) << ToString(); CHECK(state->dynamic_dimensions.empty()) << ToString(); + CHECK(state->expressions.empty()) << ToString(); CHECK(!state->layout.has_value()) << ToString(); } else if (const auto* const state = if_tuple_state()) { CHECK(state->tuple_shapes.empty()) << ToString(); @@ -509,6 +567,11 @@ bool Shape::Equal::operator()(const Shape& lhs, const Shape& rhs) { rhs.is_unbounded_dynamic_dimension(i))) { continue; } + if (i == 0 && ignore_batch_ && + (lhs.outer_multiplier() > 0 || rhs.outer_multiplier() > 0)) { + VLOG(3) << "CompareShapes: batch dimension found. Forcely compatible"; + continue; + } if (lhs.dimensions(i) != rhs.dimensions(i)) { VLOG(3) << "CompareShapes: lhs dimensions != rhs dimensions"; return false; diff --git a/third_party/xla/xla/shape.h b/third_party/xla/xla/shape.h index 8453dc17717e11..13079dab62532d 100644 --- a/third_party/xla/xla/shape.h +++ b/third_party/xla/xla/shape.h @@ -38,6 +38,7 @@ limitations under the License. #include "xla/tsl/platform/logging.h" // IWYU pragma: keep #include "xla/util.h" #include "xla/xla_data.pb.h" +#include "xla/shape_expr.h" namespace xla { @@ -218,6 +219,29 @@ class Shape { return array_state().dynamic_dimensions[dimension]; } + bool has_dynamic_expr() const { + if (auto* const state = if_array_state()) { + return absl::c_any_of(state->expressions, + [](const DExpr& e) { + return e && e->is_dynamic(); + }); + } + if (auto* const state = if_tuple_state()) { + return absl::c_any_of(state->tuple_shapes, [](Shape subshape) { + return subshape.has_dynamic_expr(); + }); + } + return false; + } + + const DExpr& expressions(int dimension) const { + if (dimension < 0) return MissingExpression(); + const auto& exprs = array_state().expressions; + const size_t dim = static_cast(dimension); + if (dim >= exprs.size()) return MissingExpression(); + return exprs[dim] ? exprs[dim] : MissingExpression(); + } + // Returns true if the given dimension is statically-sized. // Precondition: this is an array shape and `dimension` is a valid dimension // index. @@ -232,12 +256,18 @@ class Shape { // - The dimension's size is valid for the given dynamic-ness. void set_dynamic_dimension(int dimension, bool is_dynamic); + void set_expression(int dimension, DExpr e); + + void set_expressions(std::vector exprs); + // Returns a span to indicate whether each dimension is dynamic. // Precondition: this is an array shape. absl::Span dynamic_dimensions() const { return array_state().dynamic_dimensions; } + absl::Span expressions() const { return array_state().expressions; } + // Removes the given dimension from the shape. Layout, if it exists, is // adjusted to match the modified shape. // Precondition: this is an array shape, and the input dimension indices are @@ -313,7 +343,8 @@ class Shape { // - This is an array shape. // - Either `value` is >= 0, or `is_dynamic` is true and `value` is // kUnboundedSize. - void add_dimensions(int64_t value, bool is_dynamic = false); + void add_dimensions(int64_t value, bool is_dynamic = false, + DExpr expr = DExpr()); // Clears all dimensions (i.e. makes this shape a scalar). // Precondition: this is an array shape. @@ -321,6 +352,7 @@ class Shape { auto& state = array_state(); state.dimensions.clear(); state.dynamic_dimensions.clear(); + state.expressions.clear(); } // Returns a span to indicate the size of each dimension. @@ -434,6 +466,10 @@ class Shape { bool operator()(const Shape& lhs, const Shape& rhs); + Equal& IgnoreBatch(bool ignore_batch = true) { + ignore_batch_ = ignore_batch; + return *this; + } Equal& IgnoreLayout(bool ignore_layout = true) { ignore_layout_ = ignore_layout; return *this; @@ -488,6 +524,7 @@ class Shape { } private: + bool ignore_batch_ = false; bool ignore_layout_ = false; bool ignore_tiles_in_layout_ = false; bool ignore_element_size_in_layout_ = false; @@ -515,7 +552,7 @@ class Shape { } if (const auto* const state = s.if_array_state()) { h = H::combine(std::move(h), s.element_type_, state->dimensions, - state->dynamic_dimensions); + state->dynamic_dimensions, state->expressions); if (kIsLayoutSensitive) { h = H::combine(std::move(h), state->layout); } @@ -532,7 +569,11 @@ class Shape { return Shape::Hash(std::move(h), s); } + int64_t outer_multiplier() const { return outer_multiplier_; } + void set_outer_multiplier(int64_t m) { outer_multiplier_ = m; } private: + int64_t outer_multiplier_ = -1; + friend absl::Status ValidateNonLayoutProperties(const Shape& shape); // Define one state struct for each shape category. Depending on the element @@ -560,8 +601,27 @@ class Shape { // respective dimension is dynamically sized. absl::InlinedVector dynamic_dimensions; + absl::InlinedVector expressions; + // The layout of the shape. std::optional layout; + + ArrayState() = default; + ArrayState(const ArrayState& other) + : dimensions(other.dimensions), + dynamic_dimensions(other.dynamic_dimensions), + expressions(other.expressions), + layout(other.layout) {} + ArrayState& operator=(const ArrayState& other) { + if (this == &other) return *this; + dimensions = other.dimensions; + dynamic_dimensions = other.dynamic_dimensions; + expressions = other.expressions; + layout = other.layout; + return *this; + } + ArrayState(ArrayState&&) noexcept = default; + ArrayState& operator=(ArrayState&&) noexcept = default; }; struct TupleState { // The tuple element subshapes. @@ -578,12 +638,13 @@ class Shape { // CHECKs that the dimension size is valid. void CheckDimensionSize(int dim_index, int64_t size, bool is_dynamic); + static const DExpr& MissingExpression(); // Like add_dimensions(), but does not CHECK that the arguments are valid. // Instead, we rely on validation down the road to catch invalid shapes. // This is useful for code that should not crash, such as constructing a // Shape from an unvalidated proto. - void UnsafeAddDimension(int64_t value, bool is_dynamic); + void UnsafeAddDimension(int64_t value, bool is_dynamic, DExpr exp); // Convenience accessors for the state_ variant. Each if_*_state() accessor // returns a pointer to the corresponding state struct, or nullptr if the diff --git a/third_party/xla/xla/shape_expr.cc b/third_party/xla/xla/shape_expr.cc new file mode 100644 index 00000000000000..d74a42175b3380 --- /dev/null +++ b/third_party/xla/xla/shape_expr.cc @@ -0,0 +1,773 @@ +/* Copyright 2018 The OpenXLA Authors. + +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 "xla/shape.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "xla/printer.h" + +namespace xla { + +namespace { + +Constant* AsConstant(DynExpr* expr) { + return expr != nullptr && expr->kind() == DExpr::Kind::kConstant + ? static_cast(expr) + : nullptr; +} + +std::vector ExpressionChildren(DynExpr* expr) { + CHECK(expr != nullptr); + switch (expr->kind()) { + case DExpr::Kind::kUnknown: + case DExpr::Kind::kConstant: + case DExpr::Kind::kVariable: + return {}; + case DExpr::Kind::kAdd: { + auto* add = static_cast(expr); + return {add->get_lhs(), add->get_rhs()}; + } + case DExpr::Kind::kSub: { + auto* sub = static_cast(expr); + return {sub->get_lhs(), sub->get_rhs()}; + } + case DExpr::Kind::kMul: { + auto* mul = static_cast(expr); + return {mul->get_lhs(), mul->get_rhs()}; + } + case DExpr::Kind::kDiv: { + auto* div = static_cast(expr); + return {div->get_lhs(), div->get_rhs()}; + } + case DExpr::Kind::kMax: { + auto* max = static_cast(expr); + return {max->get_lhs(), max->get_rhs()}; + } + case DExpr::Kind::kGt: { + auto* gt = static_cast(expr); + return {gt->get_lhs(), gt->get_rhs()}; + } + case DExpr::Kind::kSelect: { + auto* select = static_cast(expr); + return {select->get_pred(), select->get_on_true(), + select->get_on_false()}; + } + } + return {}; +} + +DynExpr* FindSmallestCoveringSubexpression(DynExpr* expr) { + CHECK(expr != nullptr); + if (expr->kind() == DExpr::Kind::kVariable) { + return expr; + } + + DynExpr* common_core = nullptr; + for (DynExpr* child : ExpressionChildren(expr)) { + DynExpr* child_core = FindSmallestCoveringSubexpression(child); + if (child_core == nullptr) { + continue; + } + if (common_core == nullptr) { + common_core = child_core; + } else if (!DynExpr::equal(common_core, child_core)) { + return expr; + } + } + return common_core; +} + +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()); + } + case DExpr::Kind::kMax: { + auto* max = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(max->get_lhs(), target, replacement).release(), + ReplaceSubexpression(max->get_rhs(), target, replacement).release()); + } + case DExpr::Kind::kGt: { + auto* gt = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(gt->get_lhs(), target, replacement).release(), + ReplaceSubexpression(gt->get_rhs(), target, replacement).release()); + } + case DExpr::Kind::kSelect: { + auto* select = static_cast(expr); + return std::make_unique( + ReplaceSubexpression(select->get_pred(), target, replacement) + .release(), + ReplaceSubexpression(select->get_on_true(), target, replacement) + .release(), + ReplaceSubexpression(select->get_on_false(), target, replacement) + .release()); + } + } + return expr->clone(); +} + +void NormalizeFraction(int64_t* numerator, int64_t* denominator) { + CHECK(denominator != nullptr); + CHECK(*denominator != 0); + int64_t divisor = std::gcd(std::llabs(*numerator), std::llabs(*denominator)); + if (divisor > 1) { + *numerator /= divisor; + *denominator /= divisor; + } + if (*denominator < 0) { + *numerator = -*numerator; + *denominator = -*denominator; + } +} + +// We normalize affine expressions into: +// (constant + sum_i coeff_i * var_i) / denominator +// +// Concretely, the canonical forms we want to preserve are: +// p +// mX +// mX + p +// mX + nY +// (mX) / d +// (mX + p) / d +// (mX + nY + p) / d +// +// This lets us combine equivalent trees into one stable representation, e.g.: +// A/2 + A/2 -> A +// A/2 + B/2 -> (A + B) / 2 +// 4/8 * A -> A / 2 +// A + (B + 1) -> A + B + 1 (represented via coefficients + constant) +// +// The important constraint is that we only canonicalize affine expressions over +// an integer denominator. Non-affine expressions stay in tree form and only get +// a minimal local simplification in the fallback path below. +struct CanonicalAffineExpr { + int64_t denominator = 1; + int64_t constant = 0; + std::map coefficients; + + bool IsPureConstant() const { return coefficients.empty(); } +}; + +void NormalizeAffine(CanonicalAffineExpr* expr) { + CHECK(expr != nullptr); + CHECK(expr->denominator != 0); + for (auto it = expr->coefficients.begin(); it != expr->coefficients.end();) { + if (it->second == 0) { + it = expr->coefficients.erase(it); + } else { + ++it; + } + } + int64_t divisor = std::llabs(expr->constant); + for (const auto& [_, coeff] : expr->coefficients) { + divisor = std::gcd(divisor, std::llabs(coeff)); + } + divisor = std::gcd(divisor, std::llabs(expr->denominator)); + if (divisor > 1) { + expr->constant /= divisor; + expr->denominator /= divisor; + for (auto& [_, coeff] : expr->coefficients) { + coeff /= divisor; + } + } + if (expr->denominator < 0) { + expr->denominator = -expr->denominator; + expr->constant = -expr->constant; + for (auto& [_, coeff] : expr->coefficients) { + coeff = -coeff; + } + } + if (expr->constant == 0 && expr->coefficients.empty()) { + expr->denominator = 1; + } +} + +CanonicalAffineExpr MakeConstantAffine(int64_t value) { + CanonicalAffineExpr expr; + expr.constant = value; + return expr; +} + +CanonicalAffineExpr MakeVariableAffine(int id) { + CanonicalAffineExpr expr; + expr.coefficients[id] = 1; + return expr; +} + +CanonicalAffineExpr AddAffine(const CanonicalAffineExpr& lhs, + const CanonicalAffineExpr& rhs, int rhs_sign) { + CHECK(rhs_sign == 1 || rhs_sign == -1); + CanonicalAffineExpr result; + int64_t gcd = std::gcd(lhs.denominator, rhs.denominator); + int64_t lhs_scale = rhs.denominator / gcd; + int64_t rhs_scale = lhs.denominator / gcd; + result.denominator = lhs.denominator * lhs_scale; + result.constant = + lhs.constant * lhs_scale + rhs_sign * rhs.constant * rhs_scale; + for (const auto& [id, coeff] : lhs.coefficients) { + result.coefficients[id] = coeff * lhs_scale; + } + for (const auto& [id, coeff] : rhs.coefficients) { + result.coefficients[id] += rhs_sign * coeff * rhs_scale; + } + NormalizeAffine(&result); + return result; +} + +CanonicalAffineExpr MultiplyAffineByRational(const CanonicalAffineExpr& expr, + int64_t numerator, + int64_t denominator) { + CHECK(denominator != 0); + CanonicalAffineExpr result = expr; + result.constant *= numerator; + result.denominator *= denominator; + for (auto& [_, coeff] : result.coefficients) { + coeff *= numerator; + } + NormalizeAffine(&result); + return result; +} + +std::optional ToCanonicalAffine(const DynExpr* expr) { + CHECK(expr != nullptr); + switch (expr->kind()) { + case DExpr::Kind::kUnknown: + return std::nullopt; + case DExpr::Kind::kConstant: + return MakeConstantAffine(static_cast(expr)->get_val()); + case DExpr::Kind::kVariable: + return MakeVariableAffine(static_cast(expr)->get_id()); + case DExpr::Kind::kAdd: { + const auto* add = static_cast(expr); + auto lhs = ToCanonicalAffine(add->get_lhs()); + auto rhs = ToCanonicalAffine(add->get_rhs()); + if (!lhs.has_value() || !rhs.has_value()) return std::nullopt; + return AddAffine(*lhs, *rhs, /*rhs_sign=*/1); + } + case DExpr::Kind::kSub: { + const auto* sub = static_cast(expr); + auto lhs = ToCanonicalAffine(sub->get_lhs()); + auto rhs = ToCanonicalAffine(sub->get_rhs()); + if (!lhs.has_value() || !rhs.has_value()) return std::nullopt; + return AddAffine(*lhs, *rhs, /*rhs_sign=*/-1); + } + case DExpr::Kind::kMul: { + const auto* mul = static_cast(expr); + auto lhs = ToCanonicalAffine(mul->get_lhs()); + auto rhs = ToCanonicalAffine(mul->get_rhs()); + if (!lhs.has_value() || !rhs.has_value()) return std::nullopt; + if (lhs->IsPureConstant()) { + return MultiplyAffineByRational(*rhs, lhs->constant, lhs->denominator); + } + if (rhs->IsPureConstant()) { + return MultiplyAffineByRational(*lhs, rhs->constant, rhs->denominator); + } + return std::nullopt; + } + case DExpr::Kind::kDiv: { + const auto* div = static_cast(expr); + auto lhs = ToCanonicalAffine(div->get_lhs()); + auto rhs = ToCanonicalAffine(div->get_rhs()); + if (!lhs.has_value() || !rhs.has_value() || !rhs->IsPureConstant() || + rhs->constant == 0) { + return std::nullopt; + } + return MultiplyAffineByRational(*lhs, rhs->denominator, rhs->constant); + } + case DExpr::Kind::kMax: + case DExpr::Kind::kGt: + case DExpr::Kind::kSelect: + return std::nullopt; + default: + return std::nullopt; + } + return std::nullopt; +} + +bool IsNonNegativeForPositiveVariables(const DynExpr* expr) { + auto affine = ToCanonicalAffine(expr); + if (!affine.has_value()) return false; + + // Dynamic dimension variables are strictly positive. An affine expression + // has a finite lower bound only when every variable coefficient is + // non-negative; evaluate that bound with each variable set to one. + __int128 lower_bound = affine->constant; + for (const auto& [_, coefficient] : affine->coefficients) { + if (coefficient < 0) return false; + lower_bound += coefficient; + } + return lower_bound >= 0; +} + +std::unique_ptr BuildScaledVariableTerm(int id, int64_t coefficient) { + CHECK(coefficient != 0); + if (coefficient == 1) { + return std::make_unique(id); + } + return std::make_unique(DynExpr::_(coefficient), DynExpr::V(id)); +} + +std::unique_ptr BuildAffineNumerator(const CanonicalAffineExpr& expr) { + std::unique_ptr result; + for (const auto& [id, coeff] : expr.coefficients) { + auto term = BuildScaledVariableTerm(id, coeff); + if (result == nullptr) { + result = std::move(term); + } else { + result = std::make_unique(result.release(), term.release()); + } + } + if (expr.constant != 0 || result == nullptr) { + if (result == nullptr) { + result = std::make_unique(expr.constant); + } else if (expr.constant > 0) { + result = std::make_unique(result.release(), + DynExpr::_(expr.constant)); + } else { + result = std::make_unique(result.release(), + DynExpr::_(-expr.constant)); + } + } + return result; +} + +std::unique_ptr BuildCanonicalExpr(const CanonicalAffineExpr& expr) { + CanonicalAffineExpr normalized = expr; + NormalizeAffine(&normalized); + auto numerator = BuildAffineNumerator(normalized); + if (normalized.denominator == 1) { + return numerator; + } + return std::make_unique
(numerator.release(), + DynExpr::_(normalized.denominator)); +} + +std::unique_ptr SimplifyFallback(const DynExpr* expr) { + switch (expr->kind()) { + case DExpr::Kind::kUnknown: + case DExpr::Kind::kConstant: + case DExpr::Kind::kVariable: + return expr->clone(); + case DExpr::Kind::kAdd: { + const auto* add = static_cast(expr); + auto lhs = std::unique_ptr(add->get_lhs()->s()); + auto rhs = std::unique_ptr(add->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + Constant* l = AsConstant(lhs.get()); + Constant* r = AsConstant(rhs.get()); + if (l && r) return std::make_unique(l->get_val() + r->get_val()); + if (l && l->get_val() == 0) return rhs; + if (r && r->get_val() == 0) return lhs; + return std::make_unique(lhs.release(), rhs.release()); + } + case DExpr::Kind::kSub: { + const auto* sub = static_cast(expr); + auto lhs = std::unique_ptr(sub->get_lhs()->s()); + auto rhs = std::unique_ptr(sub->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + Constant* l = AsConstant(lhs.get()); + Constant* r = AsConstant(rhs.get()); + if (l && r) return std::make_unique(l->get_val() - r->get_val()); + if (r && r->get_val() == 0) return lhs; + if (*lhs == *rhs) return std::make_unique(0); + return std::make_unique(lhs.release(), rhs.release()); + } + case DExpr::Kind::kMul: { + const auto* mul = static_cast(expr); + auto lhs = std::unique_ptr(mul->get_lhs()->s()); + auto rhs = std::unique_ptr(mul->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + Constant* l = AsConstant(lhs.get()); + Constant* r = AsConstant(rhs.get()); + if (l && r) return std::make_unique(l->get_val() * r->get_val()); + if ((l && l->get_val() == 0) || (r && r->get_val() == 0)) { + return std::make_unique(0); + } + if (l && l->get_val() == 1) return rhs; + if (r && r->get_val() == 1) return lhs; + if (r != nullptr) std::swap(lhs, rhs); + return std::make_unique(lhs.release(), rhs.release()); + } + case DExpr::Kind::kDiv: { + const auto* div = static_cast(expr); + auto lhs = std::unique_ptr(div->get_lhs()->s()); + auto rhs = std::unique_ptr(div->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + Constant* l = AsConstant(lhs.get()); + Constant* r = AsConstant(rhs.get()); + if (*lhs == *rhs) { + return std::make_unique(1); + } + if (l && l->get_val() == 0 && r && r->get_val() != 0) { + return std::make_unique(0); + } + if (r && r->get_val() == 1) return lhs; + if (lhs->kind() == DExpr::Kind::kMul) { + auto* mul = static_cast(lhs.get()); + auto lhs_l = std::unique_ptr(mul->get_lhs()->s()); + auto lhs_r = std::unique_ptr(mul->get_rhs()->s()); + if (*lhs_l == *rhs) { + return lhs_r; + } + if (*lhs_r == *rhs) { + return lhs_l; + } + } + if (l && r && r->get_val() != 0) { + int64_t numerator = l->get_val(); + int64_t denominator = r->get_val(); + NormalizeFraction(&numerator, &denominator); + if (denominator == 1) return std::make_unique(numerator); + return std::make_unique
(DynExpr::_(numerator), + DynExpr::_(denominator)); + } + return std::make_unique
(lhs.release(), rhs.release()); + } + case DExpr::Kind::kMax: { + const auto* max = static_cast(expr); + auto lhs = std::unique_ptr(max->get_lhs()->s()); + auto rhs = std::unique_ptr(max->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + Constant* l = AsConstant(lhs.get()); + Constant* r = AsConstant(rhs.get()); + if (l && r) { + return std::make_unique( + std::max(l->get_val(), r->get_val())); + } + if (l && l->get_val() == 0 && + IsNonNegativeForPositiveVariables(rhs.get())) { + return rhs; + } + if (r && r->get_val() == 0 && + IsNonNegativeForPositiveVariables(lhs.get())) { + return lhs; + } + if (*lhs == *rhs) return lhs; + return std::make_unique(lhs.release(), rhs.release()); + } + case DExpr::Kind::kGt: { + const auto* gt = static_cast(expr); + auto lhs = std::unique_ptr(gt->get_lhs()->s()); + auto rhs = std::unique_ptr(gt->get_rhs()->s()); + if (lhs->kind() == DExpr::Kind::kUnknown || + rhs->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + if (lhs->is_constant() && rhs->is_constant()) { + return std::make_unique(lhs->get_val() > rhs->get_val()); + } + if (*lhs == *rhs) return std::make_unique(0); + return std::make_unique(lhs.release(), rhs.release()); + } + case DExpr::Kind::kSelect: { + const auto* select = static_cast(expr); + auto pred = std::unique_ptr(select->get_pred()->s()); + auto on_true = std::unique_ptr(select->get_on_true()->s()); + auto on_false = std::unique_ptr(select->get_on_false()->s()); + if (pred->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + if (pred->is_constant()) { + return pred->get_val() != 0 ? std::move(on_true) : std::move(on_false); + } + if (on_true->kind() == DExpr::Kind::kUnknown || + on_false->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + if (*on_true == *on_false) return on_true; + return std::make_unique(pred.release(), on_true.release(), + on_false.release()); + } + default: + return expr->clone(); + } +} + +std::unique_ptr SimplifyCanonical(const DynExpr* expr) { + if (expr->kind() == DExpr::Kind::kUnknown) { + return std::make_unique(); + } + if (auto canonical = ToCanonicalAffine(expr); canonical.has_value()) { + if (canonical->IsPureConstant()) { + CHECK_NE(canonical->denominator, 0); + return std::make_unique(canonical->constant / + canonical->denominator); + } + return BuildCanonicalExpr(*canonical); + } + return SimplifyFallback(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()); + DynExpr* result = FindSmallestCoveringSubexpression(expr_.get()); + CHECK(result != nullptr); + return DExpr::Adopt(result->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; +} + +DynExpr* operator*(DynExpr& lhs, DynExpr& rhs) { + return new Mul(lhs.clone().release(), rhs.clone().release()); +} +DynExpr* operator*(int64_t k, DynExpr& rhs) { + return new Mul(DynExpr::_(k), rhs.clone().release()); +} +DynExpr* operator*(DynExpr& lhs, int64_t k) { + return new Mul(lhs.clone().release(), DynExpr::_(k)); +} +DynExpr* operator/(DynExpr& lhs, DynExpr& rhs) { + return new Div(lhs.clone().release(), rhs.clone().release()); +} +DynExpr* operator/(DynExpr& lhs, int64_t d) { + return new Div(lhs.clone().release(), DynExpr::_(d)); +} +DynExpr* operator+(DynExpr& lhs, DynExpr& rhs) { + return new Add(lhs.clone().release(), rhs.clone().release()); +} +DynExpr* operator+(DynExpr& lhs, int64_t d) { + return new Add(lhs.clone().release(), DynExpr::_(d)); +} +DynExpr* operator-(DynExpr& lhs, DynExpr& rhs) { + return new Sub(lhs.clone().release(), rhs.clone().release()); +} +DynExpr* operator-(DynExpr& lhs, int64_t d) { + return new Sub(lhs.clone().release(), DynExpr::_(d)); +} +bool operator==(DynExpr& lhs, DynExpr& rhs) { + return DynExpr::equal(&lhs, &rhs); +} +bool operator==(DynExpr& lhs, int64_t d) { + auto rhs = std::unique_ptr(DynExpr::_(d)); + return DynExpr::equal(&lhs, rhs.get()); +} +bool operator<(DynExpr& lhs, int64_t d) { + return lhs.is_constant() && lhs.get_val() < d; +} + +DExpr DExpr::Max(const DExpr& lhs, const DExpr& rhs) { + return Adopt(new xla::MaxExpr(lhs.clone().release(), rhs.clone().release())); +} + +DExpr DExpr::Gt(const DExpr& lhs, const DExpr& rhs) { + return Adopt(new xla::GtExpr(lhs.clone().release(), rhs.clone().release())); +} + +DExpr DExpr::Select(const DExpr& pred, const DExpr& on_true, + const DExpr& on_false) { + return Adopt(new xla::SelectExpr(pred.clone().release(), + on_true.clone().release(), + on_false.clone().release())); +} + +bool DynExpr::equal(DynExpr* expr1, DynExpr* expr2) { + auto e1 = std::unique_ptr(expr1->s()); + auto e2 = std::unique_ptr(expr2->s()); + if (e1 == nullptr || e2 == nullptr) return false; + auto a1 = ToCanonicalAffine(e1.get()); + auto a2 = ToCanonicalAffine(e2.get()); + if (a1.has_value() && a2.has_value()) { + return a1->denominator == a2->denominator && + a1->constant == a2->constant && + a1->coefficients == a2->coefficients; + } + if (e1->kind() == DExpr::Kind::kConstant && + e2->kind() == DExpr::Kind::kConstant) { + return static_cast(e1.get())->get_val() == + static_cast(e2.get())->get_val(); + } + if (e1->kind() == DExpr::Kind::kVariable && + e2->kind() == DExpr::Kind::kVariable) { + return static_cast(e1.get())->get_id() == + static_cast(e2.get())->get_id(); + } + if (e1->kind() == DExpr::Kind::kUnknown && + e2->kind() == DExpr::Kind::kUnknown) { + int lhs_id = static_cast(e1.get())->get_id(); + int rhs_id = static_cast(e2.get())->get_id(); + return lhs_id != 0 && lhs_id == rhs_id; + } + if (e1->kind() == DExpr::Kind::kMul && e2->kind() == DExpr::Kind::kMul) { + auto* ab = static_cast(e1.get()); + auto* cd = static_cast(e2.get()); + auto a = ab->get_lhs(); + auto b = ab->get_rhs(); + auto c = cd->get_lhs(); + auto d = cd->get_rhs(); + return (*a == *c && *b == *d) || (*a == *d && *b == *c); + } + if (e1->kind() == DExpr::Kind::kDiv && e2->kind() == DExpr::Kind::kDiv) { + auto* ab = static_cast(e1.get()); + auto* cd = static_cast(e2.get()); + auto a = ab->get_lhs(); + auto b = ab->get_rhs(); + auto c = cd->get_lhs(); + auto d = cd->get_rhs(); + return *a == *c && *b == *d; + } + if (e1->kind() == DExpr::Kind::kAdd && e2->kind() == DExpr::Kind::kAdd) { + auto* ab = static_cast(e1.get()); + auto* cd = static_cast(e2.get()); + auto a = ab->get_lhs(); + auto b = ab->get_rhs(); + auto c = cd->get_lhs(); + auto d = cd->get_rhs(); + return (*a == *c && *b == *d) || (*a == *d && *b == *c); + } + if (e1->kind() == DExpr::Kind::kSub && e2->kind() == DExpr::Kind::kSub) { + auto* ab = static_cast(e1.get()); + auto* cd = static_cast(e2.get()); + auto* a = ab->get_lhs(); + auto* b = ab->get_rhs(); + auto* c = cd->get_lhs(); + auto* d = cd->get_rhs(); + return *a == *c && *b == *d; + } + if (e1->kind() == DExpr::Kind::kMax && e2->kind() == DExpr::Kind::kMax) { + auto* ab = static_cast(e1.get()); + auto* cd = static_cast(e2.get()); + auto* a = ab->get_lhs(); + auto* b = ab->get_rhs(); + auto* c = cd->get_lhs(); + auto* d = cd->get_rhs(); + return (*a == *c && *b == *d) || (*a == *d && *b == *c); + } + if (e1->kind() == DExpr::Kind::kGt && e2->kind() == DExpr::Kind::kGt) { + auto* lhs = static_cast(e1.get()); + auto* rhs = static_cast(e2.get()); + return *lhs->get_lhs() == *rhs->get_lhs() && + *lhs->get_rhs() == *rhs->get_rhs(); + } + if (e1->kind() == DExpr::Kind::kSelect && + e2->kind() == DExpr::Kind::kSelect) { + auto* lhs = static_cast(e1.get()); + auto* rhs = static_cast(e2.get()); + return *lhs->get_pred() == *rhs->get_pred() && + *lhs->get_on_true() == *rhs->get_on_true() && + *lhs->get_on_false() == *rhs->get_on_false(); + } + return false; +} + +DynExpr* Constant::s() { return SimplifyCanonical(this).release(); } + +DynExpr* Variable::s() { return SimplifyCanonical(this).release(); } + +DynExpr* Mul::s() { return SimplifyCanonical(this).release(); } + +DynExpr* Add::s() { return SimplifyCanonical(this).release(); } + +DynExpr* Sub::s() { return SimplifyCanonical(this).release(); } + +DynExpr* Div::s() { return SimplifyCanonical(this).release(); } + +DynExpr* MaxExpr::s() { return SimplifyCanonical(this).release(); } + +DynExpr* GtExpr::s() { return SimplifyCanonical(this).release(); } + +DynExpr* SelectExpr::s() { return SimplifyCanonical(this).release(); } + +std::ostream& operator<<(std::ostream& os, DynExpr* expr) { + auto simplified = std::unique_ptr(expr->s()); + StringPrinter printer; + simplified->print(&printer); + os << std::move(printer).ToString(); + return os; +} + +DynExpr* DynExpr::zero = new Constant(0); +DynExpr* DynExpr::one = new Constant(1); + +} // namespace xla diff --git a/third_party/xla/xla/shape_expr.h b/third_party/xla/xla/shape_expr.h new file mode 100644 index 00000000000000..babd157a56b0b9 --- /dev/null +++ b/third_party/xla/xla/shape_expr.h @@ -0,0 +1,802 @@ +/* Copyright 2018 The OpenXLA Authors. + +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. +==============================================================================*/ + +#ifndef XLA_SHAPE_EXPR_H_ +#define XLA_SHAPE_EXPR_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/hash/hash.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/types/span.h" +#include "xla/printer.h" +#include "xla/xla_data.pb.h" + +namespace xla { + +// Reserved sentinel for "missing expression". Keep this outside the normal +// expression id space so it cannot be confused with a real UnknownExpr id. +inline constexpr int kMissingExpressionSentinel = -1000001; +inline constexpr int64_t kUnknownContentSentinel = -444; + +enum class DExprKind { + kUnknown, + kConstant, + kVariable, + kAdd, + kSub, + kMul, + kDiv, + kMax, + kGt, + kSelect, +}; + +class DynExpr { + public: + virtual ~DynExpr() = default; + virtual std::unique_ptr clone() const = 0; + virtual DExprKind kind() const = 0; + virtual void print(xla::Printer* printer) const = 0; + virtual void to_proto(xla::ExpressionProto* proto) const = 0; + virtual bool is_constant() const = 0; + virtual int64_t get_val() const { return -1; } + virtual DynExpr* s() = 0; // simplify + virtual DynExpr* substitute(int id, DynExpr* v) = 0; + virtual std::set get_all_ids() = 0; + virtual std::optional solve(int64_t x) = 0; + + bool is_dynamic() { return !is_constant(); } + + static DynExpr* zero; + static DynExpr* one; + static DynExpr* _(int64_t val); + static DynExpr* V(int var_id); + static DynExpr* _s(DynExpr* expr); + static bool equal(DynExpr* expr1, DynExpr* expr2); + + friend std::ostream& operator<<(std::ostream& os, DynExpr* expr); +}; + +class DExpr { + public: + using Kind = DExprKind; + + DExpr() = default; + explicit DExpr(std::unique_ptr expr) : expr_(std::move(expr)) {} + + DExpr(const DExpr& other) { + if (other.expr_ != nullptr) { + expr_ = other.expr_->clone(); + } + } + DExpr& operator=(const DExpr& other) { + if (this == &other) return *this; + expr_.reset(); + if (other.expr_ != nullptr) { + expr_ = other.expr_->clone(); + } + return *this; + } + + DExpr(DExpr&&) noexcept = default; + DExpr& operator=(DExpr&&) noexcept = default; + + static DExpr Unknown(int id = 0); + static DExpr Adopt(DynExpr* expr) { return DExpr(std::unique_ptr(expr)); } + static DExpr Const(int64_t value) { return Adopt(DynExpr::_(value)); } + static DExpr Var(int var_id) { return Adopt(DynExpr::V(var_id)); } + static DExpr Max(const DExpr& lhs, const DExpr& rhs); + static DExpr Gt(const DExpr& lhs, const DExpr& rhs); + static DExpr Select(const DExpr& pred, const DExpr& on_true, + const DExpr& on_false); + bool is_unknown() const { + return expr_ != nullptr && expr_->kind() == DExprKind::kUnknown; + } + Kind kind() const { + CHECK(expr_ != nullptr) << "Attempted to access empty DExpr"; + return expr_->kind(); + } + + DynExpr* get() const { + CHECK(expr_ != nullptr) << "Attempted to access empty DExpr"; + return expr_.get(); + } + DynExpr& operator*() const { + CHECK(expr_ != nullptr) << "Attempted to dereference empty DExpr"; + return *expr_; + } + DynExpr* operator->() const { + CHECK(expr_ != nullptr) << "Attempted to access empty DExpr"; + return expr_.get(); + } + operator DynExpr*() const { return get(); } + explicit operator bool() const { return expr_ != nullptr && !is_unknown(); } + + std::unique_ptr clone() const { + if (expr_ == nullptr) { + return nullptr; + } + return expr_->clone(); + } + DynExpr* release() { return expr_.release(); } + + DExpr simplify() const { + return expr_ == nullptr ? DExpr() : Adopt(expr_->s()); + } + void to_proto(xla::ExpressionProto* proto) const { + CHECK(expr_ != nullptr) << "Attempted to serialize empty DExpr"; + expr_->to_proto(proto); + } + 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) { + xla::ExpressionProto proto; + if (expr.expr_ != nullptr) { + expr.expr_->to_proto(&proto); + } + return H::combine(std::move(h), proto.SerializeAsString()); + } + + private: + std::unique_ptr expr_; +}; + +class UnknownExpr : public DynExpr { + int id_; + + public: + explicit UnknownExpr(int id = 0) : id_(id) {} + std::unique_ptr clone() const override { + return std::make_unique(id_); + } + DExprKind kind() const override { return DExprKind::kUnknown; } + void print(xla::Printer* printer) const override { + if (id_ == kMissingExpressionSentinel || id_ == kUnknownContentSentinel) { + printer->Append("_"); + return; + } + printer->Append("?"); + if (id_ != 0) { + printer->Append(id_); + } + } + void to_proto(xla::ExpressionProto* proto) const override { + (void)proto; + } + bool is_constant() const override { return true; } + int get_id() const { return id_; } + DynExpr* substitute(int id, DynExpr* v) override { + (void)id; + (void)v; + return clone().release(); + } + std::set get_all_ids() override { return {}; } + std::optional solve(int64_t) override { + return std::nullopt; + } + DynExpr* s() override { return clone().release(); } +}; + +inline DExpr DExpr::Unknown(int id) { + return DExpr(std::unique_ptr(new xla::UnknownExpr(id))); +} + +// constant i +class Constant : public DynExpr { + int64_t value; + + public: + explicit Constant(int64_t v) : value(v) {} + std::unique_ptr clone() const override { + return std::make_unique(value); + } + DExprKind kind() const override { return DExprKind::kConstant; } + void print(xla::Printer* printer) const override { + if (value < 0) { + printer->Append("("); + } + printer->Append(value); + if (value < 0) { + printer->Append(")"); + } + } + void to_proto(xla::ExpressionProto* proto) const override { + proto->set_constant_value(value); + } + bool is_constant() const override { return true; } + int64_t get_val() const override { return value; } + DynExpr* substitute(int id, DynExpr* v) { return clone().release(); } + std::set get_all_ids() { return {}; } + std::optional solve(int64_t x) { return std::nullopt; } + DynExpr* s() override; +}; + +// Root variables represent strictly positive dynamic dimensions. Potentially +// signed values must be represented by expressions derived from these roots. +class Variable : public DynExpr { + int id; + + public: + explicit Variable(int identifier) : id(identifier) {} + std::unique_ptr clone() const override { + return std::make_unique(id); + } + DExprKind kind() const override { return DExprKind::kVariable; } + void print(xla::Printer* printer) const override { + if (id >= 1 && id <= 26) { + char letter = 'A' + (id - 1); + printer->Append(std::string(1, letter)); + return; + } + printer->Append("V("); + printer->Append(id); + printer->Append(")"); + } + void to_proto(xla::ExpressionProto* proto) const override { + proto->set_variable_id(id); + } + bool is_constant() const override { return false; } + int get_id() const { return id; } + DynExpr* substitute(int id, DynExpr* v) { + return get_id() == id ? v->clone().release() : clone().release(); + } + std::set get_all_ids() { return {get_id()}; } + std::optional solve(int64_t x) { return x; } + DynExpr* s() override; +}; + +// exp = exp + exp +class Add : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + Add(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique(lhs->clone().release(), rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kAdd; } + void print(xla::Printer* printer) const override { + printer->Append("("); + lhs->print(printer); + printer->Append(" + "); + rhs->print(printer); + printer->Append(")"); + } + + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + + void to_proto(xla::ExpressionProto* proto) const override { + auto* add_msg = proto->mutable_add_node(); + lhs->to_proto(add_msg->mutable_lhs()); + rhs->to_proto(add_msg->mutable_rhs()); + } + + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant(); + } + + int64_t get_val() const override { return lhs->get_val() + rhs->get_val(); } + + DynExpr* substitute(int id, DynExpr* v) { + return new Add(lhs->substitute(id, v), rhs->substitute(id, v)); + } + + std::set get_all_ids() { + auto s = lhs->get_all_ids(); + s.merge(rhs->get_all_ids()); + return s; + } + + std::optional solve(int64_t x) { + // Cannot solve if both lhs and rhs are dynamic... + if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt; + if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) { + // (A + c) = x <=> A = x - c => solve A = y with y = x - c + return lhs->solve(x - rhs->get_val()); + } + if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) { + // (c + A) = x <=> A = x - c => solve A = y with y = x - c + return rhs->solve(x - lhs->get_val()); + } + // No solution + return std::nullopt; + } + + DynExpr* s() override; + + ~Add() override = default; +}; + +// exp = exp - exp +class Sub : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + Sub(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique(lhs->clone().release(), rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kSub; } + void print(xla::Printer* printer) const override { + printer->Append("("); + lhs->print(printer); + printer->Append(" - "); + rhs->print(printer); + printer->Append(")"); + } + + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + + void to_proto(xla::ExpressionProto* proto) const override { + auto* sub_msg = proto->mutable_sub_node(); + lhs->to_proto(sub_msg->mutable_lhs()); + rhs->to_proto(sub_msg->mutable_rhs()); + } + + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant(); + } + + int64_t get_val() const override { return lhs->get_val() - rhs->get_val(); } + + DynExpr* substitute(int id, DynExpr* v) { + return new Sub(lhs->substitute(id, v), rhs->substitute(id, v)); + } + + std::set get_all_ids() { + auto s = lhs->get_all_ids(); + s.merge(rhs->get_all_ids()); + return s; + } + + std::optional solve(int64_t x) { + // Cannot solve if both lhs and rhs are dynamic... + if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt; + if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) { + // (A - c) = x <=> A = x + c => solve A = y with y = x + c + return lhs->solve(x + rhs->get_val()); + } + if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) { + // (c - A) = x <=> A = c - x => solve A = y with y = c - x + return rhs->solve(lhs->get_val() - x); + } + // No solution + return std::nullopt; + } + + DynExpr* s() override; + + ~Sub() override = default; +}; + +// exp = exp * exp +class Mul : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + Mul(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique(lhs->clone().release(), rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kMul; } + void print(xla::Printer* printer) const override { + printer->Append("("); + lhs->print(printer); + printer->Append(" * "); + rhs->print(printer); + printer->Append(")"); + } + + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + + void to_proto(xla::ExpressionProto* proto) const override { + auto* mul_msg = proto->mutable_mul_node(); + lhs->to_proto(mul_msg->mutable_lhs()); + rhs->to_proto(mul_msg->mutable_rhs()); + } + + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant(); + } + + int64_t get_val() const override { return lhs->get_val() * rhs->get_val(); } + + DynExpr* substitute(int id, DynExpr* v) { + return new Mul(lhs->substitute(id, v), rhs->substitute(id, v)); + } + + std::set get_all_ids() { + auto s = lhs->get_all_ids(); + s.merge(rhs->get_all_ids()); + return s; + } + + std::optional solve(int64_t x) { + // Cannot solve if both lhs and rhs are dynamic... + if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt; + if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) { + // (A * c) = x <=> A = x / c => solve A = y with y = x / c + int64_t c = rhs->get_val(); + if (c == 0) return x == 0 ? lhs->solve(0) : std::nullopt; + if (x % c != 0) return std::nullopt; + return lhs->solve(x / c); + } + if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) { + // (c * A) = x <=> A = x / c => solve A = y with y = x / c + int64_t c = lhs->get_val(); + if (c == 0) return x == 0 ? rhs->solve(0) : std::nullopt; + if (x % c != 0) return std::nullopt; + return rhs->solve(x / c); + } + // No solution + return std::nullopt; + } + + DynExpr* s() override; + + ~Mul() override = default; +}; + +// expr / expr +class Div : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + Div(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique
(lhs->clone().release(), rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kDiv; } + void print(xla::Printer* printer) const override { + printer->Append("("); + lhs->print(printer); + printer->Append(" / "); + rhs->print(printer); + printer->Append(")"); + } + + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + + void to_proto(xla::ExpressionProto* proto) const override { + auto* div_msg = proto->mutable_div_node(); + lhs->to_proto(div_msg->mutable_lhs()); + rhs->to_proto(div_msg->mutable_rhs()); + } + + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant() && rhs->get_val() != 0; + } + + int64_t get_val() const override { + CHECK(is_constant()) + << "Attempted to evaluate a non-constant or zero-divisor expression"; + return lhs->get_val() / rhs->get_val(); + } + + DynExpr* substitute(int id, DynExpr* v) { + return new Div(lhs->substitute(id, v), rhs->substitute(id, v)); + } + + DynExpr* s() override; + + std::set get_all_ids() { + auto s = lhs->get_all_ids(); + s.merge(rhs->get_all_ids()); + return s; + } + + std::optional solve(int64_t x) { + // Cannot solve if both lhs and rhs are dynamic... + if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt; + if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) { + // (A / c) = x <=> A = x * c => solve A = y with y = x * c + const int64_t divisor = rhs->get_val(); + if (divisor == 0) return std::nullopt; + return lhs->solve(x * divisor); + } + if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) { + // (c / A) = x <=> A = c / x => solve A = y with y = c / x + int64_t c = lhs->get_val(); + if (x == 0) return std::nullopt; + if (c % x != 0) return std::nullopt; + return rhs->solve(c / x); + } + // No solution + return std::nullopt; + } + + ~Div() override = default; +}; + +// max(lhs, rhs) +class MaxExpr : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + MaxExpr(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique(lhs->clone().release(), + rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kMax; } + void print(xla::Printer* printer) const override { + printer->Append("max("); + lhs->print(printer); + printer->Append(", "); + rhs->print(printer); + printer->Append(")"); + } + void to_proto(xla::ExpressionProto* proto) const override { + auto* max_msg = proto->mutable_max_node(); + lhs->to_proto(max_msg->mutable_lhs()); + rhs->to_proto(max_msg->mutable_rhs()); + } + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant(); + } + int64_t get_val() const override { + return std::max(lhs->get_val(), rhs->get_val()); + } + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + DynExpr* substitute(int id, DynExpr* v) override { + return new MaxExpr(lhs->substitute(id, v), rhs->substitute(id, v)); + } + std::set get_all_ids() override { + auto ids = lhs->get_all_ids(); + ids.merge(rhs->get_all_ids()); + return ids; + } + // Max is not invertible: either operand may have produced the result. + std::optional solve(int64_t x) override { + StringPrinter printer; + print(&printer); + LOG(WARNING) << "Cannot solve Max dynamic shape expression for value " << x + << ": " << std::move(printer).ToString(); + return std::nullopt; + } + DynExpr* s() override; +}; + +class GtExpr : public DynExpr { + std::unique_ptr lhs; + std::unique_ptr rhs; + + public: + GtExpr(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {} + std::unique_ptr clone() const override { + return std::make_unique(lhs->clone().release(), + rhs->clone().release()); + } + DExprKind kind() const override { return DExprKind::kGt; } + void print(xla::Printer* printer) const override { + printer->Append("("); + lhs->print(printer); + printer->Append(" > "); + rhs->print(printer); + printer->Append(")"); + } + void to_proto(xla::ExpressionProto* proto) const override { + auto* gt_msg = proto->mutable_gt_node(); + lhs->to_proto(gt_msg->mutable_lhs()); + rhs->to_proto(gt_msg->mutable_rhs()); + } + bool is_constant() const override { + return lhs->is_constant() && rhs->is_constant(); + } + int64_t get_val() const override { return lhs->get_val() > rhs->get_val(); } + DynExpr* get_lhs() const { return lhs.get(); } + DynExpr* get_rhs() const { return rhs.get(); } + DynExpr* substitute(int id, DynExpr* v) override { + return new GtExpr(lhs->substitute(id, v), rhs->substitute(id, v)); + } + std::set get_all_ids() override { + auto ids = lhs->get_all_ids(); + ids.merge(rhs->get_all_ids()); + return ids; + } + std::optional solve(int64_t x) override { + StringPrinter printer; + print(&printer); + LOG(WARNING) << "Cannot solve Gt dynamic shape expression for value " << x + << ": " << std::move(printer).ToString(); + return std::nullopt; + } + DynExpr* s() override; +}; + +class SelectExpr : public DynExpr { + std::unique_ptr pred; + std::unique_ptr on_true; + std::unique_ptr on_false; + + public: + SelectExpr(DynExpr* p, DynExpr* t, DynExpr* f) + : pred(p), on_true(t), on_false(f) {} + std::unique_ptr clone() const override { + return std::make_unique(pred->clone().release(), + on_true->clone().release(), + on_false->clone().release()); + } + DExprKind kind() const override { return DExprKind::kSelect; } + void print(xla::Printer* printer) const override { + printer->Append("select("); + pred->print(printer); + printer->Append(", "); + on_true->print(printer); + printer->Append(", "); + on_false->print(printer); + printer->Append(")"); + } + void to_proto(xla::ExpressionProto* proto) const override { + auto* select_msg = proto->mutable_select_node(); + pred->to_proto(select_msg->mutable_pred()); + on_true->to_proto(select_msg->mutable_on_true()); + on_false->to_proto(select_msg->mutable_on_false()); + } + bool is_constant() const override { + return pred->is_constant() && on_true->is_constant() && + on_false->is_constant(); + } + int64_t get_val() const override { + return pred->get_val() != 0 ? on_true->get_val() : on_false->get_val(); + } + DynExpr* get_pred() const { return pred.get(); } + DynExpr* get_on_true() const { return on_true.get(); } + DynExpr* get_on_false() const { return on_false.get(); } + DynExpr* substitute(int id, DynExpr* v) override { + return new SelectExpr(pred->substitute(id, v), on_true->substitute(id, v), + on_false->substitute(id, v)); + } + std::set get_all_ids() override { + auto ids = pred->get_all_ids(); + ids.merge(on_true->get_all_ids()); + ids.merge(on_false->get_all_ids()); + return ids; + } + std::optional solve(int64_t x) override { + StringPrinter printer; + print(&printer); + LOG(WARNING) << "Cannot solve Select dynamic shape expression for value " + << x << ": " << std::move(printer).ToString(); + return std::nullopt; + } + DynExpr* s() override; +}; + +DynExpr* operator*(DynExpr& lhs, DynExpr& rhs); +DynExpr* operator*(int64_t k, DynExpr& rhs); +DynExpr* operator*(DynExpr& lhs, int64_t k); +DynExpr* operator/(DynExpr& lhs, DynExpr& rhs); +DynExpr* operator/(DynExpr& lhs, int64_t d); +DynExpr* operator+(DynExpr& lhs, DynExpr& rhs); +DynExpr* operator+(DynExpr& lhs, int64_t d); +DynExpr* operator-(DynExpr& lhs, DynExpr& rhs); +DynExpr* operator-(DynExpr& lhs, int64_t d); +bool operator==(DynExpr& lhs, DynExpr& rhs); +bool operator==(DynExpr& lhs, int64_t d); + +inline DExpr operator*(const DExpr& lhs, const DExpr& rhs) { + return DExpr::Adopt(*lhs.get() * *rhs.get()); +} +inline DExpr operator*(int64_t lhs, const DExpr& rhs) { + return DExpr::Adopt(lhs * *rhs.get()); +} +inline DExpr operator*(const DExpr& lhs, int64_t rhs) { + return DExpr::Adopt(*lhs.get() * rhs); +} +inline DExpr operator/(const DExpr& lhs, const DExpr& rhs) { + return DExpr::Adopt(*lhs.get() / *rhs.get()); +} +inline DExpr operator/(const DExpr& lhs, int64_t rhs) { + return DExpr::Adopt(*lhs.get() / rhs); +} +inline DExpr operator+(const DExpr& lhs, const DExpr& rhs) { + return DExpr::Adopt(*lhs.get() + *rhs.get()); +} +inline DExpr operator+(const DExpr& lhs, int64_t rhs) { + return DExpr::Adopt(*lhs.get() + rhs); +} +inline DExpr operator-(const DExpr& lhs, const DExpr& rhs) { + return DExpr::Adopt(*lhs.get() - *rhs.get()); +} +inline DExpr operator-(const DExpr& lhs, int64_t rhs) { + return DExpr::Adopt(*lhs.get() - rhs); +} +inline bool operator==(const DExpr& lhs, const DExpr& rhs) { + return *lhs.get() == *rhs.get(); +} +inline bool operator==(const DExpr& lhs, int64_t rhs) { + return *lhs.get() == rhs; +} + +inline DExpr DExprFromProto(const xla::ExpressionProto& proto) { + switch (proto.node_type_case()) { + case ExpressionProto::kConstantValue: + return DExpr::Const(proto.constant_value()); + case ExpressionProto::kVariableId: + return DExpr::Var(proto.variable_id()); + case ExpressionProto::kAddNode: { + const auto& add = proto.add_node(); + return DExprFromProto(add.lhs()) + DExprFromProto(add.rhs()); + } + case ExpressionProto::kSubNode: { + const auto& sub = proto.sub_node(); + return DExprFromProto(sub.lhs()) - DExprFromProto(sub.rhs()); + } + case ExpressionProto::kMulNode: { + const auto& mul = proto.mul_node(); + return DExprFromProto(mul.lhs()) * DExprFromProto(mul.rhs()); + } + case ExpressionProto::kDivNode: { + const auto& div = proto.div_node(); + return DExprFromProto(div.lhs()) / DExprFromProto(div.rhs()); + } + case ExpressionProto::kMaxNode: { + const auto& max = proto.max_node(); + return DExpr::Max(DExprFromProto(max.lhs()), + DExprFromProto(max.rhs())); + } + case ExpressionProto::kGtNode: { + const auto& gt = proto.gt_node(); + return DExpr::Gt(DExprFromProto(gt.lhs()), DExprFromProto(gt.rhs())); + } + case ExpressionProto::kSelectNode: { + const auto& select = proto.select_node(); + return DExpr::Select(DExprFromProto(select.pred()), + DExprFromProto(select.on_true()), + DExprFromProto(select.on_false())); + } + case ExpressionProto::NODE_TYPE_NOT_SET: + default: + return DExpr::Unknown(kMissingExpressionSentinel); + } +} + +inline DynExpr* DynExpr::_(int64_t val) { + return new Constant(val); +} +inline DynExpr* DynExpr::V(int var_id) { return new Variable(var_id); } + +} // namespace xla + +#endif // XLA_SHAPE_EXPR_H_ diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc index 4f92ce19adb1f9..75970dbeceb325 100644 --- a/third_party/xla/xla/shape_util.cc +++ b/third_party/xla/xla/shape_util.cc @@ -123,6 +123,7 @@ void PrintBufferShape(Printer* printer, const Shape& shape) { // its Layout. absl::StatusOr MakeShapeWithLayoutInternal( PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions, absl::Span minor_to_major, absl::Span tiles, int64_t tail_padding_alignment_in_elements, PrimitiveType index_primitive_type, PrimitiveType pointer_primitive_type, @@ -139,7 +140,8 @@ absl::StatusOr MakeShapeWithLayoutInternal( PrimitiveType_Name(element_type)); } TF_ASSIGN_OR_RETURN(Shape shape, - ShapeUtil::MakeValidatedShape(element_type, dimensions)); + ShapeUtil::MakeValidatedShape(element_type, dimensions, + expressions)); if (element_size_in_bits == ShapeUtil::ByteSizeOfPrimitiveType(element_type) * 8) { // Only set element_size_in_bits if it's different from the default value. @@ -267,9 +269,20 @@ static std::vector MakeDynamicDimensions( return dynamic_dimensions; } +static std::vector MakeExpressions( + absl::Span dimensions) { + std::vector expressions; + expressions.reserve(dimensions.size()); + for (int64_t d : dimensions) { + expressions.push_back(DExpr::Const(d)); + } + return expressions; +} + /* static */ Shape ShapeUtil::MakeShape(PrimitiveType element_type, - absl::Span dimensions) { - return MakeValidatedShape(element_type, dimensions).value(); + absl::Span dimensions, + absl::Span expressions) { + return MakeValidatedShape(element_type, dimensions, expressions).value(); } /* static */ Shape ShapeUtil::MakeScalarShape(PrimitiveType element_type) { @@ -278,8 +291,10 @@ static std::vector MakeDynamicDimensions( /* static */ Shape ShapeUtil::MakeShape( PrimitiveType element_type, absl::Span dimensions, - const std::vector& dynamic_dimensions) { - return MakeValidatedShape(element_type, dimensions, dynamic_dimensions) + const std::vector& dynamic_dimensions, + absl::Span expressions) { + return MakeValidatedShape(element_type, dimensions, dynamic_dimensions, + expressions) .value(); } @@ -296,20 +311,37 @@ static std::vector MakeDynamicDimensions( } /* static */ absl::StatusOr ShapeUtil::MakeValidatedShape( - PrimitiveType element_type, absl::Span dimensions) { - return MakeValidatedShape(element_type, dimensions, - MakeDynamicDimensions(dimensions)); + PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions) { + if (expressions.empty() && !dimensions.empty()) { + std::vector filled_expressions = MakeExpressions(dimensions); + return MakeValidatedShape(element_type, dimensions, + MakeDynamicDimensions(dimensions), + filled_expressions); + } + return MakeValidatedShape( + element_type, dimensions, MakeDynamicDimensions(dimensions), expressions); } /* static */ absl::StatusOr ShapeUtil::MakeValidatedShape( PrimitiveType element_type, absl::Span dimensions, - const std::vector& dynamic_dimensions) { + const std::vector& dynamic_dimensions, + absl::Span expressions) { + std::vector filled_expressions; + if (expressions.empty() && !dimensions.empty()) { + filled_expressions = MakeExpressions(dimensions); + expressions = filled_expressions; + } if (dynamic_dimensions.size() != dimensions.size()) { return InvalidArgument( "dynamic dimensions size %d did not match number of dimensions %d", dynamic_dimensions.size(), dimensions.size()); } - + if (expressions.size() != dimensions.size()) { + return InvalidArgument( + "expressions size %d did not match number of dimensions %d", + expressions.size(), dimensions.size()); + } Shape shape; int64_t dense_shape_size = primitive_util::IsArrayType(element_type) ? primitive_util::ByteWidth(element_type) @@ -328,6 +360,7 @@ static std::vector MakeDynamicDimensions( for (int i = 0; i < ndims; i++) { const int64_t d = dimensions[i]; const bool is_dynamic = dynamic_dimensions[i]; + const DExpr& expression = expressions[i]; if (!Shape::IsValidDimensionSize(d, is_dynamic)) { return InvalidArgument("Invalid dimension size %d, is_dynamic=%s", d, is_dynamic ? "true" : "false"); @@ -339,7 +372,7 @@ static std::vector MakeDynamicDimensions( any_overflows |= overflow; } - shape.add_dimensions(d, is_dynamic); + shape.add_dimensions(d, is_dynamic, expression); minor_to_major->push_back(ndims - 1 - i); } @@ -352,11 +385,29 @@ static std::vector MakeDynamicDimensions( /* static */ Shape ShapeUtil::MakeShapeWithDenseLayout( PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions, absl::Span minor_to_major, absl::Span tiles, int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits, int64_t memory_space, absl::Span split_configs) { auto ret = MakeShapeWithLayoutInternal( - element_type, dimensions, minor_to_major, tiles, + element_type, dimensions, expressions, minor_to_major, tiles, + tail_padding_alignment_in_elements, + /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID, + /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, element_size_in_bits, + memory_space, split_configs, + /*physical_shape=*/std::nullopt); + TF_CHECK_OK(ret.status()); + return *ret; +} + +/* static */ Shape ShapeUtil::MakeShapeWithDenseLayout( + PrimitiveType element_type, absl::Span dimensions, + absl::Span minor_to_major, absl::Span tiles, + int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits, + int64_t memory_space, absl::Span split_configs) { + auto ret = MakeShapeWithLayoutInternal( + element_type, dimensions, MakeExpressions(dimensions), minor_to_major, + tiles, tail_padding_alignment_in_elements, /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID, /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, element_size_in_bits, @@ -373,7 +424,7 @@ static std::vector MakeDynamicDimensions( int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits, int64_t memory_space, std::optional physical_shape) { auto ret = MakeShapeWithLayoutInternal( - element_type, dimensions, minor_to_major, + element_type, dimensions, MakeExpressions(dimensions), minor_to_major, /*tiles=*/{}, tail_padding_alignment_in_elements, index_primitive_type, pointer_primitive_type, element_size_in_bits, memory_space, /*split_configs=*/{}, std::move(physical_shape)); @@ -408,10 +459,11 @@ static std::vector MakeDynamicDimensions( } /* static */ Shape ShapeUtil::MakeShapeWithDescendingLayout( - PrimitiveType element_type, absl::Span dimensions) { - std::vector layout(dimensions.size()); - std::iota(layout.rbegin(), layout.rend(), static_cast(0)); - return MakeShapeWithDenseLayout(element_type, dimensions, layout); + PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions) { + return MakeShapeWithDenseLayout( + element_type, dimensions, expressions, + LayoutUtil::MakeDescendingLayout(dimensions.size()).minor_to_major()); } /* static */ Shape @@ -425,7 +477,17 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } dims[i] = shape.dimensions(dim); } - Shape new_shape = MakeShapeWithDescendingLayout(shape.element_type(), dims); + std::vector expressions; + expressions.reserve(shape.dimensions().size()); + for (int i = 0; i < shape.dimensions().size(); ++i) { + int dim = i; + if (shape.has_layout()) { + dim = LayoutUtil::Major(shape.layout(), dim); + } + expressions.push_back(shape.expressions(dim)); + } + Shape new_shape = MakeShapeWithDescendingLayout(shape.element_type(), dims, + expressions); // Since the physical layout is kept the same, the tiles and element size are // the same also. if (shape.has_layout()) { @@ -442,6 +504,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( dim = LayoutUtil::Major(shape.layout(), dim); } new_shape.set_dynamic_dimension(i, shape.is_dynamic_dimension(dim)); + new_shape.set_expression(i, shape.expressions(dim)); } new_shape.mutable_layout()->set_memory_space(shape.layout().memory_space()); return new_shape; @@ -731,7 +794,29 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { printer->Append("?"); } } else { + // Only print constant expression if it is different than the dimension + // (i.e. it is wrong!) + const DExpr& expr = shape.expressions(i); + bool is_wrong = expr && expr->is_constant() && + expr->get_val() != shape.dimensions(i); printer->Append(shape.dimensions(i)); + if (is_wrong) { + DExpr simplified_expr = expr.simplify(); + xla::StringPrinter expr_printer; + simplified_expr->print(&expr_printer); + LOG(ERROR) << "Mismatched static shape expression at dim " << i + << ": dim=" << shape.dimensions(i) + << ", expr=" << std::move(expr_printer).ToString(); + printer->Append("print(printer); + printer->Append("!>"); + } + if (expr && expr->is_dynamic()) { + DExpr simplified_expr = expr.simplify(); + printer->Append("<"); + simplified_expr->print(printer); + printer->Append(">"); + } } }; print_dimension(0); @@ -753,6 +838,11 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { return; } PrintHumanString(printer, shape); + if (shape.outer_multiplier() > 0) { + printer->Append("(bm="); + printer->Append(shape.outer_multiplier()); + printer->Append(")"); + } if (!shape.IsArray()) return; if (!shape.has_layout()) return; if (IsScalar(shape)) { @@ -811,6 +901,10 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { const Shape& rhs) { if (!SameRank(lhs, rhs)) return false; for (int i = 0; i < lhs.dimensions().size(); ++i) { + if (i == 0 && (lhs.outer_multiplier() > 0 || rhs.outer_multiplier() > 0)) { + VLOG(3) << "CompareShapes: batch dimension found. Forcely compatible"; + continue; + } if (!lhs.is_unbounded_dynamic_dimension(i) && !rhs.is_unbounded_dynamic_dimension(i) && lhs.dimensions(i) != rhs.dimensions(i)) { @@ -826,7 +920,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { } /* static */ bool ShapeUtil::Compatible(const Shape& lhs, const Shape& rhs) { - return Shape::Equal().IgnoreDynamicDimension().IgnoreLayout()(lhs, rhs); + return Shape::Equal().IgnoreDynamicDimension().IgnoreLayout().IgnoreBatch()(lhs, rhs); } /* static */ bool ShapeUtil::CompatibleIgnoringElementType(const Shape& lhs, @@ -871,6 +965,11 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { return shape.dimensions(GetDimensionNumber(shape, dimension_number)); } +/* static */ const DExpr& ShapeUtil::GetExpression( + const Shape& shape, int64_t dimension_number) { + return shape.expressions(GetDimensionNumber(shape, dimension_number)); +} + /* static */ int64_t ShapeUtil::GetDimensionNumber(const Shape& shape, int64_t dimension_number) { if (dimension_number < 0) { @@ -1207,8 +1306,11 @@ ShapeUtil::PackedFactorFor1DInterleavedArray(const Shape& shape) { const auto permuted_dims = Permute(shape.dimensions(), permutation); const auto permuted_dynamic_dims = Permute(shape.dynamic_dimensions(), permutation); + const auto permuted_expressions = + Permute(shape.expressions(), permutation); for (int i = 0; i < permuted_dims.size(); ++i) { - new_shape.add_dimensions(permuted_dims[i], permuted_dynamic_dims[i]); + new_shape.add_dimensions(permuted_dims[i], permuted_dynamic_dims[i], + permuted_expressions[i]); } // If `shape` has a layout, by contract we choose a new layout such that the @@ -1763,7 +1865,8 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape, } Shape output_shape_with_layout = MakeShapeWithDenseLayout( - output_shape.element_type(), output_shape.dimensions(), output_layout); + output_shape.element_type(), output_shape.dimensions(), + output_shape.expressions(), output_layout); CHECK(ReshapeIsBitcast(input_shape, output_shape_with_layout)) << "reshape is not a bitcast for input_shape: " << ShapeUtil::HumanStringWithLayout(input_shape) @@ -1996,7 +2099,8 @@ struct ParallelState { } // Create the shape of the "work" which has same layout as the original shape. - Shape work_shape = ShapeUtil::MakeShape(shape.element_type(), work_dims); + Shape work_shape = ShapeUtil::MakeShape(shape.element_type(), work_dims, + shape.expressions()); *work_shape.mutable_layout() = shape.layout(); // We target one task (partition) per available thread. diff --git a/third_party/xla/xla/shape_util.h b/third_party/xla/xla/shape_util.h index 4e3602671fa153..0948c6c31df2bc 100644 --- a/third_party/xla/xla/shape_util.h +++ b/third_party/xla/xla/shape_util.h @@ -330,6 +330,11 @@ class ShapeUtil { // GetDimensionNumber(dimension_number). static int64_t GetDimension(const Shape& shape, int64_t dimension_number); + // Extracts the shape's expressions at dimension number + // GetDimensionNumber(dimension_number). + static const DExpr& GetExpression(const Shape& shape, + int64_t dimension_number); + // Resolves a dimension number, supporting negative indexing. // // Negative indexing has similar semantics to Python. For an N-dimensional @@ -406,7 +411,8 @@ class ShapeUtil { // Constructs a new shape with the given element type and sequence of // dimensions. static Shape MakeShape(PrimitiveType element_type, - absl::Span dimensions); + absl::Span dimensions, + absl::Span expressions = {}); // Make a scalar shape with given primitive type. static Shape MakeScalarShape(PrimitiveType element_type); @@ -419,7 +425,8 @@ class ShapeUtil { // the same size. static Shape MakeShape(PrimitiveType element_type, absl::Span dimensions, - const std::vector& dynamic_dimensions); + const std::vector& dynamic_dimensions, + absl::Span expressions = {}); // Constructs a new buffer shape with the given element type, and sequence of // dimensions. static Shape MakeBufferShape(PrimitiveType element_type, @@ -430,10 +437,13 @@ class ShapeUtil { // size fits in std::numeric_limits::max(), and dynamic size is not // marked static. static absl::StatusOr MakeValidatedShape( - PrimitiveType element_type, absl::Span dimensions); + PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions = {}); + static absl::StatusOr MakeValidatedShape( PrimitiveType element_type, absl::Span dimensions, - const std::vector& dynamic_dimensions); + const std::vector& dynamic_dimensions, + absl::Span expressions = {}); // Creates a Shape with element type corresponding to T and the given // dimensions @@ -443,6 +453,17 @@ class ShapeUtil { dimensions); } + // Constructs a new dense array shape with the given minor_to_major order in + // its Layout. Returns a value shape such that shape.has_layout(). + static Shape MakeShapeWithDenseLayout( + PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions, + absl::Span minor_to_major, + absl::Span tiles = {}, + int64_t tail_padding_alignment_in_elements = 1, + int64_t element_size_in_bits = 0, int64_t memory_space = 0, + absl::Span split_configs = {}); + // Constructs a new dense array shape with the given minor_to_major order in // its Layout. Returns a value shape such that shape.has_layout(). static Shape MakeShapeWithDenseLayout( @@ -476,7 +497,8 @@ class ShapeUtil { // Constructs a new shape with major-first layout (i.e. {n, n-1, ..., 0}). static Shape MakeShapeWithDescendingLayout( - PrimitiveType element_type, absl::Span dimensions); + PrimitiveType element_type, absl::Span dimensions, + absl::Span expressions = {}); // Returns a new Shape based on the given Shape with low-dimension-major // layout (i.e. {n, n-1, ..., 0}, like Fortran), and with the dimensions diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto index 98462a4f6eb83c..169e82bef57aed 100644 --- a/third_party/xla/xla/xla_data.proto +++ b/third_party/xla/xla/xla_data.proto @@ -361,6 +361,8 @@ message ShapeProto { // The layout used to back this shape. LayoutProto layout = 5; + repeated ExpressionProto expressions = 7; + // Important: if any field is added, be sure to modify ShapeUtil::Equal(), // ShapeUtil::Compatible() and Shape::Hash() appropriately to account for the // new field. @@ -1204,3 +1206,53 @@ message OriginalArrayProto { message OriginalValueProto { repeated OriginalArrayProto leaves = 1; } + +message ExpressionProto { + oneof node_type { + int32 constant_value = 1; // cons + int32 variable_id = 2; // var + AddNode add_node = 3; // exp + exp + SubNode sub_node = 4; // exp - exp + MulNode mul_node = 5; // exp * exp + DivNode div_node = 6; // exp / exp + MaxNode max_node = 7; // max(exp, exp) + GtNode gt_node = 8; // exp > exp + SelectNode select_node = 9; // select(pred, on_true, on_false) + } +} + +message AddNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message SubNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message MulNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message DivNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message MaxNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message GtNode { + ExpressionProto lhs = 1; + ExpressionProto rhs = 2; +} + +message SelectNode { + ExpressionProto pred = 1; + ExpressionProto on_true = 2; + ExpressionProto on_false = 3; +} From 5ee2428e05a8215e111528f0e9699c6cb8496135 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:03:18 +0100 Subject: [PATCH 03/21] Propagate symbolic shapes and values through TensorFlow inference Extend TensorFlow shape inference to preserve symbolic dimension identities and shape-derived tensor values. Keep unknown physical sizes separate from their symbolic expressions. --- tensorflow/compiler/jit/BUILD | 1 + tensorflow/compiler/jit/shape_inference.cc | 22 +- tensorflow/core/framework/common_shape_fns.cc | 8 +- tensorflow/core/framework/shape_inference.cc | 238 +++++-- tensorflow/core/framework/shape_inference.h | 42 +- .../core/framework/shape_inference_test.cc | 29 + .../core/grappler/costs/graph_properties.cc | 595 +++++++++++++++++- .../core/grappler/costs/graph_properties.h | 3 +- .../grappler/costs/graph_properties_test.cc | 293 ++++++++- .../core/grappler/optimizers/remapper.cc | 14 + tensorflow/core/ops/array_ops.cc | 29 +- tensorflow/core/util/strided_slice_op.cc | 115 +++- tensorflow/core/util/strided_slice_op.h | 5 + 13 files changed, 1284 insertions(+), 110 deletions(-) diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 39f93d17aa2932..951a50fd1ecfb6 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -1006,6 +1006,7 @@ cc_library( hdrs = ["shape_inference.h"], visibility = [":friends"], deps = [ + ":flags", ":shape_inference_helpers", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", diff --git a/tensorflow/compiler/jit/shape_inference.cc b/tensorflow/compiler/jit/shape_inference.cc index 7a5106aa69bbbf..b447b6802a0607 100644 --- a/tensorflow/compiler/jit/shape_inference.cc +++ b/tensorflow/compiler/jit/shape_inference.cc @@ -21,6 +21,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/shape_inference_helpers.h" #include "tensorflow/core/common_runtime/shape_refiner.h" #include "tensorflow/core/framework/function.h" @@ -49,10 +50,29 @@ absl::Status ShapeHandleToTensorShape( if (!context->RankKnown(handle)) return absl::OkStatus(); std::vector dims(context->Rank(handle)); + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + std::vector dyn_exprs; + if (flags->tf_xla_enable_dynamic_sizes) { + dyn_exprs.resize(context->Rank(handle)); + } for (int32_t i = 0, end = dims.size(); i < end; ++i) { dims[i] = context->Value(context->Dim(handle, i)); + if (flags->tf_xla_enable_dynamic_sizes) { + DimExpr* expr = context->GetDimExpr(context->Dim(handle, i)); + dyn_exprs[i] = expr != nullptr + ? *expr + : context->ValueKnown(context->Dim(handle, i)) + ? xla::DExpr::Const(dims[i]) + : xla::DExpr::Unknown( + xla::kMissingExpressionSentinel); + } + } + auto status = + PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape); + if (flags->tf_xla_enable_dynamic_sizes) { + shape->set_expressions(std::move(dyn_exprs)); } - return PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape); + return status; } absl::Status PropagateShapes( diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc index aefa2a416310e2..42cf77245b8e1b 100644 --- a/tensorflow/core/framework/common_shape_fns.cc +++ b/tensorflow/core/framework/common_shape_fns.cc @@ -2369,7 +2369,13 @@ absl::Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c, } else if (dim_y.SameHandle(dim_x)) { dims.push_back(dim_x); } else if (!c->ValueKnown(dim_x) && !c->ValueKnown(dim_y)) { - dims.push_back(c->UnknownDim()); + DimensionHandle merged; + absl::Status s = c->Merge(dim_x, dim_y, &merged); + if (s.ok()) { + dims.push_back(merged); + } else { + dims.push_back(c->UnknownDim()); + } } else { if (!incompatible_shape_error) { *out = c->UnknownShape(); diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc index b63269f68c3368..dbd0ef831871dd 100644 --- a/tensorflow/core/framework/shape_inference.cc +++ b/tensorflow/core/framework/shape_inference.cc @@ -249,6 +249,16 @@ void InferenceContext::ShapeHandleToProto(ShapeHandle handle, } else { dim_shape->set_size(-1); } + + ExpressionProto* expression = proto->add_expressions(); + if (DimExpr* expr = GetDimExpr(dim)) { + DimExprToProto(*expr, expression); + } else if (ValueKnown(dim)) { + expression->set_constant_value(Value(dim)); + } else { + DimExprToProto(DimExpr::Unknown(xla::kMissingExpressionSentinel), + expression); + } } } @@ -282,6 +292,37 @@ DimensionHandle InferenceContext::NumElements(ShapeHandle s) { } } +DimensionHandle InferenceContext::UnknownDimWithExpr( + std::unique_ptr expr) { + DimExpr* owned = shape_manager_.OwnExpr(std::move(expr)); + return shape_manager_.MakeDim(kUnknownDim, owned); +} + +DimExpr* InferenceContext::GetDimExpr(DimensionHandle d) const { + if (!d.IsSet()) return nullptr; + return d->expr_; +} + +DimExpr* InferenceContext::MakeConstExpr(int64_t v) { + return shape_manager_.OwnExpr( + std::make_unique(DimExpr::Const(v))); +} + +DimExpr* InferenceContext::ExprForDim(DimensionHandle d) { + if (!d.IsSet()) return nullptr; + + // If already tagged with expr, use it. + if (DimExpr* e = GetDimExpr(d)) return e; + + // Known dim -> const expr. + if (ValueKnown(d)) { + return MakeConstExpr(Value(d)); + } + + // Unknown dim with no expr -> cannot form expression. + return nullptr; +} + string InferenceContext::DebugString(ShapeHandle s) { if (RankKnown(s)) { std::vector vals; @@ -923,8 +964,40 @@ absl::Status InferenceContext::MakeShapeFromShapeProto( const TensorShapeProto& proto, ShapeHandle* out) { *out = nullptr; TF_RETURN_IF_ERROR(PartialTensorShape::IsValidShape(proto)); - PartialTensorShape partial_shape(proto); - return MakeShapeFromPartialTensorShape(partial_shape, out); + + if (proto.unknown_rank()) { + *out = UnknownShape(); + return absl::OkStatus(); + } + + std::vector dims; + dims.reserve(proto.dim_size()); + for (int i = 0; i < proto.dim_size(); ++i) { + const auto& dim_proto = proto.dim(i); + if (dim_proto.size() >= 0) { + // Known dimension + dims.push_back(MakeDim(dim_proto.size())); + } else { + // Unknown dimension - check for expression. + if (i < proto.expressions_size() && + proto.expressions(i).node_type_case() != + ExpressionProto::NODE_TYPE_NOT_SET) { + // Deserialize expression + DimExpr expr = DimExprFromProto(proto.expressions(i)); + if (expr) { + DimExpr* owned = shape_manager_.OwnExpr( + std::make_unique(std::move(expr))); + dims.push_back(shape_manager_.MakeDim(kUnknownDim, owned)); + } else { + dims.push_back(UnknownDim()); + } + } else { + dims.push_back(UnknownDim()); + } + } + } + *out = MakeShape(dims); + return absl::OkStatus(); } absl::Status InferenceContext::GetScalarFromTensor(const Tensor* t, @@ -1030,24 +1103,40 @@ absl::Status InferenceContext::Divide(DimensionHandle dividend, DimensionOrConstant divisor, bool evenly_divisible, DimensionHandle* out) { - const int64_t divisor_value = Value(divisor); - if (divisor_value == 1) { + const bool dividend_known = ValueKnown(dividend); + const bool divisor_known = ValueKnown(divisor); + + // Validate divisor if known. + if (divisor_known && Value(divisor) <= 0) { + return errors::InvalidArgument("Divisor must be positive but is ", + Value(divisor)); + } + // Fast-path: x / 1 = x + if (divisor_known && Value(divisor) == 1) { *out = dividend; - } else if (!ValueKnown(dividend) || - (divisor.dim.IsSet() && !ValueKnown(divisor.dim))) { - *out = UnknownDim(); - } else { + return absl::OkStatus(); + } + // If both known, do numeric divide. + if (dividend_known && divisor_known) { const int64_t v = Value(dividend); - if (divisor_value <= 0) { - return errors::InvalidArgument("Divisor must be positive but is ", - divisor_value); - } - if (evenly_divisible && (v % divisor_value) != 0) { + const int64_t d = Value(divisor); + if (evenly_divisible && (v % d) != 0) { return errors::InvalidArgument( - "Dimension size must be evenly divisible by ", divisor_value, - " but is ", v); + "Dimension size must be evenly divisible by ", d, " but is ", v); } - *out = MakeDim(v / divisor_value); + *out = MakeDim(v / d); + return absl::OkStatus(); + } + // At least one operand unknown: try to build expression. + DimExpr* lhs = ExprForDim(dividend); + DimExpr* rhs = divisor.dim.IsSet() ? ExprForDim(divisor.dim) + : MakeConstExpr(divisor.val); + if (lhs && rhs) { + DimExpr* node = + shape_manager_.OwnExpr(std::make_unique(*lhs / *rhs)); + *out = shape_manager_.MakeDim(kUnknownDim, node); + } else { + *out = UnknownDim(); // Can't form expr. } return absl::OkStatus(); } @@ -1055,26 +1144,42 @@ absl::Status InferenceContext::Divide(DimensionHandle dividend, absl::Status InferenceContext::Add(DimensionHandle first, DimensionOrConstant second, DimensionHandle* out) { - const int64_t first_value = Value(first); - const int64_t second_value = Value(second); - // Special cases. - if (first_value == 0) { + const bool first_known = ValueKnown(first); + const bool second_known = ValueKnown(second); + + // Fast-path: x + 0 = x + if (first_known && Value(first) == 0) { *out = MakeDim(second); - } else if (second_value == 0) { + return absl::OkStatus(); + } + if (second_known && Value(second) == 0) { *out = first; - } else if (first_value == kUnknownDim || second_value == kUnknownDim) { - *out = UnknownDim(); - } else { - // Invariant: Both values are known and positive. Still in run-time we can - // get pair of values which cannot be store in output. Check below will - // report error. We still need to avoid undefined behavior of signed - // overflow and use unsigned addition. - const int64_t sum = static_cast(first_value) + second_value; + return absl::OkStatus(); + } + + // If both known, do numeric add. + if (first_known && second_known) { + const int64_t sum = static_cast(Value(first)) + + static_cast(Value(second)); if (sum < 0) { return errors::InvalidArgument("Dimension size overflow from adding ", - first_value, " and ", second_value); + Value(first), " and ", Value(second)); } *out = MakeDim(sum); + return absl::OkStatus(); + } + + // At least one operand unknown: try to build expression. + DimExpr* lhs = ExprForDim(first); + DimExpr* rhs = + second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val); + + if (lhs && rhs) { + DimExpr* node = + shape_manager_.OwnExpr(std::make_unique(*lhs + *rhs)); + *out = shape_manager_.MakeDim(kUnknownDim, node); + } else { + *out = UnknownDim(); // Can't form expr. } return absl::OkStatus(); } @@ -1082,22 +1187,35 @@ absl::Status InferenceContext::Add(DimensionHandle first, absl::Status InferenceContext::Subtract(DimensionHandle first, DimensionOrConstant second, DimensionHandle* out) { - const int64_t first_value = Value(first); - const int64_t second_value = Value(second); - // Special cases. - if (second_value == 0) { + const bool first_known = ValueKnown(first); + const bool second_known = ValueKnown(second); + // Fast-path: x - 0 = x + if (second_known && Value(second) == 0) { *out = first; - } else if (first_value == kUnknownDim || second_value == kUnknownDim) { - *out = UnknownDim(); - } else { - // Invariant: Both values are known, first_value is non-negative, and - // second_value is positive. + return absl::OkStatus(); + } + // If both known, do numeric subtract. + if (first_known && second_known) { + const int64_t first_value = Value(first); + const int64_t second_value = Value(second); if (first_value < second_value) { return errors::InvalidArgument( "Negative dimension size caused by subtracting ", second_value, " from ", first_value); } *out = MakeDim(first_value - second_value); + return absl::OkStatus(); + } + // At least one operand unknown: try to build expression. + DimExpr* lhs = ExprForDim(first); + DimExpr* rhs = + second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val); + if (lhs && rhs) { + DimExpr* node = + shape_manager_.OwnExpr(std::make_unique(*lhs - *rhs)); + *out = shape_manager_.MakeDim(kUnknownDim, node); + } else { + *out = UnknownDim(); // Can't form expr. } return absl::OkStatus(); } @@ -1105,21 +1223,31 @@ absl::Status InferenceContext::Subtract(DimensionHandle first, absl::Status InferenceContext::Multiply(DimensionHandle first, DimensionOrConstant second, DimensionHandle* out) { + const bool first_known = ValueKnown(first); + const bool second_known = ValueKnown(second); const int64_t first_value = Value(first); const int64_t second_value = Value(second); - // Special cases. - if (first_value == 0) { + + // Fast-paths for identity and zero cases. + if (first_known && first_value == 0) { *out = first; - } else if (second_value == 0) { + return absl::OkStatus(); + } + if (second_known && second_value == 0) { *out = MakeDim(second); - } else if (first_value == 1) { + return absl::OkStatus(); + } + if (first_known && first_value == 1) { *out = MakeDim(second); - } else if (second_value == 1) { + return absl::OkStatus(); + } + if (second_known && second_value == 1) { *out = first; - } else if (first_value == kUnknownDim || second_value == kUnknownDim) { - *out = UnknownDim(); - } else { - // Invariant: Both values are known and greater than 1. + return absl::OkStatus(); + } + + // If both known, do numeric multiply. + if (first_known && second_known) { const int64_t product = MultiplyWithoutOverflow(first_value, second_value); if (product < 0) { return errors::InvalidArgument( @@ -1127,6 +1255,20 @@ absl::Status InferenceContext::Multiply(DimensionHandle first, first_value, " and ", second_value); } *out = MakeDim(product); + return absl::OkStatus(); + } + + // At least one operand unknown: try to build expression. + DimExpr* lhs = ExprForDim(first); + DimExpr* rhs = + second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val); + + if (lhs && rhs) { + DimExpr* node = + shape_manager_.OwnExpr(std::make_unique(*lhs * *rhs)); + *out = shape_manager_.MakeDim(kUnknownDim, node); + } else { + *out = UnknownDim(); // Can't form expr. } return absl::OkStatus(); } diff --git a/tensorflow/core/framework/shape_inference.h b/tensorflow/core/framework/shape_inference.h index 8bfd301d860de1..d1670e348c62f7 100644 --- a/tensorflow/core/framework/shape_inference.h +++ b/tensorflow/core/framework/shape_inference.h @@ -20,6 +20,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "tensorflow/core/framework/full_type.pb.h" #include "tensorflow/core/framework/node_def_util.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/macros.h" @@ -116,13 +117,15 @@ class InferenceContext; class Dimension { private: Dimension(); - Dimension(int64_t value); + Dimension(int64_t value, DimExpr* expr = nullptr); ~Dimension() {} const int64_t value_; + DimExpr* expr_; friend class InferenceContext; friend class ShapeManager; + friend class ::tensorflow::grappler::SymbolicShapeManager; Dimension(const Dimension&) = delete; void operator=(const Dimension&) = delete; }; @@ -578,6 +581,20 @@ class InferenceContext { inline DimensionHandle UnknownDim() { return MakeDim(kUnknownDim); } + // Create a new unknown dimension (size = -1) tagged with a DimExpr. + // The expression is owned by this context's ShapeManager. + DimensionHandle UnknownDimWithExpr(std::unique_ptr expr); + // Return the expression pointer for a dimension, or nullptr if none. + DimExpr* GetDimExpr(DimensionHandle d) const; + // Creates a constant DimExpr node for the given value. + // The expression is owned by this context's ShapeManager. + DimExpr* MakeConstExpr(int64_t v); + // Returns the Expr representation for the given dimension: + // - If dim has an expr, returns it + // - If dim is known, returns a new Const expr + // - If dim is unknown with no expr, returns nullptr + DimExpr* ExprForDim(DimensionHandle d); + // Returns in a scalar value from an input tensor . The input tensor // must be a 0-dimensional int32 or int64 tensor. Caller must ensure that the // input tensor is not NULL. @@ -743,8 +760,6 @@ class InferenceContext { // Adds new outputs; useful when mutating the graph. absl::Status ExpandOutputs(int new_output_size); - - private: // Creates and stores shapes for use in InferenceContext. class ShapeManager { public: @@ -760,21 +775,32 @@ class InferenceContext { // Returns a new dimension of the given size. The returned value // is owned by this class. - inline DimensionHandle MakeDim(DimensionOrConstant d) { + inline DimensionHandle MakeDim(DimensionOrConstant d, + DimExpr* expr = nullptr) { if (d.dim.IsSet()) { return d.dim; } else { - all_dims_.push_back(new Dimension(d.val)); + all_dims_.push_back(new Dimension(d.val, expr)); return all_dims_.back(); } } + // Takes ownership of an expression and returns a raw pointer to it. + DimExpr* OwnExpr(std::unique_ptr expr) { + if (!expr) return nullptr; + DimExpr* ptr = expr.get(); + all_exprs_.push_back(std::move(expr)); + return ptr; + } private: std::vector all_shapes_; // values are owned. std::vector all_dims_; // values are owned. + std::vector> all_exprs_; // expressions are owned. }; + private: friend class ::tensorflow::grappler::GraphProperties; + friend class ::tensorflow::grappler::SymbolicShapeManager; friend class ShapeInferenceTest; // For testing Relax functions. friend class ShapeInferenceTestutil; // For testing shapes. @@ -888,8 +914,10 @@ class InferenceContext { // ----------------------------------------------------------------------------- // Template and inline method implementations, please ignore -inline Dimension::Dimension() : value_(InferenceContext::kUnknownDim) {} -inline Dimension::Dimension(int64_t value) : value_(value) { +inline Dimension::Dimension() + : value_(InferenceContext::kUnknownDim), expr_(nullptr) {} +inline Dimension::Dimension(int64_t value, DimExpr* expr) + : value_(value), expr_(expr) { DCHECK(value >= 0 || value == InferenceContext::kUnknownDim) << "Dimension must be non-negative or equal to " "InferenceContext::kUnknownDim but got " diff --git a/tensorflow/core/framework/shape_inference_test.cc b/tensorflow/core/framework/shape_inference_test.cc index c5dbc299b86540..9a10a05ed65163 100644 --- a/tensorflow/core/framework/shape_inference_test.cc +++ b/tensorflow/core/framework/shape_inference_test.cc @@ -1035,6 +1035,35 @@ TEST_F(ShapeInferenceTest, KnownShapeToProto) { EXPECT_FALSE(proto.unknown_rank()); EXPECT_EQ(3, proto.dim_size()); EXPECT_EQ(1, proto.dim(0).size()); + ASSERT_EQ(3, proto.expressions_size()); + EXPECT_EQ(1, proto.expressions(0).constant_value()); + EXPECT_EQ(2, proto.expressions(1).constant_value()); + EXPECT_EQ(3, proto.expressions(2).constant_value()); + EXPECT_FALSE(proto.dim(0).has_expr()); +} + +TEST_F(ShapeInferenceTest, DynamicExpressionShapeProtoRoundTrip) { + NodeDef def; + std::vector empty; + InferenceContext c(kVersion, def, MakeOpDef(0, 2), empty, {}, {}, {}); + + TensorShapeProto input_proto; + input_proto.add_dim()->set_size(-1); + DimExprToProto(DimExpr::Var(7), input_proto.add_expressions()); + + ShapeHandle shape; + TF_ASSERT_OK(c.MakeShapeFromShapeProto(input_proto, &shape)); + TensorShapeProto proto; + c.ShapeHandleToProto(shape, &proto); + + ASSERT_EQ(1, proto.expressions_size()); + EXPECT_EQ(DimExpr::Var(7), DimExprFromProto(proto.expressions(0))); + EXPECT_FALSE(proto.dim(0).has_expr()); + + ShapeHandle restored; + TF_ASSERT_OK(c.MakeShapeFromShapeProto(proto, &restored)); + ASSERT_NE(nullptr, c.GetDimExpr(c.Dim(restored, 0))); + EXPECT_EQ(DimExpr::Var(7), *c.GetDimExpr(c.Dim(restored, 0))); } TEST_F(ShapeInferenceTest, UnknownShapeToProto) { diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc index 613b12bb18ae3a..bee993c8a1142c 100644 --- a/tensorflow/core/grappler/costs/graph_properties.cc +++ b/tensorflow/core/grappler/costs/graph_properties.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/graph/tensor_id.h" #include "tensorflow/core/grappler/costs/utils.h" #include "tensorflow/core/grappler/mutable_graph_view.h" @@ -189,6 +190,9 @@ class DisjointSet { absl::Status Merge(Handle x, Handle y); const typename HandleToObject::Object GetMergedValue(Handle value); + // Returns a pointer that uniquely identifies the set containing `value`. + // This can be used as a stable key for associating metadata with a set. + void* RootId(Handle value) { return static_cast(Find(value)); } private: // All the handles that belong to the same set are part of the same tree, and @@ -648,11 +652,13 @@ class SymbolicShapeRefiner { explicit SymbolicShapeRefiner( const GraphView& graph, const absl::flat_hash_map>& fed_ports, - const bool aggressive_shape_inference) + const bool aggressive_shape_inference, + const bool enable_dynamic_value_inference) : graph_(graph), function_library_(OpRegistry::Global(), graph.graph()->library()), fed_ports_(fed_ports), - aggressive_shape_inference_(aggressive_shape_inference) { + aggressive_shape_inference_(aggressive_shape_inference), + enable_dynamic_value_inference_(enable_dynamic_value_inference) { graph_def_version_ = graph.graph()->versions().producer(); node_to_context_.reserve(graph.graph()->node_size()); } @@ -941,7 +947,9 @@ class SymbolicShapeRefiner { TF_RETURN_IF_ERROR(gp.InferStatically( /*assume_valid_feeds=*/true, /*aggressive_shape_inference=*/aggressive_shape_inference_, - /*include_tensor_values=*/true)); + /*include_input_tensor_values=*/true, + /*include_output_tensor_values=*/true, + /*enable_dynamic_value_inference=*/enable_dynamic_value_inference_)); // Add return nodes for output shapes. int output = 0; @@ -1148,32 +1156,32 @@ class SymbolicShapeRefiner { } struct ShapeId { - const NodeDef* node; + std::string node_name; int port_id; friend bool operator==(const ShapeId& lhs, const ShapeId& rhs) { - return lhs.node == rhs.node && lhs.port_id == rhs.port_id; + return lhs.node_name == rhs.node_name && lhs.port_id == rhs.port_id; } template friend H AbslHashValue(H h, const ShapeId& s) { - return H::combine(std::move(h), s.node, s.port_id); + return H::combine(std::move(h), s.node_name, s.port_id); } }; struct DimId { - const NodeDef* node; + std::string node_name; int port_id; int dim_index; friend bool operator==(const DimId& lhs, const DimId& rhs) { - return lhs.node == rhs.node && lhs.port_id == rhs.port_id && + return lhs.node_name == rhs.node_name && lhs.port_id == rhs.port_id && lhs.dim_index == rhs.dim_index; } template friend H AbslHashValue(H h, const DimId& d) { - return H::combine(std::move(h), d.node, d.port_id, d.dim_index); + return H::combine(std::move(h), d.node_name, d.port_id, d.dim_index); } }; @@ -1391,7 +1399,7 @@ class SymbolicShapeRefiner { // Return the one ShapeHandle used to denote a fully unknown shape for a node // output. ShapeHandle GetUnknownOutputShape(const NodeDef* node, int index) { - ShapeId id{node, index}; + ShapeId id{node->name(), index}; auto it = unknown_shapes_.find(id); if (it != unknown_shapes_.end()) { return it->second; @@ -1405,16 +1413,40 @@ class SymbolicShapeRefiner { // node output. DimensionHandle GetUnknownOutputDim(const NodeDef* node, int index, int dim_id) { - DimId id{node, index, dim_id}; + DimId id{node->name(), index, dim_id}; auto it = unknown_dims_.find(id); if (it != unknown_dims_.end()) { return it->second; } InferenceContext* c = GetContext(node); - DimensionHandle dim = c->UnknownDim(); + int var_id = GetOrCreateStableVarId(id); + DimensionHandle dim; + if (node->op() == "_Arg") { + var_id *= -1; + // var_id would be minus when it's argument. + dim = c->UnknownDimWithExpr( + std::make_unique(DimExpr::Var(var_id))); + } else { + dim = c->UnknownDimWithExpr( + std::make_unique(DimExpr::Var(var_id))); + } + VLOG(1) << "[EXPR] GetUnknownOutputDim: node=" << node->name() + << " out=" << index << " dim=" << dim_id << " -> Var(" << var_id + << ")"; + // Create an unknown dim with Var(var_id) expression. unknown_dims_[id] = dim; return dim; } + // Get or create a stable integer variable ID for a given DimId. + int GetOrCreateStableVarId(const DimId& id) { + auto it = stable_var_ids_.find(id); + if (it != stable_var_ids_.end()) { + return it->second; + } + int var_id = next_var_id_++; + stable_var_ids_[id] = var_id; + return var_id; + } // Returns true if all the output tensors have known values. bool AllOutputValuesKnown(NodeContext* c) { @@ -1712,7 +1744,62 @@ class SymbolicShapeRefiner { // but instantiate a new UnknownDim to prevent incorrect symbolic shape // inference through UnknownDim from Const. InferenceContext* ic = c->inference_context.get(); + const std::string& op = node.op(); + const bool is_bin = + (op == "Sub" || op == "Add" || op == "Mul" || op == "Div"); if (!is_fed) { + // Fall through to regular value inference when symbolic shape-value + // propagation does not apply. + if (enable_dynamic_value_inference_ && is_bin && + c->input_tensors_as_shapes_to_propagate.size() >= 2) { + auto va = c->input_tensors_as_shapes_to_propagate[0]; + auto vb = c->input_tensors_as_shapes_to_propagate[1]; + + if (!va.SameHandle(tensorflow::shape_inference::ShapeHandle()) && + !vb.SameHandle(tensorflow::shape_inference::ShapeHandle()) && + ic->RankKnown(va) && ic->RankKnown(vb) && + ic->Rank(va) == ic->Rank(vb)) { + std::vector out_elems; + out_elems.reserve(ic->Rank(va)); + const auto is_unknown_from_const = [&](DimensionHandle dim) { + return ic->ValueKnown(dim) && + ic->Value(dim) == kUnknownDimFromConst; + }; + bool symbolic_propagation_succeeded = true; + + for (int i = 0; i < ic->Rank(va); ++i) { + auto da = ic->Dim(va, i); + auto db = ic->Dim(vb, i); + if (is_unknown_from_const(da) || is_unknown_from_const(db)) { + symbolic_propagation_succeeded = false; + break; + } + + tensorflow::shape_inference::DimensionHandle r; + absl::Status status; + if (op == "Sub") + status = ic->Subtract(da, db, &r); + else if (op == "Add") + status = ic->Add(da, db, &r); + else if (op == "Mul") + status = ic->Multiply(da, db, &r); + else + status = + ic->Divide(da, db, /*evenly_divisible=*/false, &r); + if (!status.ok()) { + symbolic_propagation_succeeded = false; + break; + } + out_elems.push_back(r); + } + if (symbolic_propagation_succeeded) { + c->output_tensors_as_shapes.resize(1); + c->output_tensors_as_shapes[0] = ic->MakeShape(out_elems); + return absl::OkStatus(); + } + } + } + if (IsConstant(node)) { const TensorProto& tensor_proto = node.attr().at("value").tensor(); c->output_tensor_protos.resize(1); @@ -1720,6 +1807,18 @@ class SymbolicShapeRefiner { c->output_tensors_as_shapes.resize(1); MaybeTensorProtoToShape(ic, tensor_proto, &c->output_tensors_as_shapes[0]); + } else if (IsCast(node)) { + if (c->input_tensors_as_shapes_to_propagate.empty()) { + return absl::OkStatus(); + } + const DataType src_type = node.attr().at("SrcT").type(); + const DataType dst_type = node.attr().at("DstT").type(); + if ((src_type == DT_INT32 || src_type == DT_INT64) && + (dst_type == DT_INT32 || dst_type == DT_INT64)) { + c->output_tensors_as_shapes.resize(1); + c->output_tensors_as_shapes[0] = + c->input_tensors_as_shapes_to_propagate[0]; + } } else if (IsRank(node)) { if (ic->RankKnown(ic->input(0))) { // Propagate rank value. @@ -1731,6 +1830,19 @@ class SymbolicShapeRefiner { } } else if (IsSize(node)) { DimensionHandle size = ic->NumElements(ic->input(0)); + if (enable_dynamic_value_inference_ && !ic->ValueKnown(size) && + ic->RankKnown(ic->input(0))) { + size = ic->MakeDim(1); + for (int i = 0; i < ic->Rank(ic->input(0)); ++i) { + TF_RETURN_IF_ERROR( + ic->Multiply(size, ic->Dim(ic->input(0), i), &size)); + } + } + if (enable_dynamic_value_inference_ && + (ic->ValueKnown(size) || ic->GetDimExpr(size) != nullptr)) { + c->output_tensors_as_shapes.resize(1); + c->output_tensors_as_shapes[0] = ic->MakeShape({size}); + } if (ic->ValueKnown(size)) { // Propagate size value. int64_t sz = ic->Value(size); @@ -1751,6 +1863,48 @@ class SymbolicShapeRefiner { c->output_tensor_protos[0] = &const_tensors_to_propagate_.back(); } } + } else if (enable_dynamic_value_inference_ && op == "Range") { + auto scalar_int_value = [&](int input, int64_t* value) { + const Tensor* tensor = ic->input_tensor(input); + if (tensor == nullptr || tensor->dims() != 0) { + return false; + } + if (tensor->dtype() == DT_INT32) { + *value = tensor->scalar()(); + return true; + } + if (tensor->dtype() == DT_INT64) { + *value = tensor->scalar()(); + return true; + } + return false; + }; + + int64_t start; + int64_t delta; + const bool has_positive_constant_delta = + scalar_int_value(0, &start) && scalar_int_value(2, &delta) && + start == 0 && delta > 0; + if (has_positive_constant_delta && + c->input_tensors_as_shapes_to_propagate.size() > 1) { + const ShapeHandle& limit = + c->input_tensors_as_shapes_to_propagate[1]; + if (ic->RankKnown(limit) && ic->Rank(limit) >= 1) { + DimensionHandle length = ic->Dim(limit, 0); + if (ic->ValueKnown(length) || ic->GetDimExpr(length) != nullptr) { + DimensionHandle range_length = length; + if (delta > 1) { + DimensionHandle adjusted_length; + TF_RETURN_IF_ERROR( + ic->Add(length, delta - 1, &adjusted_length)); + TF_RETURN_IF_ERROR(ic->Divide( + adjusted_length, delta, /*evenly_divisible=*/false, + &range_length)); + } + ic->set_output(0, ic->Vector(range_length)); + } + } + } } else if (IsShape(node)) { c->output_tensors_as_shapes.resize(1); c->output_tensors_as_shapes[0] = c->inference_context->input(0); @@ -1759,7 +1913,126 @@ class SymbolicShapeRefiner { for (int i = 0; i < c->inference_context->num_inputs(); ++i) { c->output_tensors_as_shapes[i] = c->inference_context->input(i); } - } else if (node.op() == "ConcatV2") { + } else if (IsGather(node)) { + if (c->input_tensors_as_shapes_to_propagate.empty()) { + return absl::OkStatus(); + } + const ShapeHandle& params = c->input_tensors_as_shapes_to_propagate[0]; + // Shape-like tensors are propagated here by storing each tensor value + // as one dimension of a ShapeHandle. Gather(axis=0) therefore becomes + // "pick those propagated entries by index". + const Tensor* indices = ic->input_tensor(1); + const bool has_axis_input = ic->num_inputs() >= 3; + const Tensor* axis = has_axis_input ? ic->input_tensor(2) : nullptr; + bool valid = ic->RankKnown(params) && indices != nullptr; + int64_t axis_value = 0; + if (valid && has_axis_input) { + valid = axis != nullptr && axis->NumElements() == 1; + } + if (valid && has_axis_input) { + axis_value = axis->dtype() == DT_INT32 ? axis->scalar()() + : axis->scalar()(); + valid = (axis_value == 0); + } + if (valid) { + std::vector dims; + auto add_index = [&](int64_t raw_index) -> bool { + int64_t index = raw_index; + if (index < 0) { + // Match Gather's negative-index semantics. + index += ic->Rank(params); + } + if (index < 0 || index >= ic->Rank(params)) { + return false; + } + // In this side channel, shape-tensor contents live in the dims. + dims.push_back(ic->Dim(params, index)); + return true; + }; + auto add_indices = [&](const auto& values) -> bool { + for (int i = 0; i < values.size(); ++i) { + if (!add_index(values(i))) { + return false; + } + } + return true; + }; + if (indices->dims() == 0) { + valid = + add_index(indices->dtype() == DT_INT32 ? indices->scalar()() + : indices->scalar()()); + } else if (indices->dims() == 1) { + if (indices->dtype() == DT_INT32) { + valid = add_indices(indices->vec()); + } else if (indices->dtype() == DT_INT64) { + valid = add_indices(indices->vec()); + } else { + valid = false; + } + } else { + valid = false; + } + if (valid) { + c->output_tensors_as_shapes.resize(1); + c->output_tensors_as_shapes[0] = ic->MakeShape(dims); + } + } + } else if (IsProd(node)) { + if (c->input_tensors_as_shapes_to_propagate.empty()) { + return absl::OkStatus(); + } + const ShapeHandle& input = c->input_tensors_as_shapes_to_propagate[0]; + // Here `input` represents a propagated shape-vector value, so its + // entries live in the dimensions of the ShapeHandle rather than in a + // normal tensor buffer. + const Tensor* reduction_indices = ic->input_tensor(1); + bool keep_dims = false; + TF_RETURN_IF_ERROR(GetNodeAttr(node, "keep_dims", &keep_dims)); + bool valid = + ic->RankKnown(input) && reduction_indices != nullptr && !keep_dims; + if (valid) { + std::vector axes; + auto read_axes = [&](const auto& values) { + axes.reserve(values.size()); + for (int i = 0; i < values.size(); ++i) { + axes.push_back(values(i)); + } + }; + if (reduction_indices->dtype() == DT_INT32) { + if (reduction_indices->dims() == 0) { + axes.push_back(reduction_indices->scalar()()); + } else if (reduction_indices->dims() == 1) { + read_axes(reduction_indices->vec()); + } else { + valid = false; + } + } else if (reduction_indices->dtype() == DT_INT64) { + if (reduction_indices->dims() == 0) { + axes.push_back(reduction_indices->scalar()()); + } else if (reduction_indices->dims() == 1) { + read_axes(reduction_indices->vec()); + } else { + valid = false; + } + } else { + valid = false; + } + if (valid) { + // For shape-vector contents, Prod(axis=0) means multiply the + // propagated entries together. + valid = axes.size() == 1 && axes[0] == 0 && ic->Rank(input) > 0; + } + if (valid) { + DimensionHandle product = ic->Dim(input, 0); + for (int i = 1; i < ic->Rank(input); ++i) { + TF_RETURN_IF_ERROR(ic->Multiply(product, ic->Dim(input, i), + &product)); + } + c->output_tensors_as_shapes.resize(1); + c->output_tensors_as_shapes[0] = ic->MakeShape({product}); + } + } + } else if (IsConcat(node)) { bool valid = true; ShapeHandle result; for (int i = 0; i < ic->num_inputs() - 1; ++i) { @@ -1798,8 +2071,12 @@ class SymbolicShapeRefiner { // possible. const ShapeHandle& shape_handle = c->input_tensors_as_shapes_to_propagate[i]; - if (ic->RankKnown(shape_handle) && ic->Rank(shape_handle) >= 1 && - ic->ValueKnown(ic->Dim(shape_handle, 0))) { + const bool has_value = + ic->RankKnown(shape_handle) && ic->Rank(shape_handle) >= 1 && + (ic->ValueKnown(ic->Dim(shape_handle, 0)) || + (enable_dynamic_value_inference_ && + ic->GetDimExpr(ic->Dim(shape_handle, 0)) != nullptr)); + if (has_value) { dims.push_back(ic->Dim(shape_handle, 0)); } else { // This is not from Const, but as it shouldn'be used as symbolic @@ -1812,6 +2089,26 @@ class SymbolicShapeRefiner { c->output_tensors_as_shapes.resize(1); c->output_tensors_as_shapes[0] = ic->MakeShape(dims); } + } else if (IsUnpack(node)) { + if (c->input_tensors_as_shapes_to_propagate.empty()) { + return absl::OkStatus(); + } + const ShapeHandle& input = c->input_tensors_as_shapes_to_propagate[0]; + bool valid = ic->RankKnown(input); + int axis = 0; + if (valid) { + TF_RETURN_IF_ERROR(GetNodeAttr(node, "axis", &axis)); + if (axis < 0) { + axis += ic->Rank(input); + } + valid = (axis == 0 && ic->num_outputs() == ic->Rank(input)); + } + if (valid) { + c->output_tensors_as_shapes.resize(ic->num_outputs()); + for (int i = 0; i < ic->num_outputs(); ++i) { + c->output_tensors_as_shapes[i] = ic->MakeShape({ic->Dim(input, i)}); + } + } } else if (IsIdentity(node) || IsIdentityNSingleInput(node)) { c->output_tensors_as_shapes.resize(1); c->output_tensors_as_shapes[0] = @@ -1919,6 +2216,99 @@ class SymbolicShapeRefiner { return absl::OkStatus(); } + absl::Status CanonicalizeOutputDims(const NodeDef* node) { + NodeContext* ctx = GetNodeContext(node); + if (!ctx) return absl::OkStatus(); + + InferenceContext* ic = ctx->inference_context.get(); + for (int out = 0; out < ic->num_outputs(); ++out) { + ShapeHandle s = ic->output(out); + + if (!ic->RankKnown(s)) { + bool recovered_rank = false; + auto it = node->attr().find("_output_shapes"); + if(it == node->attr().end()){ + it = node->attr().find("shape"); + } + if (it != node->attr().end() && out < it->second.list().shape_size()) { + const TensorShapeProto& proto = it->second.list().shape(out); + if (!proto.unknown_rank()) { + std::vector dims; + dims.reserve(proto.dim_size()); + + for (int d = 0; d < proto.dim_size(); ++d) { + int64_t size = proto.dim(d).size(); + if (size >= 0) { + dims.push_back(ic->MakeDim(size)); + } else { + dims.push_back(GetUnknownOutputDim(node, out, d)); + } + } + s = ic->MakeShape(dims); + ic->set_output(out, s); + recovered_rank = true; + } + } + if (!recovered_rank && node->op() == "_Arg") { + DimensionHandle d0 = GetUnknownOutputDim(node, out, /*dim_id=*/0); + ShapeHandle vec = ic->MakeShape({d0}); + ic->set_output(out, vec); + s = vec; + recovered_rank = true; + } + + if (!recovered_rank) { + VLOG(1) << "RANK still unknown. " << node->name(); + continue; + } + } + + if (!ic->RankKnown(s)) { + continue; + } + + bool changed = false; + std::vector dims; + dims.reserve(ic->Rank(s)); + for (int d = 0; d < ic->Rank(s); ++d) { + DimensionHandle dim = ic->Dim(s, d); + const int64_t v = ic->Value(dim); + // Keep concrete dims. + if (v >= 0) { + dims.push_back(dim); + continue; + } + // If already tagged with expr, keep it. + auto* dim_expr = ic->GetDimExpr(dim); + if (dim_expr != nullptr) { + dims.push_back(dim); + continue; + } + auto output_shapes_it = node->attr().find("_output_shapes"); + if (output_shapes_it != node->attr().end() && + out < output_shapes_it->second.list().shape_size() && + d < output_shapes_it->second.list().shape(out).dim_size()) { + const int64_t annotated_size = + output_shapes_it->second.list().shape(out).dim(d).size(); + if (annotated_size >= 0) { + changed = true; + dims.push_back(ic->MakeDim(annotated_size)); + continue; + } + } + // Canonicalize ALL unknown dims. + DimensionHandle canon = GetUnknownOutputDim(node, out, d); + changed |= !dim.SameHandle(canon); + dims.push_back(canon); + } + if (changed) { + ShapeHandle new_s = ic->MakeShape(dims); + ic->set_output(out, new_s); + } + } + return absl::OkStatus(); + } + absl::Status InferShapes(const NodeDef& node, NodeContext* c) { // Infer the shapes of output tensors. if (!c->op_data || c->op_data->shape_inference_fn == nullptr || @@ -1940,8 +2330,8 @@ class SymbolicShapeRefiner { status.Update(SetUnknownShape(&node, output_port)); } } - // Update NodeContext output fields after shape inference function runs. + status.Update(CanonicalizeOutputDims(&node)); status.Update(MaybeUpdateNodeContextOutput(node, is_fed, c)); return status; @@ -2048,6 +2438,9 @@ class SymbolicShapeRefiner { absl::flat_hash_map node_to_context_; absl::flat_hash_map unknown_shapes_; absl::flat_hash_map unknown_dims_; + // Stable variable IDs for canonical dimension symbols. + absl::flat_hash_map stable_var_ids_; + int next_var_id_ = 1; // Store function instantiations only for valid function. If function // instantiation failed it will have an `absl::nullopt`. absl::flat_hash_map> @@ -2062,6 +2455,7 @@ class SymbolicShapeRefiner { // For more aggressive shape and value inference. bool aggressive_shape_inference_; + bool enable_dynamic_value_inference_; ResourceMgr resource_mgr_; }; @@ -2080,8 +2474,9 @@ 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(dims_.Merge(InferenceContext::DimKnownRank(s1, i), - InferenceContext::DimKnownRank(s2, i))); + TF_RETURN_IF_ERROR( + MergeDimsWithExpr(InferenceContext::DimKnownRank(s1, i), + InferenceContext::DimKnownRank(s2, i))); } } return absl::OkStatus(); @@ -2090,7 +2485,7 @@ class SymbolicShapeManager { if (!d1.IsSet() || !d2.IsSet()) { return absl::OkStatus(); } - return dims_.Merge(d1, d2); + return MergeDimsWithExpr(d1, d2); } void AsTensorProperties(const ShapeHandle& shape, const DataType& type, @@ -2104,7 +2499,26 @@ class SymbolicShapeManager { shape_inference::DimensionHandle dim = InferenceContext::DimKnownRank(actual_shape, j); int64_t d = dims_.GetMergedValue(dim); - properties->mutable_shape()->add_dim()->set_size(d); + TensorShapeProto* output_shape = properties->mutable_shape(); + auto* out_dim = output_shape->add_dim(); + out_dim->set_size(d < 0 ? -1 : d); + void* root = dims_.RootId(dim); + DimExpr* expr = nullptr; + if (auto it = dim_root_expr_.find(root); it != dim_root_expr_.end()) { + expr = it->second; + } else { + expr = ExprForDim(dim); + } + ExpressionProto* output_expr = output_shape->add_expressions(); + if (expr != nullptr) { + DimExprToProto(*expr, output_expr); + // TODO: Apply simplification? + } else if (d >= 0) { + output_expr->set_constant_value(d); + } else { + DimExprToProto(DimExpr::Unknown(xla::kMissingExpressionSentinel), + output_expr); + } } } } @@ -2132,7 +2546,142 @@ class SymbolicShapeManager { } private: + // Get the variable ID from an expression, or -1 if not a variable. + static int32_t GetVarId(const DimExpr* e) { + if (!e || e->kind() != DimExpr::Kind::kVariable) return -1; + return static_cast(e->get())->get_id(); + } + + static bool IsConst(const DimExpr* e) { + return e && e->kind() == DimExpr::Kind::kConstant; + } + + static bool IsVar(const DimExpr* e) { + return e && e->kind() == DimExpr::Kind::kVariable; + } + + static bool IsPlaceHolder(const DimExpr* e) { + if (!e) return false; + if (e->kind() != DimExpr::Kind::kVariable) return false; + return static_cast(e->get())->get_id() < 0; + } + + static bool IsCompound(const DimExpr* e) { + if (!e) return false; + switch (e->kind()) { + case DimExpr::Kind::kAdd: + case DimExpr::Kind::kSub: + case DimExpr::Kind::kMul: + case DimExpr::Kind::kDiv: + case DimExpr::Kind::kMax: + case DimExpr::Kind::kGt: + case DimExpr::Kind::kSelect: + return true; + default: + return false; + } + } + + // Ranking: Const > Arg_ > Compound > Var > null + static int InfoScore(const DimExpr* e) { + if (!e) return 0; + if (IsConst(e)) return 4; + if (IsPlaceHolder(e)) return 3; + if (IsCompound(e)) return 2; + if (IsVar(e)) return 1; + return 1; // fallback (shouldn't happen) + } + + // Prefer "more informative" but keep deterministic tie-break. + static DimExpr* PreferMoreInformative(DimExpr* a, DimExpr* b) { + if (a == b) return a; + const int sa = InfoScore(a); + const int sb = InfoScore(b); + if (sa > sb) return a; + if (sb > sa) return b; + // Same score: keep stable choice. + return a; + } + + // Get the expr pointer from a dimension handle (accesses private member). + static DimExpr* GetExprFromDimHandle(const DimensionHandle& d) { + if (!d.IsSet()) return nullptr; + return d->expr_; + } + + DimExpr* ExprForDim(const DimensionHandle& d) { + if (!d.IsSet()) return nullptr; + if (DimExpr* expr = GetExprFromDimHandle(d)) { + return expr; + } + if (!InferenceContext::ValueKnown(d)) { + return nullptr; + } + const int64_t value = InferenceContext::Value(d); + auto it = const_exprs_.find(value); + if (it != const_exprs_.end()) { + return it->second.get(); + } + auto expr = std::make_unique(DimExpr::Const(value)); + DimExpr* expr_ptr = expr.get(); + const_exprs_.emplace(value, std::move(expr)); + return expr_ptr; + } + + absl::Status MergeDimsWithExpr(DimensionHandle d1, DimensionHandle d2) { + if (!d1.IsSet() || !d2.IsSet()) return absl::OkStatus(); + + void* r1 = dims_.RootId(d1); + void* r2 = dims_.RootId(d2); + + // Fetch best-known expr for each set. + auto get_best = [&](void* r, DimensionHandle d) -> DimExpr* { + auto it = dim_root_expr_.find(r); + if (it != dim_root_expr_.end()) return it->second; + return ExprForDim(d); // may be null + }; + + DimExpr* e1 = get_best(r1, d1); + DimExpr* e2 = get_best(r2, d2); + + // If already in same UF set, just keep the most informative expr. + if (r1 == r2) { + DimExpr* existing = nullptr; + if (auto it = dim_root_expr_.find(r1); it != dim_root_expr_.end()) { + existing = it->second; + } + DimExpr* chosen = PreferMoreInformative(existing, + PreferMoreInformative(e1, e2)); + if (chosen) dim_root_expr_[r1] = chosen; // keep or upgrade + return absl::OkStatus(); + } + + // Perform UF merge (rank + path compression inside DisjointSet). + TF_RETURN_IF_ERROR(dims_.Merge(d1, d2)); + + // New root after merge. + void* new_root = dims_.RootId(d1); + + // Choose best expr across both sets. + DimExpr* chosen = PreferMoreInformative(e1, e2); + + // Remove stale root keys (only the old roots). + dim_root_expr_.erase(r1); + dim_root_expr_.erase(r2); + + // Preserve any expr already stored at new_root (rare but safe). + if (auto it = dim_root_expr_.find(new_root); it != dim_root_expr_.end()) { + chosen = PreferMoreInformative(it->second, chosen); + } + + if (chosen) dim_root_expr_[new_root] = chosen; + + return absl::OkStatus(); + } DisjointSet shapes_; + absl::flat_hash_map> const_exprs_; + // Map from union-find root pointer to the best expression for that set. + absl::flat_hash_map dim_root_expr_; DisjointSet dims_; }; @@ -2518,7 +3067,8 @@ absl::Status GraphProperties::UpdateEnqueue( absl::Status GraphProperties::InferStatically( bool assume_valid_feeds, bool aggressive_shape_inference, - bool include_input_tensor_values, bool include_output_tensor_values) { + bool include_input_tensor_values, bool include_output_tensor_values, + bool enable_dynamic_value_inference) { FunctionLibraryDefinition function_library(OpRegistry::Global(), item_.graph.library()); absl::flat_hash_map> fed_ports; @@ -2607,7 +3157,8 @@ absl::Status GraphProperties::InferStatically( // Heap-allocate SymbolicShapeRefiner in order to not consume a large amount // of stack space. auto refiner = std::make_unique( - graph_view, fed_ports, aggressive_shape_inference); + graph_view, fed_ports, aggressive_shape_inference, + enable_dynamic_value_inference); TopoQueue new_shapes(topo_order); // Also seed the propagation of shapes in the fanout of primary inputs. diff --git a/tensorflow/core/grappler/costs/graph_properties.h b/tensorflow/core/grappler/costs/graph_properties.h index 1d9575e1e5c805..2d3a51cd2975f0 100644 --- a/tensorflow/core/grappler/costs/graph_properties.h +++ b/tensorflow/core/grappler/costs/graph_properties.h @@ -99,7 +99,8 @@ class GraphProperties { absl::Status InferStatically(bool assume_valid_feeds, bool aggressive_shape_inference, bool include_input_tensor_values, - bool include_output_tensor_values); + bool include_output_tensor_values, + bool enable_dynamic_value_inference = false); absl::Status InferStatically(bool assume_valid_feeds, bool aggressive_shape_inference, bool include_tensor_values) { diff --git a/tensorflow/core/grappler/costs/graph_properties_test.cc b/tensorflow/core/grappler/costs/graph_properties_test.cc index 53fe8bef1f6f85..b24ad629910909 100644 --- a/tensorflow/core/grappler/costs/graph_properties_test.cc +++ b/tensorflow/core/grappler/costs/graph_properties_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/tensor.pb.h" // NOLINT #include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/framework/versions.pb.h" @@ -42,6 +43,24 @@ namespace tensorflow { namespace grappler { namespace { +DimExpr ShapeDimExpr(const TensorShapeProto& shape, int dim) { + if (dim < shape.expressions_size()) { + return DimExprFromProto(shape.expressions(dim)); + } + return DimExpr(); +} + +bool DimExprEqual(const DimExpr& lhs_expr, const DimExpr& rhs_expr) { + if (!lhs_expr || !rhs_expr) return !lhs_expr && !rhs_expr; + return lhs_expr == rhs_expr; +} + +bool ShapeDimExprEqual(const TensorShapeProto& lhs, int lhs_dim, + const TensorShapeProto& rhs, int rhs_dim) { + return DimExprEqual(ShapeDimExpr(lhs, lhs_dim), + ShapeDimExpr(rhs, rhs_dim)); +} + using shape_inference::InferenceContext; using shape_inference::ShapeAndType; using shape_inference::ShapeHandle; @@ -776,9 +795,11 @@ TEST_F(GraphPropertiesTest, WhileLoop) { // since we concatenated along the batch dim. auto shape_in = properties.GetOutputProperties("ones").at(0).shape(); auto shape_out = properties.GetOutputProperties("while/Exit_1").at(0).shape(); - EXPECT_GE(-2, shape_in.dim(0).size()); - EXPECT_GE(-2, shape_out.dim(0).size()); - EXPECT_NE(shape_in.dim(0).size(), shape_out.dim(0).size()); + EXPECT_EQ(-1, shape_in.dim(0).size()); + EXPECT_EQ(-1, shape_out.dim(0).size()); + EXPECT_TRUE(ShapeDimExpr(shape_in, 0)); + EXPECT_TRUE(ShapeDimExpr(shape_out, 0)); + EXPECT_FALSE(ShapeDimExprEqual(shape_in, 0, shape_out, 0)); } TEST_F(GraphPropertiesTest, NestedLoop) { @@ -1969,10 +1990,10 @@ TEST_F(GraphPropertiesTest, SymbolicShapes) { const auto shape_c = properties.GetOutputProperties("c").at(0).shape(); EXPECT_EQ(2, shape_a.dim_size()); EXPECT_EQ(shape_a.dim_size(), shape_c.dim_size()); - EXPECT_GE(-2, shape_a.dim(0).size()); - EXPECT_EQ(shape_a.dim(0).size(), shape_c.dim(0).size()); - EXPECT_GE(-2, shape_a.dim(1).size()); - EXPECT_EQ(shape_a.dim(1).size(), shape_c.dim(1).size()); + EXPECT_EQ(-1, shape_a.dim(0).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_a, 0, shape_c, 0)); + EXPECT_EQ(-1, shape_a.dim(1).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_a, 1, shape_c, 1)); PartialTensorShape shape(shape_a); EXPECT_FALSE(shape.IsFullyDefined()); @@ -1982,29 +2003,29 @@ TEST_F(GraphPropertiesTest, SymbolicShapes) { const auto shape_d = properties.GetOutputProperties("d").at(0).shape(); EXPECT_EQ(1, shape_b.dim_size()); EXPECT_EQ(shape_b.dim_size(), shape_d.dim_size()); - EXPECT_GE(-2, shape_b.dim(0).size()); - EXPECT_NE(shape_a.dim(0).size(), shape_b.dim(0).size()); - EXPECT_EQ(shape_b.dim(0).size(), shape_d.dim(0).size()); + EXPECT_EQ(-1, shape_b.dim(0).size()); + EXPECT_FALSE(ShapeDimExprEqual(shape_a, 0, shape_b, 0)); + EXPECT_TRUE(ShapeDimExprEqual(shape_b, 0, shape_d, 0)); const auto shape_e = properties.GetOutputProperties("e").at(0).shape(); ASSERT_EQ(2, shape_e.dim_size()); - EXPECT_EQ(shape_e.dim(0).size(), shape_c.dim(0).size()); - EXPECT_NE(shape_e.dim(1).size(), shape_c.dim(1).size()); - EXPECT_NE(shape_e.dim(0).size(), shape_d.dim(0).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_e, 0, shape_c, 0)); + EXPECT_TRUE(ShapeDimExprEqual(shape_e, 1, shape_c, 1)); + EXPECT_FALSE(ShapeDimExprEqual(shape_e, 0, shape_d, 0)); const auto shape_f = properties.GetOutputProperties("f").at(0).shape(); ASSERT_EQ(2, shape_f.dim_size()); - EXPECT_EQ(shape_f.dim(0).size(), shape_a.dim(0).size()); - EXPECT_EQ(shape_f.dim(1).size(), shape_a.dim(1).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_f, 0, shape_a, 0)); + EXPECT_TRUE(ShapeDimExprEqual(shape_f, 1, shape_a, 1)); const auto shape_h = properties.GetOutputProperties("h").at(0).shape(); ASSERT_EQ(2, shape_f.dim_size()); - EXPECT_EQ(shape_h.dim(0).size(), shape_c.dim(0).size()); - EXPECT_EQ(shape_h.dim(1).size(), shape_c.dim(1).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_h, 0, shape_c, 0)); + EXPECT_TRUE(ShapeDimExprEqual(shape_h, 1, shape_c, 1)); const auto shape_j = properties.GetOutputProperties("j").at(0).shape(); ASSERT_EQ(1, shape_j.dim_size()); - EXPECT_EQ(shape_j.dim(0).size(), shape_a.dim(1).size()); + EXPECT_TRUE(ShapeDimExprEqual(shape_j, 0, shape_a, 1)); } TEST_F(GraphPropertiesTest, DoNotValidateColocationConstraints) { @@ -2167,6 +2188,74 @@ TEST_F(GraphPropertiesTest, StridedSlicesOfShapes) { EXPECT_EQ(shape_a.dim(1).size(), shape_o2.dim(0).size()); } +TEST_F(GraphPropertiesTest, SizeContentsPropagateToFillOutput) { + tensorflow::Scope scope = tensorflow::Scope::NewRootScope(); + Output input = ops::Placeholder( + scope.WithOpName("input"), DT_FLOAT, + ops::Placeholder::Shape(PartialTensorShape({-1, 24}))); + Output input_size = ops::Size(scope.WithOpName("input_size"), input); + Output fill_shape = + ops::Stack(scope.WithOpName("fill_shape"), {input_size}); + Output zero = ops::Const(scope.WithOpName("zero"), 0.0f, {}); + Output filled = ops::Fill(scope.WithOpName("filled"), fill_shape, zero); + + GrapplerItem item; + TF_ASSERT_OK(scope.ToGraphDef(&item.graph)); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true)); + + const TensorShapeProto& inferred_input_shape = + properties.GetOutputProperties("input").at(0).shape(); + const TensorShapeProto& inferred_fill_shape = + properties.GetOutputProperties("filled").at(0).shape(); + + ASSERT_EQ(1, inferred_fill_shape.dim_size()); + ASSERT_EQ(2, inferred_input_shape.expressions_size()); + DimExpr expected = + DimExprFromProto(inferred_input_shape.expressions(0)) * 24; + EXPECT_TRUE(DimExprEqual(expected, ShapeDimExpr(inferred_fill_shape, 0))); +} + +TEST_F(GraphPropertiesTest, SizeContentsPropagateToRangeOutput) { + tensorflow::Scope scope = tensorflow::Scope::NewRootScope(); + Output input = ops::Placeholder( + scope.WithOpName("input"), DT_FLOAT, + ops::Placeholder::Shape(PartialTensorShape({-1, 1}))); + Output input_size = ops::Size(scope.WithOpName("input_size"), input); + Output zero = ops::Const(scope.WithOpName("zero"), 0, {}); + Output three = ops::Const(scope.WithOpName("three"), 3, {}); + Output range = + ops::Range(scope.WithOpName("range"), zero, input_size, three); + + GrapplerItem item; + TF_ASSERT_OK(scope.ToGraphDef(&item.graph)); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true)); + + const TensorShapeProto& inferred_input_shape = + properties.GetOutputProperties("input").at(0).shape(); + const TensorShapeProto& inferred_range_shape = + properties.GetOutputProperties("range").at(0).shape(); + + ASSERT_EQ(1, inferred_range_shape.dim_size()); + ASSERT_EQ(2, inferred_input_shape.expressions_size()); + DimExpr expected = + (DimExprFromProto(inferred_input_shape.expressions(0)) + 2) / 3; + EXPECT_TRUE(DimExprEqual(expected, ShapeDimExpr(inferred_range_shape, 0))); +} + TEST_F(GraphPropertiesTest, StridedSliceOfShapeWithShrinkAxisMask) { tensorflow::Scope scope = tensorflow::Scope::NewRootScope(); Output placeholder = @@ -2211,6 +2300,143 @@ TEST_F(GraphPropertiesTest, StridedSliceOfShapeWithShrinkAxisMask) { } } +TEST_F(GraphPropertiesTest, ShapeTensorContentsThroughGatherProdAndUnpack) { + GrapplerItem item; + + TF_ASSERT_OK(NodeDefBuilder("input", "Placeholder") + .Attr("dtype", DT_FLOAT) + .Attr("shape", PartialTensorShape({-1, 26, 8})) + .Finalize(item.graph.add_node())); + TF_ASSERT_OK(NodeDefBuilder("shape", "Shape") + .Input("input", 0, DT_FLOAT) + .Attr("T", DT_FLOAT) + .Attr("out_type", DT_INT32) + .Finalize(item.graph.add_node())); + + Tensor gather_indices(DT_INT32, TensorShape({2})); + gather_indices.vec()(0) = 0; + gather_indices.vec()(1) = 1; + TF_ASSERT_OK(NodeDefBuilder("gather_indices", "Const") + .Attr("dtype", DT_INT32) + .Attr("value", gather_indices) + .Finalize(item.graph.add_node())); + + Tensor zero_scalar(DT_INT32, TensorShape({})); + zero_scalar.scalar()() = 0; + TF_ASSERT_OK(NodeDefBuilder("zero_axis", "Const") + .Attr("dtype", DT_INT32) + .Attr("value", zero_scalar) + .Finalize(item.graph.add_node())); + + Tensor eight_scalar(DT_INT32, TensorShape({})); + eight_scalar.scalar()() = 8; + TF_ASSERT_OK(NodeDefBuilder("eight", "Const") + .Attr("dtype", DT_INT32) + .Attr("value", eight_scalar) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("gather", "GatherV2") + .Input("shape", 0, DT_INT32) + .Input("gather_indices", 0, DT_INT32) + .Input("zero_axis", 0, DT_INT32) + .Attr("Tparams", DT_INT32) + .Attr("Tindices", DT_INT32) + .Attr("Taxis", DT_INT32) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("flat_dim", "Prod") + .Input("gather", 0, DT_INT32) + .Input("zero_axis", 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("keep_dims", false) + .Finalize(item.graph.add_node())); + + std::vector gather_pack_inputs(2); + gather_pack_inputs[0] = NodeDefBuilder::NodeOut{"flat_dim", 0, DT_INT32}; + gather_pack_inputs[1] = NodeDefBuilder::NodeOut{"eight", 0, DT_INT32}; + TF_ASSERT_OK(NodeDefBuilder("gather_shape", "Pack") + .Input(gather_pack_inputs) + .Attr("N", 2) + .Attr("T", DT_INT32) + .Attr("axis", 0) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("reshape_from_gather", "Reshape") + .Input("input", 0, DT_FLOAT) + .Input("gather_shape", 0, DT_INT32) + .Attr("T", DT_FLOAT) + .Attr("Tshape", DT_INT32) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("unpack", "Unpack") + .Input("shape", 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("axis", 0) + .Attr("num", 3) + .Finalize(item.graph.add_node())); + + std::vector tail_inputs(2); + tail_inputs[0] = NodeDefBuilder::NodeOut{"unpack", 1, DT_INT32}; + tail_inputs[1] = NodeDefBuilder::NodeOut{"unpack", 2, DT_INT32}; + TF_ASSERT_OK(NodeDefBuilder("tail_shape", "Pack") + .Input(tail_inputs) + .Attr("N", 2) + .Attr("T", DT_INT32) + .Attr("axis", 0) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("tail_prod", "Prod") + .Input("tail_shape", 0, DT_INT32) + .Input("zero_axis", 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("keep_dims", false) + .Finalize(item.graph.add_node())); + + std::vector unpack_pack_inputs(2); + unpack_pack_inputs[0] = NodeDefBuilder::NodeOut{"unpack", 0, DT_INT32}; + unpack_pack_inputs[1] = NodeDefBuilder::NodeOut{"tail_prod", 0, DT_INT32}; + TF_ASSERT_OK(NodeDefBuilder("unpack_shape", "Pack") + .Input(unpack_pack_inputs) + .Attr("N", 2) + .Attr("T", DT_INT32) + .Attr("axis", 0) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("reshape_from_unpack", "Reshape") + .Input("input", 0, DT_FLOAT) + .Input("unpack_shape", 0, DT_INT32) + .Attr("T", DT_FLOAT) + .Attr("Tshape", DT_INT32) + .Finalize(item.graph.add_node())); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically(true)); + + const auto& input_shape = properties.GetOutputProperties("input")[0].shape(); + const auto& gather_reshape_shape = + properties.GetOutputProperties("reshape_from_gather")[0].shape(); + const auto& unpack_reshape_shape = + properties.GetOutputProperties("reshape_from_unpack")[0].shape(); + + ASSERT_EQ(3, input_shape.dim_size()); + ASSERT_EQ(2, gather_reshape_shape.dim_size()); + ASSERT_EQ(2, unpack_reshape_shape.dim_size()); + ASSERT_TRUE(ShapeDimExpr(input_shape, 0)); + ASSERT_TRUE(ShapeDimExpr(gather_reshape_shape, 0)); + ASSERT_TRUE(ShapeDimExpr(unpack_reshape_shape, 0)); + + DimExpr input_expr = ShapeDimExpr(input_shape, 0); + auto expected_gather_dim0 = input_expr * DimExpr::Const(26); + EXPECT_TRUE(DimExprEqual(expected_gather_dim0, + ShapeDimExpr(gather_reshape_shape, 0))); + EXPECT_EQ(8, gather_reshape_shape.dim(1).size()); + + EXPECT_TRUE(ShapeDimExprEqual(input_shape, 0, unpack_reshape_shape, 0)); + EXPECT_EQ(208, unpack_reshape_shape.dim(1).size()); +} + TEST_F(GraphPropertiesTest, ValuePropagationThroughArithmeticOps) { tensorflow::Scope s = tensorflow::Scope::NewRootScope(); Output a = ops::Const(s.WithOpName("a"), {5, 7}, {2}); @@ -2255,6 +2481,37 @@ TEST_F(GraphPropertiesTest, ValuePropagationThroughArithmeticOps) { ExpectTensorValues({20, 24}, c_plus_b_plus_2a_prop.value()); } +TEST_F(GraphPropertiesTest, + DynamicArithmeticFallsBackToTensorValueInference) { + tensorflow::Scope scope = tensorflow::Scope::NewRootScope(); + Output one = ops::Const(scope.WithOpName("one"), {1}, {1}); + Output two = ops::Const(scope.WithOpName("two"), {2}, {1}); + Output unknown = ops::Const(scope.WithOpName("unknown"), {-1}, {1}); + Output negative = ops::Sub(scope.WithOpName("negative"), one, two); + Output unknown_plus_one = + ops::Add(scope.WithOpName("unknown_plus_one"), unknown, one); + + GrapplerItem item; + TF_ASSERT_OK(scope.ToGraphDef(&item.graph)); + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/true, + /*include_input_tensor_values=*/true, + /*include_output_tensor_values=*/true, + /*enable_dynamic_value_inference=*/true)); + + const auto& negative_prop = + properties.GetOutputProperties("negative").at(0); + ASSERT_TRUE(negative_prop.has_value()); + ExpectTensorValues({-1}, negative_prop.value()); + + const auto& unknown_plus_one_prop = + properties.GetOutputProperties("unknown_plus_one").at(0); + ASSERT_TRUE(unknown_plus_one_prop.has_value()); + ExpectTensorValues({0}, unknown_plus_one_prop.value()); +} + TEST_F(GraphPropertiesTest, ShapeAnnotation) { GrapplerItem item; TF_ASSERT_OK(NodeDefBuilder("Input", "Placeholder") diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index a5a48347f07517..c468801a6a74db 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -741,6 +741,16 @@ bool IsBiasSemanticAdd(const RemapperContext& ctx, return false; } +bool MaybeCopyOutputShapesAttr(const NodeDef& from, NodeDef* fused_op) { + const string output_shape_attr_name = "_output_shapes"; + if (from.attr().count(output_shape_attr_name) > 0) { + auto shape_attrs = from.attr().at(output_shape_attr_name); + AddNodeAttr(output_shape_attr_name, shape_attrs, fused_op); + return true; + } + return false; +} + void AddInputShapesAttr(const RemapperContext& ctx, int node_index) { auto mutable_node = ctx.graph_view.graph()->mutable_node(node_index); @@ -3334,6 +3344,7 @@ absl::Status AddFusedContractionNode(RemapperContext* ctx, } SetFusedOpAttributes(&fused_op, {"BiasAdd"}); + MaybeCopyOutputShapesAttr(bias_add, &fused_op); utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); absl::Status status; mutation->AddNode(std::move(fused_op), &status); @@ -3439,6 +3450,7 @@ absl::Status AddFusedContractionNode( } SetFusedOpAttributes(&fused_op, {"BiasAdd", activation.op()}); + MaybeCopyOutputShapesAttr(activation, &fused_op); utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); absl::Status status; @@ -3630,6 +3642,7 @@ absl::Status AddFusedContractionNode( } SetFusedOpAttributes(&contraction_node, {"BiasAdd", "Add"}, 2); + MaybeCopyOutputShapesAttr(add, &contraction_node); utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); absl::Status status; @@ -3729,6 +3742,7 @@ absl::Status AddFusedContractionNode( } SetFusedOpAttributes(&fused_conv, {"BiasAdd", "Add", activation.op()}, 2); + MaybeCopyOutputShapesAttr(add, &fused_conv); utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); absl::Status status; diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc index 8d53c6dbb38425..b570286261f75b 100644 --- a/tensorflow/core/ops/array_ops.cc +++ b/tensorflow/core/ops/array_ops.cc @@ -266,15 +266,38 @@ absl::Status SetOutputShapeForReshape(InferenceContext* c) { true /* evenly_divisible */, &inferred_dim)); DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx); + // Record the inferred scale factor on the input-side unknown. TF_RETURN_IF_ERROR( c->Merge(unknown_in_dim, inferred_dim, &unknown_in_dim)); } else if (in_unknown_idx >= 0 && out_unknown_idx >= 0) { // Exactly one unknown dimension in both input and output. These 2 are - // equal iff the known elements are equal. + // related through the ratio of known elements. + DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx); if (c->Value(known_in_elems) == c->Value(known_out_elems)) { - DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx); + // Example: [?,8] -> [?,8], so the unknown is unchanged. TF_RETURN_IF_ERROR( c->ReplaceDim(out, out_unknown_idx, unknown_in_dim, &out)); + } else if (c->Value(known_out_elems) > 0 && + c->Value(known_in_elems) % c->Value(known_out_elems) == 0) { + DimensionHandle inferred_out_dim; + // Example: [?,26,8] -> [?,8], so the output unknown absorbs *26. + TF_RETURN_IF_ERROR(c->Multiply( + unknown_in_dim, + c->Value(known_in_elems) / c->Value(known_out_elems), + &inferred_out_dim)); + TF_RETURN_IF_ERROR( + c->ReplaceDim(out, out_unknown_idx, inferred_out_dim, &out)); + } else if (c->Value(known_in_elems) > 0 && + c->Value(known_out_elems) % c->Value(known_in_elems) == 0) { + DimensionHandle inferred_out_dim; + // Example: [?,8] -> [?,26,8], so the output unknown is the + // input-side unknown divided by 26. + TF_RETURN_IF_ERROR(c->Divide( + unknown_in_dim, + c->Value(known_out_elems) / c->Value(known_in_elems), + true /* evenly_divisible */, &inferred_out_dim)); + TF_RETURN_IF_ERROR( + c->ReplaceDim(out, out_unknown_idx, inferred_out_dim, &out)); } } } @@ -719,6 +742,8 @@ REGISTER_OP("SplitV") REGISTER_OP("Const") .Output("output: dtype") .Attr("value: tensor") + .Attr("user_inferred_shape: shape = {}") + .Attr("has_dynamic: bool = false") .Attr("dtype: type") .SetShapeFn([](InferenceContext* c) { const TensorProto* proto = nullptr; diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc index 93c5a7e9818ae2..a2bb25f9c2e923 100644 --- a/tensorflow/core/util/strided_slice_op.cc +++ b/tensorflow/core/util/strided_slice_op.cc @@ -22,6 +22,8 @@ limitations under the License. #include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/lib/core/status.h" +#include "xla/shape.h" + namespace tensorflow { namespace { @@ -55,6 +57,8 @@ struct StridedSliceDenseSpec { bool end_valid; absl::InlinedVector& begin; absl::InlinedVector& end; + absl::InlinedVector& begin_expr; + absl::InlinedVector& end_expr; absl::InlinedVector& strides; // This vector helps construct the final shape of the slice. // The final tensor is reduced in rank whenever a single index e.g. foo[3] @@ -97,6 +101,8 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse, // to remove any ellipsis dense->begin.resize(dense->dims); dense->end.resize(dense->dims); + dense->begin_expr.resize(dense->dims); + dense->end_expr.resize(dense->dims); dense->strides.resize(dense->dims); dense->input_shape_gather_indices_sparse.resize(dense->dims); // What indices to get the final shape from. @@ -127,6 +133,8 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse, for (; full_index < next_index; full_index++) { // new_axis' aren't real axis so you have to skip dense->begin[full_index] = dense->end[full_index] = 0; + dense->begin_expr[full_index] = dense->end_expr[full_index] = + xla::DExpr::Const(0); dense->strides[full_index] = 1; dense->begin_mask |= (1 << full_index); dense->end_mask |= (1 << full_index); @@ -150,9 +158,13 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse, // Gather slicing spec into appropriate index if (begin_flat != nullptr) { dense->begin[full_index] = internal::SubtleMustCopy(begin_flat[i]); + dense->begin_expr[full_index] = + xla::DExpr::Const(dense->begin[full_index]); } if (end_flat != nullptr) { dense->end[full_index] = internal::SubtleMustCopy(end_flat[i]); + dense->end_expr[full_index] = + xla::DExpr::Const(dense->end[full_index]); } dense->strides[full_index] = internal::SubtleMustCopy(strides_flat[i]); @@ -193,10 +205,29 @@ absl::Status ValidateStridedSliceOp( absl::InlinedVector* begin, absl::InlinedVector* end, absl::InlinedVector* strides, + absl::InlinedVector* begin_expr, + absl::InlinedVector* end_expr, StridedSliceShapeSpec* shape_spec) { + absl::InlinedVector b; + absl::InlinedVector e; + + // HACK + if (begin_expr == nullptr) { + for (int i : *begin) { + b.push_back(xla::DExpr::Const(i)); + } + begin_expr = &b; + } + if (end_expr == nullptr) { + for (int i : *end) { + e.push_back(xla::DExpr::Const(i)); + } + end_expr = &e; + } + if (input_shape.unknown_rank()) { // Note: If the rank is unknown, "input_shape.dims()" is -1. - return errors::InvalidArgument("Unexpected input_shape with unknown rank"); + return errors::InvalidArgument("Unexpected input_shape with unknown rank"); } const bool begin_is_wrong = @@ -271,6 +302,7 @@ absl::Status ValidateStridedSliceOp( // we need to produce the missing begin_mask for the first two // dimensions i.e. from begin_mask_spec=0, end_mask_spec=2 // we achieve begin_mask=6, end_mask=7 + StridedSliceDenseSpec dense_spec = {input_shape.dims(), 0 /* begin_mask */, 0 /* end_mask */, @@ -278,6 +310,8 @@ absl::Status ValidateStridedSliceOp( false /* end_valid */, *begin, *end, + *begin_expr, + *end_expr, *strides}; if (strides_tensor.dtype() == DT_INT32) { @@ -296,17 +330,24 @@ absl::Status ValidateStridedSliceOp( *slice_dim0 = true; *is_simple_slice = true; processing_shape->Clear(); + auto dim_exprs = input_shape.get_filled_expressions(); for (int i = 0; i < input_shape.dims(); ++i) { int64_t& begin_i = (*begin)[i]; int64_t& end_i = (*end)[i]; int64_t& stride_i = (*strides)[i]; int64_t dim_i = input_shape.dim_size(i); + + xla::DExpr dim_i_expr = + i < dim_exprs.size() ? dim_exprs[i] : xla::DExpr::Const(dim_i); + if (stride_i == 0) { return errors::InvalidArgument("strides[", i, "] must be non-zero"); } bool shrink_i = (dense_spec.shrink_axis_mask & (1 << i)); if (dim_i == -1) { processing_shape->AddDim(shrink_i ? 1 : -1); + processing_shape->AddExpression(shrink_i ? xla::DExpr::Const(1) + : xla::DExpr::Const(-1)); continue; } @@ -315,15 +356,37 @@ absl::Status ValidateStridedSliceOp( const std::array valid_range = { {stride_i > 0 ? 0 : -1, stride_i > 0 ? dim_i : dim_i - 1}}; + const std::array valid_range_expr = { + {stride_i > 0 ? xla::DExpr::Const(0) : xla::DExpr::Const(-1), + stride_i > 0 ? dim_i_expr + : (dim_i_expr - xla::DExpr::Const(1)).simplify()}}; + auto canonical = [stride_i, dim_i, masks, valid_range](int64_t x, int c) { if (masks[c]) { return stride_i > 0 ? valid_range[c] : valid_range[(c + 1) & 1]; } else { int64_t x_fwd = x < 0 ? dim_i + x : x; // make negative indices positive - return x_fwd < valid_range[0] - ? valid_range[0] - : x_fwd > valid_range[1] ? valid_range[1] : x_fwd; + return x_fwd < valid_range[0] ? valid_range[0] + : x_fwd > valid_range[1] ? valid_range[1] + : x_fwd; + } + }; + auto canonical_expr = [stride_i, dim_i, masks, valid_range, + valid_range_expr, dim_i_expr](int64_t x, int c) { + if (masks[c]) { + return stride_i > 0 ? valid_range_expr[c] + : valid_range_expr[(c + 1) & 1]; + } else { + int64_t x_fwd = + x < 0 ? dim_i + x : x; // make negative indices positive + xla::DExpr x_expr = xla::DExpr::Const(x); + xla::DExpr x_fwd_expr = + x < 0 ? (dim_i_expr + x_expr) + : x_expr; // make negative indices positive + return x_fwd < valid_range[0] ? valid_range_expr[0] + : x_fwd > valid_range[1] ? valid_range_expr[1] + : x_fwd_expr; } }; if (shrink_i && stride_i <= 0) { @@ -341,15 +404,30 @@ absl::Status ValidateStridedSliceOp( // and canonical puts these to n-1 and 0, which implies a degenerate // interval. Fortunately, it is now safe to re-create end as begin+1. int64_t x_fwd = begin_i < 0 ? dim_i + begin_i : begin_i; + xla::DExpr x_fwd_expr = begin_i < 0 + ? (dim_i_expr + (*begin_expr)[i]).simplify() + : (*begin_expr)[i]; begin_i = x_fwd; end_i = begin_i + 1; + + (*begin_expr)[i] = x_fwd_expr; + (*end_expr)[i] = ((*begin_expr)[i] + xla::DExpr::Const(1)).simplify(); + if (x_fwd < 0 || x_fwd >= dim_i) { return errors::InvalidArgument( "slice index ", begin_i, " of dimension ", i, " out of bounds."); } } else { - begin_i = canonical(begin_i, 0); - end_i = canonical(end_i, 1); + const int64_t begin_raw = begin_i; + const int64_t end_raw = end_i; + begin_i = canonical(begin_raw, 0); + end_i = canonical(end_raw, 1); + if (begin_expr) { + (*begin_expr)[i] = canonical_expr(begin_raw, 0).simplify(); + } + if (end_expr) { + (*end_expr)[i] = canonical_expr(end_raw, 1).simplify(); + } } // Update optimization values bool take_all_in_dimension = @@ -362,14 +440,17 @@ absl::Status ValidateStridedSliceOp( } // Compute the processing shape (the intermediate Eigen will produce) int64_t interval_length; + xla::DExpr interval_length_expr; bool known_interval = false; if (dense_spec.begin_valid && dense_spec.end_valid) { interval_length = end_i - begin_i; + interval_length_expr = ((*end_expr)[i] - (*begin_expr)[i]).simplify(); known_interval = true; } else if (shrink_i) { // The dimension is still known as 1 for the processing_shape, but will be // discarded for the final shape. interval_length = 1; + interval_length_expr = xla::DExpr::Const(1); known_interval = true; } else if (begin_and_end_masked) { // Even if we don't have values for begin or end, we do know that this @@ -378,25 +459,35 @@ absl::Status ValidateStridedSliceOp( if (dim_i >= 0) { if (stride_i < 0) { interval_length = -dim_i; + interval_length_expr = (xla::DExpr::Const(-1) * dim_i_expr).simplify(); } else { interval_length = dim_i; + interval_length_expr = dim_i_expr; } known_interval = true; } } if (known_interval) { int64_t size_i; + xla::DExpr size_i_expr; // Hold zero if the interval is degenerate, otherwise account for // remainder if (interval_length == 0 || ((interval_length < 0) != (stride_i < 0))) { size_i = 0; + size_i_expr = xla::DExpr::Const(0); } else { size_i = interval_length / stride_i + (interval_length % stride_i != 0 ? 1 : 0); + size_i_expr = + (interval_length_expr / xla::DExpr::Const(stride_i)) + + (interval_length % stride_i != 0 ? xla::DExpr::Const(1) + : xla::DExpr::Const(0)); } processing_shape->AddDim(size_i); + processing_shape->AddExpression(size_i_expr.simplify()); } else { processing_shape->AddDim(-1); + processing_shape->AddExpression(xla::DExpr::Const(-1)); } } @@ -425,12 +516,15 @@ absl::Status ValidateStridedSliceOp( dense_spec.final_shape_gather_indices_sparse[dense_dim]; if (gather_index >= 0) { final_shape->AddDim(processing_shape->dim_size(gather_index)); + final_shape->AddExpression( + processing_shape->get_filled_expression(gather_index)); if (shape_spec != nullptr) { shape_spec->output_to_sparse_mapping.push_back(sparse_index); shape_spec->output_to_processing_mapping.push_back(gather_index); } } else if (gather_index == kNewAxis) { final_shape->AddDim(1); + final_shape->AddExpression(xla::DExpr::Const(1)); if (shape_spec != nullptr) { shape_spec->output_to_sparse_mapping.push_back(-1); shape_spec->output_to_processing_mapping.push_back(-1); @@ -441,6 +535,7 @@ absl::Status ValidateStridedSliceOp( return absl::OkStatus(); } + absl::Status ValidateStridedSliceOp( const Tensor* begin_tensor, const Tensor* end_tensor, const Tensor& strides_tensor, const PartialTensorShape& input_shape, @@ -451,16 +546,16 @@ absl::Status ValidateStridedSliceOp( absl::InlinedVector* begin, absl::InlinedVector* end, absl::InlinedVector* strides, + absl::InlinedVector* begin_expr, + absl::InlinedVector* end_expr, StridedSliceShapeSpec* shape_spec) { - // Validate with PartialTensorShape output PartialTensorShape partial_processing_shape, partial_final_shape; TF_RETURN_IF_ERROR(ValidateStridedSliceOp( begin_tensor, end_tensor, strides_tensor, input_shape, begin_mask_spec, end_mask_spec, ellipsis_mask, new_axis_mask, shrink_axis_mask, &partial_processing_shape, &partial_final_shape, is_identity, - is_simple_slice, slice_dim0, begin, end, strides, shape_spec)); - - // Verify that the output shapes are fully known + is_simple_slice, slice_dim0, begin, end, strides, begin_expr, end_expr, + shape_spec)); if (!partial_processing_shape.AsTensorShape(processing_shape) || !partial_final_shape.AsTensorShape(final_shape)) { return errors::Internal("ValidateStridedSliceOp returned partial shapes ", diff --git a/tensorflow/core/util/strided_slice_op.h b/tensorflow/core/util/strided_slice_op.h index 356b77a2a0f5b6..c7240d15adc6c7 100644 --- a/tensorflow/core/util/strided_slice_op.h +++ b/tensorflow/core/util/strided_slice_op.h @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" +#include "xla/shape.h" namespace tensorflow { @@ -74,6 +75,8 @@ absl::Status ValidateStridedSliceOp( absl::InlinedVector* begin, absl::InlinedVector* end, absl::InlinedVector* strides, + absl::InlinedVector* begin_expr = nullptr, + absl::InlinedVector* end_expr = nullptr, StridedSliceShapeSpec* shape_spec = nullptr); // Same as above, but the outputs are TensorShape, not PartialTensorShape @@ -87,6 +90,8 @@ absl::Status ValidateStridedSliceOp( absl::InlinedVector* begin, absl::InlinedVector* end, absl::InlinedVector* strides, + absl::InlinedVector* begin_expr = nullptr, + absl::InlinedVector* end_expr = nullptr, StridedSliceShapeSpec* shape_spec = nullptr); // Simple class for determining if it is possible to broadcast a tensor to a From 4bec8eb8acda620be9b2099975debca368bf2d28 Mon Sep 17 00:00:00 2001 From: Guillermo-Callaghan Date: Mon, 22 Jun 2026 16:50:57 +0100 Subject: [PATCH 04/21] Use the same symbolic ID for inputs Assign one configurable symbolic variable to dynamic dimensions of top-level graph inputs while keeping inferred internal dimensions distinct. Add GraphProperties coverage for shared input symbols and non-input dimensions. --- .../core/grappler/costs/graph_properties.cc | 71 ++++++-- .../grappler/costs/graph_properties_test.cc | 168 ++++++++++++++++++ 2 files changed, 223 insertions(+), 16 deletions(-) diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc index bee993c8a1142c..339ba1c7ed9353 100644 --- a/tensorflow/core/grappler/costs/graph_properties.cc +++ b/tensorflow/core/grappler/costs/graph_properties.cc @@ -39,6 +39,7 @@ limitations under the License. #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/util/env_var.h" namespace tensorflow { namespace grappler { @@ -55,11 +56,31 @@ using TensorVector = absl::InlinedVector; // Some ops treat "-1" specially, different from UnknownDim: // e.g., shape input to Reshape op. const int64_t kUnknownDimFromConst = INT64_MAX; +constexpr char kInputArgSharedSymbolIdEnvVar[] = + "TF_GRAPPLER_INPUT_ARG_SHARED_SYMBOL_ID"; // Skip const value instantiation if the number of elements in a const tensor // is greater than this threshold. const int kThresholdToSkipConstTensorInstantiation = 128; +bool IsPlaceholderOp(absl::string_view op_name) { + return op_name == "Placeholder" || op_name == "PlaceholderV2" || + op_name == "PlaceholderWithDefault"; +} + +int LoadInputArgSharedSymbolId() { + int64_t symbol_id = -1; + absl::Status status = + ReadInt64FromEnvVar(kInputArgSharedSymbolIdEnvVar, -1, &symbol_id); + if (!status.ok()) { + VLOG(1) << "Failed to read " << kInputArgSharedSymbolIdEnvVar + << ": " << status + << ". Falling back to default shared symbol id -1."; + return -1; + } + return static_cast(symbol_id); +} + template struct HashHandle { std::size_t operator()(const Handle& h) const { @@ -653,10 +674,14 @@ class SymbolicShapeRefiner { const GraphView& graph, const absl::flat_hash_map>& fed_ports, const bool aggressive_shape_inference, - const bool enable_dynamic_value_inference) + const bool enable_dynamic_value_inference, + absl::flat_hash_set top_level_input_placeholders) : graph_(graph), function_library_(OpRegistry::Global(), graph.graph()->library()), fed_ports_(fed_ports), + top_level_input_placeholders_( + std::move(top_level_input_placeholders)), + input_arg_shared_symbol_id_(LoadInputArgSharedSymbolId()), aggressive_shape_inference_(aggressive_shape_inference), enable_dynamic_value_inference_(enable_dynamic_value_inference) { graph_def_version_ = graph.graph()->versions().producer(); @@ -665,6 +690,12 @@ class SymbolicShapeRefiner { const GraphView& graph() const { return graph_; } + bool IsInputArg(const NodeDef& node) const { + return node.op() == "_Arg" && + top_level_input_placeholders_.find(node.name()) != + top_level_input_placeholders_.end(); + } + struct NodeContext { const OpRegistrationData* op_data; DataTypeVector input_types; @@ -1419,20 +1450,19 @@ class SymbolicShapeRefiner { return it->second; } InferenceContext* c = GetContext(node); - int var_id = GetOrCreateStableVarId(id); - DimensionHandle dim; - if (node->op() == "_Arg") { - var_id *= -1; - // var_id would be minus when it's argument. - dim = c->UnknownDimWithExpr( - std::make_unique(DimExpr::Var(var_id))); - } else { - dim = c->UnknownDimWithExpr( - std::make_unique(DimExpr::Var(var_id))); - } - VLOG(1) << "[EXPR] GetUnknownOutputDim: node=" << node->name() - << " out=" << index << " dim=" << dim_id << " -> Var(" << var_id - << ")"; + int var_id = + IsInputArg(*node) ? input_arg_shared_symbol_id_ + : GetOrCreateStableVarId(id); + + DimensionHandle dim = c->UnknownDimWithExpr( + std::make_unique(DimExpr::Var(var_id))); + + VLOG(1) << "[EXPR] GetUnknownOutputDim: node = " << node->name() + << " op = " << node->op() + << " out = " << index + << " dim = " << dim_id + << " -> Var(" << var_id << ")"; + // Create an unknown dim with Var(var_id) expression. unknown_dims_[id] = dim; return dim; @@ -2447,6 +2477,8 @@ class SymbolicShapeRefiner { fun_to_grappler_function_item_; FunctionLibraryDefinition function_library_; const absl::flat_hash_map>& fed_ports_; + const absl::flat_hash_set top_level_input_placeholders_; + const int input_arg_shared_symbol_id_; // Store TensorProtos for tensor value propagation. Note that we use deque, // not vector, as we use pointers to the TensorProtos in this container. // Vector may resize and copy the objects into a new buffer, then the existing @@ -3156,9 +3188,16 @@ absl::Status GraphProperties::InferStatically( // Heap-allocate SymbolicShapeRefiner in order to not consume a large amount // of stack space. + absl::flat_hash_set top_level_input_placeholders; + for (const NodeDef& node : item_.graph.node()) { + if (IsPlaceholderOp(node.op()) || node.op() == "_Arg") { + top_level_input_placeholders.insert(node.name()); + } + } auto refiner = std::make_unique( graph_view, fed_ports, aggressive_shape_inference, - enable_dynamic_value_inference); + enable_dynamic_value_inference, + std::move(top_level_input_placeholders)); TopoQueue new_shapes(topo_order); // Also seed the propagation of shapes in the fanout of primary inputs. diff --git a/tensorflow/core/grappler/costs/graph_properties_test.cc b/tensorflow/core/grappler/costs/graph_properties_test.cc index b24ad629910909..d88f9c0cf36033 100644 --- a/tensorflow/core/grappler/costs/graph_properties_test.cc +++ b/tensorflow/core/grappler/costs/graph_properties_test.cc @@ -61,6 +61,22 @@ bool ShapeDimExprEqual(const TensorShapeProto& lhs, int lhs_dim, ShapeDimExpr(rhs, rhs_dim)); } +void ExpectVariableId(const TensorShapeProto& shape, int dim, + int expected_var_id) { + ASSERT_GT(shape.dim_size(), dim); + const auto& expr = shape.dim(dim).expr(); + EXPECT_EQ(expr.node_type_case(), ExpressionProto::kVariableId); + EXPECT_EQ(expr.variable_id(), expected_var_id); +} + +void ExpectNotVariableId(const TensorShapeProto& shape, int dim, + int unexpected_var_id) { + ASSERT_GT(shape.dim_size(), dim); + const auto& expr = shape.dim(dim).expr(); + ASSERT_EQ(expr.node_type_case(), ExpressionProto::kVariableId); + EXPECT_NE(expr.variable_id(), unexpected_var_id); +} + using shape_inference::InferenceContext; using shape_inference::ShapeAndType; using shape_inference::ShapeHandle; @@ -2028,6 +2044,158 @@ TEST_F(GraphPropertiesTest, SymbolicShapes) { EXPECT_TRUE(ShapeDimExprEqual(shape_j, 0, shape_a, 1)); } +TEST_F(GraphPropertiesTest, SymbolicInputShapesShareSameSymbol) { + GrapplerItem item; + + AttrValue shape_a_attr; + TensorShapeProto* shape_a_proto = shape_a_attr.mutable_list()->add_shape(); + shape_a_proto->add_dim()->set_size(-1); + shape_a_proto->add_dim()->set_size(-1); + TF_ASSERT_OK(NodeDefBuilder("a", "_Arg") + .Attr("T", DT_FLOAT) + .Attr("index", 0) + .Attr("_output_shapes", shape_a_attr) + .Finalize(item.graph.add_node())); + + AttrValue shape_b_attr; + TensorShapeProto* shape_b_proto = shape_b_attr.mutable_list()->add_shape(); + shape_b_proto->add_dim()->set_size(-1); + TF_ASSERT_OK(NodeDefBuilder("b", "_Arg") + .Attr("T", DT_FLOAT) + .Attr("index", 1) + .Attr("_output_shapes", shape_b_attr) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("c", "Identity") + .Input(NodeDefBuilder::NodeOut("a", 0, DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Finalize(item.graph.add_node())); + TF_ASSERT_OK(NodeDefBuilder("d", "Identity") + .Input(NodeDefBuilder::NodeOut("b", 0, DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Finalize(item.graph.add_node())); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically(false)); + + const auto shape_a = properties.GetOutputProperties("a").at(0).shape(); + const auto shape_b = properties.GetOutputProperties("b").at(0).shape(); + const auto shape_c = properties.GetOutputProperties("c").at(0).shape(); + const auto shape_d = properties.GetOutputProperties("d").at(0).shape(); + + ASSERT_EQ(2, shape_a.dim_size()); + ASSERT_EQ(1, shape_b.dim_size()); + ExpectVariableId(shape_a, 0, -1); + ExpectVariableId(shape_a, 1, -1); + ExpectVariableId(shape_b, 0, -1); + ExpectVariableId(shape_c, 0, -1); + ExpectVariableId(shape_c, 1, -1); + ExpectVariableId(shape_d, 0, -1); +} + +TEST_F(GraphPropertiesTest, SymbolicNonInputShapesDoNotShareInputSymbol) { + GrapplerItem item; + + AttrValue input_shape_attr; + TensorShapeProto* input_shape_proto = + input_shape_attr.mutable_list()->add_shape(); + input_shape_proto->add_dim()->set_size(-1); + input_shape_proto->add_dim()->set_size(-1); + TF_ASSERT_OK(NodeDefBuilder("input", "_Arg") + .Attr("T", DT_FLOAT) + .Attr("index", 0) + .Attr("_output_shapes", input_shape_attr) + .Finalize(item.graph.add_node())); + + AttrValue reshape_shape_attr; + TensorShapeProto* reshape_shape_proto = + reshape_shape_attr.mutable_list()->add_shape(); + reshape_shape_proto->add_dim()->set_size(2); + TF_ASSERT_OK(NodeDefBuilder("target_shape", "_Arg") + .Attr("T", DT_INT32) + .Attr("index", 1) + .Attr("_output_shapes", reshape_shape_attr) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("reshaped", "Reshape") + .Input(NodeDefBuilder::NodeOut("input", 0, DT_FLOAT)) + .Input( + NodeDefBuilder::NodeOut("target_shape", 0, DT_INT32)) + .Attr("T", DT_FLOAT) + .Attr("Tshape", DT_INT32) + .Finalize(item.graph.add_node())); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically(false)); + + const auto input_shape = properties.GetOutputProperties("input").at(0).shape(); + const auto reshaped_shape = + properties.GetOutputProperties("reshaped").at(0).shape(); + + ASSERT_EQ(2, input_shape.dim_size()); + ASSERT_EQ(2, reshaped_shape.dim_size()); + ExpectVariableId(input_shape, 0, -1); + ExpectVariableId(input_shape, 1, -1); + ExpectNotVariableId(reshaped_shape, 0, -1); + ExpectNotVariableId(reshaped_shape, 1, -1); +} + +TEST_F(GraphPropertiesTest, SymbolicInputShapesShareOnlyDynamicDims) { + GrapplerItem item; + + AttrValue shape_a_attr; + TensorShapeProto* shape_a_proto = shape_a_attr.mutable_list()->add_shape(); + shape_a_proto->add_dim()->set_size(4); + shape_a_proto->add_dim()->set_size(-1); + TF_ASSERT_OK(NodeDefBuilder("a", "_Arg") + .Attr("T", DT_FLOAT) + .Attr("index", 0) + .Attr("_output_shapes", shape_a_attr) + .Finalize(item.graph.add_node())); + + AttrValue shape_b_attr; + TensorShapeProto* shape_b_proto = shape_b_attr.mutable_list()->add_shape(); + shape_b_proto->add_dim()->set_size(-1); + shape_b_proto->add_dim()->set_size(8); + TF_ASSERT_OK(NodeDefBuilder("b", "_Arg") + .Attr("T", DT_FLOAT) + .Attr("index", 1) + .Attr("_output_shapes", shape_b_attr) + .Finalize(item.graph.add_node())); + + TF_ASSERT_OK(NodeDefBuilder("c", "Identity") + .Input(NodeDefBuilder::NodeOut("a", 0, DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Finalize(item.graph.add_node())); + TF_ASSERT_OK(NodeDefBuilder("d", "Identity") + .Input(NodeDefBuilder::NodeOut("b", 0, DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Finalize(item.graph.add_node())); + + GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically(false)); + + const auto shape_a = properties.GetOutputProperties("a").at(0).shape(); + const auto shape_b = properties.GetOutputProperties("b").at(0).shape(); + const auto shape_c = properties.GetOutputProperties("c").at(0).shape(); + const auto shape_d = properties.GetOutputProperties("d").at(0).shape(); + + ASSERT_EQ(2, shape_a.dim_size()); + ASSERT_EQ(2, shape_b.dim_size()); + ASSERT_EQ(2, shape_c.dim_size()); + ASSERT_EQ(2, shape_d.dim_size()); + + EXPECT_EQ(shape_a.dim(0).size(), 4); + EXPECT_EQ(shape_b.dim(1).size(), 8); + EXPECT_EQ(shape_c.dim(0).size(), 4); + EXPECT_EQ(shape_d.dim(1).size(), 8); + + ExpectVariableId(shape_a, 1, -1); + ExpectVariableId(shape_b, 0, -1); + ExpectVariableId(shape_c, 1, -1); + ExpectVariableId(shape_d, 0, -1); +} + TEST_F(GraphPropertiesTest, DoNotValidateColocationConstraints) { tensorflow::Scope s = tensorflow::Scope::NewRootScope(); Output a = ops::Const(s.WithOpName("a"), 1.0f, {1}); From 077f22d323f1a7eaf39f7f7033aebc25e498a2d8 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:03:29 +0100 Subject: [PATCH 05/21] Preserve dynamic expressions through HLO construction and inference Carry dynamic expressions through XLA shape construction, HLO instructions, and shape inference. This keeps symbolic dimensions available while HLO graphs are built and transformed. --- third_party/xla/xla/hlo/builder/BUILD | 1 + .../xla/xla/hlo/builder/lib/approx_topk.cc | 8 +- .../xla/xla/hlo/builder/lib/arithmetic.cc | 5 +- .../xla/xla/hlo/builder/lib/broadcast.cc | 41 +- .../xla/xla/hlo/builder/lib/broadcast.h | 5 +- .../xla/xla/hlo/builder/lib/constants.cc | 3 +- third_party/xla/xla/hlo/builder/lib/matrix.cc | 15 +- third_party/xla/xla/hlo/builder/lib/prng.cc | 20 +- .../xla/xla/hlo/builder/lib/slicing.cc | 44 +- third_party/xla/xla/hlo/builder/lib/svd.cc | 23 +- .../xla/xla/hlo/builder/value_inference.cc | 38 + .../xla/xla/hlo/builder/xla_builder.cc | 326 ++++-- third_party/xla/xla/hlo/builder/xla_builder.h | 103 +- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 116 +- third_party/xla/xla/hlo/ir/hlo_instruction.h | 23 + .../xla/xla/hlo/ir/hlo_instructions.cc | 42 +- third_party/xla/xla/hlo/ir/hlo_instructions.h | 34 + .../xla/xla/hlo/pass/hlo_pass_pipeline.cc | 1 + .../collectives/all_gather_combiner.cc | 4 +- .../collectives/collective_quantizer.cc | 1 + .../expanders/bitcast_dtypes_expander.cc | 11 +- .../transforms/expanders/cholesky_expander.cc | 6 +- .../transforms/expanders/dot_decomposer.cc | 60 +- .../hlo/transforms/expanders/eigh_expander.cc | 15 +- .../hlo/transforms/expanders/qr_expander.cc | 24 +- .../transforms/expanders/reduce_decomposer.cc | 1 + .../expanders/rng_bit_generator_expander.cc | 6 +- .../service/gpu/transforms/gemm_rewriter.cc | 2 + .../gpu/transforms/windowed_einsum_handler.cc | 2 + third_party/xla/xla/service/hlo.proto | 7 +- .../xla/xla/service/hlo_creation_utils.cc | 18 +- third_party/xla/xla/service/hlo_cse.cc | 33 + .../xla/xla/service/layout_assignment.cc | 4 +- .../xla/service/reduce_scatter_combiner.cc | 6 +- .../xla/xla/service/shape_inference.cc | 276 ++++- third_party/xla/xla/service/shape_inference.h | 16 +- .../xla/xla/service/shape_inference_test.cc | 989 +++++++++++++++++- .../xla/service/triangular_solve_expander.cc | 9 +- third_party/xla/xla/shape_test.cc | 177 +++- third_party/xla/xla/shape_util_test.cc | 30 +- 40 files changed, 2276 insertions(+), 269 deletions(-) diff --git a/third_party/xla/xla/hlo/builder/BUILD b/third_party/xla/xla/hlo/builder/BUILD index 43dcce69c83579..0a508f54385cad 100644 --- a/third_party/xla/xla/hlo/builder/BUILD +++ b/third_party/xla/xla/hlo/builder/BUILD @@ -174,6 +174,7 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@local_tsl//tsl/platform:errors", + "@local_tsl//tsl/platform:protobuf", "@local_tsl//tsl/platform:stacktrace", "@local_tsl//tsl/platform:statusor", ], diff --git a/third_party/xla/xla/hlo/builder/lib/approx_topk.cc b/third_party/xla/xla/hlo/builder/lib/approx_topk.cc index 5b0f223302f4ed..9cb5fa5eefa2b8 100644 --- a/third_party/xla/xla/hlo/builder/lib/approx_topk.cc +++ b/third_party/xla/xla/hlo/builder/lib/approx_topk.cc @@ -131,10 +131,10 @@ XlaOp AggregateToTopKBuilder(XlaBuilder* builder, reduction_computation, {reduction_dim}); Shape op_shape = operands_shapes[0]; op_shape.set_dimensions(reduction_dim, 1); - auto top1_vals = - Reshape(GetTupleElement(val_args, 0), op_shape.dimensions()); - auto top1_args = - Reshape(GetTupleElement(val_args, 1), op_shape.dimensions()); + auto top1_vals = Reshape(GetTupleElement(val_args, 0), + op_shape.dimensions(), op_shape.expressions()); + auto top1_args = Reshape(GetTupleElement(val_args, 1), + op_shape.dimensions(), op_shape.expressions()); return Tuple(builder, {top1_vals, top1_args}); } diff --git a/third_party/xla/xla/hlo/builder/lib/arithmetic.cc b/third_party/xla/xla/hlo/builder/lib/arithmetic.cc index 84908846dae705..7bb8694351ce92 100644 --- a/third_party/xla/xla/hlo/builder/lib/arithmetic.cc +++ b/third_party/xla/xla/hlo/builder/lib/arithmetic.cc @@ -154,8 +154,9 @@ XlaOp ArgMinMax(XlaOp input, PrimitiveType output_type, int axis, bool is_min) { int64_t dimension_size = input_shape.dimensions(axis); auto index_type = dimension_size <= INT32_MAX ? S32 : output_type; XlaOp index_init_value = Zero(builder, index_type); - auto iota_shape = - ShapeUtil::MakeShape(index_type, input_shape.dimensions()); + auto iota_shape = ShapeUtil::MakeShape(index_type, input_shape.dimensions(), + input_shape.expressions()); + XlaOp iota = Iota(builder, iota_shape, axis); XlaComputation reducer = CreateMinMaxComputation( diff --git a/third_party/xla/xla/hlo/builder/lib/broadcast.cc b/third_party/xla/xla/hlo/builder/lib/broadcast.cc index 1baec363b53a35..46f5410f2a817a 100644 --- a/third_party/xla/xla/hlo/builder/lib/broadcast.cc +++ b/third_party/xla/xla/hlo/builder/lib/broadcast.cc @@ -31,8 +31,9 @@ limitations under the License. namespace xla { -absl::StatusOr BroadcastTo(XlaOp input, - absl::Span output_dims) { +absl::StatusOr BroadcastTo( + XlaOp input, absl::Span output_dims, + absl::Span output_exprs) { XlaBuilder* builder = input.builder(); TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); absl::Span input_dims = input_shape.dimensions(); @@ -48,6 +49,12 @@ absl::StatusOr BroadcastTo(XlaOp input, absl::StrJoin(output_dims, ","), "]"); } + if (!output_exprs.empty() && output_exprs.size() != output_dims.size()) { + return tsl::errors::InvalidArgument( + "output_exprs must be empty or have the same rank as output_dims: ", + output_exprs.size(), " vs ", output_dims.size()); + } + std::vector broadcast_dims; std::vector broadcast_shape; auto input_it = input_dims.rbegin(); @@ -79,15 +86,41 @@ absl::StatusOr BroadcastTo(XlaOp input, } TF_RET_CHECK(input_it == input_dims.rend()); + absl::Span input_exprs = input_shape.expressions(); + std::vector broadcast_exprs; + auto input_dim_et = input_dims.rbegin(); + auto input_et = input_exprs.rbegin(); + auto output_dim_et = output_dims.rbegin(); + for (auto output_et = output_exprs.rbegin(); output_et != output_exprs.rend(); + ++output_et, ++output_dim_et) { + if (input_et != input_exprs.rend()) { + if (*output_dim_et == *input_dim_et || *input_dim_et == 1 || + **output_et == **input_et || + (input_et->get()->is_constant() && input_et->get()->get_val() == 1)) { + broadcast_exprs.push_back(*output_et); + } else if (!(**output_et == **input_et)) { + broadcast_exprs.push_back(*input_et); + broadcast_exprs.push_back(*output_et / *input_et); + } + ++input_dim_et; + ++input_et; + } else { + broadcast_exprs.push_back(*output_et); + } + } + absl::c_reverse(broadcast_dims); int broadcast_shape_size = broadcast_shape.size(); for (int64_t& broadcast_dim : broadcast_dims) { broadcast_dim = broadcast_shape_size - broadcast_dim - 1; } absl::c_reverse(broadcast_shape); - XlaOp output = BroadcastInDim(input, broadcast_shape, broadcast_dims); + absl::c_reverse(broadcast_exprs); + + XlaOp output = + BroadcastInDim(input, broadcast_shape, broadcast_dims, broadcast_exprs); if (broadcast_shape != output_dims) { - output = Reshape(output, output_dims); + output = Reshape(output, output_dims, output_exprs); } return output; } diff --git a/third_party/xla/xla/hlo/builder/lib/broadcast.h b/third_party/xla/xla/hlo/builder/lib/broadcast.h index 86cf39f64ddc82..dd5535f219b971 100644 --- a/third_party/xla/xla/hlo/builder/lib/broadcast.h +++ b/third_party/xla/xla/hlo/builder/lib/broadcast.h @@ -27,8 +27,9 @@ namespace xla { // Broadcasts 'input' up to shape 'output_dims', using TensorFlow broadcasting // rules. Supports broadcasting a dimension of size x to size x*y, i.e., tiling. -absl::StatusOr BroadcastTo(XlaOp input, - absl::Span output_dims); +absl::StatusOr BroadcastTo( + XlaOp input, absl::Span output_dims, + absl::Span output_exprs = {}); } // namespace xla diff --git a/third_party/xla/xla/hlo/builder/lib/constants.cc b/third_party/xla/xla/hlo/builder/lib/constants.cc index acfa2fe0b66e2c..da6a97f4088520 100644 --- a/third_party/xla/xla/hlo/builder/lib/constants.cc +++ b/third_party/xla/xla/hlo/builder/lib/constants.cc @@ -33,7 +33,8 @@ XlaOp Zero(XlaBuilder* builder, PrimitiveType type) { } XlaOp Zeros(XlaBuilder* builder, const Shape& shape) { - return Broadcast(Zero(builder, shape.element_type()), shape.dimensions()); + return Broadcast(Zero(builder, shape.element_type()), shape.dimensions(), + shape.expressions()); } XlaOp ZerosLike(XlaOp prototype) { diff --git a/third_party/xla/xla/hlo/builder/lib/matrix.cc b/third_party/xla/xla/hlo/builder/lib/matrix.cc index e9fe29ea83ee6b..077a81c1ce38b4 100644 --- a/third_party/xla/xla/hlo/builder/lib/matrix.cc +++ b/third_party/xla/xla/hlo/builder/lib/matrix.cc @@ -209,7 +209,8 @@ XlaOp SetMatrixDiagonal(XlaOp matrix, XlaOp diag, int k) { } return Select(GetDiagonalMask(matrix, k), - BroadcastInDim(diag, shape.dimensions(), broadcast_dims), + BroadcastInDim(diag, shape.dimensions(), broadcast_dims, + shape.expressions()), matrix); }); } @@ -341,18 +342,22 @@ xla::XlaOp EinsumInverseDiagonal(XlaOp x, absl::Span config) { } TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x)); std::vector broadcast_sizes; + std::vector broadcast_exprs; int64_t x_dim = 0; for (auto label = config.begin(); label != config.end(); ++label) { auto first_label = absl::c_find(config, *label); if (first_label == label) { broadcast_sizes.push_back(x_shape.dimensions(x_dim)); + broadcast_exprs.push_back(x_shape.expressions(x_dim)); ++x_dim; } else { broadcast_sizes.push_back( broadcast_sizes[first_label - config.begin()]); + broadcast_exprs.push_back( + broadcast_exprs[first_label - config.begin()]); } } - x = BroadcastInDim(x, broadcast_sizes, labels->at(2)); + x = BroadcastInDim(x, broadcast_sizes, labels->at(2), broadcast_exprs); return EinsumDiagonalMask(x, config); }); } @@ -568,16 +573,20 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, int64_t dot_dim = 0; std::vector new_dims; + std::vector new_exprs; new_dims.reserve(output_rank); + new_exprs.reserve(output_rank); TF_ASSIGN_OR_RETURN(Shape dot_shape, builder->GetShape(dot)); for (auto d : output_config) { if (is_output_only(d)) { new_dims.push_back(1); + new_exprs.push_back(DExpr::Const(1)); } else { new_dims.push_back(dot_shape.dimensions(dot_dim)); + new_exprs.push_back(dot_shape.expressions(dot_dim)); } } - return Reshape(dot, new_dims); + return Reshape(dot, new_dims, new_exprs); }); } diff --git a/third_party/xla/xla/hlo/builder/lib/prng.cc b/third_party/xla/xla/hlo/builder/lib/prng.cc index 661981f19c17ea..7bb1c619f824d5 100644 --- a/third_party/xla/xla/hlo/builder/lib/prng.cc +++ b/third_party/xla/xla/hlo/builder/lib/prng.cc @@ -41,8 +41,9 @@ namespace xla { xla::XlaOp ConcatScalars(xla::XlaBuilder* builder, absl::Span scalars) { std::vector vectors; - absl::c_transform(scalars, std::back_inserter(vectors), - [](xla::XlaOp x) { return xla::Reshape(x, {1}); }); + absl::c_transform(scalars, std::back_inserter(vectors), [](xla::XlaOp x) { + return xla::Reshape(x, {1}, {xla::DExpr::Const(1)}); + }); return ConcatInDim(builder, vectors, 0); } @@ -154,7 +155,8 @@ std::pair GetThreeFryInputsAndUpdatedState( XlaBuilder* builder = initial_state.builder(); auto u64_shape = ShapeUtil::MakeShape(U64, shape.dimensions()); // initial_state is an R1, so reshape it to a scalar. - auto input_u64 = Broadcast(Reshape(initial_state, {}), shape.dimensions()); + auto input_u64 = Broadcast(Reshape(initial_state, {}), shape.dimensions(), + shape.expressions()); int64_t trailing_dims_product = 1; for (int64_t i = shape.dimensions().size() - 1; i >= 0; --i) { if (shape.dimensions(i) < 2) { @@ -245,8 +247,12 @@ XlaOp CombineShapePair(absl::Span pair, original_shape.dimensions(shape_pair.split_dim); std::vector reshape_dims(original_shape.dimensions().begin(), original_shape.dimensions().end()); + std::vector reshape_exprs(original_shape.expressions().begin(), + original_shape.expressions().end()); reshape_dims[shape_pair.split_dim] = RoundUpTo(pre_split_size, 2); - result = Reshape(result, reshape_dims); + reshape_exprs[shape_pair.split_dim] = + DExpr::Const(RoundUpTo(pre_split_size, 2)); + result = Reshape(result, reshape_dims, reshape_exprs); if (reshape_dims[shape_pair.split_dim] != pre_split_size) { result = Slice(result, std::vector(original_shape.dimensions().size(), 0), @@ -453,11 +459,11 @@ RngOutput PhiloxRngBit32(XlaOp op_key, XlaOp initial_state, } XlaOp numbers = ConcatInDim(builder, {bits[0], bits[1], bits[2], bits[3]}, /*dimension=*/1); - numbers = Reshape(numbers, {bits_len * 4}); + numbers = Reshape(numbers, {bits_len * 4}, {}); numbers = Slice(numbers, /*start_indices=*/{0}, /*limit_indices=*/{num_elems}, /*strides=*/{1}); - return {Reshape(numbers, shape.dimensions()), new_state}; + return {Reshape(numbers, shape.dimensions(), shape.expressions()), new_state}; } // Generates an array of primitive type U16 with the given shape containing @@ -507,7 +513,7 @@ RngOutput PhiloxRngBit64(XlaOp op_key, XlaOp initial_state, numbers = Slice(numbers, /*start_indices=*/{0}, /*limit_indices=*/{num_elems}, /*strides=*/{1}); - return {Reshape(numbers, shape.dimensions()), new_state}; + return {Reshape(numbers, shape.dimensions(), shape.expressions()), new_state}; } XlaOp ConvertRandomBitsToUniformFloatingPoint(XlaOp bits, XlaOp minval, diff --git a/third_party/xla/xla/hlo/builder/lib/slicing.cc b/third_party/xla/xla/hlo/builder/lib/slicing.cc index ae0f6b987497a1..553936d5790a95 100644 --- a/third_party/xla/xla/hlo/builder/lib/slicing.cc +++ b/third_party/xla/xla/hlo/builder/lib/slicing.cc @@ -168,13 +168,16 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) { std::vector index_broadcast_dims; std::vector input_broadcast_dims; std::vector sizes; + std::vector expressions; sizes.reserve(index_shape.dimensions().size()); + expressions.reserve(index_shape.expressions().size()); for (int64_t i = 0; i < index_shape.dimensions().size(); ++i) { if (i < dim) { input_broadcast_dims.push_back(i); index_broadcast_dims.push_back(i); } else if (i == dim) { sizes.push_back(input_shape.dimensions(i)); + expressions.push_back(input_shape.expressions(i)); input_broadcast_dims.push_back(i); index_broadcast_dims.push_back(i + 1); } else { @@ -182,15 +185,18 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) { index_broadcast_dims.push_back(i + 1); } sizes.push_back(index_shape.dimensions(i)); + expressions.push_back(index_shape.expressions(i)); } - auto mask = Eq( - BroadcastInDim(index, sizes, index_broadcast_dims), - Iota(builder, ShapeUtil::MakeShape(index_shape.element_type(), sizes), - dim)); + auto mask = + Eq(BroadcastInDim(index, sizes, index_broadcast_dims, expressions), + Iota(builder, + ShapeUtil::MakeShape(index_shape.element_type(), sizes, + expressions), + dim)); auto masked_input = Select( - mask, BroadcastInDim(input, sizes, input_broadcast_dims), - Zeros(builder, - ShapeUtil::MakeShape(input_shape.element_type(), sizes))); + mask, BroadcastInDim(input, sizes, input_broadcast_dims, expressions), + Zeros(builder, ShapeUtil::MakeShape(input_shape.element_type(), sizes, + expressions))); return Reduce(masked_input, Zero(builder, input_shape.element_type()), CreateScalarIdentityWithZeroComputation( input_shape.element_type(), builder), @@ -203,7 +209,8 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) { to_concat.reserve(input_shape.dimensions().size()); for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) { if (i == dim) { - to_concat.push_back(Reshape(index, index_shape.dimensions())); + to_concat.push_back(Reshape(index, index_shape.dimensions(), + index_shape.expressions())); } else { to_concat.push_back(Iota(builder, index_shape, i)); } @@ -229,27 +236,33 @@ XlaOp TorchScatterDense(XlaOp input, XlaOp index, XlaOp src, int64_t dim, TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); std::vector index_broadcast_dims; std::vector sizes; + std::vector expressions; const auto rank = index_shape.dimensions().size(); sizes.reserve(rank + 1); + expressions.reserve(rank + 1); for (int64_t i = 0; i < index_shape.dimensions().size(); ++i) { if (i < dim) { index_broadcast_dims.push_back(i); } else { if (i == dim) { sizes.push_back(input_shape.dimensions(i)); + expressions.push_back(input_shape.expressions(i)); } index_broadcast_dims.push_back(i + 1); } sizes.push_back(index_shape.dimensions(i)); + expressions.push_back(index_shape.expressions(i)); } auto mask = - Eq(BroadcastInDim(index, sizes, index_broadcast_dims), + Eq(BroadcastInDim(index, sizes, index_broadcast_dims, expressions), Iota(builder, - ShapeUtil::MakeShape(index_shape.element_type(), sizes), dim)); - auto masked_src = - Select(mask, BroadcastInDim(src, sizes, index_broadcast_dims), - Zeros(builder, - ShapeUtil::MakeShape(input_shape.element_type(), sizes))); + ShapeUtil::MakeShape(index_shape.element_type(), sizes, + expressions), + dim)); + auto masked_src = Select( + mask, BroadcastInDim(src, sizes, index_broadcast_dims, expressions), + Zeros(builder, ShapeUtil::MakeShape(input_shape.element_type(), sizes, + expressions))); return combiner( input, @@ -287,7 +300,8 @@ XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64_t dim, for (int64_t batch_dim = 0; batch_dim < batch_dims; ++batch_dim) { to_concat.push_back(Iota(builder, iota_shape, batch_dim)); } - to_concat.push_back(Reshape(index, index_shape.dimensions())); + to_concat.push_back( + Reshape(index, index_shape.dimensions(), index_shape.expressions())); index = ConcatInDim(builder, to_concat, gather_dnums.index_vector_dim()); } for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) { diff --git a/third_party/xla/xla/hlo/builder/lib/svd.cc b/third_party/xla/xla/hlo/builder/lib/svd.cc index 561c107ba8085a..9254085e69a85e 100644 --- a/third_party/xla/xla/hlo/builder/lib/svd.cc +++ b/third_party/xla/xla/hlo/builder/lib/svd.cc @@ -152,9 +152,9 @@ absl::StatusOr HouseRow( auto beta = Div(ScalarLike(v_0j, 2.0), (Square(Div(sigma, v_0j, broadcast_dims)) + one)); - v = Select( - BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, - v / v_0j); + v = Select(BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), + broadcast_dims, x_shape.expressions()), + v, v / v_0j); v = Select(Eq(idx, j), zeros + one, v); beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps), @@ -219,9 +219,9 @@ absl::StatusOr HouseCol( auto beta = Div(ScalarLike(v_0i, 2.0), (Square(Div(sigma, v_0i, broadcast_dims)) + one)); - v = Select( - BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v, - v / v_0i); + v = Select(BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), + broadcast_dims, x_shape.expressions()), + v, v / v_0i); v = Select(Eq(idx, i), zeros + one, v); beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps), @@ -582,11 +582,11 @@ absl::StatusOr ComputeToleranceComparison(XlaOp w, XlaOp epsilon) { diag = Select(Lt(diag, ZerosLike(diag)), -diag, diag); std::vector broadcasted_dims(num_dims - 1); std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0); - auto broadcast_to_rows = - BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); + auto broadcast_to_rows = BroadcastInDim( + diag, shape.dimensions(), broadcasted_dims, shape.expressions()); broadcasted_dims.back() = num_dims - 1; - auto broadcast_to_columns = - BroadcastInDim(diag, shape.dimensions(), broadcasted_dims); + auto broadcast_to_columns = BroadcastInDim( + diag, shape.dimensions(), broadcasted_dims, shape.expressions()); // Compute tolerance = w_{i,i} * w_{j,j} * epsilon^2 // Use at least F32 precision to avoid precision issues with small denormal. XlaOp tolerance; @@ -745,6 +745,7 @@ absl::StatusOr SortBySingularValuesAndPostProcessing( TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(result.d)); const int64_t num_dims = shape.dimensions().size(); auto dimensions = shape.dimensions(); + auto expressions = shape.expressions(); const int64_t m = ShapeUtil::GetDimension(shape, -2); const int64_t n = ShapeUtil::GetDimension(shape, -1); @@ -763,7 +764,7 @@ absl::StatusOr SortBySingularValuesAndPostProcessing( d = Select(Ge(d, zeros), d, -d); result.v = Mul(result.v, sign, broadcast_dims); - d = BroadcastInDim(d, dimensions, broadcast_dims); + d = BroadcastInDim(d, dimensions, broadcast_dims, expressions); // As m >= n, only first n column vectors need to be permuted, and the rest of // m - n vectors are appended after the sorting is done. diff --git a/third_party/xla/xla/hlo/builder/value_inference.cc b/third_party/xla/xla/hlo/builder/value_inference.cc index 169e8284aa96f5..ba72ae387e7b00 100644 --- a/third_party/xla/xla/hlo/builder/value_inference.cc +++ b/third_party/xla/xla/hlo/builder/value_inference.cc @@ -795,6 +795,26 @@ absl::StatusOr PostorderDFSVisitor::AnalyzeUpperBound( return Literal::CreateFromProto(root->literal()); } }); + } else if (root->custom_call_target() == "GetExpressionValue") { + return PostorderDFSNode() + .AddDependency(root->operand_ids(0), + PostorderDFSNodeType::kConstantUpperBound, + InferenceContext({}, {})) + .AddVisit([](Literal carrier) -> absl::StatusOr { + if (carrier.shape().dimensions_size() != 1 || + carrier.shape().dimensions(0) != 1) { + return InvalidArgument( + "GetExpressionValue carrier must be rank-1 with one " + "element, got %s", + carrier.shape().ToString()); + } + if (carrier.shape().element_type() != S32) { + return InvalidArgument( + "GetExpressionValue carrier must be s32, got %s", + carrier.shape().ToString()); + } + return LiteralUtil::CreateR0(carrier.Get({0})); + }); } else if (root->custom_call_target() == "Sharding") { return PostorderDFSNode() .AddDependency(root->operand_ids(0), @@ -975,6 +995,14 @@ absl::StatusOr PostorderDFSVisitor::AnalyzeConstant( return Literal::CreateFromProto(root->literal()); } }); + } else if (root->custom_call_target() == "GetExpressionValue") { + return PostorderDFSNode().AddVisit( + [root, context](absl::Span) -> absl::StatusOr { + TF_ASSIGN_OR_RETURN(Shape root_shape, + Shape::FromProto(root->shape())); + return CreateGarbageLiteral( + ShapeUtil::GetSubshape(root_shape, context.shape_index)); + }); } else if (root->custom_call_target() == "Sharding") { return PostorderDFSNode() .AddDependency(root->operand_ids(0), @@ -1507,6 +1535,16 @@ absl::StatusOr PostorderDFSVisitor::AnalyzeIsDynamic( } } }); + } else if (root->custom_call_target() == "GetExpressionValue") { + return PostorderDFSNode().AddVisit( + [type, root]() -> absl::StatusOr { + TF_ASSIGN_OR_RETURN(Shape root_shape, + Shape::FromProto(root->shape())); + if (type == PostorderDFSNodeType::kValueIsDynamic) { + return CreatePredLiteral(true, root_shape); + } + return CreatePredLiteral(false, root_shape); + }); } else if (root->custom_call_target() == "Sharding") { return result.AddVisit([](Literal operand) { return operand; }); } else { diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index 93d7782de50e03..7c868a41610406 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -38,6 +38,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/match.h" +#include "tsl/platform/protobuf.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/str_split.h" @@ -56,6 +57,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/permutation_util.h" +#include "xla/printer.h" #include "xla/primitive_util.h" #include "xla/service/hlo.pb.h" #include "xla/service/shape_inference.h" @@ -518,6 +520,20 @@ static std::string ShapeToString(const ShapeProto& shape) { return absl::StrCat("[", absl::StrJoin(shape.dimensions(), ", "), "]"); } +static std::vector ContentsToProto( + absl::Span contents) { + std::vector protos; + protos.reserve(contents.size()); + for (const DExpr& expr : contents) { + ExpressionProto proto; + if (expr) { + expr.to_proto(&proto); + } + protos.push_back(std::move(proto)); + } + return protos; +} + void XlaBuilder::ToStringHelper(std::string* out, int ident, int64_t op_handle) const { const HloInstructionProto& instr = @@ -707,6 +723,37 @@ absl::Status XlaBuilder::SetInstructionFrontendAttribute(const XlaOp op, return absl::OkStatus(); } +absl::Status XlaBuilder::SetInstructionContents(XlaOp op, + std::vector contents) { + auto it = handle_to_index_.find(op.handle()); + if (it == handle_to_index_.end()) { + return InvalidArgument("No XlaOp with handle %d", op.handle()); + } + const bool has_dynamic_content = + absl::c_any_of(contents, [](const DExpr& expr) { + return expr && expr->is_dynamic(); + }); + if (!has_dynamic_content) { + VLOG(1) << "SetInstructionContents clearing non-dynamic contents for op " + << op.handle() << " size=" << contents.size(); + contents.clear(); + } else { + VLOG(1) << "SetInstructionContents keeping dynamic contents for op " + << op.handle() << " size=" << contents.size(); + } + instruction_contents_.at(it->second) = std::move(contents); + return absl::OkStatus(); +} + +absl::StatusOr*> XlaBuilder::GetInstructionContents( + XlaOp op) const { + auto it = handle_to_index_.find(op.handle()); + if (it == handle_to_index_.end()) { + return InvalidArgument("No XlaOp with handle %d", op.handle()); + } + return &instruction_contents_.at(it->second); +} + absl::Status XlaBuilder::SetInstructionSharding( XlaOp op, const std::optional& sharding) { TF_ASSIGN_OR_RETURN(auto instr_proto, LookUpMutableInstruction(op)); @@ -784,7 +831,13 @@ absl::StatusOr XlaBuilder::Build( *entry.mutable_program_shape() = program_shape.ToProto(); entry.set_root_id(root_id); - for (auto& instruction : instructions_) { + for (size_t index = 0; index < instructions_.size(); ++index) { + auto& instruction = instructions_[index]; + if (!instruction_contents_[index].empty()) { + for (const auto& content : ContentsToProto(instruction_contents_[index])) { + *instruction.add_contents() = content; + } + } // Ensures that the instruction names are unique among the whole graph. instruction.set_name( GetFullName(instruction.name(), kNameSeparator, instruction.id())); @@ -810,6 +863,7 @@ absl::StatusOr XlaBuilder::Build( // Clear data held by this builder. this->instructions_.clear(); this->instruction_shapes_.clear(); + this->instruction_contents_.clear(); this->handle_to_index_.clear(); this->embedded_.clear(); this->parameter_numbers_.clear(); @@ -1017,22 +1071,24 @@ absl::StatusOr XlaBuilder::AddBroadcastSequence( Shape broadcast_shape = ShapeUtil::ChangeElementType(output_shape, operand_shape->element_type()); - // Do explicit broadcast for scalar. - if (ShapeUtil::IsScalar(*operand_shape)) { - return InDimBroadcast(ShapeUtil::MakeStaticShape(broadcast_shape), operand, - {}); - } + // Do explicit broadcast for scalar. + if (ShapeUtil::IsScalar(*operand_shape)) { + return InDimBroadcast(ShapeUtil::MakeStaticShape(broadcast_shape), + operand, {}); + } // Do explicit broadcast for degenerate broadcast. std::vector broadcast_dimensions; std::vector reshaped_dimensions; std::vector reshaped_dynamic_dimensions; + std::vector reshaped_expressions; for (int i = 0; i < operand_shape->dimensions().size(); i++) { if (operand_shape->dimensions(i) == output_shape.dimensions(i)) { broadcast_dimensions.push_back(i); reshaped_dimensions.push_back(operand_shape->dimensions(i)); reshaped_dynamic_dimensions.push_back( operand_shape->is_dynamic_dimension(i)); + reshaped_expressions.push_back(operand_shape->expressions(i)); } else { TF_RET_CHECK(operand_shape->dimensions(i) == 1 && operand_shape->is_static_dimension(i)) @@ -1046,7 +1102,8 @@ absl::StatusOr XlaBuilder::AddBroadcastSequence( Shape reshaped_shape = ShapeUtil::MakeShape(operand_shape->element_type(), reshaped_dimensions, - reshaped_dynamic_dimensions); + reshaped_dynamic_dimensions, + reshaped_expressions); // Eliminate the size one dimensions. // The added reshape reduces the rank of the tensor. Hence we cannot directly @@ -1094,15 +1151,20 @@ absl::StatusOr BroadcastToTargetRank( return origin; } - // Update target_size with origin sizes using broadcast_dimensions + // Update target_size and target_exp with origin sizes and expressions using + // broadcast_dimensions absl::Span target_dimensions = target_shape.dimensions(); + absl::Span target_expressions = target_shape.expressions(); std::vector target_size{target_dimensions.begin(), target_dimensions.end()}; + std::vector target_exp(target_expressions.begin(), + target_expressions.end()); for (int64_t origin_dim = 0; origin_dim < origin_rank; origin_dim++) { int64_t target_dim = broadcast_dimensions[origin_dim]; target_size[target_dim] = origin_shape.dimensions(origin_dim); + target_exp[target_dim] = origin_shape.expressions(origin_dim); } - return BroadcastInDim(origin, target_size, broadcast_dimensions); + return BroadcastInDim(origin, target_size, broadcast_dimensions, target_exp); } // Extract the `num_dims` counts of dimension sizes from the `op`. First, @@ -1120,7 +1182,7 @@ absl::StatusOr> ExtractDimensionSizesAndPadOnesToLeft( ? ConstantR1( /*builder=*/builder, /*values=*/{static_cast(op_shape->dimensions(i))}) - : Reshape(GetDimensionSize(op, i), {1})); + : Reshape(GetDimensionSize(op, i), {1}, {xla::DExpr::Const(1)})); } return op_dims; } @@ -1142,7 +1204,8 @@ absl::StatusOr BroadcastScalarToOutputShapeWithUnbounded( ? ConstantR1( /*builder=*/builder, /*values=*/{static_cast(output_shape.dimensions(i))}) - : Reshape(GetDimensionSize(output, i), {1}); + : Reshape(GetDimensionSize(output, i), {1}, + {xla::DExpr::Const(1)}); } return MhloDynamicBroadcastInDim( scalar, /*output_dimensions=*/ConcatInDim(builder, output_sizes, 0), {}, @@ -1525,20 +1588,13 @@ XlaOp XlaBuilder::Parameter( } XlaOp XlaBuilder::Broadcast(XlaOp operand, - absl::Span broadcast_sizes) { + absl::Span broadcast_sizes, + absl::Span broadcast_exprs) { return ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); - TF_ASSIGN_OR_RETURN( - const Shape& shape, - ShapeInference::InferBroadcastShape(*operand_shape, broadcast_sizes)); - - // The client-level broadcast op just appends dimensions on the left (adds - // lowest numbered dimensions). The HLO broadcast instruction is more - // flexible and can add new dimensions anywhere. The instruction's - // dimensions field maps operand dimensions to dimensions in the broadcast - // output, so to append dimensions on the left the instruction's dimensions - // should just be the n highest dimension numbers of the output shape where - // n is the number of input dimensions. + TF_ASSIGN_OR_RETURN(const Shape& shape, + ShapeInference::InferBroadcastShape( + *operand_shape, broadcast_sizes, broadcast_exprs)); const int64_t operand_rank = operand_shape->dimensions().size(); std::vector dimensions(operand_rank); for (int i = 0; i < operand_rank; ++i) { @@ -1550,14 +1606,13 @@ XlaOp XlaBuilder::Broadcast(XlaOp operand, XlaOp XlaBuilder::BroadcastInDim( XlaOp operand, absl::Span out_dim_size, - absl::Span broadcast_dimensions) { + absl::Span broadcast_dimensions, + absl::Span out_dim_exp) { return ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); - // Output shape, in the case of degenerate broadcast, the out_dim_size is - // not necessarily the same as the dimension sizes of the output shape. - TF_ASSIGN_OR_RETURN(auto output_shape, - ShapeUtil::MakeValidatedShape( - operand_shape->element_type(), out_dim_size)); + TF_ASSIGN_OR_RETURN(auto output_shape, ShapeUtil::MakeValidatedShape( + operand_shape->element_type(), + out_dim_size, out_dim_exp)); TF_RET_CHECK(!output_shape.is_unbounded_dynamic()) << "BroadcastInDim output must shape be static or bounded dynamic " << ShapeUtil::HumanString(output_shape); @@ -1584,6 +1639,15 @@ XlaOp XlaBuilder::BroadcastInDim( .status()); std::vector in_dim_size(out_dim_size.begin(), out_dim_size.end()); std::vector in_dim_dynamic(out_dim_size.size(), false); + std::vector in_expressions(out_dim_exp.begin(), out_dim_exp.end()); + + if (out_dim_exp.empty()) { + in_expressions.reserve(out_dim_size.size()); + std::transform(out_dim_size.begin(), out_dim_size.end(), + std::back_inserter(in_expressions), + [](int d) { return DExpr::Const(d); }); + } + for (int i = 0; i < broadcast_rank; i++) { in_dim_size[broadcast_dimensions[i]] = (operand_shape->is_unbounded_dynamic_dimension(i)) @@ -1591,19 +1655,17 @@ XlaOp XlaBuilder::BroadcastInDim( : operand_shape->dimensions(i); in_dim_dynamic[broadcast_dimensions[i]] = operand_shape->is_bounded_dynamic_dimension(i); + in_expressions[broadcast_dimensions[i]] = operand_shape->expressions(i); } - const auto& in_dim_shape = ShapeUtil::MakeShape( - operand_shape->element_type(), in_dim_size, in_dim_dynamic); + const auto& in_dim_shape = + ShapeUtil::MakeShape(operand_shape->element_type(), in_dim_size, + in_dim_dynamic, in_expressions); TF_ASSIGN_OR_RETURN( XlaOp in_dim_broadcast, InDimBroadcast(in_dim_shape, operand, broadcast_dimensions)); - - // If broadcast is not degenerate, return broadcasted result. if (ShapeUtil::Equal(in_dim_shape, output_shape)) { return in_dim_broadcast; } - - // Otherwise handle degenerate broadcast case. return AddBroadcastSequence(output_shape, in_dim_broadcast); }); } @@ -1637,10 +1699,46 @@ XlaOp XlaBuilder::Slice(XlaOp operand, absl::Span start_indices, }); } +XlaOp XlaBuilder::Slice(XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides) { + return ReportErrorOrReturn([&]() -> absl::StatusOr { + TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); + TF_ASSIGN_OR_RETURN( + Shape shape, ShapeInference::InferSliceShape(*operand_shape, + start_indices, + limit_indices, strides, + start_exprs, limit_exprs)); + return SliceInternal(shape, operand, start_indices, limit_indices, + start_exprs, limit_exprs, strides); + }); +} + +absl::StatusOr XlaBuilder::SliceInternal( + const Shape& shape, XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span strides) { + HloInstructionProto instr; + *instr.mutable_shape() = shape.ToProto(); + for (int i = 0, end = start_indices.size(); i < end; i++) { + auto* slice_config = instr.add_slice_dimensions(); + slice_config->set_start(start_indices[i]); + slice_config->set_limit(limit_indices[i]); + slice_config->set_stride(strides[i]); + } + return AddInstruction(std::move(instr), HloOpcode::kSlice, {operand}); +} + absl::StatusOr XlaBuilder::SliceInternal( const Shape& shape, XlaOp operand, absl::Span start_indices, absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, absl::Span strides) { + CHECK(start_exprs.empty() || start_exprs.size() == start_indices.size()); + CHECK(limit_exprs.empty() || limit_exprs.size() == start_indices.size()); HloInstructionProto instr; *instr.mutable_shape() = shape.ToProto(); for (int i = 0, end = start_indices.size(); i < end; i++) { @@ -1648,6 +1746,14 @@ absl::StatusOr XlaBuilder::SliceInternal( slice_config->set_start(start_indices[i]); slice_config->set_limit(limit_indices[i]); slice_config->set_stride(strides[i]); + if (i < start_exprs.size() && start_exprs[i] && + start_exprs[i]->is_dynamic()) { + start_exprs[i].simplify().to_proto(slice_config->mutable_start_expr()); + } + if (i < limit_exprs.size() && limit_exprs[i] && + limit_exprs[i]->is_dynamic()) { + limit_exprs[i].simplify().to_proto(slice_config->mutable_limit_expr()); + } } return AddInstruction(std::move(instr), HloOpcode::kSlice, {operand}); } @@ -1668,9 +1774,32 @@ XlaOp XlaBuilder::SliceInDim(XlaOp operand, int64_t start_index, }); } +XlaOp XlaBuilder::SliceInDim(XlaOp operand, int64_t start_index, + int64_t limit_index, const DExpr& start_expr, + const DExpr& limit_expr, int64_t stride, + int64_t dimno) { + return ReportErrorOrReturn([&]() -> absl::StatusOr { + TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(operand)); + std::vector starts(shape->dimensions().size(), 0); + std::vector limits(shape->dimensions().begin(), + shape->dimensions().end()); + std::vector start_exprs(shape->dimensions().size(), DExpr::Const(0)); + std::vector limit_exprs(shape->expressions().begin(), + shape->expressions().end()); + std::vector strides(shape->dimensions().size(), 1); + starts[dimno] = start_index; + limits[dimno] = limit_index; + start_exprs[dimno] = start_expr; + limit_exprs[dimno] = limit_expr; + strides[dimno] = stride; + return Slice(operand, starts, limits, start_exprs, limit_exprs, strides); + }); +} + XlaOp XlaBuilder::DynamicSlice(XlaOp operand, absl::Span start_indices, - absl::Span slice_sizes) { + absl::Span slice_sizes, + absl::Span slice_exprs) { return ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); std::vector start_indices_shape_ptrs; @@ -1679,9 +1808,9 @@ XlaOp XlaBuilder::DynamicSlice(XlaOp operand, absl::c_transform(start_indices_shapes, std::back_inserter(start_indices_shape_ptrs), [](const Shape& shape) { return &shape; }); - TF_ASSIGN_OR_RETURN(Shape shape, - ShapeInference::InferDynamicSliceShape( - *operand_shape, start_indices_shapes, slice_sizes)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDynamicSliceShape( + *operand_shape, start_indices_shapes, + slice_sizes, slice_exprs)); return DynamicSliceInternal(shape, operand, start_indices, slice_sizes); }); } @@ -1698,6 +1827,7 @@ absl::StatusOr XlaBuilder::DynamicSliceInternal( std::vector operands = {operand}; operands.insert(operands.end(), start_indices.begin(), start_indices.end()); + return AddInstruction(std::move(instr), HloOpcode::kDynamicSlice, operands); } @@ -1794,9 +1924,22 @@ XlaOp XlaBuilder::Reshape(XlaOp operand, absl::Span dimensions, int64_t inferred_dimension) { return ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); - TF_ASSIGN_OR_RETURN(const Shape shape, - ShapeInference::InferReshapeShape( - *operand_shape, dimensions, inferred_dimension)); + TF_ASSIGN_OR_RETURN(const Shape shape, ShapeInference::InferReshapeShape( + *operand_shape, dimensions, + inferred_dimension, {})); + return ReshapeInternal(shape, operand, inferred_dimension); + }); +} + +XlaOp XlaBuilder::Reshape(XlaOp operand, absl::Span dimensions, + absl::Span expressions, + int64_t inferred_dimension) { + return ReportErrorOrReturn([&]() -> absl::StatusOr { + TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); + TF_ASSIGN_OR_RETURN( + const Shape shape, + ShapeInference::InferReshapeShape(*operand_shape, dimensions, + inferred_dimension, expressions)); return ReshapeInternal(shape, operand, inferred_dimension); }); } @@ -1811,19 +1954,19 @@ XlaOp XlaBuilder::Reshape(const Shape& shape, XlaOp operand, XlaOp XlaBuilder::DynamicReshape(XlaOp operand, absl::Span dim_sizes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic) { + const std::vector& dims_are_dynamic, + absl::Span expressions) { return ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); std::vector dim_size_shape_ptrs; - TF_ASSIGN_OR_RETURN(const auto& dim_size_shapes, - GetOperandShapes(dim_sizes)); - + TF_ASSIGN_OR_RETURN(const auto& dim_size_shapes, GetOperandShapes(dim_sizes)); absl::c_transform(dim_size_shapes, std::back_inserter(dim_size_shape_ptrs), [](const Shape& shape) { return &shape; }); - TF_ASSIGN_OR_RETURN(const Shape shape, - ShapeInference::InferDynamicReshapeShape( - *operand_shape, dim_size_shape_ptrs, - new_size_bounds, dims_are_dynamic)); + TF_ASSIGN_OR_RETURN( + const Shape shape, + ShapeInference::InferDynamicReshapeShape( + *operand_shape, dim_size_shape_ptrs, new_size_bounds, + dims_are_dynamic, expressions)); TF_RETURN_IF_ERROR(first_error_); std::vector operands; operands.reserve(1 + dim_sizes.size()); @@ -1865,17 +2008,20 @@ XlaOp XlaBuilder::Collapse(XlaOp operand, VLOG(3) << "dims to collapse: " << absl::StrJoin(dimensions, ","); std::vector new_sizes; + std::vector new_exprs; for (int i = 0; i < original_shape->dimensions().size(); ++i) { if (i <= dimensions.front() || i > dimensions.back()) { new_sizes.push_back(original_shape->dimensions(i)); + new_exprs.push_back(original_shape->expressions(i)); } else { new_sizes.back() *= original_shape->dimensions(i); + new_exprs.back() = + new_exprs.back() * original_shape->expressions(i); } } VLOG(3) << "new sizes: [" << absl::StrJoin(new_sizes, ",") << "]"; - - return Reshape(operand, new_sizes); + return Reshape(operand, new_sizes, new_exprs); }); } @@ -3912,12 +4058,14 @@ XlaOp XlaBuilder::AllToAllArray( return all_to_all; } DimensionVector sizes; + std::vector expressions; const bool is_unbounded = operand_shape->is_unbounded_dynamic(); std::vector dynamic_sizes; auto GetR1DimensionSizeOrConstant = [&](XlaOp operand, int64_t dimension) -> XlaOp { if (operand_shape->is_unbounded_dynamic_dimension(dimension)) { - return Reshape(GetDimensionSize(operand, dimension), {1}); + return Reshape(GetDimensionSize(operand, dimension), {1}, + {DExpr::Const(1)}); } return ConstantR1( this, {static_cast(operand_shape->dimensions(dimension))}); @@ -3927,15 +4075,18 @@ XlaOp XlaBuilder::AllToAllArray( for (int64_t i = 0; i < operand_shape->dimensions().size(); ++i) { if (i != split_dimension) { sizes.push_back(operand_shape->dimensions(i)); + expressions.push_back(operand_shape->expressions(i)); if (is_unbounded) { dynamic_sizes.push_back(GetR1DimensionSizeOrConstant(operand, i)); } continue; } sizes.push_back(split_count); + expressions.push_back(DExpr::Const(split_count)); sizes.push_back(operand_shape->is_unbounded_dynamic_dimension(i) ? Shape::kUnboundedSize : operand_shape->dimensions(i) / split_count); + expressions.push_back(operand_shape->expressions(i) / split_count); if (is_unbounded) { dynamic_sizes.push_back(r1_split_count); @@ -3955,11 +4106,11 @@ XlaOp XlaBuilder::AllToAllArray( TF_ASSIGN_OR_RETURN( const Shape shape, ShapeUtil::MakeValidatedShape(all_to_all_shape.element_type(), sizes, - dynamic_dimensions)); + dynamic_dimensions, expressions)); all_to_all = MhloDynamicReshape(all_to_all, ConcatInDim(dynamic_sizes, 0), shape); } else { - all_to_all = Reshape(all_to_all, sizes); + all_to_all = Reshape(all_to_all, sizes, expressions); } std::vector permutation; @@ -4489,6 +4640,14 @@ XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64_t dimension) { TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGetDimensionSizeShape( *operand_shape, dimension)); + const DExpr& dim_expr = operand_shape->expressions(dimension); + if (dim_expr && dim_expr->is_dynamic()) { + DExpr simplified_expr = dim_expr.simplify(); + XlaOp dim_bound = + ConstantR0(this, operand_shape->dimensions(dimension)); + TF_RETURN_IF_ERROR(SetInstructionContents(dim_bound, {simplified_expr})); + return dim_bound; + } // Calling GetDimensionSize on a static dimension returns a constant // instruction. if (operand_shape->is_static_dimension(dimension)) { @@ -4866,6 +5025,7 @@ absl::StatusOr XlaBuilder::AddInstruction( TF_ASSIGN_OR_RETURN(Shape shape, Shape::FromProto(instructions_.back().shape())); instruction_shapes_.push_back(std::make_unique(std::move(shape))); + instruction_contents_.push_back({}); XlaOp op(handle, this); return op; @@ -4986,16 +5146,18 @@ XlaOp ConstantLiteral(XlaBuilder* builder, const LiteralSlice& literal) { return builder->ConstantLiteral(literal); } -XlaOp Broadcast(const XlaOp operand, - absl::Span broadcast_sizes) { - return operand.builder()->Broadcast(operand, broadcast_sizes); +XlaOp Broadcast(const XlaOp operand, absl::Span broadcast_sizes, + absl::Span broadcast_exprs) { + return operand.builder()->Broadcast(operand, broadcast_sizes, + broadcast_exprs); } XlaOp BroadcastInDim(const XlaOp operand, absl::Span out_dim_size, - absl::Span broadcast_dimensions) { + absl::Span broadcast_dimensions, + absl::Span out_dim_exp) { return operand.builder()->BroadcastInDim(operand, out_dim_size, - broadcast_dimensions); + broadcast_dimensions, out_dim_exp); } XlaOp MhloDynamicReshape(const XlaOp operand, const XlaOp output_shape, @@ -5030,21 +5192,38 @@ XlaOp Reshape(const XlaOp operand, absl::Span dimensions) { return operand.builder()->Reshape(operand, dimensions); } +XlaOp Reshape(const XlaOp operand, absl::Span dimensions, + absl::Span expressions) { + return operand.builder()->Reshape(operand, dimensions, expressions); +} + XlaOp Reshape(const Shape& shape, XlaOp operand) { return operand.builder()->Reshape(shape, operand); } XlaOp DynamicReshape(XlaOp operand, absl::Span dim_sizes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic) { + const std::vector& dims_are_dynamic, + absl::Span expressions) { return operand.builder()->DynamicReshape(operand, dim_sizes, new_size_bounds, - dims_are_dynamic); + dims_are_dynamic, expressions); +} + +XlaOp Slice(XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides) { + return operand.builder()->Slice(operand, start_indices, limit_indices, + start_exprs, limit_exprs, strides); } XlaOp ReshapeWithInferredDimension(XlaOp operand, absl::Span new_sizes, + absl::Span new_exprs, int64_t inferred_dimension) { - return operand.builder()->Reshape(operand, new_sizes, inferred_dimension); + return operand.builder()->Reshape(operand, new_sizes, new_exprs, + inferred_dimension); } XlaOp Collapse(const XlaOp operand, absl::Span dimensions) { @@ -5064,9 +5243,18 @@ XlaOp SliceInDim(const XlaOp operand, int64_t start_index, int64_t limit_index, stride, dimno); } +XlaOp SliceInDim(const XlaOp operand, int64_t start_index, int64_t limit_index, + const DExpr& start_expr, const DExpr& limit_expr, + int64_t stride, int64_t dimno) { + return operand.builder()->SliceInDim(operand, start_index, limit_index, + start_expr, limit_expr, stride, dimno); +} + XlaOp DynamicSlice(const XlaOp operand, absl::Span start_indices, - absl::Span slice_sizes) { - return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes); + absl::Span slice_sizes, + absl::Span slice_exprs) { + return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes, + slice_exprs); } XlaOp DynamicUpdateSlice(const XlaOp operand, const XlaOp update, @@ -6085,6 +6273,14 @@ XlaOp GetDimensionSize(const XlaOp operand, int64_t dimension) { return operand.builder()->GetDimensionSize(operand, dimension); } +XlaOp GetExpressionValue(XlaOp operand) { + XlaBuilder* builder = operand.builder(); + return CustomCall(builder, "GetExpressionValue", {operand}, + ShapeUtil::MakeShape(S32, {}), "", false, {}, + nullptr, CustomCallSchedule::SCHEDULE_NONE, + CustomCallApiVersion::API_VERSION_ORIGINAL); +} + XlaOp SetDimensionSize(const XlaOp operand, const XlaOp val, int64_t dimension) { return operand.builder()->SetDimensionSize(operand, val, dimension); diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 8f479090d86b19..b2c636c6eee530 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -484,6 +484,13 @@ class XlaBuilder { absl::Status SetInstructionFrontendAttribute(XlaOp op, std::string attribute, std::string value); + // Associates symbolic contents metadata with a specific instruction. + absl::Status SetInstructionContents(XlaOp op, std::vector contents); + + // Returns symbolic contents metadata attached to an instruction, if any. + absl::StatusOr*> GetInstructionContents( + XlaOp op) const; + // Looks up the HloInstruction and sets the sharding. If the sharding already // existed, then its value is updated. // @@ -519,10 +526,12 @@ class XlaBuilder { virtual XlaOp ConstantLiteral(const LiteralSlice& literal); - XlaOp Broadcast(XlaOp operand, absl::Span broadcast_sizes); + XlaOp Broadcast(XlaOp operand, absl::Span broadcast_sizes, + absl::Span broadcast_exprs = {}); XlaOp BroadcastInDim(XlaOp operand, absl::Span out_dim_size, - absl::Span broadcast_dimensions); + absl::Span broadcast_dimensions, + absl::Span out_dim_exp = {}); // This is an experimental API for creating the mhlo.dynamic_broadcast_in_dim // op from the XlaBuilder. This is only intended for export to MHLO or @@ -545,12 +554,17 @@ class XlaBuilder { XlaOp Reshape(XlaOp operand, absl::Span dimensions, int64_t inferred_dimension = -1); + XlaOp Reshape(XlaOp operand, absl::Span dimensions, + absl::Span expressions, + int64_t inferred_dimension = -1); + XlaOp Reshape(const Shape& shape, XlaOp operand, int64_t inferred_dimension = -1); XlaOp DynamicReshape(XlaOp operand, absl::Span dim_sizes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic); + const std::vector& dims_are_dynamic, + absl::Span expressions = {}); XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape, const Shape& shape); @@ -560,16 +574,36 @@ class XlaBuilder { XlaOp Slice(XlaOp operand, absl::Span start_indices, absl::Span limit_indices, absl::Span strides); + + XlaOp Slice(XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides); + virtual absl::StatusOr SliceInternal( const Shape& shape, XlaOp operand, absl::Span start_indices, absl::Span limit_indices, absl::Span strides); + virtual absl::StatusOr SliceInternal( + const Shape& shape, XlaOp operand, + absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides); virtual XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index, int64_t stride, int64_t dimno); + virtual XlaOp SliceInDim(XlaOp operand, int64_t start_index, + int64_t limit_index, const DExpr& start_expr, + const DExpr& limit_expr, int64_t stride, + int64_t dimno); + XlaOp DynamicSlice(XlaOp operand, absl::Span start_indices, - absl::Span slice_sizes); + absl::Span slice_sizes, + absl::Span slice_exprs = {}); virtual absl::StatusOr DynamicSliceInternal( const Shape& shape, XlaOp operand, absl::Span start_indices, absl::Span slice_sizes); @@ -1181,6 +1215,7 @@ class XlaBuilder { // A cache for the HloInstructionProto shapes, to avoid recreating Shape // objects from protos and to support the GetShapePtr() API. std::vector> instruction_shapes_; + std::vector> instruction_contents_; // Dynamic parameter configuration of this computation. DynamicParameterBinding dynamic_parameter_binding_; @@ -1244,11 +1279,13 @@ class XlaBuilder { const LiteralSlice& literal); friend XlaOp Broadcast(XlaOp operand, - absl::Span broadcast_sizes); + absl::Span broadcast_sizes, + absl::Span broadcast_expressions); friend XlaOp BroadcastInDim(XlaOp operand, absl::Span out_dim_size, - absl::Span broadcast_dimensions); + absl::Span broadcast_dimensions, + absl::Span out_dim_exp); friend XlaOp MhloDynamicBroadcastInDim( XlaOp operand, XlaOp output_dimensions, @@ -1265,18 +1302,22 @@ class XlaBuilder { friend XlaOp Reshape(XlaOp operand, absl::Span dimensions); + friend XlaOp Reshape(XlaOp operand, absl::Span dimensions, + absl::Span expressions); + friend XlaOp Reshape(const Shape& shape, XlaOp operand); friend XlaOp DynamicReshape(XlaOp operand, absl::Span dim_sizes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic); + const std::vector& dims_are_dynamic, + absl::Span expressions); friend XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape, const Shape& shape); - friend XlaOp ReshapeWithInferredDimension(XlaOp operand, - absl::Span new_sizes, - int64_t inferred_dimension); + friend XlaOp ReshapeWithInferredDimension( + XlaOp operand, absl::Span new_sizes, + absl::Span new_exprs, int64_t inferred_dimension); friend XlaOp Collapse(XlaOp operand, absl::Span dimensions); @@ -1284,12 +1325,24 @@ class XlaBuilder { absl::Span limit_indices, absl::Span strides); + friend XlaOp Slice(XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides); + friend XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index, int64_t stride, int64_t dimno); + friend XlaOp SliceInDim(XlaOp operand, int64_t start_index, + int64_t limit_index, const DExpr& start_expr, + const DExpr& limit_expr, int64_t stride, + int64_t dimno); + friend XlaOp DynamicSlice(XlaOp operand, absl::Span start_indices, - absl::Span slice_sizes); + absl::Span slice_sizes, + absl::Span slice_exprs); friend XlaOp DynamicUpdateSlice(XlaOp operand, XlaOp update, absl::Span start_indices); @@ -1975,7 +2028,8 @@ XlaOp ConstantR1(XlaBuilder* builder, int64_t length, NativeT value); // The new dimensions index into copies of the operand, i.e. // // output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM] -XlaOp Broadcast(XlaOp operand, absl::Span broadcast_sizes); +XlaOp Broadcast(XlaOp operand, absl::Span broadcast_sizes, + absl::Span broadcast_exprs = {}); // This op broadcasts the `operand` to an output with the given `shape`. // `broadcast_dimensions` are the dimensions to be broadcasting into, i.e., the @@ -1993,7 +2047,8 @@ XlaOp Broadcast(XlaOp operand, absl::Span broadcast_sizes); // {{1 , 1}, // {2 , 2}} XlaOp BroadcastInDim(XlaOp operand, absl::Span out_dim_size, - absl::Span broadcast_dimensions); + absl::Span broadcast_dimensions, + absl::Span out_dim_exp = {}); // This is an experimental API for creating the mhlo.dynamic_broadcast_in_dim // op from the XlaBuilder. This is only intended for export to MHLO or @@ -2038,7 +2093,8 @@ XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno, // dimension dimension if dims_are_dynamic[i] is true. XlaOp DynamicReshape(XlaOp operand, absl::Span dim_sizes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic); + const std::vector& dims_are_dynamic, + absl::Span expressions = {}); // This is an experimental API for creating the mhlo.dynamic_reshape op from the // XlaBuilder. This is only intended for export to MHLO or StableHLO, and cannot @@ -2050,6 +2106,9 @@ XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape, const Shape& shape); // dimension sizes. Conceptually, this is a limited form of "shape casting". XlaOp Reshape(XlaOp operand, absl::Span dimensions); +XlaOp Reshape(XlaOp operand, absl::Span dimensions, + absl::Span expressions); + // Enqueues a Reshape op that uses an explicit target shape. XlaOp Reshape(const Shape& shape, XlaOp operand); @@ -2059,6 +2118,7 @@ XlaOp Reshape(const Shape& shape, XlaOp operand); // is a dynamic dimension in the output, it must be the inferred dimension. XlaOp ReshapeWithInferredDimension(XlaOp operand, absl::Span new_sizes, + absl::Span new_exprs, int64_t inferred_dimension); // Wrapper for Reshape. @@ -2096,6 +2156,12 @@ XlaOp Slice(XlaOp operand, absl::Span start_indices, absl::Span limit_indices, absl::Span strides); +XlaOp Slice(XlaOp operand, absl::Span start_indices, + absl::Span limit_indices, + absl::Span start_exprs, + absl::Span limit_exprs, + absl::Span strides); + // Enqueues a slice operation in a given dimension, taking all other // dimensions as they are; e.g. if dimno is 1 from start_index 2 to // limit_index 4 by 1, and the shape is f32[7,8,9], this call is short-hand @@ -2105,6 +2171,10 @@ XlaOp Slice(XlaOp operand, absl::Span start_indices, XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index, int64_t stride, int64_t dimno); +XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index, + const DExpr& start_expr, const DExpr& limit_expr, + int64_t stride, int64_t dimno); + // Enqueues a slice operation onto the computation that slices the 'operand' // from dynamic start indices which are passed in 'start_indices'. // The size of the slice in each dimension is passed in 'slice_sizes', @@ -2116,7 +2186,8 @@ XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index, // Slice index calculations are computed modulo input dimension sizes to // prevent dynamic start indices from generating out-of-bound array accesses. XlaOp DynamicSlice(XlaOp operand, absl::Span start_indices, - absl::Span slice_sizes); + absl::Span slice_sizes, + absl::Span slice_exprs = {}); // Enqueues a dynamic update slice operation onto the computation, which // updates a slice of 'operand' with 'update' at dynamic 'start_indices'. @@ -3055,6 +3126,8 @@ XlaOp BatchNormGrad(XlaOp operand, XlaOp scale, XlaOp batch_mean, // array shaped. XlaOp GetDimensionSize(XlaOp operand, int64_t dimension); +XlaOp GetExpressionValue(XlaOp operand); + // Sets the size of the given dimension of the operand. The operand must be // array shaped. The result will have the same shape as the operand, but the // given dimension will be dynamic (if not already). diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 0dc7d619afc370..86b7bb0facc892 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -108,6 +108,64 @@ absl::Status EraseElementFromVector(PtrVec* container, T value) { container->erase(it); return absl::OkStatus(); } + +DynExpr* DynExprFromProtoForPrint(const ExpressionProto& proto) { + switch (proto.node_type_case()) { + case ExpressionProto::kConstantValue: + return DynExpr::_(proto.constant_value()); + case ExpressionProto::kVariableId: + return DynExpr::V(proto.variable_id()); + case ExpressionProto::kAddNode: { + const auto& add = proto.add_node(); + return new Add(DynExprFromProtoForPrint(add.lhs()), + DynExprFromProtoForPrint(add.rhs())); + } + case ExpressionProto::kSubNode: { + const auto& sub = proto.sub_node(); + return new Sub(DynExprFromProtoForPrint(sub.lhs()), + DynExprFromProtoForPrint(sub.rhs())); + } + case ExpressionProto::kMulNode: { + const auto& mul = proto.mul_node(); + return new Mul(DynExprFromProtoForPrint(mul.lhs()), + DynExprFromProtoForPrint(mul.rhs())); + } + case ExpressionProto::kDivNode: { + const auto& div = proto.div_node(); + return new Div(DynExprFromProtoForPrint(div.lhs()), + DynExprFromProtoForPrint(div.rhs())); + } + case ExpressionProto::kMaxNode: { + const auto& max = proto.max_node(); + return new MaxExpr(DynExprFromProtoForPrint(max.lhs()), + DynExprFromProtoForPrint(max.rhs())); + } + case ExpressionProto::kGtNode: { + const auto& gt = proto.gt_node(); + return new GtExpr(DynExprFromProtoForPrint(gt.lhs()), + DynExprFromProtoForPrint(gt.rhs())); + } + case ExpressionProto::kSelectNode: { + const auto& select = proto.select_node(); + return new SelectExpr(DynExprFromProtoForPrint(select.pred()), + DynExprFromProtoForPrint(select.on_true()), + DynExprFromProtoForPrint(select.on_false())); + } + case ExpressionProto::NODE_TYPE_NOT_SET: + default: + return nullptr; + } +} + +std::string ContentsExprToString(const ExpressionProto& proto) { + std::unique_ptr expr(DynExprFromProtoForPrint(proto)); + if (expr == nullptr) { + return "_"; + } + StringPrinter printer; + expr->print(&printer); + return std::move(printer).ToString(); +} } // namespace HloInstruction::Users::~Users() = default; @@ -636,14 +694,32 @@ absl::StatusOr> HloInstruction::CreateFromProto( break; case HloOpcode::kSlice: { std::vector slice_starts, slice_limits, slice_strides; + std::vector slice_start_exprs, slice_limit_exprs; + bool has_symbolic_slice_bounds = false; for (const HloInstructionProto::SliceDimensions& slice_dimensions : proto.slice_dimensions()) { slice_starts.push_back(slice_dimensions.start()); slice_limits.push_back(slice_dimensions.limit()); slice_strides.push_back(slice_dimensions.stride()); + if (slice_dimensions.has_start_expr() || + slice_dimensions.has_limit_expr()) { + has_symbolic_slice_bounds = true; + } + slice_start_exprs.push_back( + slice_dimensions.has_start_expr() + ? DExprFromProto(slice_dimensions.start_expr()) + : DExpr::Unknown(kMissingExpressionSentinel)); + slice_limit_exprs.push_back( + slice_dimensions.has_limit_expr() + ? DExprFromProto(slice_dimensions.limit_expr()) + : DExpr::Unknown(kMissingExpressionSentinel)); } - instruction = CreateSlice(shape, operands(0), slice_starts, slice_limits, - slice_strides); + instruction = has_symbolic_slice_bounds + ? CreateSlice(shape, operands(0), slice_starts, + slice_limits, slice_strides, + slice_start_exprs, slice_limit_exprs) + : CreateSlice(shape, operands(0), slice_starts, + slice_limits, slice_strides); break; } case HloOpcode::kConstant: { @@ -1373,6 +1449,14 @@ absl::StatusOr> HloInstruction::CreateFromProto( if (proto.has_frontend_attributes()) { instruction->set_frontend_attributes(proto.frontend_attributes()); } + if (proto.contents_size() > 0) { + std::vector contents; + contents.reserve(proto.contents_size()); + for (const auto& content : proto.contents()) { + contents.push_back(content); + } + instruction->set_contents(std::move(contents)); + } if (proto.has_statistics_viz()) { instruction->set_statistics_viz(proto.statistics_viz()); @@ -2046,6 +2130,18 @@ HloInstruction::CreateAddDependency(HloInstruction* data_operand, limit_indices, strides); } +/* static */ std::unique_ptr HloInstruction::CreateSlice( + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span limit_indices, + absl::Span strides, + absl::Span start_exprs, + absl::Span limit_exprs) { + return std::make_unique(shape, operand, start_indices, + limit_indices, strides, + start_exprs, limit_exprs); +} + /* static */ std::unique_ptr HloInstruction::CreateDynamicSlice( const Shape& shape, HloInstruction* operand, absl::Span start_indices, @@ -2812,6 +2908,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( SetupDerivedInstruction(clone.get()); clone->backend_config_ = BackendConfigWrapper(backend_config_); clone->set_frontend_attributes(frontend_attributes()); + clone->set_contents(contents()); // The new instruction's name will be uniquified when it's added to a // computation. clone->SetAndSanitizeName(name()); @@ -4235,6 +4332,18 @@ void HloInstruction::PrintExtraAttributes( FrontendAttributesToString(frontend_attributes())); }); } + if (has_contents()) { + printer.Next([this](Printer* printer) { + printer->Append("contents=["); + for (int64_t i = 0; i < contents().size(); ++i) { + if (i > 0) { + printer->Append(", "); + } + printer->Append(ContentsExprToString(contents()[i])); + } + printer->Append("]"); + }); + } if (opcode() != HloOpcode::kCall) { CHECK(!is_composite()) @@ -4356,6 +4465,9 @@ HloInstructionProto HloInstruction::ToProto() const { } *proto.mutable_frontend_attributes() = frontend_attributes(); + for (const auto& content : contents()) { + *proto.add_contents() = content; + } proto.set_is_composite(is_composite()); *proto.mutable_statistics_viz() = statistics_viz(); diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index 59f542b6ec3b5b..8feebcc1647a3c 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -865,6 +865,13 @@ class HloInstruction { absl::Span start_indices, absl::Span limit_indices, absl::Span strides); + static std::unique_ptr CreateSlice( + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span limit_indices, + absl::Span strides, + absl::Span start_exprs, + absl::Span limit_exprs); // Creates a slice instruction, where the first operand is sliced by // start indices specified in the second operand, and by size specified in @@ -1883,6 +1890,19 @@ class HloInstruction { return rare()->frontend_attributes; } + void set_contents(std::vector contents) { + if (!has_rare() && contents.empty()) { + return; + } + mutable_rare()->contents = std::move(contents); + } + + const std::vector& contents() const { + return rare()->contents; + } + + bool has_contents() const { return has_rare() && !rare()->contents.empty(); } + std::optional get_frontend_attribute( absl::string_view key) const { auto it = rare()->frontend_attributes.map().find(key); @@ -2506,6 +2526,9 @@ class HloInstruction { // z' = const(20), frontend_attributes={?} FrontendAttributes frontend_attributes; + // Structured symbolic contents attached to this instruction. + std::vector contents; + // Used by kCall to determine if the Call instruction is a composite. bool is_composite; diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 0557a52c67fd1d..e4853f8c56a4fb 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -74,6 +74,15 @@ namespace { using absl::CEscape; using absl::StrCat; +std::string DExprToString(const DExpr& expr) { + if (!expr) { + return "_"; + } + StringPrinter printer; + expr.simplify()->print(&printer); + return std::move(printer).ToString(); +} + bool IsInstructionElementwiseOnOperand(const HloInstruction* instruction, const HloInstruction* operand) { const auto operand_indices = instruction->OperandIndices(operand); @@ -1651,10 +1660,20 @@ HloSliceInstruction::HloSliceInstruction( const Shape& shape, HloInstruction* operand, absl::Span start_indices, absl::Span limit_indices, absl::Span strides) + : HloSliceInstruction(shape, operand, start_indices, limit_indices, strides, + /*start_exprs=*/{}, /*limit_exprs=*/{}) {} + +HloSliceInstruction::HloSliceInstruction( + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span limit_indices, absl::Span strides, + absl::Span start_exprs, absl::Span limit_exprs) : HloInstruction(HloOpcode::kSlice, shape), slice_starts_(start_indices.begin(), start_indices.end()), slice_limits_(limit_indices.begin(), limit_indices.end()), - slice_strides_(strides.begin(), strides.end()) { + slice_strides_(strides.begin(), strides.end()), + slice_start_exprs_(start_exprs.begin(), start_exprs.end()), + slice_limit_exprs_(limit_exprs.begin(), limit_exprs.end()) { AppendOperand(operand); // For backward compatibility with old serialized computations: if there are // no strides, assume all strides are 1. @@ -1671,6 +1690,14 @@ HloInstructionProto HloSliceInstruction::ToProto() const { slice_dimension->set_start(slice_starts_[i]); slice_dimension->set_limit(slice_limits_[i]); slice_dimension->set_stride(slice_strides_[i]); + const DExpr& start_expr = slice_start_exprs(i); + const DExpr& limit_expr = slice_limit_exprs(i); + if (start_expr) { + start_expr.to_proto(slice_dimension->mutable_start_expr()); + } + if (limit_expr) { + limit_expr.to_proto(slice_dimension->mutable_limit_expr()); + } } return proto; } @@ -1688,6 +1715,12 @@ void HloSliceInstruction::PrintExtraAttributesImpl( if (!omit_stride) { AppendCat(printer, ":", slice_strides_[i]); } + const DExpr& start_expr = slice_start_exprs(i); + const DExpr& limit_expr = slice_limit_exprs(i); + if (start_expr || limit_expr) { + AppendCat(printer, " start_expr=", DExprToString(start_expr), + " limit_expr=", DExprToString(limit_expr)); + } printer->Append("]"); }); printer->Append("}"); @@ -1701,7 +1734,9 @@ bool HloSliceInstruction::IdenticalSlowPath( const auto& other_slice = static_cast(other); return slice_starts_ == other_slice.slice_starts_ && slice_limits_ == other_slice.slice_limits_ && - slice_strides_ == other_slice.slice_strides_; + slice_strides_ == other_slice.slice_strides_ && + slice_start_exprs_ == other_slice.slice_start_exprs_ && + slice_limit_exprs_ == other_slice.slice_limit_exprs_; } std::unique_ptr HloSliceInstruction::CloneWithNewOperandsImpl( @@ -1709,7 +1744,8 @@ std::unique_ptr HloSliceInstruction::CloneWithNewOperandsImpl( HloCloneContext* context) const { CHECK_EQ(new_operands.size(), 1); return std::make_unique( - shape, new_operands[0], slice_starts_, slice_limits_, slice_strides_); + shape, new_operands[0], slice_starts_, slice_limits_, slice_strides_, + slice_start_exprs_, slice_limit_exprs_); } HloConstantInstruction::HloConstantInstruction(Literal literal) diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 22dd6ebe0e7b50..e90399a66a4ba9 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -1250,6 +1250,12 @@ class HloSliceInstruction : public HloInstruction { absl::Span start_indices, absl::Span limit_indices, absl::Span strides); + explicit HloSliceInstruction(const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span limit_indices, + absl::Span strides, + absl::Span start_exprs, + absl::Span limit_exprs); HloInstructionProto ToProto() const override; @@ -1275,6 +1281,32 @@ class HloSliceInstruction : public HloInstruction { const std::vector& slice_strides() const { return slice_strides_; } std::vector* mutable_slice_strides() { return &slice_strides_; } + const DExpr& slice_start_exprs(int64_t dimension) const { + static const DExpr kMissingSliceExpr = + DExpr::Unknown(kMissingExpressionSentinel); + if (dimension < 0 || + dimension >= static_cast(slice_start_exprs_.size())) { + return kMissingSliceExpr; + } + return slice_start_exprs_[dimension]; + } + const std::vector& slice_start_exprs() const { + return slice_start_exprs_; + } + + const DExpr& slice_limit_exprs(int64_t dimension) const { + static const DExpr kMissingSliceExpr = + DExpr::Unknown(kMissingExpressionSentinel); + if (dimension < 0 || + dimension >= static_cast(slice_limit_exprs_.size())) { + return kMissingSliceExpr; + } + return slice_limit_exprs_[dimension]; + } + const std::vector& slice_limit_exprs() const { + return slice_limit_exprs_; + } + static bool ClassOf(const HloInstruction* hlo) { return hlo->opcode() == HloOpcode::kSlice; } @@ -1295,6 +1327,8 @@ class HloSliceInstruction : public HloInstruction { std::vector slice_starts_; std::vector slice_limits_; std::vector slice_strides_; + std::vector slice_start_exprs_; + std::vector slice_limit_exprs_; }; class HloConstantInstruction : public HloInstruction { diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc index 07cf9a8c4bce2c..11eeefa948fa1b 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc +++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc @@ -250,6 +250,7 @@ absl::StatusOr HloPassPipeline::RunPassesInternal( } TF_RETURN_IF_ERROR(status); } + if (!pass->IsPassPipeline()) { compilation_stats_->EndPass(pass_name); } diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc index 61b261248e0e2c..aa781b082c8823 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc @@ -124,9 +124,9 @@ absl::Status CombineAllGathers(absl::Span to_combine, (*perm)[ag->all_gather_dimension()]); // Bitcast operand and update output shape. + auto sh = ShapeUtil::PermuteDimensions(*perm, operand_shape); operands.back() = - computation.AddInstruction(HloInstruction::CreateBitcast( - ShapeUtil::PermuteDimensions(*perm, operand_shape), operand)); + computation.AddInstruction(HloInstruction::CreateBitcast(sh, operand)); output_shapes.back() = ShapeUtil::PermuteDimensions(*perm, hlo->shape()); } } diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc index b3c2ffe79ec00c..e038b7338a1bc2 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc @@ -121,6 +121,7 @@ HloInstruction* ApplyUnaries(HloInstruction* instr, instr = instr->AddInstruction(unary->CloneWithNewOperands( ShapeUtil::MakeShapeWithDenseLayout( instr->shape().element_type(), unary->shape().dimensions(), + unary->shape().expressions(), unary->shape().layout().minor_to_major()), {instr})); } diff --git a/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc index ddb505801d92ca..ecdc47c0877542 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc @@ -80,10 +80,13 @@ absl::StatusOr BitcastDtypesExpander::ExpandInstruction( broadcasted_input_shape.push_back(input_bit_width / output_bit_width); reshaped_input_shape.push_back(1); int64_t output_bit_width_mask = (int64_t{1} << output_bit_width) - 1; - - TF_ASSIGN_OR_RETURN(input, - BroadcastTo(Reshape(input, reshaped_input_shape), - broadcasted_input_shape)); + std::vector reshaped_input_exprs( + from_shape.expressions().begin(), from_shape.expressions().end()); + reshaped_input_exprs.push_back(DExpr::Const(1)); + TF_ASSIGN_OR_RETURN( + input, BroadcastTo( + Reshape(input, reshaped_input_shape, reshaped_input_exprs), + broadcasted_input_shape)); input = BitcastConvertType(input, input_logical_type); TF_ASSIGN_OR_RETURN(Shape input_shape, b.GetShape(input)); XlaOp iota = Iota(&b, input_shape, input_shape.dimensions().size() - 1); diff --git a/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc index b5df50e1956c90..7b11a211ee3e02 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc @@ -218,9 +218,9 @@ XlaOp CholeskyExpander::BuildCholesky(XlaOp a, int64_t block_size, l = UpdateSliceInMinorDims(l, update, {i + k, i}); } } - return Select( - BroadcastInDim(seen_error, a_shape.dimensions(), error_dim_indices), - FullLike(l, std::numeric_limits::quiet_NaN()), l); + return Select(BroadcastInDim(seen_error, a_shape.dimensions(), + error_dim_indices, a_shape.expressions()), + FullLike(l, std::numeric_limits::quiet_NaN()), l); }); } diff --git a/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc b/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc index a3787d88ebbd93..e1aa51a390e2a0 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc @@ -80,26 +80,42 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) { lhs_non_contracting_dims.reserve(num_lhs_non_contracting_dims); int64_t lhs_contracting_size = 1; bool lhs_contracting_dynamic = false; + int64_t lhs_contracting_multiplier_accu = 1; + DExpr lhs_contracting_expression = DExpr::Const(1); int64_t lhs_non_contracting_size = 1; bool lhs_non_contracting_dynamic = false; + int64_t lhs_non_contracting_multiplier_accu = 1; + DExpr lhs_non_contracting_expression = DExpr::Const(1); std::vector batch_dim_sizes; batch_dim_sizes.reserve(num_batch_dims); std::vector batch_dynamic_dims; batch_dynamic_dims.reserve(num_batch_dims); + std::vector batch_expressions; + batch_expressions.reserve(num_batch_dims); + + bool lhs_contracting_is_static = true; + bool lhs_non_contracting_is_static = true; + for (int64_t i = 0; i < lhs_rank; ++i) { if (absl::c_linear_search(original_dnums.lhs_contracting_dimensions(), i)) { lhs_contracting_size *= lhs_shape.dimensions(i); lhs_contracting_dynamic |= lhs_shape.is_dynamic_dimension(i); + lhs_contracting_expression = + lhs_contracting_expression * lhs_shape.expressions(i); } else if (absl::c_linear_search(original_dnums.lhs_batch_dimensions(), i)) { batch_dim_sizes.push_back(lhs_shape.dimensions(i)); batch_dynamic_dims.push_back(lhs_shape.is_dynamic_dimension(i)); + batch_expressions.push_back(lhs_shape.expressions(i)); } else { lhs_non_contracting_dims.push_back(i); lhs_non_contracting_size *= lhs_shape.dimensions(i); lhs_non_contracting_dynamic |= lhs_shape.is_dynamic_dimension(i); + lhs_non_contracting_expression = + lhs_non_contracting_expression * lhs_shape.expressions(i); } } + // The canonical form of the lhs is // [BatchDims, NonContractingDimsProduct, ContractingsDimsProduct] // If NonContractingDimsProduct is 1, it is omitted. @@ -123,18 +139,21 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) { std::vector lhs_reshape_dims = batch_dim_sizes; std::vector lhs_reshape_dynamic_dims = batch_dynamic_dims; + std::vector lhs_reshape_expressions = batch_expressions; if (lhs_non_contracting_size > 1) { lhs_reshape_dims.push_back(lhs_non_contracting_size); lhs_reshape_dynamic_dims.push_back(lhs_non_contracting_dynamic); + lhs_reshape_expressions.push_back(lhs_non_contracting_expression); } lhs_reshape_dims.push_back(lhs_contracting_size); lhs_reshape_dynamic_dims.push_back(lhs_contracting_dynamic); + lhs_reshape_expressions.push_back(lhs_contracting_expression); // Reshape the contracting and non-contracting dimensions together. + auto sh_lhs = ShapeUtil::MakeShape(lhs_shape.element_type(), lhs_reshape_dims, + lhs_reshape_dynamic_dims, + lhs_reshape_expressions); HloInstruction* reshaped_lhs = computation->AddInstruction( - HloInstruction::CreateReshape( - ShapeUtil::MakeShape(lhs_shape.element_type(), lhs_reshape_dims, - lhs_reshape_dynamic_dims), - transposed_lhs), + HloInstruction::CreateReshape(sh_lhs, transposed_lhs), &transposed_lhs->metadata()); const auto& rhs_shape = original_dot->operand(1)->shape(); @@ -145,17 +164,29 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) { rhs_non_contracting_dims.reserve(num_rhs_non_contracting_dims); int64_t rhs_non_contracting_size = 1; bool rhs_non_contracting_dynamic = false; + int64_t rhs_non_contracting_multiplier_accu = 1; + DExpr rhs_non_contracting_expression = DExpr::Const(1); int64_t rhs_contracting_size = 1; bool rhs_contracting_dynamic = false; + int64_t rhs_contracting_multiplier_accu = 1; + DExpr rhs_contracting_expression = DExpr::Const(1); + + bool rhs_contracting_is_static = true; + bool rhs_non_contracting_is_static = true; + for (int64_t i = 0; i < rhs_rank; ++i) { if (absl::c_linear_search(original_dnums.rhs_contracting_dimensions(), i)) { rhs_contracting_size *= rhs_shape.dimensions(i); rhs_contracting_dynamic |= rhs_shape.is_dynamic_dimension(i); + rhs_contracting_expression = + rhs_contracting_expression * rhs_shape.expressions(i); } else if (!absl::c_linear_search(original_dnums.rhs_batch_dimensions(), i)) { rhs_non_contracting_dims.push_back(i); rhs_non_contracting_size *= rhs_shape.dimensions(i); rhs_non_contracting_dynamic |= rhs_shape.is_dynamic_dimension(i); + rhs_non_contracting_expression = + rhs_non_contracting_expression * rhs_shape.expressions(i); } } @@ -184,27 +215,35 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) { rhs_reshape_dims.push_back(rhs_contracting_size); std::vector rhs_reshape_dynamic_dims = batch_dynamic_dims; rhs_reshape_dynamic_dims.push_back(rhs_contracting_dynamic); + std::vector rhs_reshape_expressions = batch_expressions; + rhs_reshape_expressions.push_back(rhs_contracting_expression); if (rhs_non_contracting_size > 1) { rhs_reshape_dims.push_back(rhs_non_contracting_size); rhs_reshape_dynamic_dims.push_back(rhs_non_contracting_dynamic); + rhs_reshape_expressions.push_back(rhs_non_contracting_expression); } // Reshape the contracting and non-contracting dimensions together. + auto sh_rhs = ShapeUtil::MakeShape(rhs_shape.element_type(), rhs_reshape_dims, + rhs_reshape_dynamic_dims, + rhs_reshape_expressions); HloInstruction* reshaped_rhs = computation->AddInstruction( HloInstruction::CreateReshape( - ShapeUtil::MakeShape(rhs_shape.element_type(), rhs_reshape_dims, - rhs_reshape_dynamic_dims), + sh_rhs, transposed_rhs), &transposed_rhs->metadata()); std::vector dot_dims = batch_dim_sizes; std::vector dot_dynamic_dims = batch_dynamic_dims; + std::vector dot_expressions = batch_expressions; if (lhs_non_contracting_size > 1) { dot_dims.push_back(lhs_non_contracting_size); dot_dynamic_dims.push_back(lhs_non_contracting_dynamic); + dot_expressions.push_back(lhs_non_contracting_expression); } if (rhs_non_contracting_size > 1) { dot_dims.push_back(rhs_non_contracting_size); dot_dynamic_dims.push_back(rhs_non_contracting_dynamic); + dot_expressions.push_back(rhs_non_contracting_expression); } DotDimensionNumbers dot_dnums; @@ -251,12 +290,13 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) { HloInstruction::CreateReshape(result_shape, meta), &meta->metadata()); sparse_meta.push_back(meta); } + auto sh_dot = + ShapeUtil::MakeShape(original_dot->shape().element_type(), dot_dims, + dot_dynamic_dims, dot_expressions); HloInstruction* dot = computation->AddInstruction(HloInstruction::CreateDot( - ShapeUtil::MakeShape(original_dot->shape().element_type(), dot_dims, - dot_dynamic_dims), - reshaped_lhs, reshaped_rhs, dot_dnums, original_dot->precision_config(), - sparsity, sparse_meta)); + sh_dot, reshaped_lhs, reshaped_rhs, dot_dnums, + original_dot->precision_config(), sparsity, sparse_meta)); original_dot->SetupDerivedInstruction(dot); std::unique_ptr replacement = diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc index 33752d60cae8ce..6ddf023ffd77a7 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc @@ -158,8 +158,10 @@ void ApplyJacobiRotationOverRows(Eigh2x2 rotation, XlaOp& tl, XlaOp& tr, Shape shape = tl.builder()->GetShape(tl).value(); std::vector broadcast_dims(shape.dimensions().size() - 1); absl::c_iota(broadcast_dims, 0); - auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims); - auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims); + auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims, + shape.expressions()); + auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims, + shape.expressions()); auto s_conj = MaybeConjugate(s, true); std::tie(tl, tr, bl, br) = @@ -179,8 +181,10 @@ void ApplyJacobiRotationOverCols(Eigh2x2 rotation, XlaOp& tl, XlaOp& tr, std::vector broadcast_dims(shape.dimensions().size() - 1); absl::c_iota(broadcast_dims, 0); broadcast_dims.back() = shape.dimensions().size() - 1; - auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims); - auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims); + auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims, + shape.expressions()); + auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims, + shape.expressions()); auto s_conj = MaybeConjugate(s, true); std::tie(tl, tr, bl, br) = @@ -365,11 +369,12 @@ absl::Status EighExpander::SortByEigenvalues(XlaOp& v, XlaOp& w) { TF_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w)); const int64_t num_dims = v_shape.dimensions().size(); auto dimensions = v_shape.dimensions(); + auto expressions = v_shape.expressions(); std::vector broadcast_dims(num_dims - 1); std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); broadcast_dims[num_dims - 2] = num_dims - 1; - w = BroadcastInDim(w, dimensions, broadcast_dims); + w = BroadcastInDim(w, dimensions, broadcast_dims, expressions); XlaOp sort_result = Sort({w, v}, diff --git a/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc index dcf329f7c6d8b9..17775e0f132dfe 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc @@ -58,6 +58,15 @@ std::vector ConcatVectors(absl::Span xs, return output; } +std::vector ConcatEVectors(absl::Span xs, + absl::Span ys) { + std::vector output; + output.reserve(xs.size() + ys.size()); + std::copy(xs.begin(), xs.end(), std::back_inserter(output)); + std::copy(ys.begin(), ys.end(), std::back_inserter(output)); + return output; +} + // Computes sqrt(x^2 + y^2 + ...), avoiding overflow/underflow. // e.g. for 3 arguments: // def norm(x, y, z): @@ -220,11 +229,15 @@ absl::StatusOr QrExpander::QrBlock( const int64_t m = ShapeUtil::GetDimension(a_shape, -2); const int64_t n = ShapeUtil::GetDimension(a_shape, -1); + const DExpr& m_exp = ShapeUtil::GetExpression(a_shape, -2); + const DExpr& n_exp = ShapeUtil::GetExpression(a_shape, -1); const int64_t num_batch_dims = num_dims - 2; std::vector batch_dims(num_batch_dims); + std::vector batch_exprs(num_batch_dims); for (int i = 0; i < num_batch_dims; ++i) { batch_dims[i] = ShapeUtil::GetDimension(a_shape, i); + batch_exprs[i] = ShapeUtil::GetExpression(a_shape, i); } std::vector batch_dim_indices(num_batch_dims); @@ -248,9 +261,12 @@ absl::StatusOr QrExpander::QrBlock( minor_dim + 1); std::vector shape = batch_dims; + std::vector exprs = batch_exprs; shape.push_back(1); shape.push_back(m); - auto v_broadcast = Reshape(v, shape); + exprs.push_back(DExpr::Const(1)); + exprs.push_back(m_exp); + auto v_broadcast = Reshape(v, shape, exprs); // a[:, j+1:] -= np.conj(tau) * (v[:, np.newaxis] @ // (np.conj(v[np.newaxis, :]) @ a[:, j+1:])) // We use masking rather than a loop-variant shape to handle the j+1: @@ -263,7 +279,8 @@ absl::StatusOr QrExpander::QrBlock( // a[j, j] = beta // a[j+1:,j] = v[j+1:] - auto iota = Reshape(Iota(a.builder(), S32, m), {m, 1}); + auto iota = + Reshape(Iota(a.builder(), S32, m), {m, 1}, {m_exp, DExpr::Const(1)}); auto predecessor_mask = ConvertElementType(Lt(iota, j), type); auto mask = Broadcast(ConvertElementType(Eq(iota, j), type), std::vector(batch_dims.size(), 1)); @@ -279,7 +296,8 @@ absl::StatusOr QrExpander::QrBlock( std::vector dim_ids(num_dims); std::iota(dim_ids.begin(), dim_ids.end(), 0); new_x = BroadcastInDim(new_x, ConcatVectors(batch_dims, {m, n}), - /*broadcast_dimensions=*/dim_ids); + /*broadcast_dimensions=*/dim_ids, + ConcatEVectors(batch_exprs, {m_exp, n_exp})); a = Select(Eq(iota_mn, j), new_x, a); // taus[j] = tau diff --git a/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc b/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc index 2fe502429287b4..ca7eebb9cfb7e8 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc @@ -47,6 +47,7 @@ class VariadicReductionLayoutEqualizer : public DfsHloRewriteVisitor { if (first_input_s.layout() != input_s.layout()) { Shape new_input_s = ShapeUtil::MakeShapeWithDenseLayout( input_s.element_type(), input_s.dimensions(), + input_s.expressions(), first_input_s.layout().minor_to_major()); auto copy = MakeCopyHlo(input, new_input_s); changed = true; diff --git a/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc index d41aaf49be5d32..2e84d07f076078 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc @@ -66,7 +66,7 @@ RngBitGeneratorExpander::GetGeneratorComputation(const Shape& data_shape, XlaBuilder builder("rng"); XlaOp state_param = Parameter(&builder, 0, state_shape, "state"); - XlaOp key_op = Reshape(Slice(state_param, {0}, {1}, {1}), {}); + XlaOp key_op = Reshape(Slice(state_param, {0}, {1}, {1}), {}, {}); RngOutput output; switch (algorithm) { case RandomAlgorithm::RNG_THREE_FRY: @@ -83,8 +83,8 @@ RngBitGeneratorExpander::GetGeneratorComputation(const Shape& data_shape, RandomAlgorithm_Name(algorithm)); } - XlaOp final_state = - ConcatInDim(&builder, {Reshape(key_op, {1}), output.state}, 0); + XlaOp final_state = ConcatInDim( + &builder, {Reshape(key_op, {1}, {DExpr::Const(1)}), output.state}, 0); Tuple(&builder, {final_state, output.value}); TF_ASSIGN_OR_RETURN(XlaComputation xla_computation, builder.Build()); TF_ASSIGN_OR_RETURN(HloComputation * new_computation, diff --git a/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc index 5054a440778105..484a1e75ac0aeb 100644 --- a/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc +++ b/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc @@ -1316,6 +1316,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor { x = instr->AddInstruction(op.first->CloneWithNewOperands( ShapeUtil::MakeShapeWithDenseLayout( x->shape().element_type(), op.first->shape().dimensions(), + op.first->shape().expressions(), op.first->shape().layout().minor_to_major()), operands)); } @@ -1378,6 +1379,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor { instr->AddInstruction(HloInstruction::CreateCustomCall( ShapeUtil::MakeShapeWithDenseLayout( instr->shape().element_type(), new_output_shape.dimensions(), + new_output_shape.expressions(), instr->shape().layout().minor_to_major()), operands_list, kCublasLtMatmulF8CallTarget)); TF_RETURN_IF_ERROR(new_custom_call->set_backend_config(gpu_backend_config)); diff --git a/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc b/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc index ce454624144803..4585af7203946b 100644 --- a/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc +++ b/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc @@ -183,11 +183,13 @@ absl::StatusOr ShiftDequantizationF8( for (HloInstruction* unary : unaries[k]) { Shape new_shape = ShapeUtil::MakeShapeWithDenseLayout( operands[k]->shape().element_type(), unary->shape().dimensions(), + unary->shape().expressions(), unary->shape().layout().minor_to_major()); operands[k] = unary->AddInstruction(unary->CloneWithNewOperands( ShapeUtil::MakeShapeWithDenseLayout( operands[k]->shape().element_type(), unary->shape().dimensions(), + unary->shape().expressions(), unary->shape().layout().minor_to_major()), {operands[k]})); } diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index 713db5886c6f4e..de4c27b9367843 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -113,7 +113,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 92 +// Next ID: 93 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -174,6 +174,8 @@ message HloInstructionProto { int64 start = 1; int64 limit = 2; int64 stride = 3; + xla.ExpressionProto start_expr = 4; + xla.ExpressionProto limit_expr = 5; } repeated SliceDimensions slice_dimensions = 17; @@ -328,6 +330,9 @@ message HloInstructionProto { // Frontend attributes to pass to the XLA backend. xla.FrontendAttributes frontend_attributes = 68; + // Structured symbolic contents attached to this instruction. + repeated xla.ExpressionProto contents = 92; + // Specifies if all elements updated are guaranteed to be unique by // the caller. bool unique_indices = 69; diff --git a/third_party/xla/xla/service/hlo_creation_utils.cc b/third_party/xla/xla/service/hlo_creation_utils.cc index b815b1ffc79627..15c2c8417c9f28 100644 --- a/third_party/xla/xla/service/hlo_creation_utils.cc +++ b/third_party/xla/xla/service/hlo_creation_utils.cc @@ -184,7 +184,8 @@ absl::StatusOr MakeDynamicSliceHlo( TF_ASSIGN_OR_RETURN( Shape dynamic_slice_shape, ShapeInference::InferDynamicSliceShape( - operand->shape(), scalar_start_indices_shapes, slice_sizes)); + operand->shape(), scalar_start_indices_shapes, slice_sizes, + operand->shape().expressions())); return computation->AddInstruction( HloInstruction::CreateDynamicSlice(dynamic_slice_shape, operand, start_indices, slice_sizes), @@ -655,9 +656,12 @@ absl::StatusOr CollapseFirstNDims(HloInstruction* operand, CHECK_GE(operand_shape.dimensions_size(), n); int64_t new_shape_leading_bound = 1; bool new_shape_leading_is_dynamic = false; + DExpr new_shape_leading_expression = DExpr::Const(1); for (int64_t i = 0; i < n; i++) { new_shape_leading_bound *= operand_shape.dimensions(i); new_shape_leading_is_dynamic |= operand_shape.is_dynamic_dimension(i); + new_shape_leading_expression = + new_shape_leading_expression * operand_shape.expressions(i); } std::vector new_shape_dims; @@ -675,8 +679,16 @@ absl::StatusOr CollapseFirstNDims(HloInstruction* operand, operand_shape.dynamic_dimensions().end(), std::back_inserter(new_shape_dynamic_dims)); - Shape output_shape = ShapeUtil::MakeShape( - operand_shape.element_type(), new_shape_dims, new_shape_dynamic_dims); + std::vector new_shape_expressions; + new_shape_expressions.reserve(operand_shape.dimensions_size() - n + 1); + new_shape_expressions.push_back(new_shape_leading_expression); + auto exprs = operand_shape.expressions(); + std::copy(exprs.begin() + n, exprs.end(), + std::back_inserter(new_shape_expressions)); + + Shape output_shape = + ShapeUtil::MakeShape(operand_shape.element_type(), new_shape_dims, + new_shape_dynamic_dims, new_shape_expressions); return MakeReshapeHlo(output_shape, operand); } diff --git a/third_party/xla/xla/service/hlo_cse.cc b/third_party/xla/xla/service/hlo_cse.cc index 570ffa03de0edf..4b3ad32340cd44 100644 --- a/third_party/xla/xla/service/hlo_cse.cc +++ b/third_party/xla/xla/service/hlo_cse.cc @@ -37,6 +37,26 @@ namespace xla { namespace { +bool HasDynamicConstantContents(const HloInstruction* instruction) { + if (!instruction->has_contents()) { + return false; + } + for (const auto& content : instruction->contents()) { + if (content.node_type_case() != ExpressionProto::kConstantValue && + content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET) { + return true; + } + } + return false; +} + +bool HasDynamicConstantFrontendAttributes(const HloInstruction* instruction) { + const auto& attrs = instruction->frontend_attributes().map(); + return HasDynamicConstantContents(instruction) || + attrs.contains("dynamic_constant_index") || + attrs.contains("dynamic_constant_expr"); +} + template struct ConstantKey { template @@ -92,6 +112,13 @@ absl::StatusOr CombineConstants(HloComputation* computation, continue; } + // Dynamic constant frontend attrs carry semantic metadata for a later + // rewrite pass, so these constants are not interchangeable with otherwise + // identical plain constants. + if (HasDynamicConstantFrontendAttributes(instruction)) { + continue; + } + HloInstruction* match = nullptr; if (auto* constant_inst = DynCast(instruction)) { auto insert_result = constants.insert(ConstantKey{ @@ -310,6 +337,12 @@ absl::StatusOr HloCSE::RunOnComputation(HloComputation* computation) { continue; } + // These frontend attrs are semantic, not just decorative, so do not CSE + // across them. + if (HasDynamicConstantFrontendAttributes(instruction)) { + continue; + } + auto pair = representatives.insert(CseKey{instruction}); if (!pair.second) { HloInstruction* equivalent_instruction = pair.first->hlo; diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index b7bea1e55704c5..001a873118ab7f 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -1401,6 +1401,7 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( const Shape& output_shape = instruction->shape(); Shape output_shape_with_layout = ShapeUtil::MakeShapeWithDenseLayout( output_shape.element_type(), output_shape.dimensions(), + output_shape.expressions(), LayoutUtil::MinorToMajor(output_layout)); Shape operand_shape = operand->shape(); *operand_shape.mutable_layout() = @@ -1539,6 +1540,7 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( } Shape operand_shape_with_layout = ShapeUtil::MakeShapeWithDenseLayout( operand->shape().element_type(), operand->shape().dimensions(), + operand->shape().expressions(), LayoutUtil::MinorToMajor(operand_layout)); Shape output_shape = user->shape(); *output_shape.mutable_layout() = @@ -2571,7 +2573,7 @@ absl::Status LayoutAssignment::PropagateComputationLayouts( *result_layout = computed_computation_layout.result_layout(); } else { TF_RET_CHECK( - Shape::Equal().IgnoreDynamicDimension().MinorToMajorOnlyInLayout()( + Shape::Equal().IgnoreDynamicDimension().MinorToMajorOnlyInLayout().IgnoreBatch()( computed_computation_layout.result_layout().shape(), result_layout->shape())); } diff --git a/third_party/xla/xla/service/reduce_scatter_combiner.cc b/third_party/xla/xla/service/reduce_scatter_combiner.cc index a86271036f6baa..361629b34a6d5e 100644 --- a/third_party/xla/xla/service/reduce_scatter_combiner.cc +++ b/third_party/xla/xla/service/reduce_scatter_combiner.cc @@ -131,9 +131,9 @@ absl::Status CombineReduceScatters( std::swap((*perm)[most_frequent_dim], (*perm)[rs->scatter_dimension()]); // Bitcast operand and update output shape. - operands.back() = - computation.AddInstruction(HloInstruction::CreateBitcast( - ShapeUtil::PermuteDimensions(*perm, operand_shape), operand)); + auto sh = ShapeUtil::PermuteDimensions(*perm, operand_shape); + operands.back() = computation.AddInstruction( + HloInstruction::CreateBitcast(sh, operand)); output_shapes.back() = ShapeUtil::PermuteDimensions(*perm, hlo->shape()); } } diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc index 7985c930d812b5..ce006f2a132dd3 100644 --- a/third_party/xla/xla/service/shape_inference.cc +++ b/third_party/xla/xla/service/shape_inference.cc @@ -77,6 +77,15 @@ bool CompatibleDimensionSizes(int64_t size_a, int64_t size_b) { size_a == size_b; } +DExpr SymbolicElementsIn(const Shape& shape) { + DExpr product = DExpr::Const(1); + for (int64_t i = 0; i < shape.dimensions_size(); ++i) { + const DExpr& expr = shape.expressions(i); + product = product * (expr ? expr : DExpr::Const(shape.dimensions(i))); + } + return product.simplify(); +} + absl::Status ExpectArray(const Shape& shape, absl::string_view op_type) { if (!shape.IsArray()) { return InvalidArgument("Expected array argument for %s, but got %s.", @@ -195,6 +204,7 @@ absl::StatusOr InferWindowOutputShape(const Shape& base_shape, std::vector output_dimensions(window.dimensions_size()); std::vector output_is_dynamic(window.dimensions_size()); + std::vector output_expressions(window.dimensions_size()); for (int64_t i = 0; i < window.dimensions_size(); ++i) { const auto& dim = window.dimensions(i); if (dim.size() <= 0) { @@ -216,24 +226,45 @@ absl::StatusOr InferWindowOutputShape(const Shape& base_shape, window.DebugString()); } - if (IsUnboundedDynamicSize(ShapeUtil::GetDimension(base_shape, i))) { + const int64_t input_dimension = ShapeUtil::GetDimension(base_shape, i); + const DExpr& input_expression = base_shape.expressions(i); + const int64_t dilated_window = + window_util::DilatedBound(dim.size(), dim.window_dilation()); + if (IsUnboundedDynamicSize(input_dimension)) { output_dimensions[i] = Shape::kUnboundedSize; } else { const int64_t dilated_base = window_util::DilatedBound( - ShapeUtil::GetDimension(base_shape, i), dim.base_dilation()); + input_dimension, dim.base_dilation()); const int64_t padded_dilated_base = dim.padding_low() + dilated_base + dim.padding_high(); - const int64_t dilated_window = - window_util::DilatedBound(dim.size(), dim.window_dilation()); - output_dimensions[i] = window_util::StridedBound( padded_dilated_base, dilated_window, dim.stride()); } output_is_dynamic[i] = base_shape.is_dynamic_dimension(i); + if (input_expression && input_expression->is_constant()) { + output_expressions[i] = DExpr::Const(output_dimensions[i]); + continue; + } + + DExpr dilated_base_expr = input_expression; + if (dim.base_dilation() != 1) { + dilated_base_expr = + DExpr::Max((dim.base_dilation() * (input_expression - 1)) + 1, + DExpr::Const(0)) + .simplify(); + } + DExpr padded_dilated_base_expr = + (dilated_base_expr + dim.padding_low() + dim.padding_high()).simplify(); + DExpr strided_bound_expr = + (padded_dilated_base_expr - dilated_window + 1 + dim.stride() - 1) / + dim.stride(); + output_expressions[i] = + DExpr::Max(strided_bound_expr, DExpr::Const(0)).simplify(); } return ShapeUtil::MakeValidatedShape(element_type, output_dimensions, - output_is_dynamic); + output_is_dynamic, + output_expressions); } // Encapsulates inferred dimension size and bound size. @@ -472,6 +503,7 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, int64_t last_dim = operand_shape.dimensions_size() - 1; std::vector is_dynamic(operand_shape.dimensions_size()); std::vector dimensions(operand_shape.dimensions_size()); + std::vector expressions(operand_shape.dimensions_size()); TF_RET_CHECK(operand_shape.dimensions(last_dim) >= k) << "k=" << k << " is larger than the last dimension of size=" @@ -480,10 +512,13 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, is_dynamic[i] = i == last_dim ? false : operand_shape.is_dynamic_dimension(i); dimensions[i] = i == last_dim ? k : operand_shape.dimensions(i); + expressions[i] = + i == last_dim ? xla::DExpr::Const(k) : operand_shape.expressions(i); } - Shape out = ShapeUtil::MakeShape(operand_shape.element_type(), dimensions, - is_dynamic); + Shape out = + ShapeUtil::MakeShape(operand_shape.element_type(), dimensions, is_dynamic, + expressions); Shape idxs_shape = ShapeUtil::ChangeElementType(out, PrimitiveType::S32); return ShapeUtil::MakeTupleShape({out, idxs_shape}); } @@ -542,9 +577,11 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, int64_t rank = arg_shape->dimensions_size(); std::vector inferred_sizes(rank, Shape::kUnboundedSize); std::vector inferred_bounds(rank, Shape::kUnboundedSize); + std::vector inferred_expressions(rank, xla::DExpr::Const(0)); // Note: for the concatenate dimension, 0 should be the identity element: // Any dim size can keep unchanged when concatenated with 0 inferred_sizes[dimension] = 0; + inferred_expressions[dimension] = xla::DExpr::Const(0); for (const Shape* shape : arg_shapes) { for (int dim = 0; dim < rank; ++dim) { @@ -554,24 +591,32 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, int64_t leftSize = inferred_sizes[dim]; int64_t rightSize = dimension_size; int64_t leftBound = inferred_bounds[dim]; + const xla::DExpr& left_expression = inferred_expressions[dim]; int64_t rightBound = shape->is_dynamic_dimension(dim) ? dimension_size : Shape::kUnboundedSize; + const xla::DExpr& right_expression = shape->expressions(dim); + xla::DExpr inferred_expression = xla::DExpr::Const(0); + if (dim == dimension) { inferred_dim_and_bound = InferConcatenatedDimAndBound( leftSize, rightSize, leftBound, rightBound); + inferred_expression = left_expression + right_expression; } else { TF_ASSIGN_OR_RETURN( inferred_dim_and_bound, InferMostSpecificDimAndBound(dim, leftSize, rightSize, leftBound, rightBound)); + inferred_expression = right_expression; } inferred_sizes[dim] = inferred_dim_and_bound.dimension; inferred_bounds[dim] = inferred_dim_and_bound.bound; + inferred_expressions[dim] = inferred_expression; } } - Shape result = ShapeUtil::MakeShape(element_type, inferred_sizes); + Shape result = + ShapeUtil::MakeShape(element_type, inferred_sizes, inferred_expressions); for (int64_t i = 0; i < inferred_bounds.size(); ++i) { if (!IsUnboundedDynamicSize(inferred_bounds[i]) || IsUnboundedDynamicSize(inferred_sizes[i])) { @@ -756,6 +801,7 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, std::vector dimensions(operand_shape.dimensions_size()); std::vector is_dynamic(operand_shape.dimensions_size()); + std::vector expressions(operand_shape.dimensions_size()); for (int64_t i = 0; i < operand_shape.dimensions_size(); ++i) { const auto& p = padding_config.dimensions(i); if (operand_shape.is_unbounded_dynamic_dimension(i)) { @@ -771,11 +817,13 @@ absl::StatusOr InferMostSpecificDimAndBound(int64_t dim, } } is_dynamic[i] = operand_shape.is_dynamic_dimension(i); + auto diff = dimensions[i] - operand_shape.dimensions(i); + expressions[i] = operand_shape.expressions(i) + diff; } return ShapeUtil::MakeShape( ShapeUtil::HigherPrecisionElementType(operand_shape, padding_value_shape), - dimensions, is_dynamic); + dimensions, is_dynamic, expressions); } // Current DotDimensionNumbers Requirements: @@ -920,7 +968,9 @@ absl::Status CheckDotDimensionConstraints( void GenerateDotResultDimensions( const Shape& lhs, const Shape& rhs, const DotDimensionNumbers& dimension_numbers, - std::vector& dimensions, std::vector& is_dynamic, + std::vector& dimensions, + std::vector& expressions, + std::vector& is_dynamic, std::vector rhs_group_dimensions = {}) { const auto& lhs_batch_dimensions = dimension_numbers.lhs_batch_dimensions(); const auto lhs_batch_dimensions_size = @@ -930,9 +980,11 @@ void GenerateDotResultDimensions( dimension_numbers.rhs_contracting_dimensions().size() - dimension_numbers.rhs_batch_dimensions().size(); dimensions.reserve(lhs_batch_dimensions_size); + expressions.reserve(lhs_batch_dimensions_size); is_dynamic.reserve(lhs_batch_dimensions_size); for (const int64_t lhs_dim : lhs_batch_dimensions) { dimensions.push_back(lhs.dimensions(lhs_dim)); + expressions.push_back(lhs.expressions(lhs_dim)); is_dynamic.push_back(lhs.is_dynamic_dimension(lhs_dim)); } for (int64_t i = 0; i < lhs.dimensions_size(); i++) { @@ -940,6 +992,7 @@ void GenerateDotResultDimensions( i) && !absl::c_linear_search(dimension_numbers.lhs_batch_dimensions(), i)) { dimensions.push_back(lhs.dimensions(i)); + expressions.push_back(lhs.expressions(i)); is_dynamic.push_back(lhs.is_dynamic_dimension(i)); } } @@ -949,6 +1002,7 @@ void GenerateDotResultDimensions( !absl::c_linear_search(dimension_numbers.rhs_batch_dimensions(), i) && !absl::c_linear_search(rhs_group_dimensions, i)) { dimensions.push_back(rhs.dimensions(i)); + expressions.push_back(rhs.expressions(i)); is_dynamic.push_back(rhs.is_dynamic_dimension(i)); } } @@ -990,12 +1044,14 @@ void GenerateDotResultDimensions( std::vector dimensions; std::vector is_dynamic; + std::vector expressions; GenerateDotResultDimensions(lhs, rhs, dimension_numbers, dimensions, - is_dynamic); + expressions, is_dynamic); PrimitiveType type = preferred_element_type.value_or( ShapeUtil::HigherPrecisionElementType(lhs, rhs)); - Shape result = ShapeUtil::MakeShape(type, dimensions, is_dynamic); + Shape result = + ShapeUtil::MakeShape(type, dimensions, is_dynamic, expressions); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result)); VLOG(2) << "inferred dot shape: " << ShapeUtil::HumanString(result); @@ -1193,16 +1249,24 @@ void GenerateDotResultDimensions( PrimitiveType type = preferred_element_type.value_or( ShapeUtil::HigherPrecisionElementType(lhs, rhs)); std::vector dimensions; + std::vector expressions; std::vector is_dynamic; // Add the group dimension to the result shape in case of ragged contracting. if (mode == kContracting) { dimensions.push_back(num_groups); is_dynamic.push_back(is_dynamic_group_sizes); + const DExpr& group_expr = + group_sizes.expressions(group_sizes.dimensions_size() - 1); + expressions.push_back(group_expr ? group_expr + : (is_dynamic_group_sizes + ? DExpr::Unknown(70) + : DExpr::Const(num_groups))); } GenerateDotResultDimensions(lhs, rhs, dimension_numbers, dimensions, - is_dynamic, rhs_group_dimensions); + expressions, is_dynamic, rhs_group_dimensions); - Shape result = ShapeUtil::MakeShape(type, dimensions, is_dynamic); + Shape result = + ShapeUtil::MakeShape(type, dimensions, is_dynamic, expressions); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result)); VLOG(2) << "inferred ragged dot shape: " << ShapeUtil::HumanString(result); return result; @@ -1248,12 +1312,15 @@ void GenerateDotResultDimensions( // Build the resulting shape dimensions. std::vector dimensions; std::vector is_dynamic; + std::vector expressions; for (int64_t i = 0; i < operand_shape.dimensions_size(); ++i) { dimensions.push_back(i != sparsity.dimension() ? operand_shape.dimensions(i) : metadata_dimension_size); is_dynamic.push_back(operand_shape.is_dynamic_dimension(i)); + expressions.push_back(operand_shape.expressions(i)); } - return ShapeUtil::MakeShape(element_type, dimensions, is_dynamic); + return ShapeUtil::MakeShape(element_type, dimensions, is_dynamic, + expressions); } /* static */ absl::StatusOr @@ -1267,6 +1334,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, // from the lhs/rhs pair in every index. std::vector output_dimensions(lhs.dimensions_size()); std::vector output_dimensions_is_dynamic(lhs.dimensions_size()); + std::vector output_dimensions_expressions(lhs.dimensions_size()); for (int64_t i = 0; i < lhs.dimensions_size(); ++i) { if (lhs.dimensions(i) == 1 || rhs.dimensions(i) == 1) { // For the unbounded case, the operand with 1 should be broadcasted to the @@ -1283,6 +1351,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, output_dimensions_is_dynamic[i] = lhs.dimensions(i) == 1 ? rhs.is_dynamic_dimension(i) : lhs.is_dynamic_dimension(i); + output_dimensions_expressions[i] = lhs.dimensions(i) == 1 + ? rhs.expressions(i) + : lhs.expressions(i); } else if (lhs.dimensions(i) == rhs.dimensions(i)) { // LHS | RHS | Result // X | X | X @@ -1293,6 +1364,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, output_dimensions[i] = lhs.dimensions(i); output_dimensions_is_dynamic[i] = lhs.is_dynamic_dimension(i) || rhs.is_dynamic_dimension(i); + output_dimensions_expressions[i] = lhs.expressions(i); } else if (lhs.is_unbounded_dynamic_dimension(i) || rhs.is_unbounded_dynamic_dimension(i)) { // For the last two rows, consider when <=X turns out to be 1 and ? turns @@ -1309,6 +1381,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, output_dimensions_is_dynamic[i] = lhs.is_unbounded_dynamic_dimension(i) ? rhs.is_dynamic_dimension(i) : lhs.is_dynamic_dimension(i); + output_dimensions_expressions[i] = lhs.is_unbounded_dynamic_dimension(i) + ? rhs.expressions(i) + : lhs.expressions(i); } else { return InvalidArgument("Binary op with incompatible shapes: %s and %s.", ShapeUtil::HumanString(lhs), @@ -1317,7 +1392,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, } return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), - output_dimensions, output_dimensions_is_dynamic); + output_dimensions, output_dimensions_is_dynamic, + output_dimensions_expressions); } /* static */ absl::StatusOr ShapeInference::InferInDimBroadcastShape( @@ -1398,6 +1474,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, dimension_to_match, larger_shape.dimensions_size())); } int64_t small_dimension_size = smaller_shape.dimensions(i); + const DExpr& small_dimension_exp = smaller_shape.expressions(i); int64_t large_dimension_size = larger_shape.dimensions(dimension_to_match); bool small_is_dynamic = smaller_shape.is_dynamic_dimension(i); bool large_is_dynamic = @@ -1436,6 +1513,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs, output_shape.set_dimensions(dimension_to_match, small_dimension_size, small_is_dynamic); + output_shape.set_expression(dimension_to_match, small_dimension_exp); } return output_shape; @@ -1484,6 +1562,7 @@ ShapeInference::InferElementwiseBinaryOpShape( for (int64_t i = 0; i < rhs.dimensions_size(); ++i) { if (rhs.is_dynamic_dimension(i)) { result.set_dynamic_dimension(i, true); + result.set_expression(i, rhs.expressions(i)); } } @@ -1779,7 +1858,9 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { output_shape.element_type(), arg_shape->dimensions(), /*dynamic_dimensions=*/ std::vector(arg_shape->dynamic_dimensions().begin(), - arg_shape->dynamic_dimensions().end())); + arg_shape->dynamic_dimensions().end()), + /*expressions=*/ + arg_shape->expressions()); } /* static */ absl::StatusOr ShapeInference::InferBatchNormTrainingShape( @@ -1864,8 +1945,12 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { const int64_t feature_count = operand_shape.dimensions(feature_index); bool dynamic_feature = operand_shape.is_dynamic_dimension(feature_index); - Shape output_shape_for_mean_and_var = ShapeUtil::MakeShape( - operand_shape.element_type(), {feature_count}, {dynamic_feature}); + const DExpr& expression_feature = operand_shape.expressions(feature_index); + std::array feature_expressions = {expression_feature}; + + Shape output_shape_for_mean_and_var = + ShapeUtil::MakeShape(operand_shape.element_type(), {feature_count}, + {dynamic_feature}, feature_expressions); if (!CompatibleDimensionSizes(ShapeUtil::GetDimension(offset_shape, 0), feature_count)) { @@ -2148,8 +2233,12 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { const int64_t feature_count = operand_shape.dimensions(feature_index); bool dynamic_feature = operand_shape.is_dynamic_dimension(feature_index); + const DExpr& expression_feature = operand_shape.expressions(feature_index); + std::array feature_expressions = {expression_feature}; + Shape feature_shape = ShapeUtil::MakeShape( - operand_shape.element_type(), {feature_count}, {dynamic_feature}); + operand_shape.element_type(), {feature_count}, {dynamic_feature}, + feature_expressions); if (!CompatibleDimensionSizes(ShapeUtil::GetDimension(mean_shape, 0), feature_count)) { @@ -2398,13 +2487,18 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { } std::vector dynamic_dimensions(input_spatial_dims.size()); - for (auto it = input_spatial_dims.begin(); it != input_spatial_dims.end(); - ++it) { - dynamic_dimensions[it - input_spatial_dims.begin()] = - IsUnboundedDynamicSize(*it); + std::vector expressions(input_spatial_dims.size()); + for (int i = 0; i < input_spatial_dims.size(); ++i) { + const int64_t input_spatial_dimension = + dnums.input_spatial_dimensions(i); + dynamic_dimensions[i] = IsUnboundedDynamicSize(input_spatial_dims[i]); + expressions[i] = lhs.expressions(input_spatial_dimension) + ? lhs.expressions(input_spatial_dimension) + : DExpr::Const(input_spatial_dims[i]); } Shape base_shape = ShapeUtil::MakeShape( - lhs.element_type(), input_spatial_dims, dynamic_dimensions); + lhs.element_type(), input_spatial_dims, dynamic_dimensions, + expressions); TF_ASSIGN_OR_RETURN( Shape window_output_shape, InferWindowOutputShape(base_shape, window, lhs.element_type())); @@ -2418,6 +2512,15 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { window_output_shape.dimensions(i); } std::vector is_dynamic(num_dims); + std::vector output_expressions(num_dims); + output_expressions[dnums.output_batch_dimension()] = + lhs.expressions(dnums.input_batch_dimension()) / batch_group_count; + output_expressions[dnums.output_feature_dimension()] = + DExpr::Const(kernel_output_features); + for (int i = 0; i < num_spatial_dims; ++i) { + output_expressions[dnums.output_spatial_dimensions(i)] = + window_output_shape.expressions(i); + } for (int i = 0; i < num_dims; i++) { if (lhs.is_dynamic_dimension(i)) { if (i == dnums.input_batch_dimension()) { @@ -2457,7 +2560,8 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { } PrimitiveType type = preferred_element_type.value_or( ShapeUtil::HigherPrecisionElementType(lhs, rhs)); - return ShapeUtil::MakeShape(type, dimensions, is_dynamic); + return ShapeUtil::MakeShape(type, dimensions, is_dynamic, + output_expressions); } /* static */ absl::StatusOr ShapeInference::InferFftShape( @@ -2780,8 +2884,10 @@ ShapeInference::InferScalarBroadcastShape(absl::Span shapes) { const std::vector dynamic_dimensions(shape.dynamic_dimensions().begin(), shape.dynamic_dimensions().end()); + auto exprs = shape.expressions(); + std::vector expressions(exprs.begin(), exprs.end()); return ShapeUtil::MakeShape(shape.element_type(), new_dimensions, - dynamic_dimensions); + dynamic_dimensions, expressions); } /* static */ absl::StatusOr ShapeInference::InferAllToAllTupleShape( @@ -2923,23 +3029,28 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { std::vector new_dimensions; std::vector new_is_dynamic; + std::vector new_expressions; for (int i = 0; i < arg.dimensions_size(); ++i) { if (dimensions_to_reduce_set.find(i) == dimensions_to_reduce_set.end()) { new_dimensions.push_back(arg.dimensions(i)); new_is_dynamic.push_back(arg.is_dynamic_dimension(i)); + new_expressions.push_back(arg.expressions(i)); } } if (ShapeUtil::IsScalar(to_apply.result())) { - return ShapeUtil::MakeShape(to_apply.result().element_type(), - new_dimensions, new_is_dynamic); + return ShapeUtil::MakeShape( + to_apply.result().element_type(), new_dimensions, new_is_dynamic, + new_expressions); } else { std::vector result_subshapes; const auto& tuple_shapes = to_apply.result().tuple_shapes(); result_subshapes.reserve(tuple_shapes.size()); for (const Shape& subshape : tuple_shapes) { - result_subshapes.push_back(ShapeUtil::MakeShape( - subshape.element_type(), new_dimensions, new_is_dynamic)); + auto new_shape = ShapeUtil::MakeShape( + subshape.element_type(), new_dimensions, new_is_dynamic, + new_expressions); + result_subshapes.push_back(new_shape); } return ShapeUtil::MakeTupleShape(result_subshapes); } @@ -3172,7 +3283,9 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { /* static */ absl::StatusOr ShapeInference::InferSliceShape( const Shape& arg, absl::Span starts, - absl::Span limits, absl::Span strides) { + absl::Span limits, absl::Span strides, + absl::Span start_exprs, + absl::Span limit_exprs) { auto error = [&](const std::string& message) { return InvalidArgument( "%s in slice operation; argument shape: %s; starts: {%s}; limits: " @@ -3202,8 +3315,10 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { } std::vector sizes; + std::vector expressions; const auto starts_size = starts.size(); sizes.reserve(starts_size); + expressions.reserve(starts_size); for (int64_t dimension = 0; dimension < starts_size; ++dimension) { int64_t start_index = starts[dimension]; int64_t limit_index = limits[dimension]; @@ -3231,6 +3346,15 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { return InvalidArgument("Stride (%d) must be positive.", stride); } sizes.push_back((limit_index - start_index + stride - 1) / stride); + + DExpr limit_expr = + limit_exprs.empty() ? DExpr::Const(limit_index) : limit_exprs[dimension]; + DExpr start_expr = + start_exprs.empty() ? DExpr::Const(start_index) : start_exprs[dimension]; + + auto new_expr = + limit_expr - start_expr + DExpr::Const(stride) - DExpr::Const(1); + expressions.push_back((new_expr / DExpr::Const(stride)).simplify()); } std::vector is_dynamic(arg.dimensions_size()); @@ -3242,12 +3366,14 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { is_dynamic[i] = arg.is_bounded_dynamic_dimension(i); } - return ShapeUtil::MakeShape(arg.element_type(), sizes, is_dynamic); + return ShapeUtil::MakeShape(arg.element_type(), sizes, is_dynamic, + expressions); } /* static */ absl::StatusOr ShapeInference::InferDynamicSliceShape( const Shape& operand_shape, absl::Span start_index_shapes, - absl::Span slice_sizes, bool allow_scalar_indices) { + absl::Span slice_sizes, + absl::Span slice_exprs, bool allow_scalar_indices) { TF_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of dynamic slice")); auto number_of_indices = start_index_shapes.size(); // TODO(b/118437727): Remove this path. @@ -3346,8 +3472,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { VLOG(2) << StrFormat("slice_sizes[%d] = %d", dim, slice_dim_size); } - Shape result = - ShapeUtil::MakeShape(operand_shape.element_type(), slice_sizes); + Shape result = ShapeUtil::MakeShape(operand_shape.element_type(), slice_sizes, + slice_exprs); for (int64_t dimension = 0; dimension < operand_shape.dimensions_size(); ++dimension) { @@ -3646,7 +3772,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { } /* static */ absl::StatusOr ShapeInference::InferBroadcastShape( - const Shape& operand, absl::Span broadcast_sizes) { + const Shape& operand, absl::Span broadcast_sizes, + absl::Span broadcast_exprs) { // This method is used to infer shape for xla::BroadcastInDim. TF_RETURN_IF_ERROR(ExpectArray(operand, "operand of broadcast")); TF_RET_CHECK(!operand.is_unbounded_dynamic()); @@ -3666,12 +3793,14 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { std::copy(operand.dimensions().begin(), operand.dimensions().end(), dimensions.begin() + broadcast_sizes.size()); - TF_ASSIGN_OR_RETURN(Shape result, ShapeUtil::MakeValidatedShape( - operand.element_type(), dimensions)); + TF_ASSIGN_OR_RETURN( + Shape result, ShapeUtil::MakeValidatedShape(operand.element_type(), + dimensions, broadcast_exprs)); for (int64_t i = 0; i < operand.dimensions_size(); ++i) { result.set_dynamic_dimension(broadcast_sizes.size() + i, operand.is_dynamic_dimension(i)); } + return result; } @@ -3735,7 +3864,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { /* static */ absl::StatusOr ShapeInference::InferDynamicReshapeShape( const Shape& operand, absl::Span dim_size_shapes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic) { + const std::vector& dims_are_dynamic, + absl::Span expressions) { if (new_size_bounds.size() != dims_are_dynamic.size()) { return InvalidArgument( "DynamicReshape has to have the same number of elements in new_sizes " @@ -3751,9 +3881,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { dim_size_shape->ToString()); } } - Shape inferred_shape = ShapeUtil::MakeShape( - operand.element_type(), new_size_bounds, dims_are_dynamic); + operand.element_type(), new_size_bounds, dims_are_dynamic, expressions); if (ShapeUtil::ElementsIn(operand) != ShapeUtil::ElementsIn(inferred_shape)) { return InvalidArgument( "Reshape operation has mismatched element counts: from=%d (%s) " @@ -3767,10 +3896,21 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { /* static */ absl::StatusOr ShapeInference::InferReshapeShape( const Shape& operand, absl::Span dimensions, - int64_t inferred_dimension) { + int64_t inferred_dimension, absl::Span expressions) { TF_RETURN_IF_ERROR(ExpectArray(operand, "reshape")); Shape inferred_shape = - ShapeUtil::MakeShape(operand.element_type(), dimensions); + ShapeUtil::MakeShape(operand.element_type(), dimensions, expressions); + + if (expressions.empty() && operand.expressions().size() > 0 && + operand.expressions(0)->is_dynamic()) { + return InvalidArgument("Expressions is empty but operand is dynamic"); + } + + // if (!expressions.empty() && expressions[0]->is_constant() && + // expressions[0]->get_val() == 977) { + // return InvalidArgument("Expressions[0] is the magic number (977)."); + // } + VLOG(3) << "Reshape inferred shape: " << ShapeUtil::HumanString(inferred_shape); @@ -3790,6 +3930,16 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { ShapeUtil::ElementsIn(inferred_shape), ShapeUtil::HumanString(inferred_shape)); } + if (!expressions.empty()) { + DExpr input_elements = SymbolicElementsIn(operand); + DExpr output_elements = SymbolicElementsIn(inferred_shape); + if (!DynExpr::equal(input_elements.get(), output_elements.get())) { + return InvalidArgument( + "Reshape operation has mismatched symbolic element counts: " + "from=%s to=%s.", + ShapeUtil::HumanString(operand), ShapeUtil::HumanString(inferred_shape)); + } + } std::vector indices(operand.dimensions_size()); std::iota(indices.begin(), indices.end(), 0); @@ -4017,6 +4167,22 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { on_true.is_dynamic_dimension(dimension) || on_false.is_dynamic_dimension(dimension)); } + const DExpr& on_true_expr = on_true.expressions(dimension); + const DExpr& on_false_expr = on_false.expressions(dimension); + const bool has_dynamic_expression = + (on_true_expr && on_true_expr->is_dynamic()) || + (on_false_expr && on_false_expr->is_dynamic()); + if (has_dynamic_expression) { + if (!on_true_expr || !on_false_expr || + !DynExpr::equal(on_true_expr, on_false_expr)) { + return InvalidArgument( + "Select operands have mismatched expressions in dimension %d: " + "on_true=%s, on_false=%s.", + dimension, ShapeUtil::HumanString(on_true), + ShapeUtil::HumanString(on_false)); + } + result.set_expression(dimension, on_true_expr); + } } if (result.has_layout()) { result.mutable_layout()->set_element_size_in_bits( @@ -4250,18 +4416,25 @@ static absl::Status ValidateGatherDimensionNumbers( std::vector expanded_start_indices_shape; // Also tracks if an output dimension is dynamic. std::vector expanded_start_indices_shape_dynamic_dimensions; + std::vector expanded_start_indices_shape_expressions; expanded_start_indices_shape.reserve(start_indices_shape.dimensions_size()); expanded_start_indices_shape_dynamic_dimensions.reserve( start_indices_shape.dimensions_size()); + expanded_start_indices_shape_expressions.reserve( + start_indices_shape.dimensions_size()); absl::c_copy(start_indices_shape.dimensions(), std::back_inserter(expanded_start_indices_shape)); absl::c_copy( start_indices_shape.dynamic_dimensions(), std::back_inserter(expanded_start_indices_shape_dynamic_dimensions)); + absl::c_copy( + start_indices_shape.expressions(), + std::back_inserter(expanded_start_indices_shape_expressions)); if (expanded_start_indices_shape.size() == gather_dim_numbers.index_vector_dim()) { expanded_start_indices_shape.push_back(1); expanded_start_indices_shape_dynamic_dimensions.push_back(false); + expanded_start_indices_shape_expressions.push_back(DExpr::Const(1)); } TF_RETURN_IF_ERROR(ValidateGatherDimensionNumbers( @@ -4328,10 +4501,12 @@ static absl::Status ValidateGatherDimensionNumbers( output_dim_bounds.reserve(result_rank); std::vector output_dim_is_dynamic; + std::vector output_expressions; output_dim_is_dynamic.reserve(result_rank); for (int64_t i = 0; i < result_rank; i++) { int64_t current_bound; bool dim_dynamic = false; + DExpr expression = DExpr::Unknown(80); bool is_window_index = absl::c_binary_search(gather_dim_numbers.offset_dims(), i); if (is_window_index) { @@ -4353,6 +4528,9 @@ static absl::Status ValidateGatherDimensionNumbers( if (slice_sizes[offset_dims_seen] == input_shape.dimensions(offset_dims_seen)) { dim_dynamic = input_shape.is_dynamic_dimension(offset_dims_seen); + expression = input_shape.expressions(offset_dims_seen); + } else { + expression = DExpr::Const(slice_sizes[offset_dims_seen]); } current_bound = slice_sizes[offset_dims_seen++]; } else { @@ -4362,15 +4540,17 @@ static absl::Status ValidateGatherDimensionNumbers( // Forward dynamic dimensions from indices. dim_dynamic = expanded_start_indices_shape_dynamic_dimensions[gather_dims_seen]; - + expression = expanded_start_indices_shape_expressions[gather_dims_seen]; current_bound = expanded_start_indices_shape[gather_dims_seen++]; } output_dim_is_dynamic.push_back(dim_dynamic); + output_expressions.push_back(expression); output_dim_bounds.push_back(current_bound); } - return ShapeUtil::MakeShape(input_shape.element_type(), output_dim_bounds, - output_dim_is_dynamic); + auto s = ShapeUtil::MakeShape(input_shape.element_type(), output_dim_bounds, + output_dim_is_dynamic, output_expressions); + return s; } namespace { diff --git a/third_party/xla/xla/service/shape_inference.h b/third_party/xla/xla/service/shape_inference.h index 15818f01cdfa7e..cb0ee9a20f83c6 100644 --- a/third_party/xla/xla/service/shape_inference.h +++ b/third_party/xla/xla/service/shape_inference.h @@ -245,13 +245,17 @@ class ShapeInference { // e.g. slice f32[32x32] 0:16 0:16 -> f32[16x16] static absl::StatusOr InferSliceShape( const Shape& arg, absl::Span starts, - absl::Span limits, absl::Span strides); + absl::Span limits, absl::Span strides, + absl::Span start_exprs = {}, + absl::Span limit_exprs = {}); // Infers the shape produced by a dynamic slice operation of size specified // in 'slice_sizes', with dynamic start indices shape 'start_indices_shape'. static absl::StatusOr InferDynamicSliceShape( const Shape& operand_shape, absl::Span start_index_shapes, - absl::Span slice_sizes, bool allow_scalar_indices = true); + absl::Span slice_sizes, + absl::Span slice_exprs = {}, + bool allow_scalar_indices = true); // Infers the shape produced by a dynamic update slice operation based // on the shape of operand and update. @@ -283,7 +287,8 @@ class ShapeInference { // Infers the shape produced by a broadcast operation. static absl::StatusOr InferBroadcastShape( - const Shape& operand, absl::Span broadcast_sizes); + const Shape& operand, absl::Span broadcast_sizes, + absl::Span broadcast_exprs = {}); // Checks whether the given parameters can form a broadcast. Returns the same // output_shape if it's legal. @@ -295,7 +300,7 @@ class ShapeInference { // its operand and the new dimension sizes specified. static absl::StatusOr InferReshapeShape( const Shape& operand, absl::Span dimensions, - int64_t inferred_dimension); + int64_t inferred_dimension, absl::Span expressions); // Infers the shape produced by a dynamic reshape operation from the element // type of its operand and the new dimension sizes specified. The result shape @@ -304,7 +309,8 @@ class ShapeInference { static absl::StatusOr InferDynamicReshapeShape( const Shape& operand, absl::Span dim_size_shapes, absl::Span new_size_bounds, - const std::vector& dims_are_dynamic); + const std::vector& dims_are_dynamic, + absl::Span expressions); // Infers the shape produced by a transpose operation from the element type of // its operand and its dimensions field. diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc index 87841d476712d7..0b6e4bf00650b2 100644 --- a/third_party/xla/xla/service/shape_inference_test.cc +++ b/third_party/xla/xla/service/shape_inference_test.cc @@ -205,6 +205,47 @@ TEST_F(ShapeInferenceTest, SelectArrayPredBetweenArrays) { ASSERT_TRUE(ShapeUtil::Equal(matrix_64_48_, *inferred_shape)); } +TEST_F(ShapeInferenceTest, SelectPreservesExpressionsFromOperands) { + const Shape pred = ShapeUtil::MakeShape(PRED, {8, 5}); + const Shape on_true = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(33), DExpr::Const(5)}); + const Shape on_false = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(33), DExpr::Const(5)}); + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true, + on_false)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(33))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5))); +} + +TEST_F(ShapeInferenceTest, SelectRejectsMismatchedOperandExpressions) { + const Shape pred = ShapeUtil::MakeShape(PRED, {}); + const Shape on_true = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(33), DExpr::Const(5)}); + const Shape on_false = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(34), DExpr::Const(5)}); + const absl::StatusOr inferred_shape = + ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true, + on_false); + ASSERT_FALSE(inferred_shape.ok()); + EXPECT_THAT(inferred_shape.status().message(), + HasSubstr("mismatched expressions in dimension 0")); +} + +TEST_F(ShapeInferenceTest, SelectRejectsMissingDynamicOperandExpression) { + const Shape pred = ShapeUtil::MakeShape(PRED, {}); + const Shape on_true = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(33), DExpr::Const(5)}); + const Shape on_false = ShapeUtil::MakeShape(F32, {8, 5}); + const absl::StatusOr inferred_shape = + ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true, + on_false); + ASSERT_FALSE(inferred_shape.ok()); + EXPECT_THAT(inferred_shape.status().message(), + HasSubstr("mismatched expressions in dimension 0")); +} + TEST_F(ShapeInferenceTest, SelectBadShapes) { const absl::StatusOr inferred_shape_error1 = ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred_, @@ -448,6 +489,63 @@ TEST_F(ShapeInferenceTest, ReduceWindowInHalf) { ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {4, 4}), *inferred_shape)); } +TEST_F(ShapeInferenceTest, ReduceWindowBuildsWindowedExpressions) { + const Shape matrix_shape = ShapeUtil::MakeShape( + F32, {8, 8}, std::vector{DExpr::Var(1), DExpr::Var(2)}); + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(2); + dim.set_padding_low(0); + dim.set_padding_high(0); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + const Shape init_value_shape = ShapeUtil::MakeShape(F32, {}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReduceWindowShape(matrix_shape, init_value_shape, + window)); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, ShapeUtil::MakeShape( + F32, {4, 4}, + std::vector{((DExpr::Var(1) - 2) / 2) + 1, + ((DExpr::Var(2) - 2) / 2) + 1}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + ((DExpr::Var(1) - 2) / 2) + 1)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), + ((DExpr::Var(2) - 2) / 2) + 1)); +} + +TEST_F(ShapeInferenceTest, ReduceWindowBuildsPaddedStridedExpressions) { + const Shape vector_shape = + ShapeUtil::MakeShape(F32, {11}, std::vector{DExpr::Var(3)}); + Window window; + WindowDimension dim; + dim.set_size(3); + dim.set_stride(2); + dim.set_padding_low(1); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + const Shape init_value_shape = ShapeUtil::MakeShape(F32, {}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReduceWindowShape(vector_shape, init_value_shape, + window)); + + const DExpr expected = (((DExpr::Var(3) - 1) / 2) + 1).simplify(); + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, + ShapeUtil::MakeShape(F32, {6}, std::vector{expected}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), expected)); +} + TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterProperShapes) { const absl::StatusOr inferred_shape_ok = ShapeInference::InferSelectAndScatterShape( @@ -721,6 +819,58 @@ TEST_F(ShapeInferenceTest, ConvolveWithBaseDilation) { *inferred_shape)); } +TEST_F(ShapeInferenceTest, ConvolveBuildsBatchAndSpatialExpressions) { + ConvolutionDimensionNumbers dnums; + const Shape lhs_shape = ShapeUtil::MakeShape( + F32, {5, 11, 13, 3}, + std::vector{DExpr::Var(4), DExpr::Var(5), DExpr::Var(6), + DExpr::Const(3)}); + dnums.set_input_batch_dimension(0); + dnums.set_output_batch_dimension(0); + dnums.add_input_spatial_dimensions(1); + dnums.add_output_spatial_dimensions(1); + dnums.add_input_spatial_dimensions(2); + dnums.add_output_spatial_dimensions(2); + dnums.set_input_feature_dimension(3); + dnums.set_output_feature_dimension(3); + + const Shape rhs_shape = ShapeUtil::MakeShape(F32, {3, 3, 3, 7}); + dnums.add_kernel_spatial_dimensions(0); + dnums.add_kernel_spatial_dimensions(1); + dnums.set_kernel_input_feature_dimension(2); + dnums.set_kernel_output_feature_dimension(3); + + Window window; + auto* dim0 = window.add_dimensions(); + dim0->set_size(3); + dim0->set_stride(2); + dim0->set_padding_low(1); + dim0->set_padding_high(1); + dim0->set_window_dilation(1); + dim0->set_base_dilation(1); + auto* dim1 = window.add_dimensions(); + *dim1 = *dim0; + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferConvolveShape( + lhs_shape, rhs_shape, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, + /*preferred_element_type=*/std::nullopt)); + + const DExpr expected_h = (((DExpr::Var(5) - 1) / 2) + 1).simplify(); + const DExpr expected_w = (((DExpr::Var(6) - 1) / 2) + 1).simplify(); + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, + ShapeUtil::MakeShape(F32, {5, 6, 7, 7}, + std::vector{DExpr::Var(4), expected_h, + expected_w, DExpr::Const(7)}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(4))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), expected_h)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), expected_w)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(3), DExpr::Const(7))); +} + TEST_F(ShapeInferenceTest, ConvolveDimensionNumbersOverlapError) { // Dimension order for this test: batch, feature, x0, x1 const Shape lhs_shape = ShapeUtil::MakeShape(F32, {10, 11, 3, 4}); @@ -1028,7 +1178,8 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanks) { TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanksBounded) { FftType type = FftType::FFT; - const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true}); + const Shape shape = + ShapeUtil::MakeShape(C64, {16, 8}, {false, true}, {}); fft::Fail(shape, type, {}, fft::unsupported_rank); fft::Pass(shape, type, {8}, shape); fft::Pass(shape, type, {16, 8}, shape); @@ -1056,7 +1207,8 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanks) { TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanksBounded) { FftType type = FftType::IFFT; - const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true}); + const Shape shape = + ShapeUtil::MakeShape(C64, {16, 8}, {false, true}, {}); fft::Fail(shape, type, {}, fft::unsupported_rank); fft::Pass(shape, type, {8}, shape); fft::Pass(shape, type, {16, 8}, shape); @@ -1103,9 +1255,9 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestRfftDimensions) { fft::Pass(odd_shape_in, type, {16, 9}, shape_out); const Shape bounded_shape_in = - ShapeUtil::MakeShape(F32, {16, 8}, {false, true}); + ShapeUtil::MakeShape(F32, {16, 8}, {false, true}, {}); const Shape bounded_shape_out = - ShapeUtil::MakeShape(C64, {16, 5}, {false, true}); + ShapeUtil::MakeShape(C64, {16, 5}, {false, true}, {}); fft::Pass(bounded_shape_in, type, {16, 8}, bounded_shape_out); } @@ -1147,9 +1299,9 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIrfftDimensions) { fft::Pass(shape, type, {16, 9}, odd_shape_out); const Shape bounded_shape_in = - ShapeUtil::MakeShape(C64, {16, 5}, {false, true}); + ShapeUtil::MakeShape(C64, {16, 5}, {false, true}, {}); const Shape bounded_shape_out = - ShapeUtil::MakeShape(F32, {16, 9}, {false, true}); + ShapeUtil::MakeShape(F32, {16, 9}, {false, true}, {}); fft::Pass(bounded_shape_in, type, {16, 9}, bounded_shape_out); } @@ -1283,6 +1435,18 @@ TEST_F(ShapeInferenceTest, MapWithDifferentInputTypes) { EXPECT_TRUE(ShapeUtil::Equal(expected, *inferred_shape)); } +TEST_F(ShapeInferenceTest, MapPreservesExpressions) { + const Shape arg = ShapeUtil::MakeShape( + F32, {20, 7}, std::vector{true, false}, + std::vector{DExpr::Var(11), DExpr::Const(7)}); + ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, f32_); + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape, + ShapeInference::InferMapShape({&arg}, to_apply, + {0, 1})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(11))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(7))); +} + TEST_F(ReduceShapeInferenceTest, ReduceVectorToScalar) { ExpectInferredReduceShape(f32_, ShapeUtil::MakeShape(F32, {128}), /*dimensions_to_reduce=*/{0}); @@ -1367,6 +1531,58 @@ TEST_F(ReduceShapeInferenceTest, ReduceWindowMultiOutput) { *inferred_shape)); } +TEST_F(ReduceShapeInferenceTest, + ReduceWindowPreservesDynamicStridedBoundExpression) { + Shape operand = ShapeUtil::MakeShape( + F32, {101}, std::vector{true}, {DExpr::Var(1)}); + Window window; + WindowDimension* dimension = window.add_dimensions(); + dimension->set_size(3); + dimension->set_stride(2); + dimension->set_padding_low(1); + dimension->set_padding_high(1); + dimension->set_base_dilation(1); + dimension->set_window_dilation(1); + + TF_ASSERT_OK_AND_ASSIGN( + Shape inferred, + ShapeInference::InferReduceWindowShape(operand, f32_, window)); + EXPECT_EQ(51, inferred.dimensions(0)); + EXPECT_TRUE(inferred.expressions(0) == + DExpr::Max((DExpr::Var(1) + 1) / 2, DExpr::Const(0))); + + DExpr runtime_expression = + inferred.expressions(0).substitute(1, DExpr::Const(100)).simplify(); + ASSERT_TRUE(runtime_expression->is_constant()); + EXPECT_EQ(50, runtime_expression->get_val()); +} + +TEST_F(ReduceShapeInferenceTest, + ReduceWindowClampsDynamicStridedBoundAtZero) { + Shape operand = ShapeUtil::MakeShape( + F32, {101}, std::vector{true}, {DExpr::Var(2)}); + Window window; + WindowDimension* dimension = window.add_dimensions(); + dimension->set_size(5); + dimension->set_stride(1); + dimension->set_padding_low(0); + dimension->set_padding_high(0); + dimension->set_base_dilation(1); + dimension->set_window_dilation(1); + + TF_ASSERT_OK_AND_ASSIGN( + Shape inferred, + ShapeInference::InferReduceWindowShape(operand, f32_, window)); + EXPECT_EQ(97, inferred.dimensions(0)); + EXPECT_TRUE(inferred.expressions(0) == + DExpr::Max(DExpr::Var(2) - 4, DExpr::Const(0))); + + DExpr runtime_expression = + inferred.expressions(0).substitute(2, DExpr::Const(2)).simplify(); + ASSERT_TRUE(runtime_expression->is_constant()); + EXPECT_EQ(0, runtime_expression->get_val()); +} + TEST_F(ReduceShapeInferenceTest, ErrorMultiOutputBadReducerInput1) { const Shape f32_arg_shape = ShapeUtil::MakeShape(F32, {5, 3}); const Shape s32_arg_shape = ShapeUtil::MakeShape(S32, {5, 3}); @@ -1520,12 +1736,31 @@ TEST_F(ShapeInferenceTest, InferSliceShapeRank2) { } TEST_F(ShapeInferenceTest, InferSliceWithDynamicDimensions) { - const Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}, {true, true}); + const Shape matrix_shape = + ShapeUtil::MakeShape(F32, {128, 64}, {true, true}, {}); const absl::StatusOr inferred_shape = ShapeInference::InferSliceShape(matrix_shape, {32, 0}, {33, 64}, {1, 1}); ASSERT_IS_OK(inferred_shape.status()); ASSERT_TRUE(ShapeUtil::Equal( - ShapeUtil::MakeShape(F32, {1, 64}, {false, true}), *inferred_shape)); + ShapeUtil::MakeShape(F32, {1, 64}, {false, true}, {}), + *inferred_shape)); +} + +TEST_F(ShapeInferenceTest, InferSliceBuildsExpressionFromSymbolicBounds) { + const Shape vector_shape = + ShapeUtil::MakeShape(F32, {16}, std::vector{DExpr::Var(1)}); + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferSliceShape( + vector_shape, /*starts=*/{3}, /*limits=*/{8}, /*strides=*/{1}, + /*start_exprs=*/{DExpr::Const(3)}, + /*limit_exprs=*/{DExpr::Var(2)})); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, ShapeUtil::MakeShape( + F32, {5}, std::vector{DExpr::Var(2) - 3}))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(2) - 3)); } TEST_F(ShapeInferenceTest, InferSliceShapeRank2WithStrides) { @@ -1583,6 +1818,26 @@ TEST_F(ShapeInferenceTest, InferConstIndexShape) { ASSERT_TRUE(ShapeUtil::Equal(s32_, *inferred1_status)); } +TEST_F(ShapeInferenceTest, InferConstIndexPreservesExpressions) { + const Shape lhs = ShapeUtil::MakeShape( + F32, {8, 5}, std::vector{DExpr::Var(24), DExpr::Const(5)}); + const Shape rhs = ShapeUtil::MakeShape( + S32, {3, 7}, std::vector{DExpr::Const(3), DExpr::Var(25)}); + const Shape tuple_shape = ShapeUtil::MakeTupleShape({lhs, rhs}); + + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred0, + ShapeInference::InferGetTupleElementShape( + tuple_shape, /*index=*/0)); + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred1, + ShapeInference::InferGetTupleElementShape( + tuple_shape, /*index=*/1)); + + EXPECT_TRUE(DynExpr::equal(inferred0.expressions(0), DExpr::Var(24))); + EXPECT_TRUE(DynExpr::equal(inferred0.expressions(1), DExpr::Const(5))); + EXPECT_TRUE(DynExpr::equal(inferred1.expressions(0), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(inferred1.expressions(1), DExpr::Var(25))); +} + TEST_F(ShapeInferenceTest, InferTupleElementShapeOutOfBound) { const Shape tuple_shape = ShapeUtil::MakeTupleShape({f32_, s32_}); const absl::StatusOr inferredNegative_status = @@ -1622,11 +1877,13 @@ TEST_F(ShapeInferenceTest, InferReshapeDegenerateCombine) { // [<=1] // // Both output dimension can be dynamic, use inferred_dimension to tie-break. - const Shape operand = ShapeUtil::MakeShape(F32, {1, 1}, {false, true}); + const Shape operand = ShapeUtil::MakeShape(F32, {1, 1}, {false, true}, {}); const auto status = ShapeInference::InferReshapeShape(operand, {1}, - /*inferred_dimension=*/-1); - ASSERT_EQ(ShapeUtil::MakeShape(F32, {1}, {true}), *status); + /*inferred_dimension=*/-1, + /*expressions=*/{}); + ASSERT_EQ(ShapeUtil::MakeShape(F32, {1}, {true}, {}), + *status); } TEST_F(ShapeInferenceTest, InferReshapeSplit) { @@ -1635,48 +1892,218 @@ TEST_F(ShapeInferenceTest, InferReshapeSplit) { // [1, 10] // // Both output dimension can be dynamic, use inferred_dimension to tie-break. - const Shape operand = ShapeUtil::MakeShape(F32, {10}, {true}); + const Shape operand = ShapeUtil::MakeShape(F32, {10}, {true}, {}); const auto status = ShapeInference::InferReshapeShape(operand, {1, 10}, - /*inferred_dimension=*/0); - ASSERT_EQ(ShapeUtil::MakeShape(F32, {1, 10}, {true, false}), *status); + /*inferred_dimension=*/0, + /*expressions=*/{}); + ASSERT_EQ( + ShapeUtil::MakeShape(F32, {1, 10}, {true, false}, {}), + *status); } TEST_F(ShapeInferenceTest, InferReshapeCombine) { // [6, <=10] // | reshape // [<=60] - const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}); + const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}, {}); const auto status = ShapeInference::InferReshapeShape(operand, {60}, - /*inferred_dimension=*/-11); - ASSERT_EQ(ShapeUtil::MakeShape(F32, {60}, {true}), *status); + /*inferred_dimension=*/-11, + /*expressions=*/{}); + ASSERT_EQ(ShapeUtil::MakeShape(F32, {60}, {true}, {}), + *status); } TEST_F(ShapeInferenceTest, UnchangedDimension) { // [6, <=10] // | reshape // [2, 3, <=10] - const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}); + const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}, {}); const auto status = ShapeInference::InferReshapeShape(operand, {2, 3, 10}, - /*inferred_dimension=*/-11); - ASSERT_EQ(ShapeUtil::MakeShape(F32, {2, 3, 10}, {false, false, true}), + /*inferred_dimension=*/-11, + /*expressions=*/{}); + ASSERT_EQ(ShapeUtil::MakeShape(F32, {2, 3, 10}, {false, false, true}, {}), *status); } +TEST_F(ShapeInferenceTest, ReshapePreservesProvidedExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {6, 10}, std::vector{DExpr::Const(6), DExpr::Var(1)}); + const Shape expected = ShapeUtil::MakeShape( + F32, {2, 3, 10}, + std::vector{DExpr::Const(2), DExpr::Const(3), DExpr::Var(1)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReshapeShape( + operand, expected.dimensions(), + /*inferred_dimension=*/-1, expected.expressions())); + + EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Const(2))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(1))); +} + +TEST_F(ShapeInferenceTest, ReshapeCombinesLeadingSymbolicWithStaticFactor) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 16, 32}, + std::vector{DExpr::Var(1), DExpr::Const(16), DExpr::Const(32)}); + const Shape expected = ShapeUtil::MakeShape( + F32, {80, 32}, + std::vector{16 * DExpr::Var(1), DExpr::Const(32)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReshapeShape( + operand, expected.dimensions(), + /*inferred_dimension=*/-1, expected.expressions())); + + EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + 16 * DExpr::Var(1))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(32))); +} + +TEST_F(ShapeInferenceTest, ReshapeCollapsesTwoStaticDimsIntoSymbolicExtent) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 4, 8}, + std::vector{DExpr::Var(1), DExpr::Const(4), DExpr::Const(8)}); + const Shape expected = + ShapeUtil::MakeShape(F32, {160}, std::vector{32 * DExpr::Var(1)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReshapeShape( + operand, expected.dimensions(), + /*inferred_dimension=*/-1, expected.expressions())); + + EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + 32 * DExpr::Var(1))); +} + +TEST_F(ShapeInferenceTest, ReshapeSplitsSymbolicExtentByStaticFactor) { + const Shape operand = ShapeUtil::MakeShape( + F32, {80, 8}, std::vector{DExpr::Var(1), DExpr::Const(8)}); + const Shape expected = ShapeUtil::MakeShape( + F32, {40, 16}, + std::vector{DExpr::Var(1) / 2, DExpr::Const(16)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReshapeShape( + operand, expected.dimensions(), + /*inferred_dimension=*/-1, expected.expressions())); + + EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + DExpr::Var(1) / 2)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(16))); +} + +TEST_F(ShapeInferenceTest, ReshapeSplitsAndCollapsesSymbolicExtent) { + const Shape operand = ShapeUtil::MakeShape( + F32, {20, 8, 4}, + std::vector{DExpr::Var(1), DExpr::Const(8), DExpr::Const(4)}); + const Shape expected = ShapeUtil::MakeShape( + F32, {10, 64}, + std::vector{DExpr::Var(1) / 2, DExpr::Const(64)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReshapeShape( + operand, expected.dimensions(), + /*inferred_dimension=*/-1, expected.expressions())); + + EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + DExpr::Var(1) / 2)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(64))); +} + +TEST_F(ShapeInferenceTest, ReshapeRejectsIncorrectCollapsedExpression) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 16, 32}, + std::vector{DExpr::Var(1), DExpr::Const(16), DExpr::Const(32)}); + const Shape incorrect = ShapeUtil::MakeShape( + F32, {80, 32}, + std::vector{8 * DExpr::Var(1), DExpr::Const(32)}); + + const absl::StatusOr status = ShapeInference::InferReshapeShape( + operand, incorrect.dimensions(), + /*inferred_dimension=*/-1, incorrect.expressions()); + + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.status().message(), + HasSubstr("Reshape operation has mismatched symbolic element " + "counts")); +} + +TEST_F(ShapeInferenceTest, + ReshapeRejectsIncorrectSplitAndCollapseExpression) { + const Shape operand = ShapeUtil::MakeShape( + F32, {20, 8, 4}, + std::vector{DExpr::Var(1), DExpr::Const(8), DExpr::Const(4)}); + const Shape incorrect = ShapeUtil::MakeShape( + F32, {10, 64}, + std::vector{DExpr::Var(1), DExpr::Const(64)}); + + const absl::StatusOr status = ShapeInference::InferReshapeShape( + operand, incorrect.dimensions(), + /*inferred_dimension=*/-1, incorrect.expressions()); + + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.status().message(), + HasSubstr("Reshape operation has mismatched symbolic element " + "counts")); +} + +TEST_F(ShapeInferenceTest, ReshapeWithSymbolicOperandRequiresExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {10, 6}, std::vector{DExpr::Var(1), DExpr::Const(6)}); + + const absl::StatusOr status = + ShapeInference::InferReshapeShape(operand, {2, 5, 6}, + /*inferred_dimension=*/-1, + /*expressions=*/{}); + + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.status().message(), + HasSubstr("Expressions is empty but operand is dynamic")); +} + TEST_F(ShapeInferenceTest, InferDynamicBroadcast) { // CHECK: // %broadcast = s32[15,<=15]{1,0} broadcast(s32[<=15]{0}), dimensions={1} - const Shape operand_shape = ShapeUtil::MakeShape(F32, {15}, {true}); + const Shape operand_shape = ShapeUtil::MakeShape(F32, {15}, {true}, {}); const absl::StatusOr inferred_shape = ShapeInference::InferBroadcastShape(operand_shape, {15}); ASSERT_IS_OK(inferred_shape.status()); - ASSERT_EQ(ShapeUtil::MakeShape(F32, {15, 15}, {false, true}), + ASSERT_EQ(ShapeUtil::MakeShape(F32, {15, 15}, {false, true}, {}), *inferred_shape); } +TEST_F(ShapeInferenceTest, BroadcastInDimPreservesMappedExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {2, 4}, std::vector{true, true}, + {DExpr::Var(1), DExpr::Var(2)}); + const Shape output = ShapeUtil::MakeShape( + F32, {2, 3, 4}, std::vector{true, false, true}, + {DExpr::Var(1), DExpr::Const(3), DExpr::Var(2)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferBroadcastShape(operand, output, + /*broadcast_dimensions=*/{0, 2})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(2))); +} + TEST_F(ShapeInferenceTest, BroadcastScalar) { for (auto element_type : {F32, U32, S8}) { const Shape scalar_shape = ShapeUtil::MakeShape(element_type, {}); @@ -2172,6 +2599,27 @@ TEST_F(ShapeInferenceTest, SparseDotMetadata) { ShapeUtil::Equal(inferred_shape, ShapeUtil::MakeShape(U16, {5, 10, 2}))); } +TEST_F(ShapeInferenceTest, SparseDotMetadataPreservesNonSparseExpressions) { + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_batch_dimensions(0); + dot_dnums.add_lhs_contracting_dimensions(2); + SparsityDescriptor sparsity_descriptor; + sparsity_descriptor.set_type(SparsityType::SPARSITY_STRUCTURED_N_M); + sparsity_descriptor.set_n(2); + sparsity_descriptor.set_m(4); + sparsity_descriptor.set_index(0); + sparsity_descriptor.set_dimension(2); + + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 10, 16}, + std::vector{DExpr::Var(16), DExpr::Const(10), DExpr::Const(16)}); + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape, + ShapeInference::InferSparseDotMetadataShape( + operand, dot_dnums, sparsity_descriptor)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(16))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10))); +} + // mode 1 : [m,k], [g,k,n], [g] -> [m,n] TEST_F(ShapeInferenceTest, RaggedDotRaggedNonContracting) { const Shape lhs_shape = ShapeUtil::MakeShape(F32, {11, 5}); @@ -2221,6 +2669,31 @@ TEST_F(ShapeInferenceTest, RaggedDotRaggedContracting) { << " expected: " << ShapeUtil::HumanString(output_shape); } +TEST_F(ShapeInferenceTest, RaggedDotPreservesGroupExpression) { + const Shape lhs_shape = ShapeUtil::MakeShape( + F32, {11, 5}, std::vector{DExpr::Const(11), DExpr::Const(5)}); + const Shape rhs_shape = ShapeUtil::MakeShape( + F32, {5, 7}, std::vector{DExpr::Const(5), DExpr::Const(7)}); + const Shape group_sizes_shape = + ShapeUtil::MakeShape(U32, {3}, std::vector{DExpr::Var(17)}); + + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(0); + RaggedDotDimensionNumbers ragged_dot_dnums; + *ragged_dot_dnums.mutable_dot_dimension_numbers() = dot_dnums; + ragged_dot_dnums.add_lhs_ragged_dimensions(1); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferRaggedDotOpShape( + lhs_shape, rhs_shape, group_sizes_shape, ragged_dot_dnums, + /*preferred_element_type=*/std::nullopt)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(17))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(11))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(7))); +} + // mode 3 : [b,m,k], [b,k,n], [g] -> [b,m,n] TEST_F(ShapeInferenceTest, RaggedDotRaggedBatch) { const Shape lhs_shape = ShapeUtil::MakeShape(F32, {3, 11, 5}); @@ -2727,6 +3200,54 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastMatrixVector) { ASSERT_FALSE(inferred_shape_mismatch.ok()); } +TEST_F(ShapeInferenceTest, InDimBroadcastPreservesMappedExpressions) { + const Shape smaller = ShapeUtil::MakeShape( + F32, {2, 3}, + std::vector{DExpr::Var(18), DExpr::Const(3)}); + const Shape larger = ShapeUtil::MakeShape( + F32, {2, 4, 3}, + std::vector{DExpr::Const(2), DExpr::Const(4), DExpr::Const(3)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, larger, smaller, + {0, 2})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(18))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(4))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(3))); +} + +TEST_F(ShapeInferenceTest, DegenerateBroadcastUsesNonUnitExpressions) { + const Shape lhs = ShapeUtil::MakeShape( + F32, {1, 3}, + std::vector{DExpr::Const(1), DExpr::Const(3)}); + const Shape rhs = ShapeUtil::MakeShape( + F32, {5, 3}, + std::vector{DExpr::Var(19), DExpr::Const(3)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(19))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3))); +} + +TEST_F(ShapeInferenceTest, ElementwiseBinaryBroadcastPreservesExpressions) { + const Shape lhs = ShapeUtil::MakeShape( + F32, {2, 4, 3}, + std::vector{DExpr::Const(2), DExpr::Const(4), DExpr::Const(3)}); + const Shape rhs = ShapeUtil::MakeShape( + F32, {2, 3}, + std::vector{DExpr::Var(20), DExpr::Const(3)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {0, 2})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(20))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(4))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(3))); +} + TEST_F(ShapeInferenceTest, BinOpBroadcastCubeMatrix) { // Test variations of broadcasting a matrix for a binary add with a cube. const Shape cube = ShapeUtil::MakeShape(F32, {16, 8, 4}); @@ -2822,6 +3343,17 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) { HasSubstr("dimensions order is wrong")); } +TEST_F(ShapeInferenceTest, BinOpPreservesBroadcastedExpressionSameRank) { + Shape lhs = ShapeUtil::MakeShape(F32, {1}); + Shape rhs = + ShapeUtil::MakeShape(F32, {1}, std::vector{true}, {DExpr::Var(1)}); + + const absl::StatusOr inferred_shape = + ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {0}); + ASSERT_IS_OK(inferred_shape.status()); + EXPECT_EQ(inferred_shape->expressions(0), DExpr::Var(1)); +} + // Tests for the while instruction with proper shapes. TEST_F(ShapeInferenceTest, WhileWithCorrectShapes) { const Shape result_shape = ShapeUtil::MakeTupleShape({s32_, vector_32_}); @@ -2875,18 +3407,38 @@ TEST_F(ShapeInferenceTest, WhileWithBadShapes) { // Tests for the concatenate instruction with dynamic shapes. TEST_F(ShapeInferenceTest, ConcatenateWithDynamicShapes) { const auto dynamic_shape_1 = - ShapeUtil::MakeShape(F32, {32, 160, 10}, {true, false, false}); + ShapeUtil::MakeShape(F32, {32, 160, 10}, {true, false, false}, {}); const auto dynamic_shape_2 = - ShapeUtil::MakeShape(F32, {32, 160, 10}, {false, true, false}); + ShapeUtil::MakeShape(F32, {32, 160, 10}, {false, true, false}, {}); const absl::StatusOr inferred_shape = ShapeInference::InferConcatOpShape({&dynamic_shape_1, &dynamic_shape_2}, /*dimension=*/0); ASSERT_IS_OK(inferred_shape.status()); ASSERT_TRUE(ShapeUtil::Equal( - ShapeUtil::MakeShape(F32, {64, 160, 10}, {true, true, false}), + ShapeUtil::MakeShape(F32, {64, 160, 10}, {true, true, false}, {}), *inferred_shape)); } +TEST_F(ShapeInferenceTest, ConcatenateAddsConcatDimensionExpressions) { + const Shape lhs = ShapeUtil::MakeShape( + F32, {2, 5}, std::vector{DExpr::Var(1), DExpr::Const(5)}); + const Shape rhs = ShapeUtil::MakeShape( + F32, {3, 5}, std::vector{DExpr::Var(2), DExpr::Const(5)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferConcatOpShape({&lhs, &rhs}, /*dimension=*/0)); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, + ShapeUtil::MakeShape(F32, {5, 5}, + std::vector{DExpr::Var(1) + DExpr::Var(2), + DExpr::Const(5)}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + DExpr::Var(1) + DExpr::Var(2))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5))); +} + // Tests for the concatenate instruction with proper shapes. TEST_F(ShapeInferenceTest, ConcatenateWithCorrectShapes) { const absl::StatusOr inferred_shape_1 = @@ -2988,6 +3540,64 @@ TEST_F(ShapeInferenceTest, Pad) { HasSubstr("negative size for dimension 1")); } +TEST_F(ShapeInferenceTest, PadAddsConstantOffsetToExpressions) { + const Shape input_shape = ShapeUtil::MakeShape( + F32, {4, 5}, std::vector{DExpr::Var(1), DExpr::Var(2)}); + const Shape padding_value_shape = ShapeUtil::MakeShape(F32, {}); + PaddingConfig padding_config; + auto* dimension0 = padding_config.add_dimensions(); + dimension0->set_edge_padding_low(1); + dimension0->set_edge_padding_high(2); + dimension0->set_interior_padding(1); + auto* dimension1 = padding_config.add_dimensions(); + dimension1->set_edge_padding_low(0); + dimension1->set_edge_padding_high(4); + dimension1->set_interior_padding(0); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferPadShape(input_shape, padding_value_shape, + padding_config)); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, ShapeUtil::MakeShape( + F32, {10, 9}, + std::vector{DExpr::Var(1) + 6, + DExpr::Var(2) + 4}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1) + 6)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(2) + 4)); +} + +TEST_F(ShapeInferenceTest, PadBuildsExpressionsForTwoSymbolicDimensions) { + const Shape input_shape = ShapeUtil::MakeShape( + F32, {7, 9}, std::vector{DExpr::Var(34), DExpr::Var(35)}); + const Shape padding_value_shape = ShapeUtil::MakeShape(F32, {}); + PaddingConfig padding_config; + auto* dimension0 = padding_config.add_dimensions(); + dimension0->set_edge_padding_low(2); + dimension0->set_edge_padding_high(1); + dimension0->set_interior_padding(2); + auto* dimension1 = padding_config.add_dimensions(); + dimension1->set_edge_padding_low(3); + dimension1->set_edge_padding_high(4); + dimension1->set_interior_padding(1); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferPadShape(input_shape, padding_value_shape, + padding_config)); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, ShapeUtil::MakeShape( + F32, {22, 24}, + std::vector{DExpr::Var(34) + 15, + DExpr::Var(35) + 15}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + DExpr::Var(34) + 15)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), + DExpr::Var(35) + 15)); +} + TEST_F(ShapeInferenceTest, Reverse) { const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); @@ -2997,6 +3607,20 @@ TEST_F(ShapeInferenceTest, Reverse) { ASSERT_TRUE(ShapeUtil::Equal(input_shape, *inferred_shape)); } +TEST_F(ShapeInferenceTest, ReversePreservesExpressions) { + const Shape input_shape = ShapeUtil::MakeShape( + F32, {10, 25, 7}, + std::vector{DExpr::Var(26), DExpr::Const(25), DExpr::Var(27)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReverseShape(input_shape, {0, 2})); + + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(26))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(25))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(27))); +} + TEST_F(ShapeInferenceTest, ReverseInvalidDimension) { const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); @@ -3071,6 +3695,21 @@ TEST_F(ShapeInferenceTest, Transpose) { *inferred_shape_and_status)); } +TEST_F(ShapeInferenceTest, TransposePermutesExpressions) { + const Shape a_shape = ShapeUtil::MakeShape( + F32, {2, 3, 4, 5}, + std::vector{DExpr::Var(28), DExpr::Const(3), DExpr::Var(29), + DExpr::Const(5)}); + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape, + ShapeInference::InferTransposeShape( + a_shape, {1, 2, 3, 0})); + + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(29))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(5))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(3), DExpr::Var(28))); +} + TEST_F(ShapeInferenceTest, Rank1Transpose) { const Shape a_shape = ShapeUtil::MakeShape(F32, {5}); const absl::StatusOr inferred_shape_and_status = @@ -3260,8 +3899,8 @@ TEST_F(ShapeInferenceTest, ConditionalIndexed) { TEST_F(ShapeInferenceTest, ConditionalDynamic) { const Shape r0s32 = ShapeUtil::MakeShape(S32, {}); - const Shape static_shape = ShapeUtil::MakeShape(S32, {4}, {false}); - const Shape dynamic_shape = ShapeUtil::MakeShape(S32, {4}, {true}); + const Shape static_shape = ShapeUtil::MakeShape(S32, {4}, {false}, {}); + const Shape dynamic_shape = ShapeUtil::MakeShape(S32, {4}, {true}, {}); const absl::StatusOr inferred_shape0 = ShapeInference::InferConditionalShape( r0s32, @@ -3342,6 +3981,25 @@ TEST_F(ShapeInferenceTest, GoodTopK) { ShapeUtil::MakeShape(S32, {3, 4, 2})}))); } +TEST_F(ShapeInferenceTest, TopKPreservesLeadingExpressions) { + const Shape input = ShapeUtil::MakeShape( + F32, {3, 4, 5}, + std::vector{DExpr::Var(7), DExpr::Const(4), DExpr::Var(8)}); + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape, + ShapeInference::InferTopKShape(input, /*k=*/2)); + + ASSERT_TRUE(inferred_shape.IsTuple()); + ASSERT_EQ(inferred_shape.tuple_shapes_size(), 2); + const Shape& values = inferred_shape.tuple_shapes(0); + const Shape& indices = inferred_shape.tuple_shapes(1); + EXPECT_TRUE(DynExpr::equal(values.expressions(0), DExpr::Var(7))); + EXPECT_TRUE(DynExpr::equal(values.expressions(1), DExpr::Const(4))); + EXPECT_TRUE(DynExpr::equal(values.expressions(2), DExpr::Const(2))); + EXPECT_TRUE(DynExpr::equal(indices.expressions(0), DExpr::Var(7))); + EXPECT_TRUE(DynExpr::equal(indices.expressions(1), DExpr::Const(4))); + EXPECT_TRUE(DynExpr::equal(indices.expressions(2), DExpr::Const(2))); +} + TEST_F(ShapeInferenceTest, FailTopKLargeK) { const Shape input = ShapeUtil::MakeShape(F32, {3, 4, 5}); const absl::StatusOr statusor = @@ -3492,7 +4150,7 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherEntireDimension) { TF_ASSERT_OK_AND_ASSIGN( const Shape gather_shape, ShapeInference::InferGatherShape( - ShapeUtil::MakeShape(F32, {3, 2, 1}, {false, true, false}), + ShapeUtil::MakeShape(F32, {3, 2, 1}, {false, true, false}, {}), ShapeUtil::MakeShape(S64, {}), HloGatherInstruction::MakeGatherDimNumbers( /*offset_dims=*/{0, 1}, @@ -3501,7 +4159,8 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherEntireDimension) { /*index_vector_dim=*/0), /*slice_sizes=*/{1, 2, 1})); EXPECT_TRUE(ShapeUtil::Equal( - gather_shape, ShapeUtil::MakeShape(F32, {2, 1}, {true, false}))) + gather_shape, + ShapeUtil::MakeShape(F32, {2, 1}, {true, false}, {}))) << ShapeUtil::HumanString(gather_shape); } @@ -3509,7 +4168,7 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherCollapsedDimension) { TF_ASSERT_OK_AND_ASSIGN( const Shape gather_shape, ShapeInference::InferGatherShape( - ShapeUtil::MakeShape(F32, {3, 2, 1}, {true, false, false}), + ShapeUtil::MakeShape(F32, {3, 2, 1}, {true, false, false}, {}), ShapeUtil::MakeShape(S64, {}), HloGatherInstruction::MakeGatherDimNumbers( /*offset_dims=*/{0, 1}, @@ -3518,7 +4177,8 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherCollapsedDimension) { /*index_vector_dim=*/0), /*slice_sizes=*/{1, 2, 1})); EXPECT_TRUE(ShapeUtil::Equal( - gather_shape, ShapeUtil::MakeShape(F32, {2, 1}, {false, false}))) + gather_shape, + ShapeUtil::MakeShape(F32, {2, 1}, {false, false}, {}))) << ShapeUtil::HumanString(gather_shape); } @@ -3527,7 +4187,7 @@ TEST_F(GatherShapeInferenceTest, DynamicIndices) { const Shape gather_shape, ShapeInference::InferGatherShape( ShapeUtil::MakeShape(F32, {3, 2, 2}), - ShapeUtil::MakeShape(S64, {3, 4, 2}, {false, true, false}), + ShapeUtil::MakeShape(S64, {3, 4, 2}, {false, true, false}, {}), HloGatherInstruction::MakeGatherDimNumbers( /*offset_dims=*/{2, 3}, /*collapsed_slice_dims=*/{0}, @@ -3536,10 +4196,35 @@ TEST_F(GatherShapeInferenceTest, DynamicIndices) { /*slice_sizes=*/{1, 2, 2})); EXPECT_TRUE(ShapeUtil::Equal( gather_shape, - ShapeUtil::MakeShape(F32, {3, 4, 2, 2}, {false, true, false, false}))) + ShapeUtil::MakeShape(F32, {3, 4, 2, 2}, {false, true, false, false}, + {}))) << ShapeUtil::HumanString(gather_shape); } +TEST_F(GatherShapeInferenceTest, GatherPreservesIndexAndSliceExpressions) { + const Shape input = ShapeUtil::MakeShape( + F32, {3, 2, 2}, + std::vector{DExpr::Const(3), DExpr::Var(22), DExpr::Const(2)}); + const Shape indices = ShapeUtil::MakeShape( + S64, {3, 4, 2}, std::vector{false, true, false}, + std::vector{DExpr::Const(3), DExpr::Var(23), DExpr::Const(2)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape gather_shape, + ShapeInference::InferGatherShape( + input, indices, + HloGatherInstruction::MakeGatherDimNumbers( + /*offset_dims=*/{2, 3}, + /*collapsed_slice_dims=*/{0}, + /*start_index_map=*/{0, 1}, + /*index_vector_dim=*/2), + /*slice_sizes=*/{1, 2, 2})); + EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(0), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(1), DExpr::Var(23))); + EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(2), DExpr::Var(22))); + EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(3), DExpr::Const(2))); +} + TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_A) { TF_ASSERT_OK_AND_ASSIGN( const Shape gather_shape, @@ -4708,6 +5393,20 @@ TEST_F(ShapeInferenceTest, UnboundedAllToAll) { << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, AllToAllPreservesExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {12, 5}, + std::vector{DExpr::Var(14), DExpr::Const(5)}); + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferAllToAllShape(/*shape=*/operand, + /*split_dimension=*/0, + /*concat_dimension=*/0, + /*split_count=*/3)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(14))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5))); +} + TEST_F(ShapeInferenceTest, UnboundedAllToAllTupleUnsupported) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]")); TF_ASSERT_OK_AND_ASSIGN(const Shape expected, @@ -4787,6 +5486,30 @@ TEST_F(ShapeInferenceTest, UnboundedBatchNormGrad) { << " expected: " << ShapeUtil::HumanString(expected_tuple_shape); } +TEST_F(ShapeInferenceTest, BatchNormGradPreservesFeatureExpression) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 7, 11}, std::vector{false, true, false}, + std::vector{DExpr::Const(5), DExpr::Var(12), DExpr::Const(11)}); + const Shape scale = + ShapeUtil::MakeShape(F32, {7}, std::vector{DExpr::Var(12)}); + const Shape mean = + ShapeUtil::MakeShape(F32, {7}, std::vector{DExpr::Var(12)}); + const Shape variance = + ShapeUtil::MakeShape(F32, {7}, std::vector{DExpr::Var(12)}); + const Shape output_grad = operand; + + TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape, + ShapeInference::InferBatchNormGradShape( + operand, scale, mean, variance, output_grad, 1)); + ASSERT_TRUE(inferred_shape.IsTuple()); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1), DExpr::Var(12))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0), DExpr::Var(12))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(2).expressions(0), DExpr::Var(12))); +} + TEST_F(ShapeInferenceTest, UnboundedBatchNormInference) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, ?, 7]")); TF_ASSERT_OK_AND_ASSIGN(const Shape scale, ParseShape("f32[5]")); @@ -4819,6 +5542,27 @@ TEST_F(ShapeInferenceTest, UnboundedBatchNormTraining) { << " expected: " << ShapeUtil::HumanString(expected_tuple_shape); } +TEST_F(ShapeInferenceTest, BatchNormTrainingPreservesFeatureExpression) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 7, 11}, std::vector{false, true, false}, + std::vector{DExpr::Const(5), DExpr::Var(13), DExpr::Const(11)}); + const Shape scale = + ShapeUtil::MakeShape(F32, {7}, std::vector{DExpr::Var(13)}); + const Shape offset = + ShapeUtil::MakeShape(F32, {7}, std::vector{DExpr::Var(13)}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferBatchNormTrainingShape(operand, scale, offset, 1)); + ASSERT_TRUE(inferred_shape.IsTuple()); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1), DExpr::Var(13))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0), DExpr::Var(13))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(2).expressions(0), DExpr::Var(13))); +} + TEST_F(ShapeInferenceTest, UnboundedBroadcastUnsupportedOperand) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[<=2, ?]")); TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[1, <=2, ?]")); @@ -5208,6 +5952,29 @@ TEST_F(ShapeInferenceTest, UnboundedDotGeneral) { << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, DotGeneralPreservesBatchExpression) { + const Shape lhs = ShapeUtil::MakeShape( + F32, {2, 3, 5}, std::vector{true, false, false}, + {DExpr::Var(1), DExpr::Const(3), DExpr::Const(5)}); + const Shape rhs = ShapeUtil::MakeShape( + F32, {2, 5, 7}, std::vector{true, false, false}, + {DExpr::Var(1), DExpr::Const(5), DExpr::Const(7)}); + + DotDimensionNumbers dnums; + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + dnums.add_lhs_contracting_dimensions(2); + dnums.add_rhs_contracting_dimensions(1); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferDotOpShape(lhs, rhs, dnums, + /*preferred_element_type=*/std::nullopt)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(7))); +} + TEST_F(ShapeInferenceTest, UnboundedDynamicSlice) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]")); TF_ASSERT_OK_AND_ASSIGN(const Shape start_index, ParseShape("s32[]")); @@ -5216,12 +5983,30 @@ TEST_F(ShapeInferenceTest, UnboundedDynamicSlice) { const Shape inferred_shape, ShapeInference::InferDynamicSliceShape( operand, /*start_index_shapes=*/{start_index, start_index}, - /*slice_sizes=*/{2, 2}, /*allow_scalar_indices=*/true)); + /*slice_sizes=*/{2, 2}, /*slice_exprs=*/{}, + /*allow_scalar_indices=*/true)); EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected)) << "inferred: " << ShapeUtil::HumanString(inferred_shape) << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, DynamicSliceUsesProvidedExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {9, 10}, + std::vector{DExpr::Var(15), DExpr::Const(10)}); + const Shape start_index = ShapeUtil::MakeShape(S32, {}); + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferDynamicSliceShape( + operand, /*start_index_shapes=*/{start_index, start_index}, + /*slice_sizes=*/{4, 10}, + /*slice_exprs=*/{DExpr::Var(15) / 2, DExpr::Const(10)}, + /*allow_scalar_indices=*/true)); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(15) / 2)); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10))); +} + TEST_F(ShapeInferenceTest, UnboundedDynamicUpdateSlice) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]")); TF_ASSERT_OK_AND_ASSIGN(const Shape update, ParseShape("f32[?, 5]")); @@ -5237,6 +6022,42 @@ TEST_F(ShapeInferenceTest, UnboundedDynamicUpdateSlice) { << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, DynamicUpdateSlicePreservesOperandExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {12, 10}, + std::vector{DExpr::Var(30), DExpr::Const(10)}); + const Shape update = ShapeUtil::MakeShape( + F32, {4, 10}, + std::vector{DExpr::Const(4), DExpr::Const(10)}); + const Shape start_index = ShapeUtil::MakeShape(S32, {}); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferDynamicUpdateSliceShape( + operand, update, /*start_index_shapes=*/{start_index, start_index}, + /*allow_scalar_indices=*/true)); + + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(30))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10))); +} + +TEST_F(ShapeInferenceTest, DynamicReshapePreservesExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {5, 4, 8}, + std::vector{false, false, false}, + std::vector{DExpr::Var(21), DExpr::Const(4), DExpr::Const(8)}); + const Shape dim_size = ShapeUtil::MakeShape(S32, {}); + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferDynamicReshapeShape( + operand, /*dim_size_shapes=*/{&dim_size}, + /*new_size_bounds=*/{160}, + /*dims_are_dynamic=*/{false}, + /*expressions=*/{DExpr::Var(21) * DExpr::Const(32)})); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), + DExpr::Var(21) * DExpr::Const(32))); +} + TEST_F(ShapeInferenceTest, UnboundedFftWithFFT) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]")); const std::vector fft_length = {5, 10}; @@ -5472,6 +6293,56 @@ TEST_F(ShapeInferenceTest, UnboundedReduce) { << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, ReducePreservesRemainingExpressions) { + const Shape input = ShapeUtil::MakeShape( + F32, {5, 7, 11}, + std::vector{DExpr::Var(9), DExpr::Const(7), DExpr::Var(10)}); + ProgramShape to_apply = + ShapeUtil::MakeProgramShape({f32_, f32_}, f32_); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReduceShape({&input, &f32_}, {1}, to_apply)); + + EXPECT_TRUE(ShapeUtil::Equal( + inferred_shape, ShapeUtil::MakeShape( + F32, {5, 11}, + std::vector{DExpr::Var(9), DExpr::Var(10)}))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(9))); + EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(10))); +} + +TEST_F(ShapeInferenceTest, ReduceTupleOutputsPreserveRemainingExpressions) { + const Shape input0 = ShapeUtil::MakeShape( + F32, {5, 7, 11}, + std::vector{DExpr::Var(31), DExpr::Const(7), DExpr::Var(32)}); + const Shape input1 = ShapeUtil::MakeShape( + S32, {5, 7, 11}, + std::vector{DExpr::Var(31), DExpr::Const(7), DExpr::Var(32)}); + ProgramShape to_apply = ShapeUtil::MakeProgramShape( + {f32_, s32_, f32_, s32_}, ShapeUtil::MakeTupleShape({f32_, s32_})); + + TF_ASSERT_OK_AND_ASSIGN( + const Shape inferred_shape, + ShapeInference::InferReduceShape( + {&input0, &input1, &f32_, &s32_}, {1}, to_apply)); + + ASSERT_TRUE(inferred_shape.IsTuple()); + ASSERT_EQ(inferred_shape.tuple_shapes_size(), 2); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(0), + DExpr::Var(31))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1), + DExpr::Var(32))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0), + DExpr::Var(31))); + EXPECT_TRUE( + DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(1), + DExpr::Var(32))); +} + TEST_F(ShapeInferenceTest, UnboundedReduceInvalidReduceDimension) { TF_ASSERT_OK_AND_ASSIGN(const Shape input0, ParseShape("f32[7, 5]")); TF_ASSERT_OK_AND_ASSIGN(const Shape input1, ParseShape("f32[?, 5]")); @@ -5562,7 +6433,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshape) { TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[2,3]")); TF_ASSERT_OK_AND_ASSIGN( const Shape inferred, - ShapeInference::InferReshapeShape(operand, /*dimensions=*/{2, 3}, -1)); + ShapeInference::InferReshapeShape(operand, /*dimensions=*/{2, 3}, -1, + /*expressions=*/{})); ASSERT_TRUE(ShapeUtil::Equal(inferred, expected)) << "inferred: " << ShapeUtil::HumanString(inferred) << " expected: " << ShapeUtil::HumanString(expected); @@ -5573,7 +6445,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshapeUnsupportedOutputShape) { const absl::StatusOr inferred_shape = ShapeInference::InferReshapeShape( operand, - /*dimensions=*/{Shape::kUnboundedSize, Shape::kUnboundedSize}, -1); + /*dimensions=*/{Shape::kUnboundedSize, Shape::kUnboundedSize}, -1, + /*expressions=*/{}); EXPECT_THAT( inferred_shape.status().message(), HasSubstr("Reshaping with unbounded result shape is not supported.")); @@ -5583,7 +6456,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshapeUnsupportedMixOfDynamism) { TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, <=3]")); TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[<=3]")); const absl::StatusOr inferred_shape = - ShapeInference::InferReshapeShape(operand, /*dimensions=*/{3}, -1); + ShapeInference::InferReshapeShape(operand, /*dimensions=*/{3}, -1, + /*expressions=*/{}); ASSERT_THAT(inferred_shape.status().message(), HasSubstr("Reshape operand with bounded and unbounded dynamism " "not supported.")); @@ -5693,6 +6567,47 @@ TEST_F(ShapeInferenceTest, UnboundedSelectAndScatter) { << " expected: " << ShapeUtil::HumanString(expected); } +TEST_F(ShapeInferenceTest, SelectAndScatterPreservesOperandExpressions) { + const Shape operand = ShapeUtil::MakeShape( + F32, {11, 10}, std::vector{DExpr::Var(36), DExpr::Const(10)}); + const Shape source = ShapeUtil::MakeShape( + F32, {5, 10}, std::vector{(((DExpr::Var(36) - 1) / 2) + 1).simplify(), + DExpr::Const(10)}); + const Shape init_value = ShapeUtil::MakeShape(F32, {}); + + Window window; + WindowDimension dim0; + dim0.set_base_dilation(1); + dim0.set_size(3); + dim0.set_stride(2); + dim0.set_padding_low(0); + dim0.set_padding_high(1); + dim0.set_window_dilation(1); + + WindowDimension dim1; + dim1.set_base_dilation(1); + dim1.set_size(1); + dim1.set_stride(1); + dim1.set_padding_low(0); + dim1.set_padding_high(0); + dim1.set_window_dilation(1); + + *window.add_dimensions() = dim0; + *window.add_dimensions() = dim1; + + TF_ASSERT_OK_AND_ASSIGN( + const Shape result, + ShapeInference::InferSelectAndScatterShape( + operand, + /*select_shape=*/ShapeUtil::MakeProgramShape({f32_, f32_}, pred_), + window, source, init_value, + /*scatter_shape=*/ + ShapeUtil::MakeProgramShape({f32_, f32_}, f32_))); + + EXPECT_TRUE(DynExpr::equal(result.expressions(0), DExpr::Var(36))); + EXPECT_TRUE(DynExpr::equal(result.expressions(1), DExpr::Const(10))); +} + TEST_P(UnboundedBinaryOpShapeInferenceTest, UnboundedShiftLeft) { TF_ASSERT_OK_AND_ASSIGN(const Shape lhs, ParseShape(GetParam().lhs)); TF_ASSERT_OK_AND_ASSIGN(const Shape rhs, ParseShape(GetParam().rhs)); diff --git a/third_party/xla/xla/service/triangular_solve_expander.cc b/third_party/xla/xla/service/triangular_solve_expander.cc index 67dc0235810de8..3e51ffd4de8439 100644 --- a/third_party/xla/xla/service/triangular_solve_expander.cc +++ b/third_party/xla/xla/service/triangular_solve_expander.cc @@ -120,7 +120,12 @@ XlaOp DiagonalBlocks(XlaOp a, int64_t block_size) { auto last_blocks_dims = std::vector(ndims); std::copy(shape_dims.begin(), shape_dims.end(), last_blocks_dims.begin()); last_blocks_dims.insert(last_blocks_dims.end() - 2, 1); - last_blocks = Reshape(last_blocks, last_blocks_dims); + auto shape_exprs = blocks_shape.expressions(); + auto last_blocks_exprs = std::vector(ndims); + std::copy(shape_exprs.begin(), shape_exprs.end(), + last_blocks_exprs.begin()); + last_blocks_exprs.insert(last_blocks_exprs.end() - 2, DExpr::Const(1)); + last_blocks = Reshape(last_blocks, last_blocks_dims, last_blocks_exprs); // Concatenate with the other blocks if necessary if (n > block_size) { @@ -366,7 +371,7 @@ XlaOp TriangularSolveExpander::InvertDiagonalBlocks( /*broadcast_dimensions=*/{0, 1}); // Reshape back to original batch major dimensions - return Reshape(inv_diag_blocks, shape.dimensions()); + return Reshape(inv_diag_blocks, shape.dimensions(), shape.expressions()); }); } diff --git a/third_party/xla/xla/shape_test.cc b/third_party/xla/xla/shape_test.cc index fb51518b116eb2..69fdb816bd2f05 100644 --- a/third_party/xla/xla/shape_test.cc +++ b/third_party/xla/xla/shape_test.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/shape.h" +#include +#include #include #include @@ -22,6 +24,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "xla/hlo/testlib/test.h" #include "xla/layout.h" +#include "xla/printer.h" #include "xla/shape_util.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/test_benchmark.h" @@ -30,6 +33,12 @@ limitations under the License. namespace xla { namespace { +std::string DExprToString(const DExpr& expr) { + StringPrinter printer; + expr->print(&printer); + return std::move(printer).ToString(); +} + class ShapeTest : public ::testing::Test { protected: const Shape opaque_ = ShapeUtil::MakeOpaqueShape(); @@ -46,9 +55,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 @@ -98,8 +108,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); @@ -108,6 +118,165 @@ TEST_F(ShapeTest, DynamicShapeToString) { EXPECT_EQ("f32[?,784]", unbounded_.ToString()); } +TEST_F(ShapeTest, DExprSimplifyScalesMixedDenominators) { + DExpr expr = (DExpr::Var(1) / 2) + (DExpr::Var(1) / 3); + EXPECT_EQ("((5 * A) / 6)", DExprToString(expr.simplify())); +} + +TEST_F(ShapeTest, DExprSimplifyCombinesEqualFractions) { + DExpr expr = (DExpr::Var(1) / 2) + (DExpr::Var(1) / 2); + EXPECT_EQ("A", DExprToString(expr.simplify())); +} + +TEST_F(ShapeTest, DExprSolveRejectsZeroDivisor) { + DExpr expr = DExpr::Var(1) / DExpr::Const(0); + EXPECT_FALSE(expr->solve(7).has_value()); +} + +TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) { + DExpr expr = DExpr::Max(DExpr::Var(1), DExpr::Const(4)); + EXPECT_EQ("max(A, 4)", DExprToString(expr.simplify())); + EXPECT_FALSE(expr->solve(7).has_value()); + + DExpr clamped = DExpr::Max(DExpr::Var(1), DExpr::Const(0)); + EXPECT_EQ("A", DExprToString(clamped.simplify())); + + DExpr positive_affine = + DExpr::Max(2 * DExpr::Var(1) - 1, DExpr::Const(0)); + EXPECT_EQ("((2 * A) - 1)", + DExprToString(positive_affine.simplify())); + + DExpr unbounded_below = + DExpr::Max(DExpr::Var(1) - DExpr::Var(2), DExpr::Const(0)); + EXPECT_EQ("max((A + ((-1) * B)), 0)", + DExprToString(unbounded_below.simplify())); + + DExpr evaluated = expr.substitute(1, DExpr::Const(7)).simplify(); + EXPECT_EQ(DExpr::Kind::kConstant, evaluated.kind()); + EXPECT_EQ(7, evaluated->get_val()); + + DExpr divided = + DExpr::Max((DExpr::Var(1) + 1) / 2, DExpr::Const(0)); + DExpr divided_evaluated = + divided.substitute(1, DExpr::Const(100)).simplify(); + EXPECT_EQ(DExpr::Kind::kConstant, divided_evaluated.kind()); + EXPECT_EQ(50, divided_evaluated->get_val()); + + ExpressionProto proto; + expr.to_proto(&proto); + EXPECT_TRUE(expr == DExprFromProto(proto)); +} + +TEST_F(ShapeTest, DExprSelectUsesDynamicPredicate) { + DExpr delta = DExpr::Var(1) - 3; + DExpr expr = DExpr::Select(DExpr::Gt(delta, DExpr::Const(0)), + DExpr::Const(7), DExpr::Const(11)); + EXPECT_EQ("select(((A - 3) > 0), 7, 11)", + DExprToString(expr.simplify())); + EXPECT_EQ(7, expr.substitute(1, DExpr::Const(5))->s()->get_val()); + EXPECT_EQ(11, expr.substitute(1, DExpr::Const(1))->s()->get_val()); + + ExpressionProto proto; + expr.to_proto(&proto); + EXPECT_TRUE(expr == DExprFromProto(proto)); +} + +TEST_F(ShapeTest, DExprUnknownPropagatesThroughGtAndSelect) { + DExpr gt = DExpr::Gt(DExpr::Unknown(), DExpr::Const(0)).simplify(); + EXPECT_TRUE(gt.is_unknown()); + + DExpr select = + DExpr::Select(DExpr::Unknown(), DExpr::Const(7), DExpr::Const(11)) + .simplify(); + EXPECT_TRUE(select.is_unknown()); +} + +TEST_F(ShapeTest, DExprConstantDivisionIsNotDynamic) { + const DExpr expr = DExpr::Const(7) / DExpr::Const(3); + + EXPECT_FALSE(expr->is_dynamic()); + EXPECT_EQ(expr->get_val(), 2); +} + +TEST_F(ShapeTest, DExprSubstitutionReplacesEveryVariableOccurrence) { + const DExpr a = DExpr::Var(1); + const DExpr b = DExpr::Var(2); + const DExpr expr = (a + b) * (a - b); + + const DExpr substituted = expr.substitute(1, DExpr::Const(7)); + + EXPECT_EQ(substituted, (DExpr::Const(7) + b) * (DExpr::Const(7) - b)); + const std::set remaining_ids = substituted->get_all_ids(); + EXPECT_EQ(remaining_ids.size(), 1); + EXPECT_EQ(remaining_ids.count(2), 1); + EXPECT_EQ(expr->get_all_ids().size(), 2); +} + +TEST_F(ShapeTest, DExprSolveInvertsNestedSingleVariableArithmetic) { + const DExpr a = DExpr::Var(1); + const DExpr expr = ((3 * a) + 6) / 3; + + const std::optional solved = expr->solve(8); + + ASSERT_TRUE(solved.has_value()); + EXPECT_EQ(*solved, 6); + EXPECT_FALSE((a + DExpr::Var(2))->solve(8).has_value()); +} + +TEST_F(ShapeTest, DExprRoundTripPreservesRepeatedVariableStructure) { + const DExpr a = DExpr::Var(1); + const DExpr b = DExpr::Var(2); + const DExpr expr = (a + b) * (a - b); + ExpressionProto proto; + + expr.to_proto(&proto); + const DExpr decoded = DExprFromProto(proto); + ExpressionProto round_tripped_proto; + decoded.to_proto(&round_tripped_proto); + + EXPECT_EQ(proto.SerializeAsString(), + round_tripped_proto.SerializeAsString()); +} + +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, DExprFindsCommonCoreAcrossRepeatedBranches) { + const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2); + + EXPECT_EQ(shared_core, + (shared_core * shared_core) + .find_smallest_subexpression_covering_all_variables()); + EXPECT_EQ(shared_core, + (shared_core * (DExpr::Const(3) + shared_core)) + .find_smallest_subexpression_covering_all_variables()); +} + +TEST_F(ShapeTest, DExprUsesParentWhenChildCoresDiffer) { + const DExpr complete_expr = + (DExpr::Var(1) + DExpr::Var(2)) * + (DExpr::Var(1) - DExpr::Var(2)); + + EXPECT_EQ( + complete_expr, + complete_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}); diff --git a/third_party/xla/xla/shape_util_test.cc b/third_party/xla/xla/shape_util_test.cc index 743aa8ecf4d507..5a01598ad7ae1c 100644 --- a/third_party/xla/xla/shape_util_test.cc +++ b/third_party/xla/xla/shape_util_test.cc @@ -1049,7 +1049,8 @@ TEST(ShapeUtilTest, InvalidDynamicDimension) { TEST(ShapeUtilTest, PermuteDynamicDimensions) { Shape shape = ShapeUtil::MakeShape(F32, {10, 100, 1000}, - /*dynamic_dimensions*/ {false, true, true}); + /*dynamic_dimensions*/ {false, true, true}, + /*expressions=*/{}); SCOPED_TRACE(absl::StrCat("shape=", shape.ToString())); std::vector permutation(3); @@ -1129,6 +1130,15 @@ TEST(ShapeUtilTest, DeleteDimensions) { ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 2}, {1, 0})); } +TEST(ShapeUtilTest, MakeShapeWithDenseLayoutPreservesExpressions) { + std::vector expressions = {DExpr::Var(7), DExpr::Const(24)}; + Shape shape = ShapeUtil::MakeShapeWithDenseLayout( + F32, {10, 24}, expressions, {1, 0}); + + EXPECT_TRUE(shape.expressions(0) == DExpr::Var(7)); + EXPECT_TRUE(shape.expressions(1) == DExpr::Const(24)); +} + TEST(ShapeUtilTest, MakeShapeWithDescendingLayoutAndSamePhysicalLayout) { Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {128, 24, 4, 48, 48}, {2, 4, 3, 1, 0}); @@ -1153,6 +1163,24 @@ TEST(ShapeUtilTest, EXPECT_EQ(new_shape, expected_shape); } +TEST(ShapeUtilTest, + MakeShapeWithDescendingLayoutAndSamePhysicalLayoutPreservesExpressions) { + std::vector expressions = {DExpr::Var(1), DExpr::Const(24), + DExpr::Var(2), DExpr::Const(48), + DExpr::Const(48)}; + Shape shape = ShapeUtil::MakeShapeWithDenseLayout( + F32, {128, 24, 4, 48, 48}, expressions, {2, 4, 3, 1, 0}); + + Shape new_shape = + ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(shape); + + EXPECT_TRUE(new_shape.expressions(0) == DExpr::Var(1)); + EXPECT_TRUE(new_shape.expressions(1) == DExpr::Const(24)); + EXPECT_TRUE(new_shape.expressions(2) == DExpr::Const(48)); + EXPECT_TRUE(new_shape.expressions(3) == DExpr::Const(48)); + EXPECT_TRUE(new_shape.expressions(4) == DExpr::Var(2)); +} + TEST(ShapeUtilTest, DeduceTransposeDimensionsForBitcast) { Shape input_shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 3}, {1, 0}); Shape output_shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {3, 5}, {0, 1}); From 38515aecb710441ea06e71495fba2244d80c100e Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Wed, 19 Aug 2026 15:45:30 +0100 Subject: [PATCH 06/21] Add shared tf2xla dynamic expression plumbing Add the flags, shape conversion helpers, argument metadata, and XlaExpression APIs needed to carry symbolic dimensions and shape-derived contents through tf2xla. Keep concrete kernel propagation separate so this commit defines the common contract first. --- tensorflow/compiler/jit/flags.cc | 4 ++ tensorflow/compiler/jit/flags.h | 3 ++ tensorflow/compiler/tf2xla/BUILD | 7 +++ tensorflow/compiler/tf2xla/layout_util.cc | 5 ++ tensorflow/compiler/tf2xla/lib/broadcast.cc | 7 +-- tensorflow/compiler/tf2xla/lib/broadcast.h | 5 +- tensorflow/compiler/tf2xla/ops/BUILD | 1 + tensorflow/compiler/tf2xla/ops/xla_ops.cc | 25 +++++++-- tensorflow/compiler/tf2xla/shape_util.cc | 36 ++++++++++++- .../compiler/tf2xla/symbolic_content_util.h | 29 ++++++++++ tensorflow/compiler/tf2xla/xla_argument.cc | 9 ++++ tensorflow/compiler/tf2xla/xla_argument.h | 8 +++ tensorflow/compiler/tf2xla/xla_expression.cc | 53 ++++++++++++++++++- tensorflow/compiler/tf2xla/xla_expression.h | 17 ++++++ tensorflow/compiler/tf2xla/xla_op_kernel.cc | 9 +++- 15 files changed, 206 insertions(+), 12 deletions(-) create mode 100644 tensorflow/compiler/tf2xla/symbolic_content_util.h diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index d8e0b416b2d638..622b9a369fcb12 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -158,6 +158,9 @@ 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_enable_symbolic_content", + &mark_for_compilation_flags->tf_xla_enable_symbolic_content, + "Enable symbolic content propagation."), Flag("tf_xla_persistent_cache_directory", &mark_for_compilation_flags->tf_xla_persistent_cache_directory, "If non-empty, JIT-compiled executables are saved to and loaded " @@ -245,6 +248,7 @@ 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_enable_symbolic_content = false; mark_for_compilation_flags->tf_xla_persistent_cache_directory = ""; mark_for_compilation_flags->tf_xla_persistent_cache_device_types = ""; mark_for_compilation_flags->tf_xla_persistent_cache_read_only = false; diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index b4aba197559475..fb443e619e278f 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -96,6 +96,9 @@ struct MarkForCompilationPassFlags { // If true enables support of dynamic sizes. bool tf_xla_enable_dynamic_sizes; + // If true enables symbolic content propagation. + bool tf_xla_enable_symbolic_content; + // If non-empty, JIT-compiled executables are saved to and loaded from the // specified file system directory path. std::string tf_xla_persistent_cache_directory; diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index f2080a0752f41c..8432ed574b7b62 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -739,6 +739,7 @@ cc_library( "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:optional", + "@local_tsl//tsl/platform:protobuf", "@local_xla//xla/client", "@local_xla//xla/hlo/builder:value_inference", "@local_xla//xla/hlo/builder:xla_builder", @@ -851,6 +852,7 @@ cc_library( ":common", ":xla_argument", ":xla_helpers", + "//tensorflow/compiler/jit:flags", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", @@ -874,10 +876,12 @@ cc_library( hdrs = [ "literal_util.h", "shape_util.h", + "symbolic_content_util.h", "type_util.h", ], visibility = [":friends"], deps = [ + "//tensorflow/compiler/jit:flags", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", @@ -1092,6 +1096,8 @@ tf_cc_test( "//tensorflow/core:test_main", "//tensorflow/core:testlib", "//tensorflow/core/framework:tensor_testutil", + "//tensorflow/core/grappler:grappler_item", + "//tensorflow/core/grappler/costs:graph_properties", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest", @@ -1563,6 +1569,7 @@ cc_library( deps = [ ":xla_compiler", ":xla_expression", + "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:xla_compile_util", "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes", "//tensorflow/compiler/mlir/utils:array_container_utils", diff --git a/tensorflow/compiler/tf2xla/layout_util.cc b/tensorflow/compiler/tf2xla/layout_util.cc index b000c49f1f962e..1db86361065df9 100644 --- a/tensorflow/compiler/tf2xla/layout_util.cc +++ b/tensorflow/compiler/tf2xla/layout_util.cc @@ -21,6 +21,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_argument.h" @@ -131,8 +132,12 @@ absl::StatusOr ReshapeWithCorrectRepresentationAndSharding( hlo_sharding, fast_mem, shape_determination_fns, &to_shape)); } if (xla::ShapeUtil::Compatible(original_shape, to_shape)) { + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); for (int64_t i = 0; i < original_shape.dimensions().size(); ++i) { to_shape.set_dynamic_dimension(i, original_shape.is_dynamic_dimension(i)); + if (flags->tf_xla_enable_dynamic_sizes) { + to_shape.set_expression(i, original_shape.expressions(i)); + } } } return xla::Reshape(to_shape, original); diff --git a/tensorflow/compiler/tf2xla/lib/broadcast.cc b/tensorflow/compiler/tf2xla/lib/broadcast.cc index f815b91d04be33..07d09559b15d35 100644 --- a/tensorflow/compiler/tf2xla/lib/broadcast.cc +++ b/tensorflow/compiler/tf2xla/lib/broadcast.cc @@ -32,9 +32,10 @@ limitations under the License. namespace tensorflow { -absl::StatusOr BroadcastTo(xla::XlaOp input, - absl::Span output_dims) { - return xla::BroadcastTo(input, output_dims); +absl::StatusOr BroadcastTo( + xla::XlaOp input, absl::Span output_dims, + absl::Span output_exprs) { + return xla::BroadcastTo(input, output_dims, output_exprs); } absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs) { diff --git a/tensorflow/compiler/tf2xla/lib/broadcast.h b/tensorflow/compiler/tf2xla/lib/broadcast.h index 60630971e27466..b5903775679b30 100644 --- a/tensorflow/compiler/tf2xla/lib/broadcast.h +++ b/tensorflow/compiler/tf2xla/lib/broadcast.h @@ -29,8 +29,9 @@ namespace tensorflow { // Forwards to xla::BroadcastTo. // TODO(cheshire): Call the underlying function directly. -absl::StatusOr BroadcastTo(xla::XlaOp input, - absl::Span output_dims); +absl::StatusOr BroadcastTo( + xla::XlaOp input, absl::Span output_dims, + absl::Span output_exprs = {}); // Forwards to xla::BroadcastOpsToSame. absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs); diff --git a/tensorflow/compiler/tf2xla/ops/BUILD b/tensorflow/compiler/tf2xla/ops/BUILD index 61a329aa46f1ba..6014b73c35949c 100644 --- a/tensorflow/compiler/tf2xla/ops/BUILD +++ b/tensorflow/compiler/tf2xla/ops/BUILD @@ -16,6 +16,7 @@ cc_library( name = "xla_ops", srcs = ["xla_ops.cc"], deps = [ + "//tensorflow/compiler/jit:flags", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc index 6a67cfa237af70..516ae853d40290 100644 --- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc +++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/strings/match.h" #include "absl/strings/str_join.h" +#include "tensorflow/compiler/jit/flags.h" #include "xla/service/shape_inference.h" #include "xla/shape.h" #include "xla/xla_data.pb.h" @@ -1143,16 +1144,32 @@ xla::Shape GetShape(shape_inference::ShapeHandle shape_handle, } std::vector dims; std::vector dynamic_dims; + std::vector expressions; + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); for (int i = 0, rank = c->Rank(shape_handle); i < rank; ++i) { - bool is_dynamic = !c->ValueKnown(c->Dim(shape_handle, i)); + const auto dim = c->Dim(shape_handle, i); + const bool is_dynamic = !c->ValueKnown(dim); dynamic_dims.push_back(is_dynamic); - dims.push_back(is_dynamic ? xla::Shape::kUnboundedSize - : c->Value(c->Dim(shape_handle, i))); + if (flags->tf_xla_enable_dynamic_sizes) { + DimExpr* expr = c->GetDimExpr(dim); + expressions.push_back( + expr != nullptr + ? *expr + : c->ValueKnown(dim) + ? xla::DExpr::Const(c->Value(dim)) + : xla::DExpr::Unknown(xla::kMissingExpressionSentinel)); + } + dims.push_back(is_dynamic ? xla::Shape::kUnboundedSize : c->Value(dim)); } - return xla::Shape( + xla::Shape sh( // Type matters only for indices. S64 is the widest possible type. xla::PrimitiveType::S64, dims, absl::InlinedVector(dynamic_dims.begin(), dynamic_dims.end())); + + if (flags->tf_xla_enable_dynamic_sizes) { + sh.set_expressions(expressions); + } + return sh; } REGISTER_OP("XlaGather") diff --git a/tensorflow/compiler/tf2xla/shape_util.cc b/tensorflow/compiler/tf2xla/shape_util.cc index 0d7549d81c20f6..132daea9ee73e7 100644 --- a/tensorflow/compiler/tf2xla/shape_util.cc +++ b/tensorflow/compiler/tf2xla/shape_util.cc @@ -19,6 +19,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "xla/layout_util.h" #include "xla/shape_util.h" @@ -100,6 +101,11 @@ absl::Status XLAShapeToTensorShape(const xla::Shape& shape, for (int i = 0; i < shape.dimensions().size(); ++i) { TF_RETURN_IF_ERROR(tensor_shape->AddDimWithStatus(shape.dimensions(i))); } + if (!shape.expressions().empty()) { + std::vector dexprs(shape.expressions().begin(), + shape.expressions().end()); + tensor_shape->set_expressions(std::move(dexprs)); + } return absl::OkStatus(); } @@ -168,6 +174,11 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, int rank = tensor_shape.dims(); std::vector dimensions(rank); std::vector layout(rank); + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + std::vector expressions; + if (flags->tf_xla_enable_dynamic_sizes) { + expressions.resize(rank); + } for (int d = 0; d < rank; ++d) { dimensions[d] = tensor_shape.dim_size(d); if (dimensions[d] < 0) { @@ -175,11 +186,17 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, "shape; returning unknown sentinel value"; return xla::ShapeUtil::MakeShapeWithDenseLayout(type, {0}, {0}); } + if (flags->tf_xla_enable_dynamic_sizes) { + expressions[d] = tensor_shape.get_filled_expression(d); + } } // XLA uses minor-to-major; Tensorflow uses major-to-minor. std::iota(layout.rbegin(), layout.rend(), 0); xla::Shape result = xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout); + if (flags->tf_xla_enable_dynamic_sizes) { + result.set_expressions(expressions); + } return result; } @@ -200,18 +217,35 @@ absl::StatusOr TensorShapeToXLAShape( return out; } +inline static int var_id = 1; + xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, const TensorShape& tensor_shape) { int rank = tensor_shape.dims(); std::vector dimensions(rank); std::vector layout(rank); + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + std::vector expressions; + if (flags->tf_xla_enable_dynamic_sizes) { + expressions.resize(rank); + } + for (int d = 0; d < rank; ++d) { dimensions[d] = tensor_shape.dim_size(d); + if (flags->tf_xla_enable_dynamic_sizes) { + expressions[d] = tensor_shape.get_filled_expression(d); + } } + // XLA uses minor-to-major; Tensorflow uses major-to-minor. std::iota(layout.rbegin(), layout.rend(), 0); - return xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout); + auto shape = + xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout); + if (flags->tf_xla_enable_dynamic_sizes) { + shape.set_expressions(expressions); + } + return shape; } absl::StatusOr> GetShapeLayoutVector(const xla::Shape& shape) { diff --git a/tensorflow/compiler/tf2xla/symbolic_content_util.h b/tensorflow/compiler/tf2xla/symbolic_content_util.h new file mode 100644 index 00000000000000..3dd471a4e9ab02 --- /dev/null +++ b/tensorflow/compiler/tf2xla/symbolic_content_util.h @@ -0,0 +1,29 @@ +/* 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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_ +#define TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_ + +#include "tensorflow/compiler/jit/flags.h" + +namespace tensorflow { + +inline bool SymbolicContentEnabled() { + return GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content; +} + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/xla_argument.cc b/tensorflow/compiler/tf2xla/xla_argument.cc index 8b91dd3870b7d5..71ea93bf4fc628 100644 --- a/tensorflow/compiler/tf2xla/xla_argument.cc +++ b/tensorflow/compiler/tf2xla/xla_argument.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_argument.h" +#include "absl/algorithm/container.h" #include "llvm/ADT/STLExtras.h" namespace tensorflow { @@ -46,6 +47,14 @@ bool XlaArgument::operator==(const XlaArgument& other) const { if (constant_value.shape() != other.constant_value.shape()) { return false; } + if (!absl::c_equal( + constant_value_expressions, other.constant_value_expressions, + [](const xla::ExpressionProto& lhs, + const xla::ExpressionProto& rhs) { + return lhs.SerializeAsString() == rhs.SerializeAsString(); + })) { + return false; + } if (is_same_data_across_replicas != other.is_same_data_across_replicas) { return false; } diff --git a/tensorflow/compiler/tf2xla/xla_argument.h b/tensorflow/compiler/tf2xla/xla_argument.h index 9e2eccd29b1885..f037cb8610619e 100644 --- a/tensorflow/compiler/tf2xla/xla_argument.h +++ b/tensorflow/compiler/tf2xla/xla_argument.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_ +#include + #include "absl/types/optional.h" #include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h" @@ -23,6 +25,7 @@ limitations under the License. #include "xla/hlo/builder/xla_builder.h" #include "xla/hlo/ir/hlo_sharding.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" namespace tensorflow { @@ -76,6 +79,10 @@ struct XlaArgument { // host-memory tensor. Tensor constant_value; + // Symbolic expressions for each element of a compile-time constant. + // This is only used for shape-like integer tensors crossing cluster + // boundaries. + std::vector constant_value_expressions; // The upper bounds of the value. std::optional value_bound; @@ -116,6 +123,7 @@ struct XlaArgument { // Returns the dimension sizes for either TensorShape or xla::Shape. std::vector DimensionSizes() const; + std::vector DimensionExpressions() const; absl::InlinedVector DimensionSizesAsInlinedVector() const; // Returns the human-readable string for either TensorShape or xla::Shape. diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc index 61bd10e413ccf3..0488a5170ca211 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.cc +++ b/tensorflow/compiler/tf2xla/xla_expression.cc @@ -17,9 +17,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "xla/hlo/builder/value_inference.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/logging.h" namespace tensorflow { @@ -73,6 +76,50 @@ XlaExpression XlaExpression::Resource(XlaResource* resource) { return e; } +void XlaExpression::set_contents(std::vector contents) { + switch (kind_) { + case Kind::kXlaOp: + case Kind::kTensorList: + if (handle_.valid() && !handle_.IsUninitialized()) { + if (SymbolicContentEnabled()) { + auto status = handle_.builder()->SetInstructionContents( + handle_, std::move(contents)); + if (!status.ok()) { + VLOG(1) << "Failed to set XlaOp contents: " << status; + } + return; + } + break; + } + break; + case Kind::kInvalid: + case Kind::kConstant: + case Kind::kResource: + break; + } + local_contents_ = std::move(contents); +} + +absl::Span XlaExpression::contents() const { + switch (kind_) { + case Kind::kXlaOp: + case Kind::kTensorList: + if (SymbolicContentEnabled() && handle_.valid() && + !handle_.IsUninitialized()) { + auto contents_or = handle_.builder()->GetInstructionContents(handle_); + if (contents_or.ok()) { + return **contents_or; + } + } + break; + case Kind::kInvalid: + case Kind::kConstant: + case Kind::kResource: + break; + } + return local_contents_; +} + string XlaExpression::HumanString() const { switch (kind_) { case Kind::kInvalid: @@ -95,7 +142,11 @@ xla::XlaOp XlaExpression::AsXlaOp(xla::XlaBuilder* builder) const { xla::BorrowingLiteral literal; TF_RETURN_IF_ERROR( HostTensorToBorrowingLiteral(*constant_value_, &literal)); - return xla::ConstantLiteral(builder, literal); + xla::XlaOp op = xla::ConstantLiteral(builder, literal); + if (!local_contents_.empty()) { + TF_RETURN_IF_ERROR(builder->SetInstructionContents(op, local_contents_)); + } + return op; } case Kind::kTensorList: TF_FALLTHROUGH_INTENDED; diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h index d410b79a3da137..980ab6dba4b73a 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.h +++ b/tensorflow/compiler/tf2xla/xla_expression.h @@ -16,11 +16,15 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ +#include + #include "absl/types/optional.h" +#include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/client/client.h" #include "xla/hlo/builder/value_inference.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/shape.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/statusor.h" @@ -112,6 +116,13 @@ class XlaExpression { // Return the dynamism of the expression, if available. std::optional value_dynamism() const { return value_dynamism_; } + // Set symbolic content metadata for expressions whose values should retain + // links to symbolic dimensions across shape-tensor flows. + void set_contents(std::vector contents); + + // Return symbolic content metadata. + absl::Span contents() const; + XlaResource* resource() const { return resource_; } // Returns a human-readable summary of the expression. @@ -164,6 +175,12 @@ class XlaExpression { // Indicate whether each value inside a tensor is dynamic or not. std::optional value_dynamism_; + // For constant expressions, marks a single element that later passes may + // reinterpret as coming from a dynamic expression instead of the literal. + // Symbolic expressions describing tensor contents when this expression is + // used as a shape-like value. + std::vector local_contents_; + // The resource, if kind_ == kResource. Not owned. XlaResource* resource_ = nullptr; }; diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index 4a570827029330..a1dce933db5309 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -464,6 +464,12 @@ absl::Status XlaOpKernelContext::ConstantInputAsShape( ", result: ", num_elements); } *shape = TensorShape(dims); + const auto& contents = InputExpression(index).contents(); + for (int i = 0; i < shape->dims() && i < contents.size(); ++i) { + if (contents[i] && contents[i]->is_dynamic()) { + shape->set_expression(i, contents[i]); + } + } return absl::OkStatus(); } @@ -736,7 +742,8 @@ absl::Status AssignVariableTensor(const Tensor& tensor, DataType type, xla::Shape xla_shape; TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); if (!xla::ShapeUtil::Compatible(xla_shape, representation_shape)) { - handle = xla::Reshape(handle, representation_shape.dimensions()); + handle = xla::Reshape(handle, representation_shape.dimensions(), + representation_shape.expressions()); } variable->SetRepresentationShape(representation_shape); return variable->SetValue(handle); From d7e6a3e3e5a4ec170e2e3b40cd9d8468eb57d3df Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Wed, 19 Aug 2026 15:45:42 +0100 Subject: [PATCH 07/21] Propagate dynamic expressions through tf2xla kernels Update individual tf2xla kernel lowerings to derive and preserve symbolic shape and value expressions. Cover shape-producing, reshaping, slicing, sequence, partitioning, and related kernels using the shared expression plumbing. --- tensorflow/compiler/tf2xla/kernels/BUILD | 3 + .../tf2xla/kernels/batchtospace_op.cc | 32 ++- tensorflow/compiler/tf2xla/kernels/beta_op.cc | 9 +- .../compiler/tf2xla/kernels/bincount_op.cc | 24 +- tensorflow/compiler/tf2xla/kernels/cast_op.cc | 9 + .../tf2xla/kernels/clip_by_value_op.cc | 4 +- .../compiler/tf2xla/kernels/concat_op.cc | 86 +++++++- .../compiler/tf2xla/kernels/const_op.cc | 90 +++++++- .../tf2xla/kernels/conv_op_helpers.cc | 13 +- .../compiler/tf2xla/kernels/conv_ops.cc | 9 +- .../tf2xla/kernels/depthtospace_op.cc | 41 +++- tensorflow/compiler/tf2xla/kernels/diag_op.cc | 43 +++- .../tf2xla/kernels/dynamic_partition_op.cc | 83 ++++++- .../tf2xla/kernels/dynamic_stitch_op.cc | 13 +- .../tf2xla/kernels/fake_quantize_ops.cc | 16 +- tensorflow/compiler/tf2xla/kernels/fill_op.cc | 11 +- .../compiler/tf2xla/kernels/gather_op.cc | 3 +- .../compiler/tf2xla/kernels/image_ops.cc | 2 +- .../compiler/tf2xla/kernels/in_topk_op.cc | 6 +- .../tf2xla/kernels/lower_upper_bound_ops.cc | 6 +- .../tf2xla/kernels/matrix_diag_ops.cc | 23 +- .../kernels/matrix_triangular_solve_op.cc | 12 +- tensorflow/compiler/tf2xla/kernels/pack_op.cc | 57 ++++- .../compiler/tf2xla/kernels/pooling_ops.cc | 8 +- .../kernels/quantize_and_dequantize_op.cc | 3 +- .../compiler/tf2xla/kernels/reduction_ops.cc | 5 +- .../tf2xla/kernels/reduction_ops_common.cc | 6 +- tensorflow/compiler/tf2xla/kernels/relu_op.cc | 32 ++- .../compiler/tf2xla/kernels/resampler_ops.cc | 16 +- .../compiler/tf2xla/kernels/reshape_op.cc | 71 +++++- .../tf2xla/kernels/reverse_sequence_op.cc | 21 +- tensorflow/compiler/tf2xla/kernels/roll_op.cc | 4 +- .../compiler/tf2xla/kernels/scatter_nd_op.cc | 3 +- .../tf2xla/kernels/segment_reduction_ops.cc | 3 +- .../compiler/tf2xla/kernels/select_op.cc | 7 +- .../compiler/tf2xla/kernels/sequence_ops.cc | 191 +++++++++++++++- .../compiler/tf2xla/kernels/shape_op.cc | 167 ++++++++++++-- .../compiler/tf2xla/kernels/slice_op.cc | 86 +++++++- .../compiler/tf2xla/kernels/softmax_op.cc | 65 +++++- .../tf2xla/kernels/spacetobatch_op.cc | 23 +- .../tf2xla/kernels/spacetodepth_op.cc | 41 +++- .../tf2xla/kernels/sparse_to_dense_op.cc | 3 +- .../compiler/tf2xla/kernels/split_op.cc | 50 ++++- .../compiler/tf2xla/kernels/stack_ops.cc | 3 +- .../tf2xla/kernels/stateless_random_ops.cc | 18 +- .../tf2xla/kernels/strided_slice_op.cc | 207 +++++++++++++++--- .../tf2xla/kernels/tensor_array_ops.cc | 25 ++- .../tf2xla/kernels/tensor_list_ops.cc | 16 +- .../tf2xla/kernels/tensor_list_utils.cc | 59 ++++- .../compiler/tf2xla/kernels/tile_ops.cc | 22 +- .../compiler/tf2xla/kernels/unique_op.cc | 45 +++- .../compiler/tf2xla/kernels/unpack_op.cc | 12 +- .../compiler/tf2xla/kernels/variable_ops.cc | 20 +- .../compiler/tf2xla/kernels/where_op.cc | 41 +++- tensorflow/compiler/tf2xla/lib/data_format.cc | 10 +- tensorflow/core/kernels/strided_slice_op.cc | 4 +- 56 files changed, 1649 insertions(+), 233 deletions(-) diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index cd11a497b83514..37d68dbd23c242 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -1504,6 +1504,7 @@ tf_kernel_library( name = "slice_op", srcs = ["slice_op.cc"], deps = [ + "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compilation_device", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_context", @@ -2200,6 +2201,7 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", + "@local_tsl//tsl/platform:protobuf", "@local_xla//xla/hlo/builder:xla_builder", ], ) @@ -2210,6 +2212,7 @@ tf_kernel_library( deps = [ ":shape_util", ":tensor_list_utils", + "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compilation_device", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_context", diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc index 7a42150f3a9c19..6deeaa7ba11037 100644 --- a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc @@ -36,6 +36,8 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input, const int input_rank = input_tensor_shape.dims(); const absl::InlinedVector input_shape = input_tensor_shape.dim_sizes(); + const std::vector input_exprs = + input_tensor_shape.get_filled_expressions(); const int block_rank = block_shape.size(); OP_REQUIRES( @@ -76,11 +78,21 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input, ") is not divisible by product of block sizes (", block_num_elems, ")")); std::vector reshaped_shape(input_rank + block_rank); + std::vector reshaped_exprs(input_rank + block_rank); std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin()); + std::fill(reshaped_exprs.begin(), reshaped_exprs.begin() + block_rank, + xla::DExpr::Const(0)); + for (int i = 0; i < block_rank; ++i) { + reshaped_exprs[i] = xla::DExpr::Const(block_shape[i]); + } reshaped_shape[block_rank] = batch_size / block_num_elems; + reshaped_exprs[block_rank] = + (input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify(); std::copy(input_shape.begin() + 1, input_shape.end(), reshaped_shape.begin() + block_rank + 1); - xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape); + std::copy(input_exprs.begin() + 1, input_exprs.end(), + reshaped_exprs.begin() + block_rank + 1); + xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs); // 2. Permute dimensions of `reshaped` to produce `permuted` of shape // [batch / prod(block_shape), @@ -111,15 +123,22 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input, // ..., // input_shape[N-1]] std::vector reshaped_permuted_shape(input_rank); + std::vector reshaped_permuted_exprs(input_rank); reshaped_permuted_shape[0] = batch_size / block_num_elems; + reshaped_permuted_exprs[0] = + (input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify(); for (int i = 0; i < block_rank; ++i) { reshaped_permuted_shape[1 + i] = block_shape[i] * input_shape[1 + i]; + reshaped_permuted_exprs[1 + i] = + (xla::DExpr::Const(block_shape[i]) * input_exprs[1 + i]).simplify(); } std::copy(remainder_shape.begin(), remainder_shape.end(), reshaped_permuted_shape.begin() + 1 + block_rank); + std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(), + reshaped_permuted_exprs.begin() + 1 + block_rank); xla::XlaOp reshaped_permuted = - xla::Reshape(permuted, reshaped_permuted_shape); + xla::Reshape(permuted, reshaped_permuted_shape, reshaped_permuted_exprs); // 4. Crop the start and end of dimensions `[1, ..., M]` of // `reshaped_permuted` according to `crops` to produce the output of shape: @@ -133,6 +152,9 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input, std::vector start_indices(input_rank, 0); std::vector end_indices = reshaped_permuted_shape; std::vector strides(input_rank, 1); + std::vector start_exprs(input_rank, xla::DExpr::Const(0)); + std::vector end_exprs(reshaped_permuted_exprs.begin(), + reshaped_permuted_exprs.end()); for (int i = 0; i < block_rank; ++i) { int64_t crop_start = crops.Get({i, 0}); int64_t crop_end = crops.Get({i, 1}); @@ -140,14 +162,16 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input, errors::InvalidArgument("Crops must be non-negative")); start_indices[1 + i] = crop_start; end_indices[1 + i] -= crop_end; + start_exprs[1 + i] = xla::DExpr::Const(crop_start); + end_exprs[1 + i] = (reshaped_permuted_exprs[1 + i] - crop_end).simplify(); OP_REQUIRES( ctx, start_indices[1 + i] <= end_indices[1 + i], errors::InvalidArgument( "Cropped size must be non-negative: start: ", crop_start, " end: ", crop_end, " size ", reshaped_permuted_shape[1 + i])); } - xla::XlaOp output = - xla::Slice(reshaped_permuted, start_indices, end_indices, strides); + xla::XlaOp output = xla::Slice(reshaped_permuted, start_indices, end_indices, + start_exprs, end_exprs, strides); ctx->SetOutput(0, output); } diff --git a/tensorflow/compiler/tf2xla/kernels/beta_op.cc b/tensorflow/compiler/tf2xla/kernels/beta_op.cc index 4ead9f76fcee11..40cdd190d173b4 100644 --- a/tensorflow/compiler/tf2xla/kernels/beta_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/beta_op.cc @@ -63,11 +63,14 @@ class BetaincOp : public XlaOpKernel { auto result = builder->ReportErrorOrReturn([&]() -> absl::StatusOr { TF_ASSIGN_OR_RETURN( - auto a, BroadcastTo(ctx->Input(0), merged_shape.dim_sizes())); + auto a, BroadcastTo(ctx->Input(0), merged_shape.dim_sizes(), + merged_shape.get_filled_expressions())); TF_ASSIGN_OR_RETURN( - auto b, BroadcastTo(ctx->Input(1), merged_shape.dim_sizes())); + auto b, BroadcastTo(ctx->Input(1), merged_shape.dim_sizes(), + merged_shape.get_filled_expressions())); TF_ASSIGN_OR_RETURN( - auto x, BroadcastTo(ctx->Input(2), merged_shape.dim_sizes())); + auto x, BroadcastTo(ctx->Input(2), merged_shape.dim_sizes(), + merged_shape.get_filled_expressions())); return xla::RegularizedIncompleteBeta(a, b, x); }); ctx->SetOutput(0, result); diff --git a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc index fadd2f87219464..5e277772f40e4e 100644 --- a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc @@ -110,16 +110,24 @@ class DenseBincountOp : public XlaOpKernel { scatter_dnums.add_scatter_dims_to_operand_dims(0); if (rank == 2) { - output_shape = xla::ShapeUtil::MakeShape(dtype, {size, output_size}); + output_shape = xla::ShapeUtil::MakeShape( + dtype, {size, output_size}, + std::vector{input_shape.expressions(0), + xla::DExpr::Const(output_size)}); scatter_dnums.add_inserted_window_dims(1); scatter_dnums.add_scatter_dims_to_operand_dims(1); - auto i_shape = - xla::ShapeUtil::MakeShape(input_xla_type, {input_shape.dimensions()}); + auto i_shape = xla::ShapeUtil::MakeShape(input_xla_type, + input_shape.dimensions(), + input_shape.expressions()); auto i = xla::Iota(ctx->builder(), i_shape, 0); + xla::DExpr flattened_expr = + input_shape.expressions(0) * input_shape.expressions(1); i = xla::Reshape( - i, {input_shape.dimensions(0) * input_shape.dimensions(1), 1}); + i, {input_shape.dimensions(0) * input_shape.dimensions(1), 1}, + {flattened_expr, xla::DExpr::Const(1)}); auto j = xla::Reshape( - input, {input_shape.dimensions(0) * input_shape.dimensions(1), 1}); + input, {input_shape.dimensions(0) * input_shape.dimensions(1), 1}, + {flattened_expr, xla::DExpr::Const(1)}); std::vector iotas_to_concat; iotas_to_concat.push_back(i); iotas_to_concat.push_back(j); @@ -127,10 +135,12 @@ class DenseBincountOp : public XlaOpKernel { updates = xla::Broadcast( one, {input_shape.dimensions(0) * input_shape.dimensions(1)}); output = xla::Broadcast( - zero, {output_shape.dimensions(0), output_shape.dimensions(1)}); + zero, {output_shape.dimensions(0), output_shape.dimensions(1)}, + {output_shape.expressions(0), output_shape.expressions(1)}); if (has_weights && !binary_output_) { weights = xla::Reshape( - weights, {input_shape.dimensions(0) * input_shape.dimensions(1)}); + weights, {input_shape.dimensions(0) * input_shape.dimensions(1)}, + {flattened_expr}); updates = weights; } } else { diff --git a/tensorflow/compiler/tf2xla/kernels/cast_op.cc b/tensorflow/compiler/tf2xla/kernels/cast_op.cc index 1779cfcc1ced40..74e9c5aa4003ff 100644 --- a/tensorflow/compiler/tf2xla/kernels/cast_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cast_op.cc @@ -92,6 +92,15 @@ class CastOp : public XlaOpKernel { output = xla::ConvertElementType(input, dst_type_); } + const auto& input_contents = ctx->InputExpression(0).contents(); + if (!input_contents.empty()) { + auto output_expr = + XlaExpression::XlaOp(output, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::vector(input_contents.begin(), + input_contents.end())); + ctx->SetOutputExpression(0, output_expr); + return; + } ctx->SetOutput(0, output); } diff --git a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc index 6b4f278c72beff..928fbb6e22d76f 100644 --- a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc @@ -46,11 +46,11 @@ class ClipByValueOp : public XlaOpKernel { if (shape != min_shape) { OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(min_shape), shape_error()); - min = xla::Broadcast(min, shape.dim_sizes()); + min = xla::Broadcast(min, shape.dim_sizes(), shape.get_filled_expressions()); } if (shape != max_shape) { OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(max_shape), shape_error()); - max = xla::Broadcast(max, shape.dim_sizes()); + max = xla::Broadcast(max, shape.dim_sizes(), shape.get_filled_expressions()); } ctx->SetOutput(0, xla::Clamp(min, input, max)); } diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc index bed3479941ca41..8a71a7dc4d1128 100644 --- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc @@ -18,9 +18,11 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/log/log.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/tf2xla/kernels/shape_util.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" @@ -41,6 +43,44 @@ limitations under the License. namespace tensorflow { namespace { +bool AppendConcatInputContents(const XlaExpression& expr, + const TensorShape& shape, + std::vector* contents) { + const auto& input_contents = expr.contents(); + if (!input_contents.empty()) { + if (shape.dims() == 0) { + if (input_contents.size() != 1) { + return false; + } + contents->push_back(input_contents[0]); + return true; + } + if (shape.dims() == 1) { + contents->insert(contents->end(), input_contents.begin(), + input_contents.end()); + return true; + } + return false; + } + if (shape.dims() == 0) { + contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel)); + return true; + } + if (shape.dims() != 1) { + return false; + } + for (int64_t i = 0; i < shape.dim_size(0); ++i) { + contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel)); + } + return true; +} + +bool HasDynamicContents(absl::Span contents) { + return absl::c_any_of(contents, [](const xla::DExpr& expr) { + return expr && expr->is_dynamic(); + }); +} + // -------------------------------------------------------------------------- class ConcatBaseOp : public XlaOpKernel { public: @@ -74,6 +114,29 @@ class ConcatBaseOp : public XlaOpKernel { // Make a vector holding the XlaOp for each of the inputs that has non-zero // elements. std::vector input_data; + std::vector> input_contents; + input_contents.resize(N); + bool has_output_contents = SymbolicContentEnabled() && axis == 0; + std::vector output_contents; + if (has_output_contents) { + for (int i = 0; i < N; ++i) { + if (!AppendConcatInputContents(ctx->InputExpression(ValueInputIndex(i)), + shapes[i], + &output_contents)) { + has_output_contents = false; + output_contents.clear(); + break; + } + if (shapes[i].dims() == 0) { + input_contents[i].push_back(output_contents.back()); + } else if (shapes[i].dims() == 1) { + input_contents[i].insert(input_contents[i].end(), + output_contents.end() - shapes[i].dim_size(0), + output_contents.end()); + } + } + has_output_contents = has_output_contents && HasDynamicContents(output_contents); + } int output_concat_dim = 0; for (int i = 0; i < N; ++i) { xla::XlaOp handle = values[i]; @@ -86,7 +149,12 @@ class ConcatBaseOp : public XlaOpKernel { "] = ", in_shape.DebugString())); if (in_shape.dims() == 0) { // Inputs that come in as scalars must be reshaped to 1-vectors. - input_data.push_back(xla::Reshape(handle, {1})); + xla::XlaOp reshaped = xla::Reshape(handle, {1}); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(reshaped, input_contents[i])); + } + input_data.push_back(reshaped); } else { input_data.push_back(handle); } @@ -94,10 +162,24 @@ class ConcatBaseOp : public XlaOpKernel { } VLOG(1) << "Concat dim " << concat_dim << " equivalent to " << axis; - ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), input_data, axis)); + auto output = xla::ConcatInDim(ctx->builder(), input_data, axis); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(output, output_contents)); + auto output_expr = + XlaExpression::XlaOp(output, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } + ctx->SetOutput(0, output); } private: + int ValueInputIndex(int value_index) const { + return value_index < axis_index_ ? value_index : value_index + 1; + } + int axis_index_; }; diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc index d2463a9974b1bb..33d8b2079897d0 100644 --- a/tensorflow/compiler/tf2xla/kernels/const_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc @@ -15,19 +15,29 @@ limitations under the License. #include #include +#include +#include "xla/printer.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/shape_expr.h" +#include "tsl/platform/protobuf.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/tensor.pb.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/platform/logging.h" namespace tensorflow { namespace { +constexpr char kUserInferredValueContentsAttrName[] = + "_user_inferred_value_contents"; + template DstT CastTo(SrcT src) { return static_cast(src); @@ -100,6 +110,40 @@ xla::XlaOp GetScalarConst(const TensorProto& proto, xla::XlaBuilder* b) { return xla::XlaOp(); } +std::vector BuildShapeContentsFromTensorShapeProto( + const TensorShapeProto& shape) { + std::vector contents; + contents.reserve(shape.dim_size()); + for (int i = 0; i < shape.dim_size(); ++i) { + xla::DExpr expr; + if (i < shape.expressions_size()) { + expr = DimExprFromProto(shape.expressions(i)); + } + if (expr) { + xla::StringPrinter printer; + expr->print(&printer); + VLOG(1) << "BuildShapeContentsFromTensorShape dim=" << i + << " expr=" << std::move(printer).ToString() + << " dynamic=" << (expr->is_dynamic() ? "true" : "false"); + } else { + VLOG(1) << "BuildShapeContentsFromTensorShape dim=" << i + << " expr=_ dynamic=false"; + } + contents.push_back(expr && expr->is_dynamic() + ? std::move(expr) + : xla::DExpr::Unknown( + xla::kUnknownContentSentinel)); + } + return contents; +} + +bool CanAttachContentsFromTensorShapeProto(const TensorShape& tensor_shape, + const TensorShapeProto& contents) { + return (tensor_shape.dims() == 0 && contents.dim_size() == 1) || + (tensor_shape.dims() == 1 && + tensor_shape.dim_size(0) == contents.dim_size()); +} + class ConstOp : public XlaOpKernel { public: explicit ConstOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -117,13 +161,49 @@ class ConstOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { xla::XlaBuilder* b = ctx->builder(); + bool has_dynamic = false; + TensorShapeProto inferred_shape_proto; + TensorShapeProto inferred_value_contents_proto; + string inferred_value_contents_serialized; + if (GetNodeAttr(ctx->op_kernel().def(), "has_dynamic", &has_dynamic).ok() && + has_dynamic) { + GetNodeAttr(ctx->op_kernel().def(), "user_inferred_shape", + &inferred_shape_proto) + .IgnoreError(); + } + if (GetNodeAttr(ctx->op_kernel().def(), kUserInferredValueContentsAttrName, + &inferred_value_contents_serialized) + .ok()) { + if (!inferred_value_contents_proto.ParseFromString( + inferred_value_contents_serialized)) { + inferred_value_contents_proto.Clear(); + } + } + const bool has_contents_proto = inferred_value_contents_proto.dim_size() > 0; + const TensorShapeProto& contents_proto = + has_contents_proto ? inferred_value_contents_proto : inferred_shape_proto; + // To avoid blowups for large constants filled with the same value, // recognize that case and emit a scalar broadcast instead. TensorShape shape(proto_.tensor_shape()); if (shape.num_elements() > 1) { xla::XlaOp value = GetScalarConst(proto_, b); if (value.valid()) { - ctx->SetOutput(0, xla::Broadcast(value, shape.dim_sizes())); + if (has_dynamic) { + VLOG(1) << "ConstOp broadcast fast path shape=" + << shape.DebugString() << " inferred_rank=" + << inferred_shape_proto.dim_size(); + } + xla::XlaOp broadcast = + xla::Broadcast(value, shape.dim_sizes(), shape.get_expressions()); + XlaExpression output = + XlaExpression::XlaOp(broadcast, ctx->expected_output_dtype(0)); + if ((has_contents_proto || has_dynamic) && + CanAttachContentsFromTensorShapeProto(shape, contents_proto)) { + output.set_contents( + BuildShapeContentsFromTensorShapeProto(contents_proto)); + } + ctx->SetOutputExpression(0, output); return; } } @@ -132,7 +212,13 @@ class ConstOp : public XlaOpKernel { OP_REQUIRES(ctx, tensor.FromProto(cpu_allocator(), proto_), errors::InvalidArgument("Cannot parse tensor from proto: ", proto_.DebugString())); - ctx->SetConstantOutput(0, tensor); + XlaExpression output = XlaExpression::Constant(tensor); + if ((has_contents_proto || has_dynamic) && + CanAttachContentsFromTensorShapeProto(tensor.shape(), contents_proto)) { + output.set_contents( + BuildShapeContentsFromTensorShapeProto(contents_proto)); + } + ctx->SetOutputExpression(0, output); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc index 3fe22dcb4441e7..7889852a24c580 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc @@ -97,8 +97,11 @@ xla::XlaOp TransposeFilterForGroupConvolutionBackpropInput( CHECK_GE(num_dims, 2); // Crash OK xla::Shape new_shape = filter_shape; new_shape.set_dimensions(num_dims - 1, num_groups); - new_shape.add_dimensions(filter_shape.dimensions(num_dims - 1) / num_groups); - xla::XlaOp result = xla::Reshape(filter, new_shape.dimensions()); + new_shape.add_dimensions( + filter_shape.dimensions(num_dims - 1) / num_groups, + filter_shape.expressions(num_dims - 1) / num_groups); + xla::XlaOp result = + xla::Reshape(filter, new_shape.dimensions(), new_shape.expressions()); // 2. Transpose to [H, W, ..., G, filter_in_depth, out_depth / G] std::vector transpose_dims(num_dims + 1); @@ -118,7 +121,8 @@ xla::XlaOp ReshapeFilterForDepthwiseConvolution(const xla::Shape& filter_shape, xla::XlaOp filter) { return xla::Reshape( filter, - GroupedFilterShapeForDepthwiseConvolution(filter_shape).dimensions()); + GroupedFilterShapeForDepthwiseConvolution(filter_shape).dimensions(), + GroupedFilterShapeForDepthwiseConvolution(filter_shape).expressions()); } // Performs some basic checks on ConvOpAttrs that are true for all kinds of XLA @@ -603,7 +607,8 @@ absl::StatusOr MakeXlaBackpropFilterConvOp( } if (attrs.depthwise) { - filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions()); + filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions(), + filter_shape.expressions()); } return filter_backprop; diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc index b1da0acd61608f..f90bad207cfbd4 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc @@ -110,7 +110,8 @@ class ConvNDOp : public XlaOpKernel { expanded_input_shape.set_dimensions(i + 1, input_shape.dimensions(i)); } expanded_input_shape.set_dimensions(0, 1); - input = xla::Reshape(input, expanded_input_shape.dimensions()); + input = xla::Reshape(input, expanded_input_shape.dimensions(), + expanded_input_shape.expressions()); } else if (attrs_.batch_dims > 1) { // Flatten batch_dims. std::vector to_collapse(attrs_.batch_dims); @@ -131,7 +132,8 @@ class ConvNDOp : public XlaOpKernel { if (attrs_.batch_dims == 0) { xla::Shape no_batch_shape(out_shape); no_batch_shape.DeleteDimension(0); - out = xla::Reshape(out, no_batch_shape.dimensions()); + out = xla::Reshape(out, no_batch_shape.dimensions(), + no_batch_shape.expressions()); } else if (attrs_.batch_dims > 1) { xla::Shape expanded_out_shape(input_shape); for (int i = attrs_.batch_dims; i < input_shape.dimensions().size(); @@ -139,7 +141,8 @@ class ConvNDOp : public XlaOpKernel { expanded_out_shape.set_dimensions( i, out_shape.dimensions(i - (attrs_.batch_dims - 1))); } - out = xla::Reshape(out, expanded_out_shape.dimensions()); + out = xla::Reshape(out, expanded_out_shape.dimensions(), + expanded_out_shape.expressions()); } ctx->SetOutput(0, out); } diff --git a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc index e8e2babffd529c..3a651d1adae655 100644 --- a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc @@ -64,6 +64,7 @@ class DepthToSpaceOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, input_xla_shape.status()); absl::Span input_shape = input_xla_shape.value().dimensions(); + const xla::Shape& input_shape_with_exprs = input_xla_shape.value(); int input_rank = input_shape.size(); static const int kRequiredDims = 4; @@ -77,20 +78,31 @@ class DepthToSpaceOp : public XlaOpKernel { std::vector reshaped_shape; std::vector transpose_order; std::vector output_shape; + std::vector reshaped_exprs; + std::vector output_exprs; reshaped_shape.reserve(input_rank); transpose_order.reserve(input_rank); output_shape.reserve(input_rank); + reshaped_exprs.reserve(input_rank + num_spatial_dims); + output_exprs.reserve(input_rank); if (data_format == FORMAT_NHWC) { reshaped_shape.push_back(input_shape[0]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(0)); for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(input_shape[1 + i]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(1 + i)); } int64_t block_elems = 1; for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(block_size_); + reshaped_exprs.push_back(xla::DExpr::Const(block_size_)); block_elems *= block_size_; } reshaped_shape.push_back(input_shape[feature_dim] / block_elems); + reshaped_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) / + xla::DExpr::Const(block_elems)) + .simplify()); transpose_order.push_back(0); for (int i = 0; i < num_spatial_dims; ++i) { @@ -100,21 +112,37 @@ class DepthToSpaceOp : public XlaOpKernel { transpose_order.push_back(feature_dim + num_spatial_dims); output_shape.push_back(input_shape[0]); + output_exprs.push_back(input_shape_with_exprs.expressions(0)); for (int i = 0; i < num_spatial_dims; ++i) { output_shape.push_back(input_shape[1 + i] * block_size_); + output_exprs.push_back( + (input_shape_with_exprs.expressions(1 + i) * + xla::DExpr::Const(block_size_)) + .simplify()); } output_shape.push_back(input_shape[feature_dim] / block_elems); + output_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) / + xla::DExpr::Const(block_elems)) + .simplify()); } else { // NCHW format. reshaped_shape.push_back(input_shape[0]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(0)); int64_t block_elems = 1; for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(block_size_); + reshaped_exprs.push_back(xla::DExpr::Const(block_size_)); block_elems *= block_size_; } reshaped_shape.push_back(input_shape[feature_dim] / block_elems); + reshaped_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) / + xla::DExpr::Const(block_elems)) + .simplify()); for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(input_shape[2 + i]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(2 + i)); } transpose_order.push_back(0); @@ -125,9 +153,18 @@ class DepthToSpaceOp : public XlaOpKernel { } output_shape.push_back(input_shape[0]); + output_exprs.push_back(input_shape_with_exprs.expressions(0)); output_shape.push_back(input_shape[feature_dim] / block_elems); + output_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) / + xla::DExpr::Const(block_elems)) + .simplify()); for (int i = 0; i < num_spatial_dims; ++i) { output_shape.push_back(input_shape[2 + i] * block_size_); + output_exprs.push_back( + (input_shape_with_exprs.expressions(2 + i) * + xla::DExpr::Const(block_size_)) + .simplify()); } } @@ -148,7 +185,7 @@ class DepthToSpaceOp : public XlaOpKernel { ") is not divisible by square of the block size (", block_size_, ")")); - xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape); + xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs); // 2. Permute dimensions of `reshaped` to produce // `permuted_reshaped` of shape: @@ -169,7 +206,7 @@ class DepthToSpaceOp : public XlaOpKernel { // input_shape[2] * block_size_, // depth / (block_size_ * block_size_)] // - xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape); + xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape, output_exprs); // If this used to be a vectorized format turn it back now. if (data_format != data_format_) { diff --git a/tensorflow/compiler/tf2xla/kernels/diag_op.cc b/tensorflow/compiler/tf2xla/kernels/diag_op.cc index 404fa9f5e04e45..9644d52fa489a6 100644 --- a/tensorflow/compiler/tf2xla/kernels/diag_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/diag_op.cc @@ -29,6 +29,7 @@ limitations under the License. #include "xla/hlo/builder/lib/matrix.h" #include "xla/hlo/builder/lib/pooling.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/shape_util.h" #include "xla/util.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/op_kernel.h" @@ -38,7 +39,9 @@ namespace { // Create a diagonal / batch diagonal matrix with 'input' on the diagonal. xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size, - absl::Span other_dims) { + const xla::DExpr& last_dim_expr, + absl::Span other_dims, + absl::Span other_dim_exprs) { xla::XlaBuilder* builder = input.builder(); // Create two matrices that have the following forms, and compare them: // @@ -49,14 +52,23 @@ xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size, // // This produces a predicate matrix of the right size, with "true" on the // diagonal. - xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size); - xla::XlaOp iota_broadcast = xla::Broadcast(iota, {last_dim_size}); + xla::XlaOp iota = xla::Iota( + builder, + xla::ShapeUtil::MakeShape(xla::S32, std::vector{last_dim_size}, + std::vector{last_dim_expr}), + /*iota_dimension=*/0); + xla::XlaOp iota_broadcast = xla::Broadcast( + iota, {last_dim_size}, {last_dim_expr, last_dim_expr}); xla::XlaOp mask = xla::Eq(iota_broadcast, iota, {0}); // If this is a batched diagonal, broadcast the mask across the other // dimensions. if (!other_dims.empty()) { - mask = xla::Broadcast(mask, other_dims); + std::vector mask_exprs(other_dim_exprs.begin(), + other_dim_exprs.end()); + mask_exprs.push_back(last_dim_expr); + mask_exprs.push_back(last_dim_expr); + mask = xla::Broadcast(mask, other_dims, mask_exprs); } // Broadcast the input, and then use the mask computed above to select the @@ -69,13 +81,17 @@ xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size, std::vector out_dim_sizes(other_dims.begin(), other_dims.end()); out_dim_sizes.push_back(last_dim_size); out_dim_sizes.push_back(last_dim_size); + std::vector out_dim_exprs(other_dim_exprs.begin(), + other_dim_exprs.end()); + out_dim_exprs.push_back(last_dim_expr); + out_dim_exprs.push_back(last_dim_expr); // Broadcast into the second to last dimension. std::vector broadcast_dimensions(other_dims.size() + 1); absl::c_iota(broadcast_dimensions, 0); ++broadcast_dimensions.back(); - xla::XlaOp input_broadcast = - xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions); + xla::XlaOp input_broadcast = xla::BroadcastInDim( + input, out_dim_sizes, broadcast_dimensions, out_dim_exprs); return xla::Select(mask, input_broadcast, xla::ZerosLike(input_broadcast)); } @@ -102,17 +118,26 @@ class DiagOp : public XlaOpKernel { // [0, 0, 0, 4]] // Flattens the input to 1D. + xla::DExpr flattened_expr = xla::DExpr::Const(1); + std::vector input_exprs = input_shape.get_filled_expressions(); + for (const xla::DExpr& expr : input_exprs) { + flattened_expr = (flattened_expr * expr).simplify(); + } int64_t size = input_shape.num_elements(); - input = xla::Reshape(input, {size}); + input = xla::Reshape(input, {size}, {flattened_expr}); // Create an R2 with the R1 diagonal. - xla::XlaOp diag = CreateDiagonal(input, size, /*other_dims=*/{}); + xla::XlaOp diag = + CreateDiagonal(input, size, flattened_expr, /*other_dims=*/{}, + /*other_dim_exprs=*/{}); // Reshapes to the final shape. std::vector new_dims(dims.size() * 2); std::copy(dims.begin(), dims.end(), new_dims.begin()); std::copy(dims.begin(), dims.end(), new_dims.begin() + dims.size()); - diag = xla::Reshape(diag, new_dims); + std::vector new_exprs(input_exprs.begin(), input_exprs.end()); + new_exprs.insert(new_exprs.end(), input_exprs.begin(), input_exprs.end()); + diag = xla::Reshape(diag, new_dims, new_exprs); ctx->SetOutput(0, diag); } diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc index ceeea010ee7858..46404dc31ed23f 100644 --- a/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc @@ -43,6 +43,26 @@ limitations under the License. namespace tensorflow { namespace { +std::vector GetFilledExpressions(const xla::Shape& shape) { + std::vector expressions; + expressions.reserve(shape.dimensions_size()); + auto shape_expressions = shape.expressions(); + for (int64_t i = 0; i < shape.dimensions_size(); ++i) { + expressions.push_back(i < shape_expressions.size() && shape_expressions[i] + ? shape_expressions[i] + : xla::DExpr::Const(shape.dimensions(i))); + } + return expressions; +} + +xla::DExpr CollapseExpressions(absl::Span expressions) { + xla::DExpr collapsed = xla::DExpr::Const(1); + for (const xla::DExpr& expression : expressions) { + collapsed = collapsed * expression; + } + return collapsed; +} + class DynamicPartitionOp : public XlaOpKernel { public: explicit DynamicPartitionOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -67,6 +87,26 @@ class DynamicPartitionOp : public XlaOpKernel { xla::XlaOp partitions_1d, const xla::Shape& data_1d_shape, const xla::Shape& partition_1d_shape) { int64_t input_count = data_1d_shape.dimensions(0); + VLOG(1) << "data_1d_shape=" + << xla::ShapeUtil::HumanString(data_1d_shape) + << " partition_1d_shape=" + << xla::ShapeUtil::HumanString(partition_1d_shape) + << " input_count=" << input_count + << " num_partitions=" << num_partitions_; + // Use the smaller runtime size to mask off padded tail elements when one + // input loses dynamic shape information and falls back to its static bound. + // For valid DynamicPartition inputs the true prefix sizes should match, so + // this keeps us from treating padding on either side as real elements. + xla::XlaOp dynamic_input_count = xla::Min( + xla::GetDimensionSize(data_1d, 0), + xla::GetDimensionSize(partitions_1d, 0)); + xla::XlaOp input_index = xla::Iota(ctx->builder(), xla::S32, input_count); + xla::XlaOp valid_element = xla::Lt(input_index, dynamic_input_count); + xla::XlaOp invalid_partition = + xla::Broadcast(xla::ConstantR0(ctx->builder(), num_partitions_), + {input_count}, partition_1d_shape.expressions()); + partitions_1d = xla::Select(valid_element, partitions_1d, invalid_partition); + std::vector to_sort = {partitions_1d, data_1d}; std::vector types_to_sort = { partition_1d_shape.element_type(), data_1d_shape.element_type()}; @@ -114,11 +154,13 @@ class DynamicPartitionOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { xla::Shape data_shape = ctx->InputXlaShape(0).value(); xla::Shape partition_shape = ctx->InputXlaShape(1).value(); + xla::Shape flattened_partition_shape = partition_shape; xla::XlaOp data = ctx->Input(0); xla::XlaOp partitions = ctx->Input(1); std::vector partitions_static; bool partitions_are_static = - ctx->ConstantInputReshapedToIntVector(1, &partitions_static).ok(); + ctx->ConstantInputReshapedToIntVector(1, &partitions_static).ok() && + partition_shape.is_static(); // We know how to solve DynamicPartition on 1D inputs using // DynamicPartition1D. For other input, we do two things: // @@ -140,8 +182,11 @@ class DynamicPartitionOp : public XlaOpKernel { for (int64_t i = 0; i < rank; ++i) { broadcasted_dims.push_back(i); } - partitions = xla::BroadcastInDim(partitions, data_shape.dimensions(), - broadcasted_dims); + partitions = + xla::BroadcastInDim(partitions, data_shape.dimensions(), + broadcasted_dims, data_shape.expressions()); + flattened_partition_shape = data_shape; + flattened_partition_shape.set_element_type(partition_shape.element_type()); } // Output shape bounded is calculated by @@ -159,19 +204,39 @@ class DynamicPartitionOp : public XlaOpKernel { } int64_t input_count = xla::ShapeUtil::ElementsIn(data_shape); - auto data_1d = xla::Reshape(data, {input_count}); - auto partitions_1d = xla::Reshape(partitions, {input_count}); - xla::Shape data_1d_shape = - xla::ShapeUtil::MakeShape(data_shape.element_type(), {input_count}); + auto data_exprs = GetFilledExpressions(data_shape); + auto flattened_partition_exprs = GetFilledExpressions(flattened_partition_shape); + auto data_1d = + xla::Reshape(data, {input_count}, {CollapseExpressions(data_exprs)}); + auto partitions_1d = xla::Reshape( + partitions, {input_count}, + {CollapseExpressions(flattened_partition_exprs)}); + xla::Shape data_1d_shape = xla::ShapeUtil::MakeShape( + data_shape.element_type(), {input_count}, + std::vector{CollapseExpressions(data_exprs)}); xla::Shape partitions_1d_shape = xla::ShapeUtil::MakeShape( - partition_shape.element_type(), {input_count}); + partition_shape.element_type(), {input_count}, + std::vector{ + CollapseExpressions(flattened_partition_exprs)}); std::vector output, partition_length; std::tie(output, partition_length) = DynamicPartition1D( ctx, data_1d, partitions_1d, data_1d_shape, partitions_1d_shape); + std::vector output_exprs; + output_exprs.reserve(data_shape.dimensions().size() - + partition_shape.dimensions().size() + 1); + output_exprs.push_back(CollapseExpressions(absl::MakeConstSpan( + flattened_partition_exprs).subspan( + 0, partition_shape.dimensions().size()))); + for (int64_t i = partition_shape.dimensions().size(); + i < data_shape.dimensions().size(); ++i) { + output_exprs.push_back(data_exprs[i]); + } for (int64_t i = 0; i < num_partitions_; ++i) { - auto reshape = xla::Reshape(output[i], output_shape_bound_dims); + xla::Shape output_shape = xla::ShapeUtil::MakeShape( + data_shape.element_type(), output_shape_bound_dims, output_exprs); + auto reshape = xla::Reshape(output_shape, output[i]); if (partitions_are_static) { int64_t size = absl::c_count(partitions_static, i); ctx->SetOutput(i, xla::SliceInDim(reshape, 0, size, 1, 0)); diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc index cb7e4f6f96437e..fe45493d452d46 100644 --- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc @@ -132,13 +132,18 @@ class DynamicStitchOp : public XlaOpKernel { int64_t result_rank = 1 + data0_shape.dims() - indices0_shape.dims(); if (number_of_indices == 0) { std::vector result_shape(result_rank); + std::vector result_expressions(result_rank, + xla::DExpr::Const(0)); for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) { result_shape[d - indices0_shape.dims() + 1] = data0_shape.dim_size(d); + result_expressions[d - indices0_shape.dims() + 1] = + data0_shape.get_filled_expression(d); } xla::PrimitiveType element_type = ctx->input_xla_type(ctx->num_inputs() - 1); xla::Literal empty_literal = xla::Literal::CreateFromShape( - xla::ShapeUtil::MakeShape(element_type, result_shape)); + xla::ShapeUtil::MakeShape(element_type, result_shape, + result_expressions)); ctx->SetOutput(0, xla::ConstantLiteral(ctx->builder(), empty_literal)); return; } @@ -174,16 +179,20 @@ class DynamicStitchOp : public XlaOpKernel { TensorShape new_shape; // first reshaped dimension is the number of indices for this input. new_shape.AddDim(indices[input_num].shape().dimensions(0)); + new_shape.AddExpression( + xla::DExpr::Const(indices[input_num].shape().dimensions(0))); // Then the rest are the common extra shape. for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) { new_shape.AddDim(data0_shape.dim_size(d)); + new_shape.AddExpression(data0_shape.get_filled_expression(d)); } // Get the data, shaped appropriately. auto handle = data[input_num]; if (new_shape == data_shapes[input_num]) { input[input_num] = handle; } else { - input[input_num] = xla::Reshape(handle, new_shape.dim_sizes()); + input[input_num] = xla::Reshape(handle, new_shape.dim_sizes(), + new_shape.get_filled_expressions()); } } diff --git a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc index 2a65441eb79bf9..e7bc0890317316 100644 --- a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc @@ -127,7 +127,8 @@ class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel { xla::XlaOp between_nudged_min_max = xla::And( xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max)); xla::XlaOp zeroes = xla::Broadcast(XlaHelpers::Zero(b, data_type), - gradient_shape.dim_sizes()); + gradient_shape.dim_sizes(), + gradient_shape.get_filled_expressions()); xla::XlaOp output = xla::Select(between_nudged_min_max, gradient, zeroes); ctx->SetOutput(0, output); } @@ -213,7 +214,8 @@ class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel { xla::XlaOp between_nudged_min_max = xla::And( xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max)); xla::XlaOp zero = XlaHelpers::Zero(b, data_type); - xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes()); + xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes(), + gradient_shape.get_filled_expressions()); xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes); ctx->SetOutput(0, output0); @@ -270,7 +272,8 @@ class FakeQuantWithMinMaxVarsPerChannelOp : public XlaOpKernel { absl::Span input_dimensions = input_shape.dimensions(); auto convert_to_input_shape = [&](const xla::XlaOp op) { return xla::BroadcastInDim(op, input_dimensions, - {input_shape.dimensions_size() - 1}); + {input_shape.dimensions_size() - 1}, + input_shape.expressions()); }; input_min = convert_to_input_shape(input_min); input_max = convert_to_input_shape(input_max); @@ -323,7 +326,6 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel { xla::XlaBuilder* b = ctx->builder(); xla::Shape input_shape = b->GetShape(input).value(); absl::Span input_dimensions = input_shape.dimensions(); - std::vector reduce_axes; for (int64_t i = 0; i + 1 < input_shape.dimensions_size(); ++i) { reduce_axes.push_back(i); @@ -331,7 +333,8 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel { auto convert_to_input_shape = [&](const xla::XlaOp op) { return xla::BroadcastInDim(op, input_dimensions, - {input_shape.dimensions_size() - 1}); + {input_shape.dimensions_size() - 1}, + input_shape.expressions()); }; input_min = convert_to_input_shape(input_min); input_max = convert_to_input_shape(input_max); @@ -343,7 +346,8 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel { xla::XlaOp between_nudged_min_max = xla::And( xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max)); xla::XlaOp zero = XlaHelpers::Zero(b, data_type); - xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes()); + xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes(), + gradient_shape.get_filled_expressions()); xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes); ctx->SetOutput(0, output0); diff --git a/tensorflow/compiler/tf2xla/kernels/fill_op.cc b/tensorflow/compiler/tf2xla/kernels/fill_op.cc index 6e5a1430538365..97e4c6e85d4dcf 100644 --- a/tensorflow/compiler/tf2xla/kernels/fill_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/fill_op.cc @@ -53,11 +53,20 @@ class FillOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector( "dims", &dims, xla::ValueInferenceMode::kUpperBound)); + std::vector dim_exprs; + dim_exprs.reserve(dims.size()); + const auto& contents = ctx->InputExpression("dims").contents(); + for (int64_t i = 0; i < dims.size(); ++i) { + dim_exprs.push_back(i < contents.size() && contents[i] && + contents[i]->is_dynamic() + ? contents[i] + : xla::DExpr::Const(dims[i])); + } std::vector dynamic_dims; OP_REQUIRES_OK( ctx, ctx->ResolveInputDynamismIntoPredVector("dims", &dynamic_dims)); - auto output = xla::Broadcast(ctx->Input("value"), dims); + auto output = xla::Broadcast(ctx->Input("value"), dims, dim_exprs); for (int64_t i = 0; i < dims.size(); ++i) { // If a dimension is dynamic, call set-dimension-size on the output. if (dynamic_dims[i]) { diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 2783951e1b6b0f..f058f5a9722506 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -88,7 +88,8 @@ absl::Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape, out_shape.AppendShape(input_shape_post_axis); *gather_output = - xla::Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes()); + xla::Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes(), + out_shape.get_filled_expressions()); return absl::OkStatus(); } diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc index a8eb7bbf794268..6cc0b873a355fd 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc @@ -59,7 +59,7 @@ std::array RGBToHSV(XlaOpKernelContext* ctx, xla::XlaBuilder* b, auto minimum = xla::Min(xla::Min(red, green), blue); auto range = xla::Sub(value, minimum); - auto zeros = xla::Broadcast(zero, shape.dim_sizes()); + auto zeros = xla::Broadcast(zero, shape.dim_sizes(), shape.get_filled_expressions()); auto saturation = xla::Select(xla::Gt(value, zero), xla::Div(range, value), zeros); diff --git a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc index f357262a39c35b..3808df4937b3ec 100644 --- a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc @@ -87,9 +87,11 @@ class InTopKOp : public XlaOpKernel { // which indicates the target is in topk. xla::XlaOp gt_r2 = xla::Gt(predictions_r2, targets_values_r1, {0}); xla::XlaOp zero_r0 = xla::Zero(xla_builder, xla::S32); - xla::XlaOp zero_r2 = xla::Broadcast(zero_r0, predictions_shape.dim_sizes()); + xla::XlaOp zero_r2 = xla::Broadcast(zero_r0, predictions_shape.dim_sizes(), + predictions_shape.get_filled_expressions()); xla::XlaOp one_r0 = xla::One(xla_builder, xla::S32); - xla::XlaOp one_r2 = xla::Broadcast(one_r0, predictions_shape.dim_sizes()); + xla::XlaOp one_r2 = xla::Broadcast(one_r0, predictions_shape.dim_sizes(), + predictions_shape.get_filled_expressions()); xla::XlaOp one_hot_r2 = xla::Select(gt_r2, one_r2, zero_r2); xla::XlaOp num_gt_r1 = xla::Reduce( one_hot_r2, zero_r0, diff --git a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc index 46e46f6d8b3d32..cf2389110eefb4 100644 --- a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc @@ -49,14 +49,16 @@ void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype, // dimension of sorted_sequence. auto new_values_shape = values_shape; new_values_shape.InsertDim(/* d */ 2, /* size */ 1); - auto values_reshaped = xla::Reshape(values, new_values_shape.dim_sizes()); + auto values_reshaped = xla::Reshape(values, new_values_shape.dim_sizes(), + new_values_shape.get_filled_expressions()); // Add a new penultimate dimension to sorted_inputs, to allow broadcasting of // sorted_sequence entries for each value. auto new_sorted_inputs_shape = sorted_inputs_shape; new_sorted_inputs_shape.InsertDim(/* d */ 1, /* size */ 1); auto sorted_inputs_reshaped = - xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes()); + xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes(), + new_sorted_inputs_shape.get_filled_expressions()); // We are relying on broadcasting to compare each value against each entry in // the associated sorted_inputs row. diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc index 48e8f976cc67bb..8ef6dfbd9a1702 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc @@ -234,7 +234,8 @@ xla::XlaOp SetMatrixDiag(const xla::XlaOp input, const xla::XlaOp diag, // Broadcast and mask. xla::XlaOp diag_broadcast = xla::BroadcastInDim( - diag_slice, input_shape.dim_sizes(), broadcast_dimensions); + diag_slice, input_shape.dim_sizes(), broadcast_dimensions, + input_shape.get_filled_expressions()); const auto mask = xla::GetDiagonalMask(output, diag_index); output = xla::Select(mask, diag_broadcast, output); } @@ -327,8 +328,11 @@ class MatrixDiagOp : public XlaOpKernel { TensorShape output_shape = diag_shape; output_shape.RemoveLastDims((num_diags == 1) ? 1 : 2); output_shape.AddDim(num_rows); + output_shape.AddExpression(xla::DExpr::Const(num_rows)); output_shape.AddDim(num_cols); - xla::XlaOp output = xla::Broadcast(padding_value, output_shape.dim_sizes()); + output_shape.AddExpression(xla::DExpr::Const(num_cols)); + xla::XlaOp output = xla::Broadcast(padding_value, output_shape.dim_sizes(), + output_shape.get_filled_expressions()); xla::XlaOp diag = context->Input(0); context->SetOutput( 0, SetMatrixDiag(output, diag, output_shape, diag_rank, num_diags, @@ -404,11 +408,15 @@ class MatrixDiagPartOp : public XlaOpKernel { TensorShape output_shape = input_shape; output_shape.RemoveLastDims(2); const int num_diags = upper_diag_index - lower_diag_index + 1; - if (num_diags > 1) output_shape.AddDim(num_diags); + if (num_diags > 1) { + output_shape.AddDim(num_diags); + output_shape.AddExpression(xla::DExpr::Const(num_diags)); + } const int32_t max_diag_len = std::min(num_rows + std::min(upper_diag_index, int64_t{0}), num_cols - std::max(lower_diag_index, int64_t{0})); output_shape.AddDim(max_diag_len); + output_shape.AddExpression(xla::DExpr::Const(max_diag_len)); // Computes output. xla::XlaOp input = context->Input(0); @@ -447,7 +455,8 @@ class MatrixDiagPartOp : public XlaOpKernel { } auto concat = xla::ConcatInDim(context->builder(), diag_list, input_rank - 2); - context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes())); + context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes(), + output_shape.get_filled_expressions())); } private: @@ -519,11 +528,15 @@ class MatrixSetDiagOp : public XlaOpKernel { TensorShape expected_diag_shape = input_shape; expected_diag_shape.RemoveLastDims(2); - if (num_diags > 1) expected_diag_shape.AddDim(num_diags); + if (num_diags > 1) { + expected_diag_shape.AddDim(num_diags); + expected_diag_shape.AddExpression(xla::DExpr::Const(num_diags)); + } const int32_t max_diag_len = std::min(num_rows + std::min(upper_diag_index, int64_t{0}), num_cols - std::max(lower_diag_index, int64_t{0})); expected_diag_shape.AddDim(max_diag_len); + expected_diag_shape.AddExpression(xla::DExpr::Const(max_diag_len)); OP_REQUIRES( context, expected_diag_shape == diag_shape, errors::InvalidArgument( diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc index 17b5ae7a70375a..207bc6f7ce8984 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc @@ -96,8 +96,12 @@ MatrixTriangularSolveOp::Broadcast(xla::XlaOp lhs, const TensorShape& lhs_shape, TensorShape lhs_broadcast_shape(broadcast_helper.output_batch_shape()); lhs_broadcast_shape.AddDim(m); + lhs_broadcast_shape.AddExpression(xla::DExpr::Const(m)); lhs_broadcast_shape.AddDim(m); - auto lhs_output = BroadcastTo(lhs, lhs_broadcast_shape.dim_sizes()); + lhs_broadcast_shape.AddExpression(xla::DExpr::Const(m)); + auto lhs_output = + BroadcastTo(lhs, lhs_broadcast_shape.dim_sizes(), + lhs_broadcast_shape.get_filled_expressions()); if (!lhs_output.ok()) { xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status()); return {error, error}; @@ -105,8 +109,12 @@ MatrixTriangularSolveOp::Broadcast(xla::XlaOp lhs, const TensorShape& lhs_shape, TensorShape rhs_broadcast_shape(broadcast_helper.output_batch_shape()); rhs_broadcast_shape.AddDim(m); + rhs_broadcast_shape.AddExpression(xla::DExpr::Const(m)); rhs_broadcast_shape.AddDim(n); - auto rhs_output = BroadcastTo(rhs, rhs_broadcast_shape.dim_sizes()); + rhs_broadcast_shape.AddExpression(xla::DExpr::Const(n)); + auto rhs_output = + BroadcastTo(rhs, rhs_broadcast_shape.dim_sizes(), + rhs_broadcast_shape.get_filled_expressions()); if (!rhs_output.ok()) { xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status()); return {error, error}; diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc index ba4e8bbef7b136..0b236ae9c0c052 100644 --- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc @@ -17,6 +17,8 @@ limitations under the License. #include +#include "absl/algorithm/container.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_builder.h" @@ -28,6 +30,33 @@ limitations under the License. namespace tensorflow { namespace { +bool TryBuildPackedContents(XlaOpKernelContext* ctx, int num, int axis, + std::vector* contents) { + contents->clear(); + if (axis != 0) { + return false; + } + for (int i = 0; i < num; ++i) { + if (ctx->InputShape(i).dims() != 0) { + contents->clear(); + return false; + } + const auto& input_contents = ctx->InputExpression(i).contents(); + if (!input_contents.empty()) { + if (input_contents.size() != 1) { + contents->clear(); + return false; + } + contents->push_back(input_contents[0]); + continue; + } + contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel)); + } + return absl::c_any_of(*contents, [](const xla::DExpr& expr) { + return expr && expr->is_dynamic(); + }); +} + class PackOp : public XlaOpKernel { public: explicit PackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -64,14 +93,36 @@ class PackOp : public XlaOpKernel { std::vector reshaped_inputs(num); TensorShape child_shape(shapes[0]); + std::vector exprs = child_shape.get_filled_expressions(); child_shape.InsertDim(axis, 1); - + exprs.insert(exprs.begin() + axis, xla::DExpr::Const(1)); + std::vector output_contents; + const bool has_output_contents = + SymbolicContentEnabled() && + TryBuildPackedContents(ctx, num, axis, &output_contents); for (int i = 0; i < num; ++i) { // Reshape the inputs to have an extra dimension of size 1. - reshaped_inputs[i] = xla::Reshape(values[i], child_shape.dim_sizes()); + reshaped_inputs[i] = xla::Reshape(values[i], child_shape.dim_sizes(), + exprs); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents( + reshaped_inputs[i], + {output_contents[static_cast(i)]})); + } } - ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), reshaped_inputs, axis)); + auto output = xla::ConcatInDim(ctx->builder(), reshaped_inputs, axis); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(output, output_contents)); + auto output_expr = + XlaExpression::XlaOp(output, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } + ctx->SetOutput(0, output); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc index aa7c78b8b8f97a..fe2401fade704a 100644 --- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc @@ -264,10 +264,14 @@ class MaxPoolOp : public PoolingOp { absl::InlinedVector new_dims(result_shape->dimensions().begin(), result_shape->dimensions().end()); + absl::InlinedVector new_exprs( + result_shape->expressions().begin(), + result_shape->expressions().end()); new_dims[1] /= *vect_width; + new_exprs[1] = new_exprs[1] / xla::DExpr::Const(*vect_width); new_dims.insert(new_dims.begin() + 2, *vect_width); - pooling = - xla::Transpose(xla::Reshape(pooling, new_dims), {0, 1, 3, 4, 2}); + pooling = xla::Transpose(xla::Reshape(pooling, new_dims, new_exprs), + {0, 1, 3, 4, 2}); } ctx->SetOutput(0, pooling); diff --git a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc index cac9f8a68f234e..6b7e990f7030fd 100644 --- a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc @@ -174,7 +174,8 @@ class QuantizeAndDequantizeOp : public XlaOpKernel { xla::Shape input_shape = b->GetShape(input).value(); absl::Span input_dimensions = input_shape.dimensions(); auto convert_to_input_shape = [&](const xla::XlaOp op) { - return xla::BroadcastInDim(op, input_dimensions, {axis_}); + return xla::BroadcastInDim(op, input_dimensions, {axis_}, + input_shape.expressions()); }; min_range = convert_to_input_shape(min_range); max_range = convert_to_input_shape(max_range); diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc index 5f911018c244b5..f1342c7e187ea2 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -147,10 +148,10 @@ class MeanOp : public XlaReductionOp { xla::XlaOp result = reduce_output; xla::Shape bounded_shape = builder->GetShape(input).value(); int64_t divisor_value = bounded_shape.dimensions(dimensions_to_reduce[0]); - auto divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]); + xla::XlaOp divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]); for (int i = 1; i < dimensions_to_reduce.size(); i++) { int64_t size_value = bounded_shape.dimensions(dimensions_to_reduce[i]); - auto size = xla::GetDimensionSize(input, dimensions_to_reduce[i]); + xla::XlaOp size = xla::GetDimensionSize(input, dimensions_to_reduce[i]); if (size_value * divisor_value > std::numeric_limits::max()) { result = result / xla::ConvertElementType(divisor, xla_reduction_type_); divisor_value = size_value; diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc index 6a8a98342c1123..6a5f40441beff3 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc @@ -106,16 +106,19 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { } std::vector final_shape; + std::vector final_exprs; for (int i = 0; i < data_shape.dims(); ++i) { if (!bitmap[i]) { // If we are not reducing along dimension i. int64_t dim = data_shape.dim_size(i); final_shape.push_back(dim); + final_exprs.push_back(data_shape.get_filled_expression(i)); } else if (keep_dims_) { // We are reducing along dimension i, but we want to keep the // same number of dimensions, so we set the dimension of i to // '1'. final_shape.push_back(1); + final_exprs.push_back(xla::DExpr::Const(1)); } } @@ -139,7 +142,8 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { auto reduce = xla::Reduce(data, initial, reduction_computation, xla_axes); auto finalized = BuildFinalizer(b, data, reduce, xla_axes); - auto result = keep_dims_ ? xla::Reshape(finalized, final_shape) : finalized; + auto result = keep_dims_ ? xla::Reshape(finalized, final_shape, final_exprs) + : finalized; ctx->SetOutput(0, result); } diff --git a/tensorflow/compiler/tf2xla/kernels/relu_op.cc b/tensorflow/compiler/tf2xla/kernels/relu_op.cc index f274b271596ff5..fd932247a0f558 100644 --- a/tensorflow/compiler/tf2xla/kernels/relu_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/relu_op.cc @@ -40,8 +40,28 @@ XlaOp Relu6(XlaOp x) { namespace tensorflow { namespace { -REGISTER_XLA_OP(Name("Relu"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Relu6"), MlirXlaOpKernel); +class ReluOp : public XlaOpKernel { + public: + explicit ReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Computes the max of the scalar input x and 0. + void Compile(XlaOpKernelContext* ctx) override { + ctx->SetOutput(0, xla::Relu(ctx->Input(0))); + } +}; + +class Relu6Op : public XlaOpKernel { + public: + explicit Relu6Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Computes the max of the scalar input x and 0. + void Compile(XlaOpKernelContext* ctx) override { + ctx->SetOutput(0, xla::Relu6(ctx->Input(0))); + } +}; + +REGISTER_XLA_OP(Name("Relu"), ReluOp); +// REGISTER_XLA_OP(Name("Relu"), MlirXlaOpKernel); +// REGISTER_XLA_OP(Name("Relu6"), MlirXlaOpKernel); +REGISTER_XLA_OP(Name("Relu6"), Relu6Op); class LeakyReluOp : public XlaOpKernel { public: @@ -70,9 +90,11 @@ class Relu6GradOp : public XlaOpKernel { xla::XlaBuilder* b = ctx->builder(); const TensorShape shape = ctx->InputShape(0); const auto zero = - xla::Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes()); - const auto six = xla::Broadcast( - XlaHelpers::IntegerLiteral(b, input_type(0), 6), shape.dim_sizes()); + xla::Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes(), + shape.get_filled_expressions()); + const auto six = + xla::Broadcast(XlaHelpers::IntegerLiteral(b, input_type(0), 6), + shape.dim_sizes(), shape.get_filled_expressions()); auto out = xla::Select( xla::And(xla::Lt(ctx->Input(1), six), xla::Gt(ctx->Input(1), zero)), ctx->Input(0), zero); diff --git a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc index c54c4613d29e44..602d480b5b6541 100644 --- a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc @@ -122,12 +122,15 @@ XlaOp ConcatenateIota(xla::XlaBuilder* b, XlaOp indices, for (auto dim : warp_shape) { dimensions.push_back(dim.size); } + std::vector expressions = + warp_shape.get_filled_expressions(); // Except the last dimension, which is of size 1. dimensions.back() = 1; + expressions.back() = xla::DExpr::Const(1); - auto batch_indices = - xla::Iota(b, xla::ShapeUtil::MakeShape(xla::S32, dimensions), - /*iota_dimension=*/0); + auto batch_indices = xla::Iota( + b, xla::ShapeUtil::MakeShape(xla::S32, dimensions, expressions), + /*iota_dimension=*/0); return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1); } @@ -365,14 +368,17 @@ XlaOp CalculateGradWarp(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio, auto warp_dims = warp_shape.dim_sizes(); std::vector warp_dims_without_last_dims(warp_dims.begin(), warp_dims.end() - 1); + std::vector warp_expressions = + warp_shape.get_filled_expressions(); // With dimension [batch, dim_0, ...dim_n, 4] std::vector neighbor_broadcast_dims = warp_dims_without_last_dims; neighbor_broadcast_dims.push_back(4); + warp_expressions.back() = xla::DExpr::Const(4); // With dimension [batch, dim_0, ...dim_n, 4] - auto neighbor_broadcast_shape = - xla::ShapeUtil::MakeShape(data_type, neighbor_broadcast_dims); + auto neighbor_broadcast_shape = xla::ShapeUtil::MakeShape( + data_type, neighbor_broadcast_dims, warp_expressions); const int64_t last_warp_dim = warp_shape.dims() - 1; diff --git a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc index ba17d1b295b763..eb861e895ac03b 100644 --- a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc @@ -57,8 +57,10 @@ class ReshapeOp : public XlaOpKernel { // is one. TensorShape shape; int64_t product = 1; + xla::DExpr product_expr = xla::DExpr::Const(1); int unknown_index = -1; bool shape_has_zero_dim = false; + int ratio = 1; for (int d = 0; d < num_dims; ++d) { const int64_t size = shape_input[d]; if (size == -1) { @@ -68,23 +70,60 @@ class ReshapeOp : public XlaOpKernel { unknown_index, " and ", d)); unknown_index = d; shape.AddDim(1); + shape.AddExpression(xla::DExpr::Const(1)); + ratio = 1; } else if (size == 0) { // We don't include zero-sized dimension in product, so that we can // still calculate number of elements for non-zero-sized dimensions and // therefore infer their shapes. shape.AddDim(size); + shape.AddExpression(xla::DExpr::Const(size)); shape_has_zero_dim = true; } else { + xla::DExpr size_expr; OP_REQUIRES(ctx, size >= 0, errors::InvalidArgument( "size ", d, " must be non-negative, not ", size)); shape.AddDim(size); + if (d < input_shape.dims() && input_shape.get_filled_expression(d) && + input_shape.get_filled_expression(d)->is_dynamic()) { + int old = input_shape.dim_size(d); + bool is_split = (old > size); + int local_ratio = ratio * (is_split ? old / size : size / old); + xla::DExpr input_dexpr = input_shape.get_filled_expression(d); + xla::DExpr ratio_expr = xla::DExpr::Const(local_ratio); + xla::DExpr new_expr = + (size > old) ? input_dexpr * ratio_expr // Reduce [x,y] -> [x*y] + : input_dexpr / ratio_expr; // Split [xy] -> [x/y,y] + + // Pass ratio to next dimension if this is a split, otherwise just + // reset it to 1. + ratio = is_split ? local_ratio : 1; + size_expr = new_expr.simplify(); + + } else { + size_expr = xla::DExpr::Const(size); + if (ratio != 1) { + // A split dynamic dimension can be materialized by multiple later + // known dimensions. Any unresolved remainder is kept in `ratio` + // and may be consumed by a subsequent `-1` dimension (if present); + // otherwise, it remains unapplied. + if (ratio % size == 0) { + ratio /= size; + } else if (size % ratio == 0) { + ratio = 1; + } + } + } product *= size; + product_expr = product_expr * size_expr; + shape.AddExpression(size_expr); } } auto input = ctx->Input(0); if (unknown_index != -1) { int64_t input_num_elements = 1; + xla::DExpr input_num_elements_expr = xla::DExpr::Const(1); bool input_has_zero_dim = false; for (int dim = 0; dim < input_shape.dims(); dim++) { // For zero dimension, we don't count it into `input_num_elements` @@ -92,12 +131,18 @@ class ReshapeOp : public XlaOpKernel { // infer shapes for other dimensions. if (input_shape.dim_size(dim) > 0 || !shape_has_zero_dim) { input_num_elements *= input_shape.dim_size(dim); + input_num_elements_expr = + (input_num_elements_expr * input_shape.get_filled_expression(dim)) + .simplify(); } else { input_has_zero_dim = true; } } int64_t missing = input_num_elements / product; + input_num_elements_expr = input_num_elements_expr.simplify(); + product_expr = product_expr.simplify(); + auto missing_expr = input_num_elements_expr / product_expr; if (!input_has_zero_dim) { if (input_xla_shape->is_static() || input_xla_shape->dimensions().size() != 1) { @@ -119,10 +164,19 @@ class ReshapeOp : public XlaOpKernel { input, xla::Zero(ctx->builder(), input_xla_shape->element_type()), 0, 0, padded_input_num - input_num_elements); input_shape.set_dim(0, padded_input_num); + missing_expr = + (input_num_elements_expr + (product - 1)) / product; + missing_expr = missing_expr.simplify(); + xla::DExpr padded_input_num_expr = + (missing_expr * xla::DExpr::Const(product)) + .simplify(); + input_shape.set_expression(0, padded_input_num_expr); } } shape.set_dim(unknown_index, missing); + shape.set_expression(unknown_index, missing_expr.simplify()); } + OP_REQUIRES(ctx, shape.num_elements() == input_shape.num_elements(), errors::InvalidArgument("Input to reshape is a tensor with ", input_shape.num_elements(), @@ -131,19 +185,23 @@ class ReshapeOp : public XlaOpKernel { VLOG(2) << "Reshape from " << input_shape.DebugString() << " to " << shape.DebugString() << ", unknown_index=" << unknown_index; + if (input_xla_shape->is_static()) { - ctx->SetOutput(0, xla::Reshape(input, shape.dim_sizes())); + ctx->SetOutput( + 0, xla::Reshape(input, shape.dim_sizes(), shape.get_filled_expressions())); return; } std::vector output_dim_sizes; std::vector dims_are_dynamic; + std::vector output_dim_exprs; const auto& dims = shape.dims(); dims_are_dynamic.reserve(dims); output_dim_sizes.reserve(dims); for (int64_t i = 0; i < dims; ++i) { output_dim_sizes.push_back( xla::Reshape(xla::Slice(ctx->Input(1), {i}, {i + 1}, {1}), {})); + output_dim_exprs.push_back(xla::DExpr::Unknown(111)); } OP_REQUIRES_OK( ctx, ctx->ResolveInputDynamismIntoPredVector(1, &dims_are_dynamic)); @@ -151,7 +209,7 @@ class ReshapeOp : public XlaOpKernel { // No unknown index. ctx->SetOutput( 0, xla::DynamicReshape(input, output_dim_sizes, shape.dim_sizes(), - dims_are_dynamic)); + dims_are_dynamic, output_dim_exprs)); return; } auto common_factors = @@ -166,21 +224,27 @@ class ReshapeOp : public XlaOpKernel { // reshape(Tensor([2, 3, 3]), [3, -1, 3]) product of the group // containing -1 will be 6. xla::XlaOp product = xla::One(ctx->builder(), xla::S32); + xla::DExpr product_expr = xla::DExpr::Const(1); for (int64_t dim = start.first; dim < end.first; ++dim) { if (input_xla_shape->is_dynamic_dimension(dim)) { input_is_dynamic = true; } product = xla::Mul(product, xla::GetDimensionSize(input, dim)); + product_expr = + (product_expr * input_shape.get_filled_expression(dim)).simplify(); } bool unknown_dim_in_group = false; // The real size for the -1 dimension in a reshape. E.g., in // reshape(Tensor([2, 3, 3]), [3, -1, 3]) this will be 2. xla::XlaOp unknown_dim_size = product; + xla::DExpr unknown_dim_expr = product_expr; for (int64_t dim = start.second; dim < end.second; ++dim) { if (dim == unknown_index) { unknown_dim_in_group = true; } else { unknown_dim_size = xla::Div(unknown_dim_size, output_dim_sizes[dim]); + unknown_dim_expr = + (unknown_dim_expr / output_dim_exprs[dim]).simplify(); } } @@ -188,12 +252,13 @@ class ReshapeOp : public XlaOpKernel { // If input dim is dynamic, output dim at the -1 position must be // dynamic. Similarly, if input dim is static, output dim has to be // static at the -1 dimension. + output_dim_exprs[unknown_index] = unknown_dim_expr; dims_are_dynamic[unknown_index] = input_is_dynamic; output_dim_sizes[unknown_index] = unknown_dim_size; ctx->SetOutput( 0, xla::DynamicReshape(input, output_dim_sizes, shape.dim_sizes(), - dims_are_dynamic)); + dims_are_dynamic, output_dim_exprs)); VLOG(2) << "Reshape from " << ctx->InputXlaShape(0)->ToString() << " to " << xla::VectorString(shape.dim_sizes()) << ", dynamic_dims=" << xla::VectorString(dims_are_dynamic); diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc index 5cecbf37706283..9f3bbf333eb4eb 100644 --- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc @@ -71,6 +71,9 @@ class ReverseSequenceOp : public XlaOpKernel { xla::XlaBuilder* builder = context->builder(); const auto input = context->Input(0); const auto seq_lens = context->Input(1); + auto input_xla_shape_or = context->InputXlaShape(0); + OP_REQUIRES(context, input_xla_shape_or.ok(), input_xla_shape_or.status()); + const xla::Shape& input_xla_shape = input_xla_shape_or.value(); const int64_t batch_size = input_shape.dim_size(batch_dim_); if (batch_size == 0) { @@ -86,11 +89,21 @@ class ReverseSequenceOp : public XlaOpKernel { xla::XlaOp back = xla::Sub(seq_lens, xla::ScalarLike(seq_lens, 1)); xla::XlaOp batch_idx = xla::Iota( builder, - xla::ShapeUtil::MakeShape(seq_lens_type, {batch_size, max_seq_len, 1}), + xla::ShapeUtil::MakeShape( + seq_lens_type, {batch_size, max_seq_len, 1}, + std::vector{ + input_shape.get_filled_expression(batch_dim_), + input_shape.get_filled_expression(seq_dim_), + xla::DExpr::Const(1)}), /*iota_dimension=*/0); xla::XlaOp forward_idx = xla::Iota( builder, - xla::ShapeUtil::MakeShape(seq_lens_type, {batch_size, max_seq_len, 1}), + xla::ShapeUtil::MakeShape( + seq_lens_type, {batch_size, max_seq_len, 1}, + std::vector{ + input_shape.get_filled_expression(batch_dim_), + input_shape.get_filled_expression(seq_dim_), + xla::DExpr::Const(1)}), /*iota_dimension=*/1); xla::XlaOp reverse_idx = xla::Sub(back, forward_idx, {0}); reverse_idx = xla::Select(xla::Lt(reverse_idx, xla::ZerosLike(reverse_idx)), @@ -129,8 +142,8 @@ class ReverseSequenceOp : public XlaOpKernel { slice_sizes[batch_dim_] = 1; slice_sizes[seq_dim_] = 1; - context->SetOutput(0, - xla::Gather(input, start_indices, dnums, slice_sizes)); + xla::XlaOp gathered = xla::Gather(input, start_indices, dnums, slice_sizes); + context->SetOutput(0, xla::Reshape(input_xla_shape, gathered)); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/roll_op.cc b/tensorflow/compiler/tf2xla/kernels/roll_op.cc index 0fcc6bec56095b..49b5cd3f01b8b8 100644 --- a/tensorflow/compiler/tf2xla/kernels/roll_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/roll_op.cc @@ -94,8 +94,8 @@ class RollOp : public XlaOpKernel { std::vector start_indices( input_shape.dims(), xla::Zero(ctx->builder(), shift_type)); start_indices[cur_axis] = axis_size - offset; - output = - xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes()); + output = xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes(), + input_shape.get_filled_expressions()); } ctx->SetOutput(0, output); } diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc index 694b4eb17ef298..f5e03242666a1a 100644 --- a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc @@ -118,7 +118,8 @@ class ScatterNdOp : public XlaOpKernel { xla::XlaBuilder* builder = context->builder(); auto buffer = xla::Broadcast(XlaHelpers::Zero(builder, dtype), - buffer_shape.dim_sizes()); + buffer_shape.dim_sizes(), + buffer_shape.get_filled_expressions()); auto indices = context->Input(0); auto updates = context->Input(1); auto combine = diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc index 21eaac25f058ed..f731727f582411 100644 --- a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc @@ -96,7 +96,8 @@ class SegmentReduce : public XlaOpKernel { buffer_shape.InsertDim(0, num_segments); auto buffer = - xla::Broadcast(InitialValue(builder), buffer_shape.dim_sizes()); + xla::Broadcast(InitialValue(builder), buffer_shape.dim_sizes(), + buffer_shape.get_filled_expressions()); // Build dynamic dim sizes for buffer, as well as whether each dimension // size is dynamic or static. We build two parts: num_sgement part and diff --git a/tensorflow/compiler/tf2xla/kernels/select_op.cc b/tensorflow/compiler/tf2xla/kernels/select_op.cc index 85aaabe87076c2..f6297e9572f76a 100644 --- a/tensorflow/compiler/tf2xla/kernels/select_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/select_op.cc @@ -70,8 +70,10 @@ class SelectOp : public XlaOpKernel { // Broadcast into the dimensions on the right. std::vector broadcast_dimensions(cond_shape.dims()); absl::c_iota(broadcast_dimensions, 0); + cond_handle = xla::BroadcastInDim(cond_handle, then_shape.dim_sizes(), - broadcast_dimensions); + broadcast_dimensions, + then_shape.get_filled_expressions()); } ctx->SetOutput(0, xla::Select(cond_handle, then_handle, else_handle)); } @@ -81,7 +83,8 @@ class SelectOp : public XlaOpKernel { void operator=(const SelectOp&) = delete; }; -REGISTER_XLA_OP(Name("Select"), MlirXlaOpKernel); +// REGISTER_XLA_OP(Name("Select"), MlirXlaOpKernel); +REGISTER_XLA_OP(Name("Select"), SelectOp); class SelectOpV2 : public XlaOpKernel { public: diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc index 108bf3848aae93..cd011016a4672c 100644 --- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc @@ -17,8 +17,9 @@ limitations under the License. #include #include - +#include "absl/algorithm/container.h" #include "absl/status/statusor.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -26,6 +27,7 @@ limitations under the License. #include "xla/hlo/builder/xla_builder.h" #include "xla/literal.h" #include "xla/primitive_util.h" +#include "xla/shape_util.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" @@ -33,17 +35,90 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { namespace { +template +xla::DExpr GetScalarExpr(const XlaExpression& expression, + const xla::LiteralSlice& literal) { + const auto& contents = expression.contents(); + if (!contents.empty() && contents[0]) { + return contents[0]; + } + return xla::DExpr::Const(literal.Get({})); +} + +bool HasDynamicContent(const XlaExpression& expression) { + return absl::c_any_of(expression.contents(), [](const xla::DExpr& expr) { + return expr && expr->is_dynamic(); + }); +} + +template +std::vector BuildRangeContents(const XlaExpression& start_expr, + const XlaExpression& delta_expr, + const xla::LiteralSlice& start, + const xla::LiteralSlice& delta, + int64_t size) { + std::vector contents; + contents.reserve(size); + xla::DExpr start_symbol = GetScalarExpr(start_expr, start); + xla::DExpr delta_symbol = GetScalarExpr(delta_expr, delta); + for (int64_t i = 0; i < size; ++i) { + xla::DExpr offset = xla::DExpr::Const(static_cast(i)); + contents.push_back((start_symbol + (delta_symbol * offset).simplify()).simplify()); + } + return contents; +} + +template +xla::DExpr BuildRangeSizeExpr(const XlaExpression& start_expr, + const XlaExpression& limit_expr, + const XlaExpression& delta_expr, + const xla::LiteralSlice& start, + const xla::LiteralSlice& limit, + const xla::LiteralSlice& delta, + int64_t fallback_size) { + xla::DExpr start_symbol = GetScalarExpr(start_expr, start); + xla::DExpr limit_symbol = GetScalarExpr(limit_expr, limit); + xla::DExpr delta_symbol = GetScalarExpr(delta_expr, delta); + + const auto& start_contents = start_expr.contents(); + xla::DExpr effective_start = + (!start_contents.empty() && start_contents[0]) ? start_contents[0] + : start_symbol; + const auto& limit_contents = limit_expr.contents(); + xla::DExpr effective_limit = + (!limit_contents.empty() && limit_contents[0]) ? limit_contents[0] + : limit_symbol; + const auto& delta_contents = delta_expr.contents(); + xla::DExpr effective_delta = + (!delta_contents.empty() && delta_contents[0]) ? delta_contents[0] + : delta_symbol; + + xla::DExpr positive_diff = (effective_limit - effective_start).simplify(); + xla::DExpr positive_size = + (((positive_diff - 1) / effective_delta) + 1).simplify(); + xla::DExpr negative_step = + (xla::DExpr::Const(0) - effective_delta).simplify(); + xla::DExpr negative_diff = (effective_start - effective_limit).simplify(); + xla::DExpr negative_size = + (((negative_diff - 1) / negative_step) + 1).simplify(); + return xla::DExpr::Select(xla::DExpr::Gt(delta_symbol, xla::DExpr::Const(0)), + positive_size, negative_size) + .simplify(); +} + // The type-specific part of the implementation of Range. template absl::StatusOr CreateRangeTensor( const xla::LiteralSlice& start_literal, const xla::LiteralSlice& limit_literal, - const xla::LiteralSlice& delta_literal, xla::XlaBuilder* builder) { + const xla::LiteralSlice& delta_literal, xla::XlaBuilder* builder, + xla::DExpr size_expr = xla::DExpr()) { T start = start_literal.Get({}); T limit = limit_literal.Get({}); T delta = delta_literal.Get({}); @@ -70,10 +145,17 @@ absl::StatusOr CreateRangeTensor( : (std::abs(limit - start) - 1) / std::abs(delta) + 1) : std::ceil(std::abs((limit - start) / delta))); - return xla::ConstantR0(builder, start) + - xla::ConstantR0(builder, delta) * - xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType(), - size); + xla::XlaOp iota = + (std::is_integral::value && size_expr) + ? xla::Iota(builder, + xla::ShapeUtil::MakeShape( + xla::primitive_util::NativeToPrimitiveType(), + {size}, std::vector{size_expr}), + /*iota_dimension=*/0) + : xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType(), + size); + + return xla::ConstantR0(builder, start) + xla::ConstantR0(builder, delta) * iota; } class RangeOp : public XlaOpKernel { @@ -103,13 +185,50 @@ class RangeOp : public XlaOpKernel { DataType type = input_type(0); absl::StatusOr output; switch (type) { - case DT_INT32: - output = CreateRangeTensor(start, limit, delta, ctx->builder()); + case DT_INT32: { + int32 start_value = start.Get({}); + int32 limit_value = limit.Get({}); + int32 delta_value = delta.Get({}); + int64_t size = static_cast( + limit_value == start_value + ? 0 + : (std::abs(limit_value - start_value) - 1) / + std::abs(delta_value) + + 1); + xla::DExpr size_expr = xla::DExpr::Const(size); + if (HasDynamicContent(ctx->InputExpression(0)) || + HasDynamicContent(ctx->InputExpression(1)) || + HasDynamicContent(ctx->InputExpression(2))) { + size_expr = BuildRangeSizeExpr( + ctx->InputExpression(0), ctx->InputExpression(1), + ctx->InputExpression(2), start, limit, delta, size); + } + output = CreateRangeTensor(start, limit, delta, ctx->builder(), + size_expr); break; - case DT_INT64: - output = - CreateRangeTensor(start, limit, delta, ctx->builder()); + } + case DT_INT64: { + int64_t start_value = start.Get({}); + int64_t limit_value = limit.Get({}); + int64_t delta_value = delta.Get({}); + int64_t size = + limit_value == start_value + ? 0 + : (std::abs(limit_value - start_value) - 1) / + std::abs(delta_value) + + 1; + xla::DExpr size_expr = xla::DExpr::Const(size); + if (HasDynamicContent(ctx->InputExpression(0)) || + HasDynamicContent(ctx->InputExpression(1)) || + HasDynamicContent(ctx->InputExpression(2))) { + size_expr = BuildRangeSizeExpr( + ctx->InputExpression(0), ctx->InputExpression(1), + ctx->InputExpression(2), start, limit, delta, size); + } + output = CreateRangeTensor(start, limit, delta, ctx->builder(), + size_expr); break; + } case DT_FLOAT: output = CreateRangeTensor(start, limit, delta, ctx->builder()); break; @@ -145,7 +264,55 @@ class RangeOp : public XlaOpKernel { } } - ctx->SetOutput(0, output.value()); + const XlaExpression& start_expr = ctx->InputExpression(0); + const XlaExpression& limit_expr = ctx->InputExpression(1); + const XlaExpression& delta_expr = ctx->InputExpression(2); + const bool symbolic_enabled = SymbolicContentEnabled(); + const bool has_dynamic_content = + HasDynamicContent(start_expr) || HasDynamicContent(limit_expr) || + HasDynamicContent(delta_expr); + + if (type == DT_INT32) { + int32 start_value = start.Get({}); + int32 limit_value = limit.Get({}); + int32 delta_value = delta.Get({}); + if (symbolic_enabled && has_dynamic_content) { + int64_t size = static_cast( + limit_value == start_value + ? 0 + : (std::abs(limit_value - start_value) - 1) / + std::abs(delta_value) + + 1); + auto output_expr = + XlaExpression::XlaOp(output.value(), ctx->expected_output_dtype(0)); + output_expr.set_contents(BuildRangeContents( + start_expr, delta_expr, start, delta, size)); + ctx->SetOutputExpression(0, output_expr); + } else { + ctx->SetOutput(0, output.value()); + } + } else if (type == DT_INT64) { + int64_t start_value = start.Get({}); + int64_t limit_value = limit.Get({}); + int64_t delta_value = delta.Get({}); + if (symbolic_enabled && has_dynamic_content) { + int64_t size = static_cast( + limit_value == start_value + ? 0 + : (std::abs(limit_value - start_value) - 1) / + std::abs(delta_value) + + 1); + auto output_expr = + XlaExpression::XlaOp(output.value(), ctx->expected_output_dtype(0)); + output_expr.set_contents(BuildRangeContents( + start_expr, delta_expr, start, delta, size)); + ctx->SetOutputExpression(0, output_expr); + } else { + ctx->SetOutput(0, output.value()); + } + } else { + ctx->SetOutput(0, output.value()); + } } }; diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc index 7e8889cb2ccee6..1d624598d6ba07 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/kernels/shape_util.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/kernels/tensor_list_utils.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" @@ -46,6 +47,19 @@ limitations under the License. namespace tensorflow { namespace { +std::vector BuildShapeContents(const TensorShape& input_shape) { + std::vector contents; + contents.reserve(input_shape.dims()); + for (int64_t i = 0; i < input_shape.dims(); ++i) { + xla::DExpr expr = input_shape.get_filled_expression(i); + contents.push_back(expr && expr->is_dynamic() + ? expr + : xla::DExpr::Unknown( + xla::kUnknownContentSentinel)); + } + return contents; +} + class ShapeOp : public XlaOpKernel { public: explicit ShapeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -58,18 +72,46 @@ class ShapeOp : public XlaOpKernel { const int rank = input_shape.dims(); if (rank != 0) { for (int64_t i = 0; i < rank; ++i) { - operands.push_back(xla::Broadcast( - xla::ConvertElementType(xla::GetDimensionSize(ctx->Input(0), i), - ctx->output_xla_type(0)), - {1})); + xla::DExpr expr = input_shape.get_filled_expression(i); + std::vector content = { + expr && expr->is_dynamic() + ? expr + : xla::DExpr::Unknown(xla::kUnknownContentSentinel)}; + xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(0), i); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(dim_size, content)); + } + xla::XlaOp converted = + xla::ConvertElementType(dim_size, ctx->output_xla_type(0)); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(converted, content)); + } + xla::XlaOp broadcast = xla::Broadcast(converted, {1}); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(broadcast, content)); + } + operands.push_back(broadcast); } - ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), operands, 0)); + xla::XlaOp concat = xla::ConcatInDim(ctx->builder(), operands, 0); + XlaExpression output = + XlaExpression::XlaOp(concat, ctx->expected_output_dtype(0)); + if (SymbolicContentEnabled()) { + output.set_contents(BuildShapeContents(input_shape)); + } + ctx->SetOutputExpression(0, output); } else { // Rank 0 won't have dynamic size dimension, use constant output. Tensor shape_constant(out_dtype_, TensorShape({input_shape.dims()})); OP_REQUIRES_OK(ctx, TensorShapeToConstant(input_shape, &shape_constant)); - ctx->SetConstantOutput(0, shape_constant); + XlaExpression output = XlaExpression::Constant(shape_constant); + if (SymbolicContentEnabled()) { + output.set_contents(BuildShapeContents(input_shape)); + } + ctx->SetOutputExpression(0, output); } } @@ -196,19 +238,47 @@ class ShapeNOp : public XlaOpKernel { // Each dimension can be dynamic, so use GetDimensionSize to get the // runtime dimension. for (int64_t dim = 0; dim < rank; ++dim) { - operands.push_back(xla::Broadcast( - xla::ConvertElementType(xla::GetDimensionSize(ctx->Input(i), dim), - ctx->output_xla_type(i)), - {1})); + xla::DExpr expr = input_shape.get_filled_expression(dim); + std::vector content = { + expr && expr->is_dynamic() + ? expr + : xla::DExpr::Unknown(xla::kUnknownContentSentinel)}; + xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(i), dim); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(dim_size, content)); + } + xla::XlaOp converted = + xla::ConvertElementType(dim_size, ctx->output_xla_type(i)); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(converted, content)); + } + xla::XlaOp broadcast = xla::Broadcast(converted, {1}); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(broadcast, content)); + } + operands.push_back(broadcast); } - ctx->SetOutput(i, xla::ConcatInDim(ctx->builder(), operands, 0)); + XlaExpression output = + XlaExpression::XlaOp(xla::ConcatInDim(ctx->builder(), operands, 0), + ctx->expected_output_dtype(i)); + if (SymbolicContentEnabled()) { + output.set_contents(BuildShapeContents(input_shape)); + } + ctx->SetOutputExpression(i, output); } else { // Rank 0 won't have dynamic size dimension, use constant output. Tensor shape_constant(out_dtype_, TensorShape({input_shape.dims()})); OP_REQUIRES_OK(ctx, TensorShapeToConstant(input_shape, &shape_constant)); - ctx->SetConstantOutput(i, shape_constant); + XlaExpression output = XlaExpression::Constant(shape_constant); + if (SymbolicContentEnabled()) { + output.set_contents(BuildShapeContents(input_shape)); + } + ctx->SetOutputExpression(i, output); } } } @@ -244,6 +314,7 @@ class SizeOp : public XlaOpKernel { const TensorShape input_shape = ctx->InputShape(0); xla::XlaBuilder* builder = ctx->builder(); auto size = xla::One(builder, ctx->output_xla_type(0)); + xla::DExpr size_expr = xla::DExpr::Const(1); const int rank = input_shape.dims(); for (int64_t dim = 0; dim < rank; ++dim) { @@ -256,12 +327,39 @@ class SizeOp : public XlaOpKernel { "on all dimensions, found ", input_shape.dim_size(dim), " elements on dimension ", dim))); - size = xla::Mul(size, xla::ConvertElementType( - xla::GetDimensionSize(ctx->Input(0), dim), - ctx->output_xla_type(0))); + xla::DExpr dim_expr = input_shape.get_filled_expression(dim); + std::vector contents = { + dim_expr && dim_expr->is_dynamic() + ? dim_expr + : xla::DExpr::Unknown(xla::kUnknownContentSentinel)}; + xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(0), dim); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK(ctx, + builder->SetInstructionContents(dim_size, contents)); + } + xla::XlaOp converted = + xla::ConvertElementType(dim_size, ctx->output_xla_type(0)); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK(ctx, + builder->SetInstructionContents(converted, contents)); + } + size = xla::Mul(size, converted); + size_expr = + (size_expr * input_shape.get_filled_expression(dim)).simplify(); + if (SymbolicContentEnabled()) { + OP_REQUIRES_OK(ctx, + builder->SetInstructionContents(size, {size_expr})); + } } - ctx->SetOutput(0, size); + if (SymbolicContentEnabled()) { + XlaExpression output = + XlaExpression::XlaOp(size, ctx->expected_output_dtype(0)); + output.set_contents({size_expr}); + ctx->SetOutputExpression(0, output); + } else { + ctx->SetOutput(0, size); + } } }; @@ -290,6 +388,8 @@ class ExpandDimsOp : public XlaOpKernel { " dimensions.")); auto existing_dims = input_shape.dim_sizes(); + auto existing_exprs = input_shape.get_filled_expressions(); + // Safe - # elements in tensor dims bounded. const int existing_dims_size = static_cast(existing_dims.size()); std::vector new_shape(existing_dims_size); @@ -297,6 +397,13 @@ class ExpandDimsOp : public XlaOpKernel { new_shape[i] = existing_dims[i]; } + const int existing_exprs_size = static_cast(existing_exprs.size()); + std::vector new_exprs; + new_exprs.reserve(existing_exprs_size); + for (size_t i = 0; i < existing_exprs.size(); ++i) { + new_exprs.push_back(existing_exprs[i]); + } + // We emulate numpy's interpretation of the dim axis when // -input.dims() >= dim <= input.dims(). if (dim < 0) { @@ -306,8 +413,8 @@ class ExpandDimsOp : public XlaOpKernel { // Clamp to the end if needed. dim = std::min(dim, existing_dims_size); new_shape.emplace(new_shape.begin() + dim, 1); - - ctx->SetOutput(0, xla::Reshape(ctx->Input("input"), new_shape)); + new_exprs.emplace(new_exprs.begin() + dim, xla::DExpr::Const(1)); + ctx->SetOutput(0, xla::Reshape(ctx->Input("input"), new_shape, new_exprs)); } }; REGISTER_XLA_OP(Name("ExpandDims").CompileTimeConstantInput("dim"), @@ -331,6 +438,7 @@ class SqueezeOp : public XlaOpKernel { absl::flat_hash_set wrapped_squeeze_dims; wrapped_squeeze_dims.reserve(squeeze_dims_.size()); std::vector new_shape; + std::vector new_exprs; // Validate squeeze dims against the input. for (int32_t dim : squeeze_dims_) { OP_REQUIRES( @@ -358,6 +466,7 @@ class SqueezeOp : public XlaOpKernel { } else { // This dimension is not being squeezed. new_shape.push_back(existing_dim); + new_exprs.push_back(shape.expressions(i)); } } else { OP_REQUIRES( @@ -368,11 +477,25 @@ class SqueezeOp : public XlaOpKernel { // Copy over all non-1-length dimensions. if (existing_dim != 1) { new_shape.push_back(existing_dim); + new_exprs.push_back(shape.expressions(i)); } } } - ctx->SetOutput(0, xla::Reshape(ctx->Input(0), new_shape)); + xla::XlaOp output = xla::Reshape(ctx->Input(0), new_shape, new_exprs); + const auto& input_contents = ctx->InputExpression(0).contents(); + if (SymbolicContentEnabled() && !input_contents.empty()) { + std::vector output_contents(input_contents.begin(), + input_contents.end()); + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(output, output_contents)); + XlaExpression output_expr = + XlaExpression::XlaOp(output, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } + ctx->SetOutput(0, output); } private: @@ -430,7 +553,8 @@ class ZerosLikeOp : public XlaOpKernel { auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0)); xla::XlaOp input = ctx->Input(0); auto input_shape = ctx->InputXlaShape(0).value(); - auto result = xla::Broadcast(zero, input_shape.dimensions()); + auto result = xla::Broadcast(zero, input_shape.dimensions(), + input_shape.expressions()); // Setting up dynamic dimensions of the broadcast. for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) { @@ -455,7 +579,8 @@ class OnesLikeOp : public XlaOpKernel { const TensorShape input_shape = ctx->InputShape(0); auto one = XlaHelpers::One(ctx->builder(), input_type(0)); - ctx->SetOutput(0, xla::Broadcast(one, input_shape.dim_sizes())); + ctx->SetOutput(0, xla::Broadcast(one, input_shape.dim_sizes(), + input_shape.get_filled_expressions())); } }; diff --git a/tensorflow/compiler/tf2xla/kernels/slice_op.cc b/tensorflow/compiler/tf2xla/kernels/slice_op.cc index 844a31f97990fc..5d3d9023bf5a02 100644 --- a/tensorflow/compiler/tf2xla/kernels/slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/slice_op.cc @@ -18,8 +18,11 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/container/inlined_vector.h" #include "absl/types/span.h" +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -37,11 +40,42 @@ limitations under the License. namespace tensorflow { namespace { +bool TryBuildSlicedContents(const XlaExpression& input_expr, + const TensorShape& input_shape, + absl::Span begin, + absl::Span size, + std::vector* output_contents) { + output_contents->clear(); + const auto& input_contents = input_expr.contents(); + if (input_contents.empty() || input_shape.dims() != 1 || begin.size() != 1 || + size.size() != 1) { + return false; + } + const int64_t start = begin[0]; + const int64_t count = + size[0] == -1 ? input_shape.dim_size(0) - start : size[0]; + for (int64_t i = 0; i < count; ++i) { + const int64_t index = start + i; + if (index < 0 || index >= input_contents.size()) { + output_contents->clear(); + return false; + } + const xla::DExpr& expr = input_contents[index]; + output_contents->push_back( + expr ? expr : xla::DExpr::Unknown(xla::kUnknownContentSentinel)); + } + return absl::c_any_of(*output_contents, [](const xla::DExpr& expr) { + return expr && expr->is_dynamic(); + }); +} + class SliceOp : public XlaOpKernel { public: explicit SliceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { + const bool enable_dynamic_sizes = + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; const TensorShape input_shape = ctx->InputShape(0); const TensorShape begin_tensor_shape = ctx->InputShape(1); const TensorShape size_tensor_shape = ctx->InputShape(2); @@ -66,13 +100,17 @@ class SliceOp : public XlaOpKernel { ctx->ConstantInputAsIntVector(2, &size).ok(); if (all_begins_are_constant && all_sizes_are_constant) { std::vector wrapped_size(size.size()); + std::vector wrapped_size_exprs(size.size()); // `begin` is a compile-time constant. for (int i = 0; i < input_dims; ++i) { if (size[i] == -1) { // A size[i] of -1 means "all elements from begin[i] to dim_size(i)". wrapped_size[i] = input_shape.dim_size(i) - begin[i]; + wrapped_size_exprs[i] = + (input_shape.get_filled_expression(i) - begin[i]).simplify(); } else { wrapped_size[i] = size[i]; + wrapped_size_exprs[i] = xla::DExpr::Const(size[i]); } } @@ -97,13 +135,32 @@ class SliceOp : public XlaOpKernel { } } + std::vector begin_exprs; + for (int d : begin){ + begin_exprs.push_back(xla::DExpr::Const(d)); + } std::vector limits; + std::vector exprs; limits.reserve(begin.size()); + exprs.reserve(begin.size()); for (int i = 0; i < begin.size(); ++i) { limits.push_back(begin[i] + wrapped_size[i]); + exprs.push_back(begin_exprs[i] + wrapped_size_exprs[i]); } std::vector strides(begin.size(), 1); - auto slice = xla::Slice(ctx->Input(0), begin, limits, strides); + auto slice = enable_dynamic_sizes + ? xla::Slice(ctx->Input(0), begin, limits, begin_exprs, + exprs, strides) + : xla::Slice(ctx->Input(0), begin, limits, strides); + std::vector output_contents; + const bool has_output_contents = + SymbolicContentEnabled() && + TryBuildSlicedContents(ctx->InputExpression(0), input_shape, begin, + size, &output_contents); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(slice, output_contents)); + } // Check for slice on dynamic dimensions. std::vector size_is_dynamic; OP_REQUIRES_OK( @@ -114,13 +171,27 @@ class SliceOp : public XlaOpKernel { if (size[i] != -1) { // If there is a dynamic dimension, properly set dimension size of // the slice. - auto dynamic_size = - xla::Reshape(xla::Slice(ctx->Input(2), {i}, {i + 1}, {1}), {}); + auto dynamic_size = xla::Reshape( + xla::Slice(ctx->Input(2), {i}, {i + 1}, {xla::DExpr::Const(i)}, + {xla::DExpr::Const(i + 1)}, {1}), + {}); slice = xla::SetDimensionSize(slice, dynamic_size, i); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, + ctx->builder()->SetInstructionContents(slice, output_contents)); + } } } } + if (has_output_contents) { + auto output_expr = + XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } ctx->SetOutput(0, slice); } else { // When a size is -1, we take rest of the dimension according to @@ -153,7 +224,14 @@ class SliceOp : public XlaOpKernel { } if (all_sizes_are_constant && !constant_size_is_minus_one) { xla::XlaOp input = ctx->Input(0); - ctx->SetOutput(0, xla::DynamicSlice(input, begin_indices, size)); + std::vector output_exprs; + output_exprs.reserve(size.size()); + for (int64_t d : size) { + output_exprs.push_back(xla::DExpr::Const(d)); + } + ctx->SetOutput(0, + xla::DynamicSlice(input, begin_indices, size, + output_exprs)); } else { // Size is not constant, use input size as upperbound and then set // dimension size on it. diff --git a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc index 330479bc8d4150..b1d7e16bbda741 100644 --- a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc @@ -35,11 +35,72 @@ limitations under the License. #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/bcast.h" +#include "tensorflow/compiler/tf2xla/type_util.h" + namespace tensorflow { namespace { -REGISTER_XLA_OP(Name("Softmax"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("LogSoftmax"), MlirXlaOpKernel); +class SoftmaxOp : public XlaOpKernel { + public: + explicit SoftmaxOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + log_ = absl::StartsWith(type_string(), "Log"); + } + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape logits_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, TensorShapeUtils::IsVectorOrHigher(logits_shape), + errors::InvalidArgument("logits must have >= 1 dimension, got ", + logits_shape.DebugString())); + + // Major dimensions are batch dimensions, minor dimension is the class + // dimension. + std::vector batch_dims(logits_shape.dims() - 1); + std::iota(batch_dims.begin(), batch_dims.end(), 0); + const int kClassDim = logits_shape.dims() - 1; + + const DataType type = input_type(0); + const xla::PrimitiveType xla_type = ctx->input_xla_type(0); + auto logits = ctx->Input(0); + + xla::XlaBuilder* const b = ctx->builder(); + + const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type); + + // Find the max in each batch, resulting in a tensor of shape [batch] + auto logits_max = + xla::Reduce(logits, xla::MinValue(b, xla_type), max_func, {kClassDim}); + // Subtract the max in batch b from every element in batch b. Broadcasts + // along the batch dimension. + auto shifted_logits = xla::Sub(logits, logits_max, batch_dims); + auto exp_shifted = xla::Exp(shifted_logits); + const DataType accumulation_type = XlaHelpers::SumAccumulationType(type); + xla::PrimitiveType xla_accumulation_type; + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(accumulation_type, + &xla_accumulation_type)); + auto converted = + xla::ConvertElementType(exp_shifted, xla_accumulation_type); + auto reduce = + xla::Reduce(converted, xla::Zero(b, xla_accumulation_type), + *ctx->GetOrCreateAdd(accumulation_type), {kClassDim}); + auto sum = XlaHelpers::ConvertElementType(reduce, type); + auto softmax = + log_ + // softmax = shifted_logits - log(sum(exp(shifted_logits))) + ? xla::Sub(shifted_logits, xla::Log(sum), batch_dims) + // softmax = exp(shifted_logits) / sum(exp(shifted_logits)) + : xla::Div(exp_shifted, sum, batch_dims); + ctx->SetOutput(0, softmax); + } + + private: + bool log_; +}; + +REGISTER_XLA_OP(Name("Softmax"), SoftmaxOp); +REGISTER_XLA_OP(Name("LogSoftmax"), SoftmaxOp); std::pair CrossEntropyWithLogits( XlaOpKernelContext* ctx, DataType type, xla::PrimitiveType xla_type, diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc index d4a93e0556143d..e10c009a5b7d16 100644 --- a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc @@ -44,6 +44,8 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, const int input_rank = input_tensor_shape.dims(); const absl::InlinedVector input_shape = input_tensor_shape.dim_sizes(); + const std::vector input_exprs = + input_tensor_shape.get_filled_expressions(); const int block_rank = block_shape.size(); OP_REQUIRES( @@ -68,6 +70,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, // input according to `paddings` to produce `padded` of shape `padded_shape`. xla::PaddingConfig padding_config; std::vector padded_shape(input_shape.begin(), input_shape.end()); + std::vector padded_exprs(input_exprs.begin(), input_exprs.end()); int64_t block_num_elems = 1LL; padding_config.add_dimensions(); // Don't pad the batch dimension. for (int i = 0; i < block_rank; ++i) { @@ -83,6 +86,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, dim->set_edge_padding_low(pad_start); dim->set_edge_padding_high(pad_end); padded_shape[1 + i] += pad_start + pad_end; + padded_exprs[1 + i] = (padded_exprs[1 + i] + pad_start + pad_end).simplify(); block_num_elems = MultiplyWithoutOverflow(block_num_elems, block_shape[i]); } // Don't pad the remainder dimensions. @@ -116,7 +120,9 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, // block_shape[M-1]] + // remaining_shape std::vector reshaped_padded_shape(input_rank + block_rank); + std::vector reshaped_padded_exprs(input_rank + block_rank); reshaped_padded_shape[0] = batch_size; + reshaped_padded_exprs[0] = padded_exprs[0]; for (int i = 0; i < block_rank; ++i) { OP_REQUIRES(ctx, padded_shape[1 + i] % block_shape[i] == 0, errors::InvalidArgument("padded_shape[", 1 + i, @@ -126,11 +132,17 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, reshaped_padded_shape[1 + i * 2] = padded_shape[1 + i] / block_shape[i]; reshaped_padded_shape[1 + i * 2 + 1] = block_shape[i]; + reshaped_padded_exprs[1 + i * 2] = + (padded_exprs[1 + i] / block_shape[i]).simplify(); + reshaped_padded_exprs[1 + i * 2 + 1] = xla::DExpr::Const(block_shape[i]); } std::copy(remainder_shape.begin(), remainder_shape.end(), reshaped_padded_shape.begin() + 1 + 2 * block_rank); + std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(), + reshaped_padded_exprs.begin() + 1 + 2 * block_rank); - xla::XlaOp reshaped_padded = xla::Reshape(padded, reshaped_padded_shape); + xla::XlaOp reshaped_padded = + xla::Reshape(padded, reshaped_padded_shape, reshaped_padded_exprs); // 3. Permute dimensions of `reshaped_padded` to produce // `permuted_reshaped_padded` of shape: @@ -163,14 +175,21 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input, // Determine the length of the prefix of block dims that can be combined // into the batch dimension due to having no padding and block_shape=1. std::vector output_shape(input_rank); + std::vector output_exprs(input_rank); output_shape[0] = output_dim; + output_exprs[0] = (input_exprs[0] * xla::DExpr::Const(block_num_elems)).simplify(); for (int i = 0; i < block_rank; ++i) { output_shape[1 + i] = padded_shape[1 + i] / block_shape[i]; + output_exprs[1 + i] = + (padded_exprs[1 + i] / block_shape[i]).simplify(); } std::copy(remainder_shape.begin(), remainder_shape.end(), output_shape.begin() + 1 + block_rank); + std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(), + output_exprs.begin() + 1 + block_rank); - xla::XlaOp output = xla::Reshape(permuted_reshaped_padded, output_shape); + xla::XlaOp output = + xla::Reshape(permuted_reshaped_padded, output_shape, output_exprs); ctx->SetOutput(0, output); } diff --git a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc index ac33e0877200dc..d09fa5f4daacbb 100644 --- a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc @@ -67,6 +67,7 @@ class SpaceToDepthOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, input_xla_shape.status()); absl::Span input_shape = input_xla_shape.value().dimensions(); + const xla::Shape& input_shape_with_exprs = input_xla_shape.value(); int input_rank = input_shape.size(); static const int kRequiredDims = 4; @@ -80,9 +81,13 @@ class SpaceToDepthOp : public XlaOpKernel { std::vector reshaped_shape; std::vector transpose_order; std::vector output_shape; + std::vector reshaped_exprs; + std::vector output_exprs; reshaped_shape.reserve(input_rank); transpose_order.reserve(input_rank); output_shape.reserve(input_rank); + reshaped_exprs.reserve(input_rank + num_spatial_dims); + output_exprs.reserve(input_rank); if (data_format == FORMAT_NHWC) { int64_t block_elems = 1; for (int i = 0; i < num_spatial_dims; ++i) { @@ -94,11 +99,18 @@ class SpaceToDepthOp : public XlaOpKernel { } reshaped_shape.push_back(input_shape[0]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(0)); for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(input_shape[1 + i] / block_size_); + reshaped_exprs.push_back( + (input_shape_with_exprs.expressions(1 + i) / + xla::DExpr::Const(block_size_)) + .simplify()); reshaped_shape.push_back(block_size_); + reshaped_exprs.push_back(xla::DExpr::Const(block_size_)); } reshaped_shape.push_back(input_shape[feature_dim]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(feature_dim)); transpose_order.push_back(0); for (int i = 0; i < num_spatial_dims; ++i) { @@ -110,10 +122,19 @@ class SpaceToDepthOp : public XlaOpKernel { transpose_order.push_back(feature_dim + num_spatial_dims); output_shape.push_back(input_shape[0]); + output_exprs.push_back(input_shape_with_exprs.expressions(0)); for (int i = 0; i < num_spatial_dims; ++i) { output_shape.push_back(input_shape[1 + i] / block_size_); + output_exprs.push_back( + (input_shape_with_exprs.expressions(1 + i) / + xla::DExpr::Const(block_size_)) + .simplify()); } output_shape.push_back(input_shape[feature_dim] * block_elems); + output_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) * + xla::DExpr::Const(block_elems)) + .simplify()); } else { // FORMAT_NCHW int64_t block_elems = 1; @@ -126,10 +147,17 @@ class SpaceToDepthOp : public XlaOpKernel { } reshaped_shape.push_back(input_shape[0]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(0)); reshaped_shape.push_back(input_shape[feature_dim]); + reshaped_exprs.push_back(input_shape_with_exprs.expressions(feature_dim)); for (int i = 0; i < num_spatial_dims; ++i) { reshaped_shape.push_back(input_shape[2 + i] / block_size_); + reshaped_exprs.push_back( + (input_shape_with_exprs.expressions(2 + i) / + xla::DExpr::Const(block_size_)) + .simplify()); reshaped_shape.push_back(block_size_); + reshaped_exprs.push_back(xla::DExpr::Const(block_size_)); } transpose_order.push_back(0); @@ -142,9 +170,18 @@ class SpaceToDepthOp : public XlaOpKernel { } output_shape.push_back(input_shape[0]); + output_exprs.push_back(input_shape_with_exprs.expressions(0)); output_shape.push_back(input_shape[feature_dim] * block_elems); + output_exprs.push_back( + (input_shape_with_exprs.expressions(feature_dim) * + xla::DExpr::Const(block_elems)) + .simplify()); for (int i = 0; i < num_spatial_dims; ++i) { output_shape.push_back(input_shape[2 + i] / block_size_); + output_exprs.push_back( + (input_shape_with_exprs.expressions(2 + i) / + xla::DExpr::Const(block_size_)) + .simplify()); } } @@ -156,7 +193,7 @@ class SpaceToDepthOp : public XlaOpKernel { // input_shape[1] / block_size_, block_size_, // input_shape[2] / block_size_, block_size_, // depth] - xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape); + xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs); // 2. Permute dimensions of `reshaped` to produce // `permuted_reshaped` of shape: @@ -176,7 +213,7 @@ class SpaceToDepthOp : public XlaOpKernel { // input_shape[2] / block_size_, // block_size_ * block_size_ * depth] // - xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape); + xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape, output_exprs); // If this used to be a vectorized format turn it back now. if (data_format != data_format_) { diff --git a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc index b4d589f183108e..709d5d8c28d17c 100644 --- a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc @@ -83,7 +83,8 @@ class SparseToDenseOp : public XlaOpKernel { sparse_values = Broadcast(sparse_values, {num_elems}); } xla::XlaBuilder* builder = context->builder(); - auto buffer = Broadcast(default_value, output_shape.dim_sizes()); + auto buffer = Broadcast(default_value, output_shape.dim_sizes(), + output_shape.get_filled_expressions()); std::vector dynamic_dims; OP_REQUIRES_OK( context, context->ResolveInputDynamismIntoPredVector(1, &dynamic_dims)); diff --git a/tensorflow/compiler/tf2xla/kernels/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc index 4f7c4ae99b6b6b..dcecb71e11e7d7 100644 --- a/tensorflow/compiler/tf2xla/kernels/split_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_builder.h" @@ -35,6 +36,8 @@ class SplitOp : public XlaOpKernel { explicit SplitOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { + const bool enable_dynamic_sizes = + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; const int32_t num_split = num_outputs(); const TensorShape split_dim_shape = ctx->InputShape("split_dim"); const TensorShape input_shape = ctx->InputShape(1); @@ -81,17 +84,25 @@ class SplitOp : public XlaOpKernel { // All the slices are the same size: this is the size along the // split dimension. const int32_t slice_size = input_shape.dim_size(split_dim) / num_split; + xla::DExpr slice_expr = + input_shape.get_filled_expression(split_dim) / num_split; // The vectors we will use to define the slice. The entry for the // split dimensions varies for each output. std::vector begin(input_shape.dims(), 0); std::vector limits(input_shape.dims()); + std::vector begin_expr; + std::vector limits_expr; + begin_expr.reserve(input_shape.dims()); + limits_expr.reserve(input_shape.dims()); std::vector strides(input_shape.dims(), 1); for (int i = 0; i < input_shape.dims(); ++i) { // Initially set up the limits to be the full size of the input: // the split dimension is filled in below. int64_t dim = input_shape.dim_size(i); limits[i] = dim; + begin_expr.push_back(xla::DExpr::Const(0)); + limits_expr.push_back(input_shape.get_filled_expression(i)); } // Create each of the outputs. @@ -99,7 +110,15 @@ class SplitOp : public XlaOpKernel { // Slice out the ith split from the split dimension. begin[split_dim] = i * slice_size; limits[split_dim] = (i + 1) * slice_size; - ctx->SetOutput(i, xla::Slice(input, begin, limits, strides)); + + if (enable_dynamic_sizes) { + begin_expr[split_dim] = (i * slice_expr).simplify(); + limits_expr[split_dim] = ((i + 1) * slice_expr).simplify(); + ctx->SetOutput(i, xla::Slice(input, begin, limits, begin_expr, + limits_expr, strides)); + } else { + ctx->SetOutput(i, xla::Slice(input, begin, limits, strides)); + } } } }; @@ -111,6 +130,8 @@ class SplitVOp : public XlaOpKernel { explicit SplitVOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { + const bool enable_dynamic_sizes = + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; const int32_t num_split = num_outputs(); const TensorShape input_shape = ctx->InputShape(0); const TensorShape index_shape = ctx->InputShape(2); @@ -202,21 +223,38 @@ class SplitVOp : public XlaOpKernel { input_shape.dim_size(split_dim) - total_split_size; } - // The vectors we will use to define the slice. The entry for the - // split dimensions varies for each output. + // The vectors we will use to define the slice. The entry for the split + // dimension varies for each output. std::vector begin(input_shape.dims(), 0); auto dim_sizes = input_shape.dim_sizes(); std::vector limits(dim_sizes.begin(), dim_sizes.end()); std::vector strides(input_shape.dims(), 1); + std::vector begin_expr(input_shape.dims(), xla::DExpr::Const(0)); + auto input_exprs = input_shape.get_filled_expressions(); + std::vector limits_expr; + limits_expr.reserve(input_exprs.size()); + for (auto expr : input_exprs) { + limits_expr.push_back(expr); + } for (int i = 0; i < num_split; ++i) { - TensorShape output_shape(input_shape); int slice_size = split_sizes[i]; - output_shape.set_dim(split_dim, slice_size); + xla::DExpr slice_expr = xla::DExpr::Const(slice_size); // Slice out the ith split from the split dimension. limits[split_dim] = begin[split_dim] + slice_size; - ctx->SetOutput(i, xla::Slice(input, begin, limits, strides)); + if (enable_dynamic_sizes) { + limits_expr[split_dim] = + (begin_expr[split_dim] + slice_expr).simplify(); + ctx->SetOutput(i, + xla::Slice(input, begin, limits, begin_expr, + limits_expr, strides)); + } else { + ctx->SetOutput(i, xla::Slice(input, begin, limits, strides)); + } begin[split_dim] = limits[split_dim]; + if (enable_dynamic_sizes) { + begin_expr[split_dim] = limits_expr[split_dim]; + } } } }; diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index 3c99ad63565266..28ca29002084c4 100644 --- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc @@ -157,7 +157,8 @@ class StackPushOp : public XlaOpKernel { TensorShape slice_shape = elem_shape; slice_shape.InsertDim(0, 1LL); - auto update = xla::Reshape(value, slice_shape.dim_sizes()); + auto update = xla::Reshape(value, slice_shape.dim_sizes(), + slice_shape.get_filled_expressions()); // TODO(phawkins): We don't check the index is in bounds --- there is no // error mechanism in XLA. diff --git a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc index aa71c5c34d2e1a..9e3d5b6fa6aad5 100644 --- a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc @@ -49,8 +49,9 @@ xla::BitGeneratorTy GetBitGeneratorForDevice( device_type_string == DEVICE_CPU_XLA_JIT) { return [=](xla::XlaOp key, xla::XlaOp state, const xla::Shape& shape) { std::tie(state, key) = xla::ScramblePhiloxKey(key); - xla::XlaOp philox_state = - xla::ConcatInDim(key.builder(), {xla::Reshape(key, {1}), state}, 0); + xla::XlaOp philox_state = xla::ConcatInDim( + key.builder(), + {xla::Reshape(key, {1}, {xla::DExpr::Const(1)}), state}, 0); xla::XlaOp result = xla::RngBitGenerator(xla::RandomAlgorithm::RNG_PHILOX, philox_state, shape); return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1), @@ -420,20 +421,23 @@ class StatelessParameterizedTruncatedNormalOp : public XlaOpKernel { auto rng_dtype = MaybeConvertBF16ToF32(dtype_); xla::Shape xla_shape; OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(rng_dtype, shape, &xla_shape)); - - auto bcasted_means = BroadcastTo(ctx->Input(2), shape.dim_sizes()); + auto bcasted_means = + BroadcastTo(ctx->Input(2), shape.dim_sizes(), shape.get_filled_expressions()); OP_REQUIRES_OK(ctx, bcasted_means.status()); auto means = bcasted_means.value(); - auto bcasted_stddevs = BroadcastTo(ctx->Input(3), shape.dim_sizes()); + auto bcasted_stddevs = + BroadcastTo(ctx->Input(3), shape.dim_sizes(), shape.get_filled_expressions()); OP_REQUIRES_OK(ctx, bcasted_stddevs.status()); auto stddevs = bcasted_stddevs.value(); - auto bcasted_minvals = BroadcastTo(ctx->Input(4), shape.dim_sizes()); + auto bcasted_minvals = + BroadcastTo(ctx->Input(4), shape.dim_sizes(), shape.get_filled_expressions()); OP_REQUIRES_OK(ctx, bcasted_minvals.status()); auto minvals = bcasted_minvals.value(); - auto bcasted_maxvals = BroadcastTo(ctx->Input(5), shape.dim_sizes()); + auto bcasted_maxvals = + BroadcastTo(ctx->Input(5), shape.dim_sizes(), shape.get_filled_expressions()); OP_REQUIRES_OK(ctx, bcasted_maxvals.status()); auto maxvals = bcasted_maxvals.value(); diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc index e15196bd756462..8fd205097fe8ef 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -24,7 +24,9 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/types/span.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/literal_util.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" @@ -48,6 +50,37 @@ namespace tensorflow { namespace { using errors::InvalidArgument; +bool TryBuildSlicedContents(const XlaExpression& input_expr, + const TensorShape& input_shape, + const absl::InlinedVector& begin, + const absl::InlinedVector& strides, + const TensorShape& final_shape, + std::vector* output_contents) { + output_contents->clear(); + const auto& input_contents = input_expr.contents(); + if (input_contents.empty() || input_shape.dims() != 1 || begin.size() != 1 || + strides.size() != 1) { + return false; + } + + const int64_t output_elements = final_shape.num_elements(); + const int64_t start = begin[0]; + const int64_t stride = strides[0]; + for (int64_t i = 0; i < output_elements; ++i) { + const int64_t index = start + i * stride; + if (index < 0 || index >= input_contents.size()) { + output_contents->clear(); + return false; + } + const xla::DExpr& expr = input_contents[index]; + output_contents->push_back( + expr ? expr : xla::DExpr::Unknown(xla::kUnknownContentSentinel)); + } + return absl::c_any_of(*output_contents, [](const xla::DExpr& expr) { + return expr && expr->is_dynamic(); + }); +} + class StridedSliceOp : public XlaOpKernel { public: explicit StridedSliceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -82,6 +115,9 @@ class StridedSliceOp : public XlaOpKernel { partial_final_shape.set_dim( i, input_shape.dim_size(shape_spec.output_to_processing_mapping[i])); + partial_final_shape.set_expression( + i, input_shape.get_expression( + shape_spec.output_to_processing_mapping[i])); } } @@ -97,6 +133,8 @@ class StridedSliceOp : public XlaOpKernel { // Use input shape to update unknown dimension of partial shape -- if a // dimension is unknown, we use input shape as bound. partial_processing_shape.set_dim(i, input_shape.dim_size(i)); + partial_processing_shape.set_expression(i, + input_shape.get_expression(i)); } } TensorShape processing_shape; @@ -157,7 +195,10 @@ class StridedSliceOp : public XlaOpKernel { auto zero = xla::Zero(ctx->builder(), ctx->InputXlaType("begin")); xla::XlaOp begin_index, end_index; int64_t sparse_index = shape_spec.processing_to_sparse_mapping[i]; - bool xla_input_is_dynamic = input_xla_shape.is_dynamic_dimension(i); + const xla::DExpr& input_expr = input_xla_shape.expressions(i); + bool xla_input_is_dynamic = + input_xla_shape.is_dynamic_dimension(i) || + input_expr->is_dynamic(); xla::XlaOp dim_size; if (xla_input_is_dynamic) { dim_size = xla::GetDimensionSize(ctx->Input(0), i); @@ -215,10 +256,12 @@ class StridedSliceOp : public XlaOpKernel { } slice = - xla::DynamicSlice(slice, start_indices, processing_shape.dim_sizes()); + xla::DynamicSlice(slice, start_indices, processing_shape.dim_sizes(), + processing_shape.get_expressions()); // new_axis_mask_, ellipsis_mask_ and shrink_axis_mask_ may add or remove // size 1 dims of a shape. - slice = xla::Reshape(slice, final_shape.dim_sizes()); + slice = xla::Reshape(slice, final_shape.dim_sizes(), + final_shape.get_expressions()); for (int64_t i = 0; i < final_shape.dims(); ++i) { int64 processing_shape_dim = shape_spec.output_to_processing_mapping[i]; // If processing_shape_dim is -1, it means the output dimension was newly @@ -238,6 +281,8 @@ class StridedSliceOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { + const bool enable_dynamic_sizes = + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; const TensorShape input_shape = ctx->InputShape(0); const TensorShape begin_shape = ctx->InputShape("begin"); OP_REQUIRES( @@ -246,6 +291,8 @@ class StridedSliceOp : public XlaOpKernel { absl::InlinedVector begin; absl::InlinedVector end; + absl::InlinedVector begin_expr; + absl::InlinedVector end_expr; absl::InlinedVector strides; xla::Literal begin_literal, end_literal, strides_literal; @@ -268,14 +315,15 @@ class StridedSliceOp : public XlaOpKernel { PartialTensorShape partial_processing_shape, partial_final_shape; bool dummy = false; StridedSliceShapeSpec shape_spec; + OP_REQUIRES_OK( - ctx, - ValidateStridedSliceOp( - begin_is_constant ? &begin_tensor : nullptr, - end_is_constant ? &end_tensor : nullptr, strides_tensor, - input_shape, begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, - shrink_axis_mask_, &partial_processing_shape, &partial_final_shape, - &dummy, &dummy, &dummy, &begin, &end, &strides, &shape_spec)); + ctx, ValidateStridedSliceOp( + begin_is_constant ? &begin_tensor : nullptr, + end_is_constant ? &end_tensor : nullptr, strides_tensor, + input_shape, begin_mask_, end_mask_, ellipsis_mask_, + new_axis_mask_, shrink_axis_mask_, &partial_processing_shape, + &partial_final_shape, &dummy, &dummy, &dummy, &begin, &end, + &strides, &begin_expr, &end_expr, &shape_spec)); xla::XlaOp slice = ctx->Input(0); std::vector begins_are_dynamic; @@ -294,17 +342,30 @@ class StridedSliceOp : public XlaOpKernel { ", output shape must be a compile-time constant")); absl::InlinedVector dimensions_to_reverse; absl::InlinedVector slice_begin, slice_end, slice_strides; + absl::InlinedVector slice_begin_expr, slice_end_expr; for (int i = 0; i < begin.size(); ++i) { if (strides[i] > 0) { slice_begin.push_back(begin[i]); + slice_begin_expr.push_back(begin_expr[i]); slice_end.push_back(std::max(end[i], begin[i])); + slice_end_expr.push_back( + xla::DExpr::Max(end_expr[i], begin_expr[i]).simplify()); slice_strides.push_back(strides[i]); } else { // Negative stride: swap begin and end, add 1 because the interval // is semi-open, and mark the dimension to be reversed. + xla::DExpr input_expr = input_shape.get_filled_expression(i); slice_begin.push_back(input_shape.dim_size(i) - begin[i] - 1); + slice_begin_expr.push_back( + (input_expr - begin_expr[i] - xla::DExpr::Const(1)).simplify()); slice_end.push_back(std::max(input_shape.dim_size(i) - end[i] - 1, input_shape.dim_size(i) - begin[i] - 1)); + slice_end_expr.push_back( + xla::DExpr::Max( + (input_expr - end_expr[i] - xla::DExpr::Const(1)).simplify(), + (input_expr - begin_expr[i] - xla::DExpr::Const(1)) + .simplify()) + .simplify()); slice_strides.push_back(-strides[i]); dimensions_to_reverse.push_back(i); } @@ -312,7 +373,40 @@ class StridedSliceOp : public XlaOpKernel { if (!dimensions_to_reverse.empty()) { slice = xla::Rev(slice, dimensions_to_reverse); } - slice = xla::Slice(slice, slice_begin, slice_end, slice_strides); + for (int i = 0; i < partial_processing_shape.dims(); ++i) { + partial_processing_shape.set_expression( + i, ((slice_end_expr[i] - slice_begin_expr[i] + + xla::DExpr::Const(slice_strides[i]) - xla::DExpr::Const(1)) / + xla::DExpr::Const(slice_strides[i])) + .simplify()); + } + for (int i = 0; i < partial_final_shape.dims(); ++i) { + int64_t processing_index = shape_spec.output_to_processing_mapping[i]; + partial_final_shape.set_expression( + i, processing_index == -1 + ? xla::DExpr::Const(partial_final_shape.dim_size(i)) + : partial_processing_shape.get_filled_expression( + processing_index)); + } + OP_REQUIRES( + ctx, partial_final_shape.AsTensorShape(&final_shape), + InvalidArgument("XLA can't deduce compile time constant output " + "shape for strided slice: ", + partial_final_shape.DebugString(), + ", output shape must be a compile-time constant")); + slice = enable_dynamic_sizes + ? xla::Slice(slice, slice_begin, slice_end, slice_begin_expr, + slice_end_expr, slice_strides) + : xla::Slice(slice, slice_begin, slice_end, slice_strides); + std::vector output_contents; + const bool has_output_contents = + SymbolicContentEnabled() && + TryBuildSlicedContents(ctx->InputExpression(0), input_shape, begin, + strides, final_shape, &output_contents); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(slice, output_contents)); + } auto operand_shape_or = ctx->builder()->GetShape(ctx->Input(0)); OP_REQUIRES_OK(ctx, operand_shape_or.status()); xla::Shape xla_shape = operand_shape_or.value(); @@ -325,8 +419,20 @@ class StridedSliceOp : public XlaOpKernel { bool ends_are_static = absl::c_all_of( ends_are_dynamic, [](bool dynamic) { return !dynamic; }); // Static output shape, return a static slice. - slice = xla::Reshape(slice, final_shape.dim_sizes()); + slice = xla::Reshape(slice, final_shape.dim_sizes(), + final_shape.get_expressions()); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(slice, output_contents)); + } if (xla_shape.is_static() && ends_are_static) { + if (has_output_contents) { + auto output_expr = + XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } ctx->SetOutput(0, slice); return; } @@ -385,8 +491,20 @@ class StridedSliceOp : public XlaOpKernel { xla::Sub(operand_size, xla::ConstantR0( ctx->builder(), begin[input_index])), i); + if (has_output_contents) { + OP_REQUIRES_OK( + ctx, ctx->builder()->SetInstructionContents(slice, + output_contents)); + } } } + if (has_output_contents) { + auto output_expr = + XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } ctx->SetOutput(0, slice); return; } else { @@ -431,11 +549,15 @@ class StridedSliceGradOp : public XlaOpKernel { void CompileAsDynamicUpdateSlice(XlaOpKernelContext* ctx, const TensorShape& input_shape, const xla::Literal& strides_literal) { + const bool enable_dynamic_sizes = + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; bool dummy = false; Tensor strides_tensor; PartialTensorShape processing_shape, final_shape; absl::InlinedVector begin; absl::InlinedVector end; + absl::InlinedVector begin_expr; + absl::InlinedVector end_expr; absl::InlinedVector strides; StridedSliceShapeSpec shape_spec; OP_REQUIRES_OK(ctx, LiteralToHostTensor(strides_literal, index_type_, @@ -445,7 +567,7 @@ class StridedSliceGradOp : public XlaOpKernel { nullptr, nullptr, strides_tensor, input_shape, begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_, &processing_shape, &final_shape, &dummy, &dummy, &dummy, - &begin, &end, &strides, &shape_spec)); + &begin, &end, &strides, &begin_expr, &end_expr, &shape_spec)); for (int64_t i = 0; i < processing_shape.dims(); ++i) { OP_REQUIRES( ctx, strides[i] == 1, @@ -459,14 +581,22 @@ class StridedSliceGradOp : public XlaOpKernel { VLOG(1) << "xla final_shape" << final_shape; VLOG(1) << "input_shape" << input_shape.DebugString(); auto input_sizes = input_shape.dim_sizes(); + std::vector input_exprs; + input_exprs.reserve(input_shape.dims()); + for (int64_t i = 0; i < input_shape.dims(); ++i) { + input_exprs.push_back(input_shape.get_filled_expression(i)); + } // For unknown output dim the bound of the output shape is input. Pad and // double the size of input shape to leave enough buffer to avoid OOB // dynamic update slice. auto input_sizes_padded = input_shape.dim_sizes(); + auto input_exprs_padded = input_exprs; bool need_padding = false; for (int64_t i = 0; i < processing_shape.dims(); ++i) { if (processing_shape.dim_size(i) == -1) { input_sizes_padded[i] *= 2; + input_exprs_padded[i] = + (xla::DExpr::Const(2) * input_exprs_padded[i]).simplify(); need_padding = true; } } @@ -477,6 +607,7 @@ class StridedSliceGradOp : public XlaOpKernel { if (shape_spec.output_to_processing_mapping[i] != -1) { processing_shape.set_dim(shape_spec.output_to_processing_mapping[i], grad_shape.dimensions(i)); + // TODO Pass it back } } @@ -506,15 +637,21 @@ class StridedSliceGradOp : public XlaOpKernel { } auto zero = XlaHelpers::Zero(ctx->builder(), ctx->expected_output_dtype(0)); - zero = xla::Broadcast(zero, input_sizes_padded); - grad = xla::Reshape(grad, processing_shape.dim_sizes()); + zero = xla::Broadcast(zero, input_sizes_padded, input_exprs_padded); + grad = xla::Reshape(grad, processing_shape.dim_sizes(), + processing_shape.get_expressions()); grad = xla::DynamicUpdateSlice(zero, grad, begins); if (need_padding) { // We padded the input shape to avoid OOB when DUS. Now slice out the // padding in the final result. std::vector strides(input_shape.dims(), 1); std::vector start_indices(input_shape.dims(), 0); - grad = xla::Slice(grad, start_indices, input_sizes, strides); + std::vector start_exprs(input_shape.dims(), + xla::DExpr::Const(0)); + grad = enable_dynamic_sizes + ? xla::Slice(grad, start_indices, input_sizes, start_exprs, + input_exprs, strides) + : xla::Slice(grad, start_indices, input_sizes, strides); } ctx->SetOutput(0, grad); } @@ -522,6 +659,8 @@ class StridedSliceGradOp : public XlaOpKernel { TensorShape processing_shape, final_shape; absl::InlinedVector begin; absl::InlinedVector end; + absl::InlinedVector begin_expr; + absl::InlinedVector end_expr; absl::InlinedVector strides; TensorShape input_shape; @@ -547,11 +686,13 @@ class StridedSliceGradOp : public XlaOpKernel { bool dummy = false; OP_REQUIRES_OK( - ctx, ValidateStridedSliceOp( - &begin_tensor, &end_tensor, strides_tensor, input_shape, - begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, - shrink_axis_mask_, &processing_shape, &final_shape, &dummy, - &dummy, &dummy, &begin, &end, &strides)); + ctx, + ValidateStridedSliceOp( + &begin_tensor, &end_tensor, strides_tensor, input_shape, + begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, + shrink_axis_mask_, &processing_shape, &final_shape, &dummy, &dummy, + &dummy, &begin, &end, &strides, &begin_expr, &end_expr, + /*shape_spec=*/nullptr)); // Check to make sure dy is consistent with the original slice const TensorShape dy_shape = ctx->InputShape(4); @@ -570,7 +711,8 @@ class StridedSliceGradOp : public XlaOpKernel { xla::XlaOp grad = ctx->Input(4); // Undo any new/shrink axes. - grad = xla::Reshape(grad, processing_shape.dim_sizes()); + grad = xla::Reshape(grad, processing_shape.dim_sizes(), + processing_shape.get_expressions()); // Pad the input gradients. absl::InlinedVector dimensions_to_reverse; @@ -662,6 +804,8 @@ class StridedSliceAssignOp : public XlaOpKernel { TensorShape final_shape; absl::InlinedVector begin; absl::InlinedVector end; + absl::InlinedVector begin_expr; + absl::InlinedVector end_expr; absl::InlinedVector strides; xla::Literal begin_literal, end_literal, strides_literal; @@ -690,12 +834,14 @@ class StridedSliceAssignOp : public XlaOpKernel { TensorShape dummy_processing_shape; bool dummy = false; - OP_REQUIRES_OK(ctx, - ValidateStridedSliceOp( - &begin_tensor, &end_tensor, strides_tensor, lhs_shape, - begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, - shrink_axis_mask_, &dummy_processing_shape, &final_shape, - &dummy, &dummy, &dummy, &begin, &end, &strides)); + OP_REQUIRES_OK( + ctx, + ValidateStridedSliceOp( + &begin_tensor, &end_tensor, strides_tensor, lhs_shape, begin_mask_, + end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_, + &dummy_processing_shape, &final_shape, &dummy, &dummy, &dummy, + &begin, &end, &strides, &begin_expr, &end_expr, + /*shape_spec=*/nullptr)); if (final_shape.num_elements() == 0 && rhs_shape.num_elements() == 0) { // DynamicUpdateSlice does not allow 0-element updates. We should probably @@ -717,6 +863,7 @@ class StridedSliceAssignOp : public XlaOpKernel { absl::InlinedVector dimensions_to_reverse; absl::InlinedVector slice_begin; absl::InlinedVector slice_dims; + absl::InlinedVector slice_exprs; for (int i = 0; i < begin.size(); ++i) { // TODO(b/121179231): implement strides != 1 OP_REQUIRES( @@ -726,12 +873,14 @@ class StridedSliceAssignOp : public XlaOpKernel { slice_begin.push_back( xla::ConstantR0(ctx->builder(), begin[i])); slice_dims.push_back(end[i] - begin[i]); + slice_exprs.push_back(xla::DExpr::Const(end[i] - begin[i])); } else { // Negative stride: swap begin and end, add 1 because the interval // is semi-open, and mark the dimension to be reversed. slice_begin.push_back( xla::ConstantR0(ctx->builder(), end[i] + 1)); slice_dims.push_back(begin[i] - end[i]); + slice_exprs.push_back(xla::DExpr::Const(begin[i] - end[i])); dimensions_to_reverse.push_back(i); } } @@ -739,7 +888,7 @@ class StridedSliceAssignOp : public XlaOpKernel { if (!dimensions_to_reverse.empty()) { rhs = xla::Rev(rhs, dimensions_to_reverse); } - rhs = xla::Reshape(rhs, slice_dims); + rhs = xla::Reshape(rhs, slice_dims, slice_exprs); lhs = xla::DynamicUpdateSlice(lhs, rhs, slice_begin); diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 888908e30b2331..9bc5d0d95bfd8c 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -165,9 +165,11 @@ class TensorArrayOp : public XlaOpKernel { CHECK(element_shape_.AsTensorShape(&shape)); TensorShape ta_shape; ta_shape.AddDim(size); + ta_shape.AddExpression(xla::DExpr::Const(size)); ta_shape.AppendShape(shape); xla::XlaOp zero = XlaHelpers::Zero(b, dtype_); - value = xla::Broadcast(zero, ta_shape.dim_sizes()); + value = xla::Broadcast(zero, ta_shape.dim_sizes(), + ta_shape.get_filled_expressions()); } XlaResource* var = @@ -223,7 +225,8 @@ class TensorArrayWriteOp : public XlaOpKernel { TensorShape slice_shape = elem_shape; slice_shape.InsertDim(0, 1LL); - auto update = xla::Reshape(value, slice_shape.dim_sizes()); + auto update = xla::Reshape(value, slice_shape.dim_sizes(), + slice_shape.get_filled_expressions()); xla::XlaOp written; if (resource->tensor_array_multiple_writes_aggregate()) { @@ -274,9 +277,12 @@ class TensorArrayReadOp : public XlaOpKernel { start_indices[0] = index; auto slice_shape = ta_shape.dim_sizes(); + auto slice_exprs = ta_shape.get_filled_expressions(); slice_shape[0] = 1LL; + slice_exprs[0] = xla::DExpr::Const(1LL); - xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape); + xla::XlaOp read = + xla::DynamicSlice(ta, start_indices, slice_shape, slice_exprs); // Remove the leading '1' dimension. std::vector value_shape(slice_shape.begin() + 1, @@ -469,9 +475,16 @@ class TensorArrayConcatOp : public XlaOpKernel { xla::XlaOp ta = resource->value(); auto ta_dims = ta_shape.dim_sizes(); + auto ta_exprs = ta_shape.get_filled_expressions(); std::vector shape(ta_dims.begin() + 1, ta_dims.end()); + std::vector exprs; + exprs.reserve(ta_exprs.size() - 1); + for (auto it = ta_exprs.begin() + 1; it != ta_exprs.end(); ++it) { + exprs.push_back(*it); + } shape[0] *= ta_shape.dim_size(0); - ctx->SetOutput(0, xla::Reshape(ta, shape)); + exprs[0] = ta_exprs[0] * ta_shape.get_filled_expression(0); + ctx->SetOutput(0, xla::Reshape(ta, shape, exprs)); Tensor lengths(DT_INT64, {ta_dims[0]}); auto lengths_vec = lengths.vec(); @@ -526,6 +539,7 @@ class TensorArraySplitOp : public XlaOpKernel { TensorShape ta_shape; ta_shape.AddDim(resource->max_array_size()); + ta_shape.AddExpression(xla::DExpr::Const(resource->max_array_size())); ta_shape.AppendShape(elem_shape); OP_REQUIRES(ctx, lengths.size() == resource->max_array_size(), @@ -541,7 +555,8 @@ class TensorArraySplitOp : public XlaOpKernel { value_shape.DebugString(), " vs. ", ta_shape.DebugString())); - const xla::XlaOp reshape = xla::Reshape(value, ta_shape.dim_sizes()); + const xla::XlaOp reshape = + xla::Reshape(value, ta_shape.dim_sizes(), ta_shape.get_filled_expressions()); if (dtype_ == DT_BOOL) { ta = xla::Or(ta, reshape); } else { diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index a1f58d5ae9b40e..92128c0dc19873 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -131,7 +131,8 @@ absl::Status TryGetElementShapeFromInput(XlaOpKernelContext* ctx, return absl::OkStatus(); } - *shape = xla::ShapeUtil::MakeShape(dtype, partial_shape.dim_sizes()); + *shape = xla::ShapeUtil::MakeShape(dtype, partial_shape.dim_sizes(), + partial_shape.get_filled_expressions()); *got_shape = true; return absl::OkStatus(); } @@ -503,6 +504,8 @@ class TensorListConcatOp : public XlaOpKernel { xla::Shape element_shape = std::move(shape_or).value(); std::vector element_dims = xla::SpanToVector(element_shape.dimensions()); + std::vector element_exprs(element_shape.expressions().begin(), + element_shape.expressions().end()); OP_REQUIRES( ctx, element_dims.size() > 1, errors::Unimplemented("TensorList of scalars is not supported")); @@ -510,12 +513,15 @@ class TensorListConcatOp : public XlaOpKernel { int64_t tensor_lengths = element_dims[1]; std::vector new_dims = {num_elements * tensor_lengths}; + std::vector new_exprs = { + xla::DExpr::Const(num_elements * tensor_lengths)}; for (int i = 2; i < element_dims.size(); i++) { new_dims.push_back(element_dims[i]); + new_exprs.push_back(element_exprs[i]); } - xla::XlaOp out = xla::Reshape(buffer, new_dims); + xla::XlaOp out = xla::Reshape(buffer, new_dims, new_exprs); ctx->SetOutput(0, out); // Second output is a tensor of lengths of returned tensors. @@ -550,6 +556,8 @@ class TensorListSplitOp : public XlaOpKernel { xla::Shape element_shape = std::move(shape_or).value(); std::vector element_dims = xla::SpanToVector(element_shape.dimensions()); + std::vector element_exprs(element_shape.expressions().begin(), + element_shape.expressions().end()); OP_REQUIRES( ctx, !element_dims.empty(), errors::Unimplemented("Element dimensions have to be non-empty")); @@ -569,11 +577,13 @@ class TensorListSplitOp : public XlaOpKernel { ctx, element_dims[0] % length == 0, errors::Unimplemented("Buffer size has to be a multiple of length")); std::vector new_dims = {element_dims[0] / length, length}; + std::vector new_exprs = {element_exprs[0] / xla::DExpr::Const(length), + xla::DExpr::Const(length)}; for (int i = 1; i < element_dims.size(); i++) { new_dims.push_back(element_dims[i]); } - xla::XlaOp reshaped = xla::Reshape(input_tensor, new_dims); + xla::XlaOp reshaped = xla::Reshape(input_tensor, new_dims, new_exprs); xla::XlaOp result; OP_REQUIRES_OK(ctx, ExecuteTensorListFromTensor(length, reshaped, &result)); diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc index 683dc4737e6dab..412c01f24dfc19 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc @@ -241,9 +241,12 @@ absl::Status GetTensorListShapeFromElementTensorListShape( const xla::Shape& shape = xla::ShapeUtil::GetTupleElementShape(element_tensor_list_shape, i); std::vector dimensions = xla::SpanToVector(shape.dimensions()); + std::vector expressions = xla::SpanToVector(shape.expressions()); dimensions.insert(dimensions.begin(), leading_dim); + expressions.insert(expressions.begin(), xla::DExpr::Const(leading_dim)); shapes.push_back( - xla::ShapeUtil::MakeShape(shape.element_type(), dimensions)); + xla::ShapeUtil::MakeShape(shape.element_type(), dimensions, + expressions)); if (leading_dim_is_dynamic) { shapes.back().set_dynamic_dimension(0, true); } @@ -267,9 +270,13 @@ absl::Status GetTensorListShapeFromElementShape(const xla::Shape& element_shape, std::vector shapes; std::vector dimensions = xla::SpanToVector(element_shape.dimensions()); + std::vector expressions = + xla::SpanToVector(element_shape.expressions()); dimensions.insert(dimensions.begin(), leading_dim); + expressions.insert(expressions.begin(), xla::DExpr::Const(leading_dim)); shapes.push_back( - xla::ShapeUtil::MakeShape(element_shape.element_type(), dimensions)); + xla::ShapeUtil::MakeShape(element_shape.element_type(), dimensions, + expressions)); shapes.back().set_dynamic_dimension(0, leading_dim_is_dynamic); shapes.push_back(xla::ShapeUtil::MakeShape(xla::PrimitiveType::S32, std::vector{})); @@ -289,7 +296,8 @@ absl::Status CreateZerosTensorListWithShape( xla::ShapeUtil::GetTupleElementShape(list_shape, i); xla::XlaOp zero = xla::ConstantLiteral(b, xla::LiteralUtil::Zero(shape.element_type())); - xla::XlaOp zeros = xla::Broadcast(zero, shape.dimensions()); + xla::XlaOp zeros = + xla::Broadcast(zero, shape.dimensions(), shape.expressions()); TF_RET_CHECK(dynamic_dims[i].size() == shape.dimensions().size()); for (int64_t dim = 0; dim < shape.dimensions().size(); ++dim) { if (shape.is_dynamic_dimension(dim)) { @@ -389,7 +397,14 @@ absl::Status ExecuteTensorListPushBack(xla::XlaOp list, xla::XlaOp element, std::vector element_part_dims = xla::SpanToVector(element_part_shape.dimensions()); element_part_dims.insert(element_part_dims.begin(), 1); - element_part = xla::Reshape(element_part, element_part_dims); + std::vector element_part_exprs; + element_part_exprs.reserve(element_part_shape.expressions().size() + 1); + element_part_exprs.push_back(xla::DExpr::Const(1)); + for (auto expr : element_part_shape.expressions()) { + element_part_exprs.push_back(expr); + } + element_part = + xla::Reshape(element_part, element_part_dims, element_part_exprs); std::vector start_indices( element_part_shape.dimensions().size() + 1, @@ -406,7 +421,13 @@ absl::Status ExecuteTensorListPushBack(xla::XlaOp list, xla::XlaOp element, std::vector element_dims = xla::SpanToVector(element_shape.dimensions()); element_dims.insert(element_dims.begin(), 1); - xla::XlaOp update = xla::Reshape(element, element_dims); + std::vector element_exprs; + element_exprs.reserve(element_shape.expressions().size() + 1); + element_exprs.push_back(xla::DExpr::Const(1)); + for (auto expr : element_shape.expressions()) { + element_exprs.push_back(expr); + } + xla::XlaOp update = xla::Reshape(element, element_dims, element_exprs); std::vector start_indices(element_shape.dimensions().size() + 1, xla::ConstantR0(b, 0)); @@ -455,11 +476,19 @@ absl::Status ExecuteTensorListPopBack(xla::XlaOp list, xla::XlaOp* list_result, xla::SpanToVector(list_part_shape.dimensions()); slice_shape[0] = 1LL; + std::vector slice_exprs; + slice_exprs.reserve(list_part_shape.expressions().size()); + for (auto expr : list_part_shape.expressions()) { + slice_exprs.push_back(expr); + } + slice_exprs[0] = xla::DExpr::Const(1LL); + xla::XlaOp list_part = xla::GetTupleElement(list, i); xla::XlaOp read = xla::DynamicSlice(list_part, start_indices, slice_shape); slice_shape.erase(slice_shape.begin()); - element_result_parts.push_back(xla::Reshape(read, slice_shape)); + slice_exprs.erase(slice_exprs.begin()); + element_result_parts.push_back(xla::Reshape(read, slice_shape, slice_exprs)); list_result_parts.push_back(list_part); } list_result_parts.push_back(push_index); @@ -493,7 +522,13 @@ absl::Status ExecuteTensorListSetItem(xla::XlaOp list, xla::XlaOp index, std::vector element_dims = xla::SpanToVector(element_shape.dimensions()); element_dims.insert(element_dims.begin(), 1); - xla::XlaOp update = xla::Reshape(element, element_dims); + std::vector element_exprs; + element_exprs.reserve(element_shape.expressions().size() + 1); + element_exprs.push_back(xla::DExpr::Const(1)); + for (auto expr : element_shape.expressions()) { + element_exprs.push_back(expr); + } + xla::XlaOp update = xla::Reshape(element, element_dims, element_exprs); std::vector start_indices(element_shape.dimensions().size() + 1, xla::ConstantR0(b, 0)); @@ -557,6 +592,13 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index, xla::SpanToVector(buffer_shape.dimensions()); slice_shape[0] = 1LL; + std::vector slice_exprs; + slice_exprs.reserve(buffer_shape.expressions().size()); + for (auto expr : buffer_shape.expressions()) { + slice_exprs.push_back(expr); + } + slice_exprs[0] = xla::DExpr::Const(1LL); + xla::XlaOp list_part = xla::GetTupleElement(list, 0); xla::XlaOp read = xla::DynamicSlice(list_part, start_indices, slice_shape); // Propagate dynamic dimensions from buffer to the sliced buffer, except for @@ -569,7 +611,8 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index, } } slice_shape.erase(slice_shape.begin()); - *result = xla::Reshape(read, slice_shape); + slice_exprs.erase(slice_exprs.begin()); + *result = xla::Reshape(read, slice_shape, slice_exprs); return absl::OkStatus(); } diff --git a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc index 6c39981ba5b937..10752ebb0f8c25 100644 --- a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc @@ -67,17 +67,35 @@ class TileOp : public XlaOpKernel { xla::ValueInferenceMode::kUpperBound)); std::vector output_dims(input_shape.dims()); + std::vector output_exprs(input_shape.dims()); + + auto expr_sizes = input_shape.get_filled_expressions(); + const auto& multiples_contents = + ctx->InputExpression("multiples").contents(); + for (int64_t i = 0; i < input_shape.dims(); ++i) { OP_REQUIRES(ctx, multiples_bounds[i] >= 0, errors::InvalidArgument("Expected multiples[", i, "] >= 0, but got ", output_dims[i])); output_dims[i] = input_shape.dim_size(i) * multiples_bounds[i]; + const xla::DExpr multiple_expr = + i < multiples_contents.size() && multiples_contents[i] && + multiples_contents[i]->is_dynamic() + ? multiples_contents[i] + : xla::DExpr::Const(multiples_bounds[i]); + output_exprs[i] = (expr_sizes[i] * multiple_expr).simplify(); } std::vector multiples_are_dynamic; OP_REQUIRES_OK(ctx, ctx->ResolveInputDynamismIntoPredVector( 1, &multiples_are_dynamic)); + for (int64_t i = 0; i < multiples_are_dynamic.size(); ++i) { + if (i < multiples_contents.size() && multiples_contents[i] && + multiples_contents[i]->is_dynamic()) { + multiples_are_dynamic[i] = true; + } + } bool all_multiples_are_static = absl::c_all_of( multiples_are_dynamic, [](bool dynamic) { return !dynamic; }); @@ -91,8 +109,8 @@ class TileOp : public XlaOpKernel { return; } } - - auto result_or = BroadcastTo(ctx->Input("input"), output_dims); + auto result_or = + BroadcastTo(ctx->Input("input"), output_dims, output_exprs); OP_REQUIRES_OK(ctx, result_or.status()); auto result = result_or.value(); diff --git a/tensorflow/compiler/tf2xla/kernels/unique_op.cc b/tensorflow/compiler/tf2xla/kernels/unique_op.cc index 46de3dd89b6115..6590dd2f55cf78 100644 --- a/tensorflow/compiler/tf2xla/kernels/unique_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unique_op.cc @@ -83,9 +83,12 @@ class UniqueOpBase : public XlaOpKernel { // // This is implemented as an hlo while loop. xla::XlaOp RollingSelectR1(XlaOpKernelContext* ctx, xla::XlaOp data, - xla::XlaOp mask, int64_t size) { + xla::XlaOp mask, int64_t size, + const xla::DExpr& expr) { xla::XlaComputation cond, body; - const xla::Shape r1_shape = xla::ShapeUtil::MakeShape(xla::S32, {size}); + xla::Shape r1_shape = xla::ShapeUtil::MakeShape(xla::S32, {size}); + r1_shape.set_expression(0, expr); + const xla::Shape counter_shape = xla::ShapeUtil::MakeScalarShape(xla::S32); const xla::Shape& single_element_shape = counter_shape; @@ -136,7 +139,7 @@ class UniqueOpBase : public XlaOpKernel { } auto zero = xla::Zero(ctx->builder(), xla::S32); - auto zero_broadcast = xla::Broadcast(zero, {size}); + auto zero_broadcast = xla::Broadcast(zero, {size}, {expr}); auto init = xla::Tuple(ctx->builder(), {zero, data, mask, zero_broadcast}); return xla::GetTupleElement(xla::While(cond, body, init), 3); } @@ -153,13 +156,19 @@ class UniqueOpBase : public XlaOpKernel { auto aux = MoveAxis(input, axis, 0, input_shape); auto aux_shape = ctx->builder()->GetShape(aux).value(); int64_t leading_size = aux_shape.dimensions(0); + auto leading_expr = aux_shape.expressions(0); int64_t product = 1; + auto product_expr = xla::DExpr::Const(1); for (int64_t i = 1; i < aux_shape.dimensions().size(); ++i) { product *= aux_shape.dimensions(i); + product_expr = product_expr * aux_shape.expressions(i); } - aux = xla::Reshape(aux, {leading_size, product}); + product_expr = product_expr.simplify(); + aux = xla::Reshape(aux, {leading_size, product}, + {leading_expr, product_expr}); if (leading_size == 0) { - auto result_data = xla::Reshape(aux, aux_shape.dimensions()); + auto result_data = + xla::Reshape(aux, aux_shape.dimensions(), aux_shape.expressions()); result_data = MoveAxis(result_data, 0, axis, aux_shape); ctx->SetOutput(0, result_data); ctx->SetOutput(1, xla::Iota(ctx->builder(), xla::S32, leading_size)); @@ -169,12 +178,18 @@ class UniqueOpBase : public XlaOpKernel { sort_keys.reserve(product + 1); std::vector sort_types; sort_types.reserve(product + 1); + xla::Shape leading_shape = xla::ShapeUtil::MakeShape( + input_shape.element_type(), {leading_size}, + std::vector{leading_expr}); for (int64_t i = 0; i < product; ++i) { xla::XlaOp slice = xla::SliceInDim(aux, i, i + 1, 1, 1); - sort_keys.push_back(xla::Reshape(slice, {leading_size})); + sort_keys.push_back(xla::Reshape(leading_shape, slice)); sort_types.push_back(input_shape.element_type()); } - auto iota = xla::Iota(ctx->builder(), xla::S32, leading_size); + xla::Shape iota_shape = xla::ShapeUtil::MakeShape( + xla::S32, {leading_size}, std::vector{leading_expr}); + iota_shape.set_expression(0, leading_expr); + auto iota = xla::Iota(ctx->builder(), iota_shape, 0); sort_keys.push_back(iota); sort_types.push_back(xla::S32); @@ -202,16 +217,20 @@ class UniqueOpBase : public XlaOpKernel { gather_dim_numbers.add_collapsed_slice_dims(0); auto permuted = xla::Gather(aux, perm, gather_dim_numbers, {1, product}); // Tail is everything except for first element. - auto tail = xla::SliceInDim(permuted, 1, leading_size, 1, 0); + auto tail = xla::SliceInDim(permuted, 1, leading_size, + xla::DExpr::Const(1), leading_expr, 1, 0); // Init is everything except for last element. - auto init = xla::SliceInDim(permuted, 0, leading_size - 1, 1, 0); + auto init = xla::SliceInDim(permuted, 0, leading_size - 1, + xla::DExpr::Const(0), + (leading_expr - xla::DExpr::Const(1)).simplify(), + 1, 0); auto ne = xla::Compare(tail, init, xla::ComparisonDirection::kNe); auto reduce = xla::Reduce(ne, xla::ConstantR0(ctx->builder(), false), CreateScalarOrComputation(xla::PRED, ctx->builder()), {1}); auto mask = xla::ConvertElementType(reduce, xla::S32); mask = xla::PadInDim(mask, xla::One(ctx->builder(), xla::S32), 0, 1, 0); - auto iperm = RollingSelectR1(ctx, perm, mask, leading_size); + auto iperm = RollingSelectR1(ctx, perm, mask, leading_size, leading_expr); auto sort_by_iperm = xla::Sort({iperm, mask, perm}, @@ -232,12 +251,14 @@ class UniqueOpBase : public XlaOpKernel { /*is_stable=*/true); auto mask_permute = xla::GetTupleElement(mask_sort, 1); permuted = xla::Gather(aux, mask_permute, gather_dim_numbers, {1, product}); - auto result_data = xla::Reshape(permuted, aux_shape.dimensions()); + auto result_data = xla::Reshape(aux_shape, permuted); result_data = MoveAxis(result_data, 0, axis, aux_shape); result_data = xla::SetDimensionSize(result_data, dynamic_size, axis); ctx->SetOutput(0, result_data); auto imask = CumSumR1(ctx, mask, leading_size); - imask = xla::Sub(imask, xla::One(ctx->builder(), xla::S32), {}); + auto one = xla::One(ctx->builder(), xla::S32); + auto one_broadcast = xla::Broadcast(one, {leading_size}, {leading_expr}); + imask = xla::Sub(imask, one_broadcast, {}); auto idx = xla::GetTupleElement( xla::Sort({perm_sort, imask}, xla::CreateScalarLtComputation({xla::S32, xla::S32}, diff --git a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc index cca29f7f585907..55899c7f7b7d95 100644 --- a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc @@ -60,16 +60,24 @@ class UnpackOp : public XlaOpKernel { std::vector start_indices(input_shape.dims(), 0); std::vector limit_indices(input_shape.dims()); std::vector strides(input_shape.dims(), 1); + std::vector start_exprs(input_shape.dims(), xla::DExpr::Const(0)); + std::vector limit_exprs; + limit_exprs.reserve(input_shape.dims()); for (int i = 0; i < input_shape.dims(); ++i) { limit_indices[i] = input_shape.dim_size(i); + limit_exprs.push_back(input_shape.get_filled_expression(i)); } for (int i = 0; i < num; ++i) { start_indices[axis] = i; limit_indices[axis] = i + 1; - auto slice = xla::Slice(input, start_indices, limit_indices, strides); + start_exprs[axis] = xla::DExpr::Const(i); + limit_exprs[axis] = xla::DExpr::Const(i + 1); + auto slice = xla::Slice(input, start_indices, limit_indices, start_exprs, + limit_exprs, strides); // Reshape to drop the 'axis' dimension. - auto result = xla::Reshape(slice, output_shape.dim_sizes()); + auto result = xla::Reshape(slice, output_shape.dim_sizes(), + output_shape.get_filled_expressions()); ctx->SetOutput(i, result); } } diff --git a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc index a7a1a438f95b9e..b588525f3831a0 100644 --- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" #include "tensorflow/compiler/tf2xla/kernels/shape_util.h" #include "tensorflow/compiler/tf2xla/lib/scatter.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" @@ -50,6 +51,19 @@ absl::Status ValidateAssignUpdateVariableOpShapes(XlaOpKernelContext* ctx) { return absl::OkStatus(); } +std::vector BuildVariableShapeContents(const TensorShape& shape) { + std::vector contents; + contents.reserve(shape.dims()); + for (int i = 0; i < shape.dims(); ++i) { + xla::DExpr expr = shape.get_filled_expression(i); + contents.push_back(expr && expr->is_dynamic() + ? expr + : xla::DExpr::Unknown( + xla::kUnknownContentSentinel)); + } + return contents; +} + class VarIsInitializedOp : public XlaOpKernel { public: explicit VarIsInitializedOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} @@ -75,7 +89,11 @@ class VariableShapeOp : public XlaOpKernel { ctx->GetVariableTypeAndShape(0, &variable_dtype, &shape)); Tensor shape_constant(out_dtype_, TensorShape({shape.dims()})); OP_REQUIRES_OK(ctx, TensorShapeToConstant(shape, &shape_constant)); - ctx->SetConstantOutput(0, shape_constant); + auto output = XlaExpression::Constant(shape_constant); + if (SymbolicContentEnabled()) { + output.set_contents(BuildVariableShapeContents(shape)); + } + ctx->SetOutputExpression(0, output); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/where_op.cc b/tensorflow/compiler/tf2xla/kernels/where_op.cc index f97e6d5077efa7..4920a106808605 100644 --- a/tensorflow/compiler/tf2xla/kernels/where_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/where_op.cc @@ -162,10 +162,16 @@ absl::StatusOr CompileWhereWithSort(XlaOpKernelContext* ctx) { TF_ASSIGN_OR_RETURN(xla::Shape input_shape, ctx->builder()->GetShape(condition)); auto iota_shape = - xla::ShapeUtil::MakeShape(xla::S32, input_shape.dimensions()); + xla::ShapeUtil::MakeShape(xla::S32, input_shape.dimensions(), + input_shape.expressions()); int64_t flattened_size = xla::Product(iota_shape.dimensions()); - XlaOp reshaped_condition = xla::Reshape(condition, {flattened_size}); + xla::DExpr flattened_expr = xla::DExpr::Const(1); + for (auto e : iota_shape.expressions()){ + flattened_expr = flattened_expr * e; + } + XlaOp reshaped_condition = + xla::Reshape(condition, {flattened_size}, {flattened_expr}); XlaOp zeros = xla::ZerosLike(reshaped_condition); XlaOp compared = xla::Ne(reshaped_condition, zeros); @@ -175,7 +181,7 @@ absl::StatusOr CompileWhereWithSort(XlaOpKernelContext* ctx) { // indices of each element. for (int64_t axis = 0; axis < iota_shape.dimensions_size(); ++axis) { XlaOp iota = xla::Iota(ctx->builder(), iota_shape, axis); - XlaOp reshaped = xla::Reshape(iota, {flattened_size}); + XlaOp reshaped = xla::Reshape(iota, {flattened_size}, {flattened_expr}); to_sort.push_back(reshaped); types_to_sort.push_back(xla::S32); } @@ -186,7 +192,9 @@ absl::StatusOr CompileWhereWithSort(XlaOpKernelContext* ctx) { std::vector to_concat; for (int64_t i = 0; i < iota_shape.dimensions_size(); ++i) { XlaOp index_single_dim = xla::GetTupleElement(sorted, i + 1); - to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1})); + to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1}, + {flattened_expr, + xla::DExpr::Const(1)})); } XlaOp result = xla::ConcatInDim(ctx->builder(), to_concat, 1); @@ -214,7 +222,12 @@ absl::StatusOr CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) { TF_ASSIGN_OR_RETURN(xla::Shape input_shape, b->GetShape(condition)); int64_t flattened_size = xla::Product(input_shape.dimensions()); - XlaOp reshaped_condition = xla::Reshape(condition, {flattened_size}); + xla::DExpr flattened_expr = xla::DExpr::Const(1); + for (auto e : input_shape.expressions()) { + flattened_expr = flattened_expr * e; + } + XlaOp reshaped_condition = + xla::Reshape(condition, {flattened_size}, {flattened_expr}); XlaOp zeros = xla::ZerosLike(reshaped_condition); XlaOp preds = xla::ConvertElementType(xla::Ne(reshaped_condition, zeros), S32); @@ -276,11 +289,16 @@ absl::StatusOr CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) { // // and then scatter iotas[out_idxs] into the output. std::vector iotas_to_concat; - auto iota_shape = xla::ShapeUtil::MakeShape(S32, input_shape.dimensions()); + auto iota_shape = xla::ShapeUtil::MakeShape( + S32, input_shape.dimensions(), input_shape.expressions()); iotas_to_concat.reserve(iota_shape.dimensions_size()); for (int64_t axis = 0; axis < iota_shape.dimensions_size(); ++axis) { - iotas_to_concat.push_back( - xla::Reshape(xla::Iota(b, iota_shape, axis), {flattened_size, 1})); + XlaOp flattened_iota = + xla::Reshape(xla::Iota(b, iota_shape, axis), {flattened_size}, + {flattened_expr}); + iotas_to_concat.push_back(xla::Reshape( + flattened_iota, {flattened_size, 1}, + {flattened_expr, xla::DExpr::Const(1)})); } XlaOp iotas = xla::ConcatInDim(b, iotas_to_concat, /*dimension=*/1); @@ -305,7 +323,12 @@ absl::StatusOr CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) { XlaOp scattered = xla::Scatter( /*input=*/xla::Zeros( b, /*shape=*/xla::ShapeUtil::MakeShape( - S32, {flattened_size, iota_shape.dimensions_size()})), + S32, + std::vector{flattened_size, + iota_shape.dimensions_size()}, + std::vector{ + flattened_expr, + xla::DExpr::Const(iota_shape.dimensions_size())})), /*scatter_indices=*/out_idxs, /*updates=*/iotas, /*update_computation=*/assn_computation, scatter_dnums, /*indices_are_sorted=*/true, /*unique_indices=*/true); diff --git a/tensorflow/compiler/tf2xla/lib/data_format.cc b/tensorflow/compiler/tf2xla/lib/data_format.cc index 2473b97af4c2bd..b59dbf59e54120 100644 --- a/tensorflow/compiler/tf2xla/lib/data_format.cc +++ b/tensorflow/compiler/tf2xla/lib/data_format.cc @@ -53,7 +53,13 @@ absl::StatusOr Contract(xla::XlaOp input, int64_t dim) { input_shape.dimensions().end() - 1); contracted_shape[dim] *= 4; - return xla::Reshape(xla::Transpose(input, permutation), contracted_shape); + std::vector contracted_exprs( + input_shape.expressions().begin(), input_shape.expressions().end() - 1); + contracted_exprs[dim] = + contracted_exprs[dim] * xla::DExpr::Const(4); + + return xla::Reshape(xla::Transpose(input, permutation), contracted_shape, + contracted_exprs); } absl::StatusOr Expand(xla::XlaOp input, int64_t dim) { @@ -85,7 +91,7 @@ absl::StatusOr Expand(xla::XlaOp input, int64_t dim) { } permutation.push_back(dim + 1); - return xla::Transpose(xla::Reshape(input, expanded_shape), permutation); + return xla::Transpose(xla::Reshape(input, expanded_shape, {}), permutation); } } // namespace diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc index 4523126f9a71bd..7a2cf9d55922af 100644 --- a/tensorflow/core/kernels/strided_slice_op.cc +++ b/tensorflow/core/kernels/strided_slice_op.cc @@ -309,6 +309,8 @@ class StridedSliceAssignOp : public OpKernel { bool is_simple_slice = true; absl::InlinedVector begin; absl::InlinedVector end; + absl::InlinedVector begin_expr; + absl::InlinedVector end_expr; absl::InlinedVector strides; Tensor* old_lhs = nullptr; @@ -353,7 +355,7 @@ class StridedSliceAssignOp : public OpKernel { old_lhs->shape(), begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask, &processing_shape, &final_shape, &is_identity, &is_simple_slice, &slice_dim0, - &begin, &end, &strides, &shape_spec)); + &begin, &end, &strides, &begin_expr, &end_expr, &shape_spec)); if (processing_shape.num_elements() > 0) { const Tensor& input = context->input(4); From 95402c00cb7e5ae10150c8936a6a49cf8f13a6a5 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:03:50 +0100 Subject: [PATCH 08/21] Preserve dynamic expressions through binary broadcasting Derive binary-op output expressions from the operand XLA shapes during broadcasting. Preserve the non-singleton expression when an operand expands from a singleton dimension. --- .../compiler/tf2xla/kernels/binary_ops.cc | 253 ++++++++++++----- .../tf2xla/kernels/broadcast_to_op.cc | 3 +- .../compiler/tf2xla/kernels/cwise_ops.cc | 264 +++++++++++++++++- .../compiler/tf2xla/kernels/cwise_ops.h | 12 +- 4 files changed, 453 insertions(+), 79 deletions(-) diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc index e9f571d830d619..34762dfeb9b545 100644 --- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc @@ -42,34 +42,46 @@ namespace { // A subclass of a XlaBinaryOp must build the computation that // describes the (tensor,tensor)->tensor function to apply to each element of // the input. -#define XLA_MAKE_BINARY(NAME, HLO) \ +#define XLA_MAKE_BINARY(NAME, HLO, SYMBOLIC_HLO) \ class NAME##Op : public XlaBinaryOp { \ public: \ explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {} \ xla::XlaOp Computation( \ XlaOpKernelContext* ctx, const xla::XlaOp& lhs, \ - const absl::Span& lhs_shape, const xla::XlaOp& rhs, \ + const absl::Span& lhs_shape, \ + const xla::XlaOp& rhs, \ const absl::Span& rhs_shape, \ const BCast& broadcast_helper, \ + const absl::Span& broadcast_output_exprs, \ const std::vector& extend_dimensions) override { \ xla::XlaBuilder* b = ctx->builder(); \ (void)b; \ (void)lhs_shape; \ (void)rhs_shape; \ + (void)broadcast_output_exprs; \ (void)extend_dimensions; \ return HLO; \ } \ + xla::DExpr SymbolicComputation(const xla::DExpr& lhs, \ + const xla::DExpr& rhs) override { \ + return SYMBOLIC_HLO; \ + } \ }; \ REGISTER_XLA_OP(Name(#NAME), NAME##Op) -XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(AddV2, xla::Add(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions), + (lhs + rhs).simplify()); +XLA_MAKE_BINARY(AddV2, xla::Add(lhs, rhs, extend_dimensions), + (lhs + rhs).simplify()); +XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions), + (lhs - rhs).simplify()); +XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions), + (lhs * rhs).simplify()); +XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions), + (lhs / rhs).simplify()); -XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions), xla::DExpr()); +XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions), xla::DExpr()); // Implementation of DivNoNan. Pseudo-code: // if (y == 0) { @@ -78,8 +90,12 @@ XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions)); // return x / y; // } static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, + const BCast& broadcast_helper) { + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto zero = XlaHelpers::Zero(b, dtype); auto y_equals_0 = xla::Eq(y, zero); auto zeros = xla::ZerosLike(x); @@ -87,7 +103,9 @@ static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, return result; } XLA_MAKE_BINARY(DivNoNan, - DivNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper)); + DivNoNanImpl(b, input_type(0), lhs, rhs, + broadcast_output_exprs, broadcast_helper), + xla::DExpr()); // Implementation of MulNoNan. Pseudo-code: // if (y == 0) { @@ -96,8 +114,12 @@ XLA_MAKE_BINARY(DivNoNan, // return x * y; // } static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, + const BCast& broadcast_helper) { + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto zero = XlaHelpers::Zero(b, dtype); auto y_equals_0 = xla::Eq(y, zero); auto zeros = xla::ZerosLike(x); @@ -105,7 +127,9 @@ static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, return result; } XLA_MAKE_BINARY(MulNoNan, - MulNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper)); + MulNoNanImpl(b, input_type(0), lhs, rhs, + broadcast_output_exprs, broadcast_helper), + xla::DExpr()); // Implementation of FloorDiv. // @@ -118,8 +142,12 @@ XLA_MAKE_BINARY(MulNoNan, // return z; // } static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, + const BCast& broadcast_helper) { + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); if (DataTypeIsFloating(dtype)) { if (dtype == DataType::DT_BFLOAT16) { // The result of a BF16 division may produce the Ceil of what was @@ -144,43 +172,61 @@ static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, return xla::Select(round_down, xla::Sub(x_div_y, one), x_div_y); } XLA_MAKE_BINARY(FloorDiv, - FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper)); + FloorDivImpl(b, input_type(0), lhs, rhs, + broadcast_output_exprs, broadcast_helper), + (lhs / rhs).simplify()); xla::XlaOp XlogyImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto zero = xla::ZerosLike(x); auto is_zero = xla::Eq(x, zero); return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y))); } -XLA_MAKE_BINARY(Xlogy, XlogyImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Xlogy, + XlogyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); xla::XlaOp Xlog1pyImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto non_zero = xla::Mul(x, xla::Log1p(y)); auto zero = xla::ZerosLike(non_zero); auto x_is_zero = xla::Eq(x, zero); return xla::Select(x_is_zero, zero, non_zero); } -XLA_MAKE_BINARY(Xlog1py, Xlog1pyImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Xlog1py, + Xlog1pyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); xla::XlaOp XdivyImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto zero = xla::ZerosLike(x); auto is_zero = xla::Eq(x, zero); return xla::Select(is_zero, zero, xla::Div(x, y)); } -XLA_MAKE_BINARY(Xdivy, XdivyImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Xdivy, + XdivyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); // Implementation of FloorMod. Pseudo-code: // T trunc_mod = std::fmod(x, y); // return trunc_mod != 0 && (y < 0 != trunc_mod < 0) ? trunc_mod + y // : trunc_mod; static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, + const BCast& broadcast_helper) { + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); auto zero = XlaHelpers::Zero(b, dtype); auto trunc_mod = xla::Rem(x, y); auto trunc_mod_not_zero = xla::Ne(trunc_mod, zero); @@ -189,43 +235,62 @@ static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, return xla::Select(do_plus, xla::Add(trunc_mod, y), trunc_mod); } XLA_MAKE_BINARY(FloorMod, - FloorModImpl(b, input_type(0), lhs, rhs, broadcast_helper)); - -XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions)); - -XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions)); + FloorModImpl(b, input_type(0), lhs, rhs, + broadcast_output_exprs, broadcast_helper), + xla::DExpr()); + +XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions), + xla::DExpr()); + +XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions), + xla::DExpr()); XLA_MAKE_BINARY(RightShift, (DataTypeIsUnsigned(ctx->input_type(0)) ? xla::ShiftRightLogical(lhs, rhs, extend_dimensions) - : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions))); - -XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs)))); + : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)), + xla::DExpr()); + +XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions), xla::DExpr()); +XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions), + (lhs / rhs).simplify()); +XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))), + xla::DExpr()); XLA_MAKE_BINARY( RsqrtGrad, xla::Mul((lhs * lhs) * lhs, xla::Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)), - extend_dimensions)); + extend_dimensions), + xla::DExpr()); XLA_MAKE_BINARY( SqrtGrad, xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)), - lhs, extend_dimensions)); + lhs, extend_dimensions), + xla::DExpr()); // Implementation of TruncateDiv. // // For floating-point values, returns trunc(x / y). For integers, simply // returns x / y. static xla::XlaOp TruncateDivImpl(xla::XlaBuilder* b, DataType dtype, - xla::XlaOp x, xla::XlaOp y, + xla::XlaOp x, + xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); if (!DataTypeIsFloating(dtype)) { return xla::Div(x, y); } @@ -235,35 +300,46 @@ static xla::XlaOp TruncateDivImpl(xla::XlaBuilder* b, DataType dtype, return xla::Select(round_up, xla::Ceil(x_div_y), xla::Floor(x_div_y)); } XLA_MAKE_BINARY(TruncateDiv, - TruncateDivImpl(b, input_type(0), lhs, rhs, broadcast_helper)); -XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions)); + TruncateDivImpl(b, input_type(0), lhs, rhs, + broadcast_output_exprs, broadcast_helper), + (lhs / rhs).simplify()); +XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions), + xla::DExpr()); // Comparison ops -XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions)); -XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions), xla::DExpr()); +XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions), + xla::DExpr()); +XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions), xla::DExpr()); +XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions), + xla::DExpr()); // Non-linear ops XLA_MAKE_BINARY(SigmoidGrad, xla::Mul(xla::Mul(rhs, lhs), - xla::Sub(XlaHelpers::One(b, input_type(0)), lhs))); + xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)), + xla::DExpr()); -XLA_MAKE_BINARY(SoftplusGrad, xla::Mul(lhs, xla::Logistic(rhs))); +XLA_MAKE_BINARY(SoftplusGrad, xla::Mul(lhs, xla::Logistic(rhs)), + xla::DExpr()); // softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2 XLA_MAKE_BINARY(SoftsignGrad, xla::Div(lhs, xla::Square(xla::Add(XlaHelpers::One(b, input_type(0)), - xla::Abs(rhs))))); + xla::Abs(rhs)))), + xla::DExpr()); XLA_MAKE_BINARY(TanhGrad, xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)), - xla::Mul(lhs, lhs)))); + xla::Mul(lhs, lhs))), + xla::DExpr()); -XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions), xla::DExpr()); xla::XlaOp SquaredDifferenceImpl( DataType dtype, xla::XlaOp x, xla::XlaOp y, @@ -277,55 +353,86 @@ xla::XlaOp SquaredDifferenceImpl( } XLA_MAKE_BINARY(SquaredDifference, SquaredDifferenceImpl(input_type(0), lhs, rhs, - extend_dimensions)); + extend_dimensions), + xla::DExpr()); xla::XlaOp IgammaImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); return xla::Igamma(x, y); } -XLA_MAKE_BINARY(Igamma, IgammaImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Igamma, + IgammaImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); xla::XlaOp IgammaGradAImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); return xla::IgammaGradA(x, y); } -XLA_MAKE_BINARY(IgammaGradA, IgammaGradAImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(IgammaGradA, + IgammaGradAImpl(lhs, rhs, broadcast_output_exprs, + broadcast_helper), + xla::DExpr()); xla::XlaOp RandomGammaGradImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& + broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); return xla::RandomGammaGrad(x, y); } XLA_MAKE_BINARY(RandomGammaGrad, - RandomGammaGradImpl(lhs, rhs, broadcast_helper)); + RandomGammaGradImpl(lhs, rhs, broadcast_output_exprs, + broadcast_helper), + xla::DExpr()); xla::XlaOp IgammacImpl(xla::XlaOp x, xla::XlaOp y, + const absl::Span& broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); + std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper, + broadcast_output_exprs); return xla::Igammac(x, y); } -XLA_MAKE_BINARY(Igammac, IgammacImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Igammac, + IgammacImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); xla::XlaOp PolygammaImpl(xla::XlaOp n, xla::XlaOp x, + const absl::Span& + broadcast_output_exprs, const BCast& broadcast_helper) { - std::tie(n, x) = XlaBinaryOp::Broadcast(n, x, broadcast_helper); + std::tie(n, x) = XlaBinaryOp::Broadcast(n, x, broadcast_helper, + broadcast_output_exprs); return xla::Polygamma(n, x); } -XLA_MAKE_BINARY(Polygamma, PolygammaImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Polygamma, + PolygammaImpl(lhs, rhs, broadcast_output_exprs, + broadcast_helper), + xla::DExpr()); -xla::XlaOp ZetaImpl(xla::XlaOp x, xla::XlaOp q, const BCast& broadcast_helper) { - std::tie(x, q) = XlaBinaryOp::Broadcast(x, q, broadcast_helper); +xla::XlaOp ZetaImpl(xla::XlaOp x, xla::XlaOp q, + const absl::Span& broadcast_output_exprs, + const BCast& broadcast_helper) { + std::tie(x, q) = XlaBinaryOp::Broadcast(x, q, broadcast_helper, + broadcast_output_exprs); return xla::Zeta(x, q); } -XLA_MAKE_BINARY(Zeta, ZetaImpl(lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Zeta, + ZetaImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper), + xla::DExpr()); #undef XLA_MAKE_BINARY diff --git a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc index 975179466bf104..d3e76ee63da1c3 100644 --- a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc @@ -37,7 +37,8 @@ class BroadcastToOp : public XlaOpKernel { context->ConstantInputAsShape( 1, &output_shape, xla::ValueInferenceMode::kUpperBound)); auto output_status_or = - BroadcastTo(context->Input(0), output_shape.dim_sizes()); + BroadcastTo(context->Input(0), output_shape.dim_sizes(), + output_shape.get_filled_expressions()); OP_REQUIRES_OK(context, output_status_or.status()); auto output = output_status_or.value(); std::vector dynamic_dims; diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc index 6c91556862d9e2..a796509e36a4ba 100644 --- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc @@ -32,10 +32,182 @@ limitations under the License. #include "xla/shape.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/util/bcast.h" namespace tensorflow { +namespace { + +bool IsSymbolicContentType(DataType type) { + type = BaseType(type); + return type == DT_INT32 || type == DT_INT64; +} + +bool TryGetIntContentsFromConstant(const Tensor& tensor, + std::vector* contents) { + contents->clear(); + if (tensor.dims() > 1) { + return false; + } + if (tensor.dtype() == DT_INT32) { + if (tensor.dims() == 0) { + contents->push_back(xla::DExpr::Const(tensor.scalar()())); + return true; + } + auto flat = tensor.flat(); + contents->reserve(flat.size()); + for (int i = 0; i < flat.size(); ++i) { + contents->push_back(xla::DExpr::Const(flat(i))); + } + return true; + } + if (tensor.dtype() == DT_INT64) { + if (tensor.dims() == 0) { + contents->push_back(xla::DExpr::Const(tensor.scalar()())); + return true; + } + auto flat = tensor.flat(); + contents->reserve(flat.size()); + for (int i = 0; i < flat.size(); ++i) { + contents->push_back(xla::DExpr::Const(flat(i))); + } + return true; + } + return false; +} + +bool TryGetIntContentsFromLiteral(const xla::LiteralSlice& literal, + std::vector* contents) { + contents->clear(); + if (literal.shape().dimensions_size() > 1) { + return false; + } + if (literal.shape().element_type() == xla::S32) { + if (literal.shape().dimensions_size() == 0) { + contents->push_back(xla::DExpr::Const(literal.Get({}))); + return true; + } + const int64_t size = literal.shape().dimensions(0); + contents->reserve(size); + for (int64_t i = 0; i < size; ++i) { + contents->push_back(xla::DExpr::Const(literal.Get({i}))); + } + return true; + } + if (literal.shape().element_type() == xla::S64) { + if (literal.shape().dimensions_size() == 0) { + contents->push_back(xla::DExpr::Const(literal.Get({}))); + return true; + } + const int64_t size = literal.shape().dimensions(0); + contents->reserve(size); + for (int64_t i = 0; i < size; ++i) { + contents->push_back(xla::DExpr::Const(literal.Get({i}))); + } + return true; + } + return false; +} + +bool TryGetInputContents(XlaOpKernelContext* ctx, const XlaExpression& expr, + const TensorShape& shape, + std::vector* contents) { + if (shape.dims() > 1) { + return false; + } + if (!expr.contents().empty()) { + if (absl::c_any_of(expr.contents(), + [](const xla::DExpr& e) { return !e; })) { + return false; + } + contents->assign(expr.contents().begin(), expr.contents().end()); + return true; + } + auto constant = expr.constant_value(); + if (!constant.has_value()) { + if (!expr.handle().valid() || expr.handle().IsUninitialized()) { + return false; + } + auto literal_or = + ctx->value_inference().AnalyzeConstant(expr.handle(), + xla::ValueInferenceMode::kValue); + if (!literal_or.ok() || !literal_or->AllValid()) { + return false; + } + auto literal = literal_or->GetValue(); + if (!literal.has_value()) { + return false; + } + return TryGetIntContentsFromLiteral(*literal, contents); + } + return TryGetIntContentsFromConstant(*constant, contents); +} + +xla::DExpr BroadcastedContentAt(absl::Span contents, + const TensorShape& shape, + int64_t output_index) { + if (shape.dims() == 0) { + return contents.empty() ? xla::DExpr() : contents[0]; + } + if (contents.empty()) { + return xla::DExpr(); + } + if (shape.dim_size(0) == 1) { + return contents[0]; + } + if (output_index >= contents.size()) { + return xla::DExpr(); + } + return contents[output_index]; +} + +bool TryBuildSymbolicBinaryContents(XlaOpKernelContext* ctx, + XlaBinaryOp* op, + const TensorShape& lhs_shape, + const TensorShape& rhs_shape, + const BCast& bcast, + std::vector* contents) { + contents->clear(); + if (!IsSymbolicContentType(ctx->input_type(0)) || + !IsSymbolicContentType(ctx->expected_output_dtype(0))) { + return false; + } + const auto& output_shape = bcast.output_shape(); + if (output_shape.size() > 1) { + return false; + } + + std::vector lhs_contents; + std::vector rhs_contents; + if (!TryGetInputContents(ctx, ctx->InputExpression(0), lhs_shape, + &lhs_contents) || + !TryGetInputContents(ctx, ctx->InputExpression(1), rhs_shape, + &rhs_contents)) { + return false; + } + + int64_t output_elements = output_shape.empty() ? 1 : output_shape[0]; + contents->reserve(output_elements); + for (int64_t i = 0; i < output_elements; ++i) { + xla::DExpr lhs_expr = BroadcastedContentAt(lhs_contents, lhs_shape, i); + xla::DExpr rhs_expr = BroadcastedContentAt(rhs_contents, rhs_shape, i); + if (!lhs_expr || !rhs_expr) { + contents->clear(); + return false; + } + xla::DExpr out_expr = op->SymbolicComputation(lhs_expr, rhs_expr); + if (!out_expr) { + contents->clear(); + return false; + } + contents->push_back(out_expr.simplify()); + } + return true; +} + +} // namespace + void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { TensorShape lhs_shape = ctx->InputShape(0); TensorShape rhs_shape = ctx->InputShape(1); @@ -68,6 +240,7 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { lhs = xla::SliceInDim(lhs, 0, rhs_xla_shape.dimensions(i), 1, /*dimno=*/i); lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i)); + lhs_tensor_shape->set_expression(i, rhs_xla_shape.expressions(i)); // Propagate dynamic dimension. lhs = xla::SetDimensionSize(lhs, size, i); } @@ -90,6 +263,7 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { lhs, xla::Zero(ctx->builder(), lhs_xla_shape.element_type()), i, 0, diff); lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i)); + lhs_tensor_shape->set_expression(i, rhs_xla_shape.expressions(i)); // Propagate dynamic dimension. lhs = xla::SetDimensionSize(lhs, size, i); } @@ -137,6 +311,10 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { lhs = xla::SetDimensionSize(lhs, size, i); lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i)); + lhs_tensor_shape->set_expression( + i, (lhs_tensor_shape->get_filled_expression(i) * + rhs_xla_shape.expressions(i)) + .simplify()); } } } @@ -162,6 +340,67 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { return; } + auto build_broadcast_output_expressions = + [&lhs_shape, &rhs_shape, &bcast]() -> std::vector { + auto merge_broadcast_dim = [](bool has_lhs, int64_t lhs_dim, + const xla::DExpr& lhs_expr, bool has_rhs, + int64_t rhs_dim, const xla::DExpr& rhs_expr, + int64_t output_dim) { + if (!has_lhs) { + return rhs_expr; + } + if (!has_rhs) { + return lhs_expr; + } + if (lhs_dim == 1 && rhs_dim != 1) { + // A broadcasted singleton usually inherits the other side's + // expression, but keep a dynamic singleton visible by folding it into + // the result. + return lhs_expr && lhs_expr->is_dynamic() + ? (lhs_expr * rhs_expr).simplify() + : rhs_expr; + } + if (rhs_dim == 1 && lhs_dim != 1) { + // Symmetric case for a singleton rhs broadcast. + return rhs_expr && rhs_expr->is_dynamic() + ? (rhs_expr * lhs_expr).simplify() + : lhs_expr; + } + // When both sides describe the same logical dimension, prefer whichever + // side still carries a dynamic symbolic expression. + if (lhs_expr && lhs_expr->is_dynamic()) { + return lhs_expr; + } + if (rhs_expr && rhs_expr->is_dynamic()) { + return rhs_expr; + } + return xla::DExpr::Const(output_dim); + }; + + const auto& output_shape = bcast.output_shape(); + std::vector output_exprs(output_shape.size()); + + for (int out_i = output_shape.size() - 1, lhs_i = lhs_shape.dims() - 1, + rhs_i = rhs_shape.dims() - 1; + out_i >= 0; --out_i, --lhs_i, --rhs_i) { + const bool has_lhs = lhs_i >= 0; + const bool has_rhs = rhs_i >= 0; + xla::DExpr lhs_expr = has_lhs ? lhs_shape.get_filled_expression(lhs_i) + : xla::DExpr::Const(1); + xla::DExpr rhs_expr = has_rhs ? rhs_shape.get_filled_expression(rhs_i) + : xla::DExpr::Const(1); + const int64_t lhs_dim = has_lhs ? lhs_shape.dim_size(lhs_i) : 1; + const int64_t rhs_dim = has_rhs ? rhs_shape.dim_size(rhs_i) : 1; + output_exprs[out_i] = merge_broadcast_dim( + has_lhs, lhs_dim, lhs_expr, has_rhs, rhs_dim, rhs_expr, + output_shape[out_i]); + } + + return output_exprs; + }; + std::vector broadcast_output_exprs = + build_broadcast_output_expressions(); + // If the ranks of the inputs don't match, TensorFlow automatically // reshapes the smaller by padding with dimensions of size 1 as a // prefix. In other words to pad a 5-vector to a 3-dimensional @@ -187,22 +426,39 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) { // Call virtual method to emit the computation. xla::XlaOp output = Computation(ctx, lhs_handle, lhs_shape.dim_sizes(), rhs_handle, - rhs_shape.dim_sizes(), bcast, extend_dimension); + rhs_shape.dim_sizes(), bcast, + broadcast_output_exprs, extend_dimension); // The TensorFlow helper computed the post-broadcast shape in // output_shape: we rely on subclassed Computations to implement the // same broadcast semantics. + std::vector output_contents; + if (TryBuildSymbolicBinaryContents(ctx, this, lhs_shape, rhs_shape, bcast, + &output_contents)) { + auto output_expr = + XlaExpression::XlaOp(output, ctx->expected_output_dtype(0)); + output_expr.set_contents(std::move(output_contents)); + ctx->SetOutputExpression(0, output_expr); + return; + } ctx->SetOutput(0, output); } /* static */ std::pair XlaBinaryOp::Broadcast( - xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper) { - auto lhs_output = BroadcastTo(lhs, broadcast_helper.output_shape()); + xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper, + absl::Span output_exprs) { + CHECK_EQ(output_exprs.size(), broadcast_helper.output_shape().size()); + for (const xla::DExpr& expr : output_exprs) { + CHECK(expr); + } + auto lhs_output = + BroadcastTo(lhs, broadcast_helper.output_shape(), output_exprs); if (!lhs_output.ok()) { xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status()); return {error, error}; } - auto rhs_output = BroadcastTo(rhs, broadcast_helper.output_shape()); + auto rhs_output = + BroadcastTo(rhs, broadcast_helper.output_shape(), output_exprs); if (!rhs_output.ok()) { xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status()); return {error, error}; diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h index d22e6eb74039b4..6e3734abe7d53f 100644 --- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h +++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h @@ -65,15 +65,25 @@ class XlaBinaryOp : public XlaOpKernel { XlaOpKernelContext* ctx, const xla::XlaOp& lhs, const absl::Span& lhs_shape, const xla::XlaOp& rhs, const absl::Span& rhs_shape, const BCast& broadcast_helper, + const absl::Span& broadcast_output_exprs, const std::vector& extend_dimensions) = 0; + // Returns a symbolic expression for one output element when content metadata + // should be propagated through this op. Returns an empty DExpr when the operation + // should not propagate symbolic contents. + virtual xla::DExpr SymbolicComputation(const xla::DExpr& lhs, + const xla::DExpr& rhs) { + return xla::DExpr(); + } + void Compile(XlaOpKernelContext* ctx) override; // Helper function that performs the broadcasting described by // 'broadcast_helper', yielding arguments 'lhs' and 'rhs' that have the same // shape. static std::pair Broadcast( - xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper); + xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper, + absl::Span output_exprs); }; } // namespace tensorflow From 193ab1a297d77d6fd165dc917c0b17d5e25cb89c Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:03:59 +0100 Subject: [PATCH 09/21] Preserve shape-derived expressions through constant folding Retain symbolic contents when shape-derived operations are rewritten or folded to constants. Add a dedicated rewrite path so optimization does not erase runtime shape dependencies. --- tensorflow/core/common_runtime/BUILD | 5 + .../core/common_runtime/constant_folding.cc | 542 +++++++++++++++++- .../common_runtime/constant_folding_test.cc | 338 +++++++++++ .../xla/xla/hlo/transforms/simplifiers/BUILD | 1 + .../simplifiers/algebraic_simplifier.cc | 103 +++- .../simplifiers/hlo_constant_folding.cc | 62 ++ third_party/xla/xla/service/BUILD | 20 + .../xla/service/dynamic_constant_rewriter.cc | 165 ++++++ .../xla/service/dynamic_constant_rewriter.h | 41 ++ 9 files changed, 1263 insertions(+), 14 deletions(-) create mode 100644 third_party/xla/xla/service/dynamic_constant_rewriter.cc create mode 100644 third_party/xla/xla/service/dynamic_constant_rewriter.h diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD index 301015eba61fe8..f99259ad3e3862 100644 --- a/tensorflow/core/common_runtime/BUILD +++ b/tensorflow/core/common_runtime/BUILD @@ -2754,6 +2754,7 @@ tf_cc_test( ":direct_session_internal", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", + "//tensorflow/cc:function_ops", "//tensorflow/cc:sendrecv_ops", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", @@ -2769,9 +2770,13 @@ tf_cc_test( "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:concat_op", "//tensorflow/core/kernels:cwise_op", + "//tensorflow/core/kernels:gather_op", "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:immutable_constant_op", "//tensorflow/core/kernels:matmul_op", + "//tensorflow/core/kernels:pack_op", + "//tensorflow/core/kernels:reshape_op", + "//tensorflow/core/kernels:slice_op", "//tensorflow/core/kernels:topk_op", "@eigen_archive//:eigen3", ], diff --git a/tensorflow/core/common_runtime/constant_folding.cc b/tensorflow/core/common_runtime/constant_folding.cc index 6820a5ddd696d3..9a6f7dc7269ed3 100644 --- a/tensorflow/core/common_runtime/constant_folding.cc +++ b/tensorflow/core/common_runtime/constant_folding.cc @@ -31,6 +31,8 @@ limitations under the License. #include "tensorflow/core/common_runtime/rendezvous_mgr.h" #include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/graph/algorithm.h" @@ -40,6 +42,7 @@ limitations under the License. #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/denormal.h" #include "tensorflow/core/platform/setround.h" #include "tensorflow/core/public/session_options.h" @@ -49,6 +52,461 @@ namespace tensorflow { namespace { const char kScopedAllocatorAttrName[] = "_scoped_allocator"; +const char kXlaShapeDerivedAttrName[] = "_xla_shape_derived"; +const char kUserInferredValueContentsAttrName[] = + "_user_inferred_value_contents"; + +bool IsShapeOp(const Node* n); + +bool GetShapeFromArgNode(const Node* node, TensorShapeProto* out_shape) { + for (const Edge* edge : node->in_edges()) { + if (edge->IsControlEdge()) continue; + + Node* input_node = edge->src(); + if (input_node->type_string() == "_Arg") { + std::vector shapes; + if (GetNodeAttr(input_node->def(), "_output_shapes", &shapes).ok() && + !shapes.empty() && HasDynamicDimExprs(shapes[0])) { + *out_shape = shapes[0]; + return true; + } + } + } + return false; +} + +bool GetShapeFromDirectDynamicSource(const Node* node, + TensorShapeProto* out_shape) { + return (IsShapeOp(node) || + node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) && + GetShapeFromArgNode(node, out_shape); +} + +bool TryGetContentsProtoAttr(const AttrSlice& attrs, + TensorShapeProto* out_contents) { + string serialized_contents; + if (!GetNodeAttr(attrs, kUserInferredValueContentsAttrName, + &serialized_contents) + .ok()) { + return false; + } + out_contents->Clear(); + return out_contents->ParseFromString(serialized_contents); +} + +bool HasTransitiveDynamicShapeContents( + const Node* node, std::unordered_map* memo, + absl::flat_hash_set* visiting) { + auto it = memo->find(node); + if (it != memo->end()) { + return it->second; + } + if (!visiting->insert(node).second) { + return false; + } + + TensorShapeProto contents_proto; + if (TryGetContentsProtoAttr(node->attrs(), &contents_proto) && + HasDynamicDimExprs(contents_proto)) { + return (*memo)[node] = true; + } + + bool has_dynamic = false; + TensorShapeProto inferred_shape_proto; + if (GetNodeAttr(node->attrs(), "has_dynamic", &has_dynamic).ok() && + has_dynamic && + GetNodeAttr(node->attrs(), "user_inferred_shape", &inferred_shape_proto) + .ok() && + HasDynamicDimExprs(inferred_shape_proto)) { + return (*memo)[node] = true; + } + + if (GetShapeFromDirectDynamicSource(node, &inferred_shape_proto) || + node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) { + return (*memo)[node] = true; + } + + for (const Edge* edge : node->in_edges()) { + if (edge->IsControlEdge()) continue; + if (HasTransitiveDynamicShapeContents(edge->src(), memo, visiting)) { + return (*memo)[node] = true; + } + } + + return (*memo)[node] = false; +} + +bool GetConstTensor(const Node* node, Tensor* tensor) { + if (node == nullptr || !node->IsConstant()) { + return false; + } + const TensorProto* tensor_proto; + if (!GetNodeAttr(node->attrs(), "value", &tensor_proto).ok()) { + return false; + } + DataType dtype; + if (!GetNodeAttr(node->attrs(), "dtype", &dtype).ok()) { + return false; + } + *tensor = Tensor(dtype); + return tensor->FromProto(cpu_allocator(), *tensor_proto); +} + +bool GetInputConstTensor(const Node* node, int input_index, Tensor* tensor) { + const Edge* edge; + if (!node->input_edge(input_index, &edge).ok()) { + return false; + } + return GetConstTensor(edge->src(), tensor); +} + +bool GetTensorIntValues(const Tensor& tensor, std::vector* values) { + values->clear(); + if (tensor.dims() == 0) { + values->reserve(1); + switch (tensor.dtype()) { + case DT_INT32: + values->push_back(tensor.scalar()()); + return true; + case DT_INT64: + values->push_back(tensor.scalar()()); + return true; + default: + return false; + } + } + if (tensor.dims() != 1) { + return false; + } + values->reserve(tensor.NumElements()); + switch (tensor.dtype()) { + case DT_INT32: { + auto flat = tensor.flat(); + for (int i = 0; i < flat.size(); ++i) values->push_back(flat(i)); + return true; + } + case DT_INT64: { + auto flat = tensor.flat(); + for (int i = 0; i < flat.size(); ++i) values->push_back(flat(i)); + return true; + } + default: + return false; + } +} + +void CopyContentAt(const TensorShapeProto& input_contents, int64_t index, + TensorShapeProto* output_contents) { + output_contents->add_dim()->CopyFrom(input_contents.dim(index)); + ExpressionProto* output_expression = output_contents->add_expressions(); + if (index < input_contents.expressions_size()) { + output_expression->CopyFrom(input_contents.expressions(index)); + } else { + output_expression->set_constant_value(input_contents.dim(index).size()); + } +} + +void AppendScalarConstantContent(int64_t value, + TensorShapeProto* output_contents) { + output_contents->add_dim()->set_size(value); + output_contents->add_expressions()->set_constant_value(value); +} + +void AppendScalarContentFromTensor(const Tensor& tensor, + TensorShapeProto* output_contents) { + if (tensor.dtype() == DT_INT32) { + AppendScalarConstantContent(tensor.scalar()(), output_contents); + } else { + AppendScalarConstantContent(tensor.scalar()(), output_contents); + } +} + +ExpressionProto MakeConstantExpressionProto(int64_t value) { + ExpressionProto expr; + expr.set_constant_value(value); + return expr; +} + +ExpressionProto GetContentExpressionProto(const TensorShapeProto& contents, + int64_t index) { + if (index < contents.expressions_size()) { + return contents.expressions(index); + } + return MakeConstantExpressionProto(contents.dim(index).size()); +} + +ExpressionProto MakeMulExpressionProto(ExpressionProto lhs, + ExpressionProto rhs) { + ExpressionProto expr; + auto* mul = expr.mutable_mul_node(); + *mul->mutable_lhs() = std::move(lhs); + *mul->mutable_rhs() = std::move(rhs); + return expr; +} + +bool TryGetFoldedValueContents(const Node* node, int output_index, + TensorShapeProto* out_contents, + absl::flat_hash_set* visiting) { + out_contents->Clear(); + if (output_index != 0) { + return false; + } + if (!visiting->insert(node).second) { + return false; + } + auto remove_from_visiting = + gtl::MakeCleanup([&] { visiting->erase(node); }); + + if (TryGetContentsProtoAttr(node->attrs(), out_contents)) { + return true; + } + + bool has_dynamic = false; + TensorShapeProto user_inferred_shape; + if (GetNodeAttr(node->attrs(), "has_dynamic", &has_dynamic).ok() && + has_dynamic && + GetNodeAttr(node->attrs(), "user_inferred_shape", &user_inferred_shape) + .ok()) { + *out_contents = user_inferred_shape; + return true; + } + + if (GetShapeFromDirectDynamicSource(node, out_contents)) { + return true; + } + + auto recurse_input = [&](int input_index, + TensorShapeProto* input_contents) -> bool { + const Edge* input_edge; + if (!node->input_edge(input_index, &input_edge).ok()) { + return false; + } + return TryGetFoldedValueContents(input_edge->src(), input_edge->src_output(), + input_contents, visiting); + }; + + if (node->IsIdentity() || node->type_string() == "Cast") { + return recurse_input(0, out_contents); + } + + TensorShapeProto input_contents; + if (node->type_string() == "Reshape" && recurse_input(0, &input_contents)) { + Tensor shape_tensor; + std::vector shape_dims; + if (!GetInputConstTensor(node, 1, &shape_tensor) || + !GetTensorIntValues(shape_tensor, &shape_dims)) { + return false; + } + if (input_contents.dim_size() == 1 && shape_dims.empty()) { + CopyContentAt(input_contents, 0, out_contents); + return true; + } + if (shape_dims.size() == 1 && input_contents.dim_size() == shape_dims[0]) { + out_contents->CopyFrom(input_contents); + return true; + } + return false; + } + + if (node->type_string() == "Pack") { + for (int i = 0; i < node->num_inputs(); ++i) { + TensorShapeProto scalar_contents; + Tensor scalar_tensor; + if (recurse_input(i, &scalar_contents)) { + if (scalar_contents.dim_size() != 1) { + return false; + } + CopyContentAt(scalar_contents, 0, out_contents); + } else if (GetInputConstTensor(node, i, &scalar_tensor) && + scalar_tensor.dims() == 0 && + (scalar_tensor.dtype() == DT_INT32 || + scalar_tensor.dtype() == DT_INT64)) { + AppendScalarContentFromTensor(scalar_tensor, out_contents); + } else { + return false; + } + } + return true; + } + + if (node->type_string() == "ConcatV2") { + Tensor axis_tensor; + std::vector axis_values; + if (!GetInputConstTensor(node, node->num_inputs() - 1, &axis_tensor) || + !GetTensorIntValues(axis_tensor, &axis_values) || + axis_values.size() != 1) { + return false; + } + int64_t axis = axis_values[0]; + if (axis != 0 && axis != -1) { + return false; + } + for (int i = 0; i < node->num_inputs() - 1; ++i) { + TensorShapeProto part_contents; + if (!recurse_input(i, &part_contents)) { + return false; + } + for (int64_t j = 0; j < part_contents.dim_size(); ++j) { + CopyContentAt(part_contents, j, out_contents); + } + } + return true; + } + + if ((node->type_string() == "Gather" || node->type_string() == "GatherV2") && + recurse_input(0, &input_contents)) { + Tensor indices_tensor; + std::vector indices; + if (!GetInputConstTensor(node, 1, &indices_tensor) || + !GetTensorIntValues(indices_tensor, &indices)) { + return false; + } + + int64_t axis = 0; + if (node->type_string() == "GatherV2") { + Tensor axis_tensor; + std::vector axis_values; + if (!GetInputConstTensor(node, 2, &axis_tensor) || + !GetTensorIntValues(axis_tensor, &axis_values) || + axis_values.size() != 1) { + return false; + } + axis = axis_values[0]; + } + + const int64_t params_rank = 1; + if (axis < 0) axis += params_rank; + if (axis != 0) { + return false; + } + + const int64_t rank = input_contents.dim_size(); + for (int64_t index : indices) { + if (index < 0) index += rank; + if (index < 0 || index >= rank) { + return false; + } + CopyContentAt(input_contents, index, out_contents); + } + return true; + } + + if (node->type_string() == "Prod" && recurse_input(0, &input_contents)) { + Tensor reduction_indices_tensor; + std::vector axes; + bool keep_dims = false; + if (!GetInputConstTensor(node, 1, &reduction_indices_tensor) || + !GetTensorIntValues(reduction_indices_tensor, &axes) || + !GetNodeAttr(node->attrs(), "keep_dims", &keep_dims).ok() || + keep_dims || axes.size() != 1 || + (axes[0] != 0 && axes[0] != -1) || input_contents.dim_size() == 0) { + return false; + } + int64_t value = 1; + ExpressionProto expr = GetContentExpressionProto(input_contents, 0); + for (int64_t i = 0; i < input_contents.dim_size(); ++i) { + value *= input_contents.dim(i).size(); + if (i > 0) { + expr = MakeMulExpressionProto(std::move(expr), + GetContentExpressionProto(input_contents, i)); + } + } + out_contents->add_dim()->set_size(value); + out_contents->add_expressions()->Swap(&expr); + return true; + } + + if (node->type_string() == "Slice" && recurse_input(0, &input_contents)) { + Tensor begin_tensor; + Tensor size_tensor; + std::vector begin; + std::vector size; + if (!GetInputConstTensor(node, 1, &begin_tensor) || + !GetInputConstTensor(node, 2, &size_tensor) || + !GetTensorIntValues(begin_tensor, &begin) || + !GetTensorIntValues(size_tensor, &size) || begin.size() != 1 || + size.size() != 1) { + return false; + } + int64_t start = begin[0]; + if (start < 0 || start > input_contents.dim_size()) { + return false; + } + int64_t length = size[0] < 0 ? input_contents.dim_size() - start : size[0]; + if (length < 0 || start + length > input_contents.dim_size()) { + return false; + } + for (int64_t i = 0; i < length; ++i) { + CopyContentAt(input_contents, start + i, out_contents); + } + return true; + } + + if (node->type_string() == "StridedSlice" && + recurse_input(0, &input_contents)) { + Tensor begin_tensor; + Tensor end_tensor; + Tensor strides_tensor; + std::vector begin; + std::vector end; + std::vector strides; + int64_t begin_mask = 0; + int64_t end_mask = 0; + int64_t ellipsis_mask = 0; + int64_t new_axis_mask = 0; + int64_t shrink_axis_mask = 0; + if (!GetInputConstTensor(node, 1, &begin_tensor) || + !GetInputConstTensor(node, 2, &end_tensor) || + !GetInputConstTensor(node, 3, &strides_tensor) || + !GetTensorIntValues(begin_tensor, &begin) || + !GetTensorIntValues(end_tensor, &end) || + !GetTensorIntValues(strides_tensor, &strides) || begin.size() != 1 || + end.size() != 1 || strides.size() != 1 || + !GetNodeAttr(node->attrs(), "begin_mask", &begin_mask).ok() || + !GetNodeAttr(node->attrs(), "end_mask", &end_mask).ok() || + !GetNodeAttr(node->attrs(), "ellipsis_mask", &ellipsis_mask).ok() || + !GetNodeAttr(node->attrs(), "new_axis_mask", &new_axis_mask).ok() || + !GetNodeAttr(node->attrs(), "shrink_axis_mask", &shrink_axis_mask).ok()) { + return false; + } + if (ellipsis_mask != 0 || new_axis_mask != 0) { + return false; + } + const int64_t rank = input_contents.dim_size(); + int64_t stride = strides[0]; + if (stride == 0) { + return false; + } + int64_t start = (begin_mask & 1) ? (stride > 0 ? 0 : rank - 1) : begin[0]; + int64_t stop = (end_mask & 1) ? (stride > 0 ? rank : -1) : end[0]; + if (start < 0) start += rank; + if (stop < 0 && !(end_mask & 1 && stride < 0)) stop += rank; + if (shrink_axis_mask & 1) { + if (start < 0 || start >= rank) { + return false; + } + CopyContentAt(input_contents, start, out_contents); + return true; + } + if (stride < 0) { + return false; + } + start = std::max(0, start); + stop = std::min(rank, stop); + for (int64_t i = start; i < stop; i += stride) { + CopyContentAt(input_contents, i, out_contents); + } + return true; + } + + return false; +} + +bool TryGetFoldedValueContents(const Node* node, int output_index, + TensorShapeProto* out_contents) { + absl::flat_hash_set visiting; + return TryGetFoldedValueContents(node, output_index, out_contents, &visiting); +} // For stateless RNGs ops, they are pure but device-dependent. Those ops are not // constant-foldable. @@ -242,8 +700,23 @@ bool IsConstantFoldable( shape_map, const std::function& consider, int64_t max_constant_size_in_bytes, - std::unordered_map>* - shape_replacement_map) { + std::unordered_map>* shape_replacement_map, + std::unordered_map* dynamic_contents_memo) { + TensorShapeProto exact_contents; + const bool has_exact_contents = + TryGetFoldedValueContents(n, 0, &exact_contents); + TensorShapeProto dynamic_shape; + const bool has_dynamic = GetShapeFromDirectDynamicSource(n, &dynamic_shape); + const bool is_shape_derived = + n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr; + absl::flat_hash_set dynamic_contents_visiting; + const bool has_transitive_dynamic_contents = + HasTransitiveDynamicShapeContents(n, dynamic_contents_memo, + &dynamic_contents_visiting); + if ((has_dynamic || is_shape_derived || has_transitive_dynamic_contents) && + (!has_exact_contents || n->num_outputs() > 1)) { + return false; + } if (n->IsConstant()) { // Skip constant folding resources as they cannot be deep copied. return n->output_type(0) != DT_RESOURCE; @@ -326,10 +799,11 @@ void ConsiderConstantFoldableNode( Node* n, const ConstantFoldingOptions& opts, std::vector* nodes, std::unordered_map>* constant_control_deps, std::unordered_map>* shape_replacement_map, + std::unordered_map* dynamic_contents_memo, bool* internal_node_inserted) { if (!IsConstantFoldable(n, opts.shape_map, opts.consider, opts.max_constant_size_in_bytes, - shape_replacement_map)) { + shape_replacement_map, dynamic_contents_memo)) { return; } // A node is constant provided all of its non-control incoming Tensors come @@ -383,16 +857,17 @@ void FindConstantFoldableNodes( const Graph* graph, const ConstantFoldingOptions& opts, std::vector* nodes, std::unordered_map>* constant_control_deps, - std::unordered_map>* - shape_replacement_map) { + std::unordered_map>* shape_replacement_map) { bool internal_node_inserted = false; + std::unordered_map dynamic_contents_memo; // Walk the nodes in data flow order. ReverseDFS( *graph, nullptr, [nodes, constant_control_deps, shape_replacement_map, - &internal_node_inserted, &opts](Node* n) { + &dynamic_contents_memo, &internal_node_inserted, &opts](Node* n) { ConsiderConstantFoldableNode(n, opts, nodes, constant_control_deps, shape_replacement_map, + &dynamic_contents_memo, &internal_node_inserted); }, NodeComparatorName()); @@ -449,13 +924,29 @@ void AddShapeNodeToConstantGraph( shape_replacement_map, std::unordered_map>* node_map, const ConstantFoldNameGenerator& generate_new_name, Graph* constant_graph) { + TensorShapeProto user_inferred_shape; + const bool has_dynamic = + GetShapeFromDirectDynamicSource(n, &user_inferred_shape); + TensorShapeProto exact_contents; + const bool has_exact_contents = + TryGetFoldedValueContents(n, 0, &exact_contents); std::vector& added = (*node_map)[n]; const string& node_name = n->name(); for (const Tensor& t : shape_replacement_map.at(n)) { + VLOG(1) << "Constant folding shape node " << node_name + << " into Const with value " << t.SummarizeValue(16); auto builder = NodeDefBuilder(generate_new_name(constant_graph, node_name), "Const") .Attr("dtype", t.dtype()) .Attr("value", t); + if (has_dynamic) { + builder.Attr("has_dynamic", has_dynamic) + .Attr("user_inferred_shape", user_inferred_shape); + } + if (has_exact_contents && HasDynamicDimExprs(exact_contents)) { + builder.Attr(kUserInferredValueContentsAttrName, + exact_contents.SerializeAsString()); + } NodeDef def; CHECK(builder.Finalize(&def).ok()); Node* constant_node; @@ -546,6 +1037,12 @@ bool ReplaceTensorWithConstant( ? DeviceType{partition_device->device_type()} : DEVICE_CPU; if (partition_device && device_type != DEVICE_CPU) { + // Constant folding replaces one output edge-set at a time. Be + // conservative for non-CPU multi-output ops, since partially replacing a + // node can violate per-output placement or memory-type assumptions. + if (tensor.first->num_outputs() > 1) { + return false; + } MemoryTypeVector input_mvec; MemoryTypeVector output_mvec; if (!MemoryTypesForNode(graph->op_registry(), device_type, @@ -574,10 +1071,40 @@ bool ReplaceTensorWithConstant( } } const string& node_name = n->name(); + TensorShapeProto user_inferred_shape; + const bool has_dynamic = + GetShapeFromDirectDynamicSource(tensor.first, &user_inferred_shape); + const bool is_shape_derived = + tensor.first->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr; + if (tensor.second != 0 && (has_dynamic || is_shape_derived)) { + VLOG(1) << "Skipping replacement of " << tensor.first->name() << " :: " + << tensor.second + << " because symbolic content preservation is only supported for " + << "single-output replacements"; + return false; + } + TensorShapeProto exact_contents; + const bool has_exact_contents = + TryGetFoldedValueContents(tensor.first, tensor.second, &exact_contents); + if ((has_dynamic || is_shape_derived) && !has_exact_contents) { + VLOG(1) << "Skipping replacement of " << tensor.first->name() << " :: " + << tensor.second + << " because constant folding could not preserve symbolic contents"; + return false; + } Node* constant_node; auto builder = NodeDefBuilder(generate_new_name(graph, node_name), "Const") .Attr("dtype", constant.dtype()) .Attr("value", constant); + if (has_dynamic) { + builder.Attr("has_dynamic", has_dynamic) + .Attr("user_inferred_shape", user_inferred_shape); + } + if (has_exact_contents && HasDynamicDimExprs(exact_contents)) { + builder.Attr("has_dynamic", true) + .Attr(kUserInferredValueContentsAttrName, + exact_contents.SerializeAsString()); + } if (partition_device) { builder.Device(partition_device->name()); } @@ -592,6 +1119,9 @@ bool ReplaceTensorWithConstant( VLOG(1) << "Replacing " << tensor.first->name() << " :: " << tensor.second << " with a constant"; + VLOG(1) << "ReplaceTensorWithConstant creating Const from " + << tensor.first->name() << " :: " << tensor.second + << " with value " << constant.SummarizeValue(16); if (!NodeBuilder(builder).Finalize(graph, &constant_node).ok()) { return false; diff --git a/tensorflow/core/common_runtime/constant_folding_test.cc b/tensorflow/core/common_runtime/constant_folding_test.cc index 481a85add4893c..fa4cf9f2b9c6a2 100644 --- a/tensorflow/core/common_runtime/constant_folding_test.cc +++ b/tensorflow/core/common_runtime/constant_folding_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include "tensorflow/cc/ops/array_ops_internal.h" +#include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/nn_ops.h" #include "tensorflow/cc/ops/sendrecv_ops.h" #include "tensorflow/cc/ops/standard_ops.h" @@ -32,6 +33,7 @@ limitations under the License. #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/node_builder.h" @@ -45,6 +47,15 @@ limitations under the License. namespace tensorflow { namespace { +TensorShapeProto MakeDynamicShapeProto977x16() { + TensorShapeProto proto; + proto.add_dim()->set_size(977); + proto.add_dim()->set_size(16); + proto.add_expressions()->set_variable_id(0); + proto.add_expressions()->set_constant_value(16); + return proto; +} + class ConstantFoldingTest : public ::testing::Test { protected: template @@ -634,6 +645,333 @@ TEST_F(ConstantFoldingTest, ConstShapeKnown) { } } +TEST_F(ConstantFoldingTest, FoldShapeFromDynamicArgPreservesContents) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto send = ops::_Send(s.WithOpName("send"), shape, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ExpectNodeEqual(folded, {977, 16}, {2}); + + // This test checks the stable contract we care about: after folding + // Shape(arg), the replacement Const still carries symbolic contents. + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 2); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_FALSE(IsDynamicDimExpr(contents_proto.expressions(1))); + EXPECT_EQ(contents_proto.dim(0).size(), 977); + EXPECT_EQ(contents_proto.dim(1).size(), 16); +} + +TEST_F(ConstantFoldingTest, FoldSliceOfDynamicShapePreservesContents) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto begin = ops::Const(s.WithOpName("begin"), {0}, {1}); + auto size = ops::Const(s.WithOpName("size"), {1}, {1}); + auto slice = ops::Slice(s.WithOpName("slice"), shape, begin, size); + auto send = ops::_Send(s.WithOpName("send"), slice, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ExpectNodeEqual(folded, {977}, {1}); + + // Folding Slice(Shape(arg), [0], [1]) should preserve the selected symbolic + // content, not just the concrete value 977. + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 1); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_EQ(contents_proto.dim(0).size(), 977); +} + +TEST_F(ConstantFoldingTest, FoldGatherOfDynamicShapePreservesContents) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto index = ops::Const(s.WithOpName("index"), 0); + auto gather = ops::GatherV2(s.WithOpName("gather"), shape, index, + ops::Const(s.WithOpName("axis"), 0)); + auto send = ops::_Send(s.WithOpName("send"), gather, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ExpectNodeEqual(folded, {977}, {}); + + // Folding Gather(Shape(arg), 0, axis=0) should preserve the selected + // symbolic content, not just the concrete scalar value 977. + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 1); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_EQ(contents_proto.dim(0).size(), 977); +} + +TEST_F(ConstantFoldingTest, FoldReshapeOfDynamicShapePreservesContents) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto begin = ops::Const(s.WithOpName("begin"), {0}, {1}); + auto size = ops::Const(s.WithOpName("size"), {1}, {1}); + auto slice = ops::Slice(s.WithOpName("slice"), shape, begin, size); + auto scalar_shape = ops::Const(s.WithOpName("scalar_shape"), {}); + auto reshape = + ops::Reshape(s.WithOpName("reshape"), slice, scalar_shape); + auto send = ops::_Send(s.WithOpName("send"), reshape, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ExpectNodeEqual(folded, {977}, {}); + + // Folding Reshape(Slice(Shape(arg), [0], [1]), []) should preserve the + // selected symbolic content when the shape-vector result is scalarized. + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 1); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_EQ(contents_proto.dim(0).size(), 977); +} + +TEST_F(ConstantFoldingTest, FoldPackWithRepeatedDynamicInputPreservesContents) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto first = ops::GatherV2(s.WithOpName("first"), shape, + ops::Const(s.WithOpName("index"), 0), + ops::Const(s.WithOpName("axis"), 0)); + // Reusing the same producer is a DAG, not a recursion cycle. Both visits must + // preserve its symbolic content. + OutputList pack_inputs = {first, first}; + auto pack = ops::Stack(s.WithOpName("pack"), pack_inputs, + ops::Stack::Axis(0)); + auto send = ops::_Send(s.WithOpName("send"), pack, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ASSERT_TRUE(folded->IsConstant()); + ExpectNodeEqual(folded, {977, 977}, {2}); + + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 2); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(1))); + EXPECT_EQ(contents_proto.dim(0).size(), 977); + EXPECT_EQ(contents_proto.dim(1).size(), 977); +} + +TEST_F(ConstantFoldingTest, FoldPackKeepsSymbolicContentsAligned) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto dynamic_value = ops::GatherV2( + s.WithOpName("dynamic_value"), shape, + ops::Const(s.WithOpName("index"), 0), + ops::Const(s.WithOpName("axis"), 0)); + auto static_value = ops::Const(s.WithOpName("static_value"), 16); + auto pack = ops::Stack(s.WithOpName("pack"), + OutputList{static_value, dynamic_value}, + ops::Stack::Axis(0)); + auto send = ops::_Send(s.WithOpName("send"), pack, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + const Edge* send_input = nullptr; + TF_ASSERT_OK(index_by_name.at("send")->input_edge(0, &send_input)); + Node* folded = send_input->src(); + ASSERT_TRUE(folded->IsConstant()); + ExpectNodeEqual(folded, {16, 977}, {2}); + + string serialized_contents_proto; + TF_ASSERT_OK(GetNodeAttr(folded->attrs(), + "_user_inferred_value_contents", + &serialized_contents_proto)); + TensorShapeProto contents_proto; + ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto)); + ASSERT_EQ(contents_proto.expressions_size(), 2); + EXPECT_FALSE(IsDynamicDimExpr(contents_proto.expressions(0))); + EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(1))); + EXPECT_EQ(contents_proto.dim(0).size(), 16); + EXPECT_EQ(contents_proto.dim(1).size(), 977); +} + +TEST_F(ConstantFoldingTest, DoNotFoldUnsupportedDynamicContentsTransform) { + Graph g(OpRegistry::Global()); + Scope s = Scope::NewRootScope(); + auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0); + auto shape = ops::Shape(s.WithOpName("shape"), arg); + auto added = ops::Add(s.WithOpName("added"), shape, shape); + auto send = ops::_Send(s.WithOpName("send"), added, "send", "sender", 0, + "receiver"); + TF_ASSERT_OK(s.ToGraph(&g)); + + std::unordered_map index_by_name = g.BuildNodeNameIndex(); + Node* arg_node = index_by_name.at("arg"); + arg_node->AddAttr("_output_shapes", + std::vector{MakeDynamicShapeProto977x16()}); + + PartialTensorShape partial_shape({977, 16}); + std::unordered_map> shape_map; + shape_map[arg_node->name()].push_back(partial_shape); + + ConstantFoldingOptions opts; + opts.shape_map = &shape_map; + bool was_mutated = false; + TF_ASSERT_OK( + ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated)); + + index_by_name = g.BuildNodeNameIndex(); + Node* send_node = index_by_name.at("send"); + const Edge* send_input = nullptr; + TF_ASSERT_OK(send_node->input_edge(0, &send_input)); + EXPECT_FALSE(send_input->src()->IsConstant()); +} + TEST_F(ConstantFoldingTest, NoReplacePartialOutput) { Graph g(OpRegistry::Global()); { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index b3f1e108ba3122..814a0d2388f457 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -944,6 +944,7 @@ cc_library( "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 6493c56fb93dee..171c1bdca2e2bf 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -87,6 +87,22 @@ namespace m = match; using primitive_util::NativeTypeOf; +static bool HasDynamicConstantFrontendAttributes( + const HloInstruction* instruction); + +static bool HasDynamicConstantContents(const HloInstruction* instruction) { + if (!instruction->has_contents()) { + return false; + } + for (const auto& content : instruction->contents()) { + if (content.node_type_case() != ExpressionProto::kConstantValue && + content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET) { + return true; + } + } + return false; +} + // Unwraps broadcasts hunting for a constant. If we find one, checks if the // constant contains only the given value. bool IsAll(const HloInstruction* op, int8_t value) { @@ -94,6 +110,9 @@ bool IsAll(const HloInstruction* op, int8_t value) { case HloOpcode::kBroadcast: return IsAll(op->operand(0), value); case HloOpcode::kConstant: + if (HasDynamicConstantFrontendAttributes(op)) { + return false; + } return op->literal().IsAll(value); default: return false; @@ -107,6 +126,9 @@ bool IsAllFloat(const HloInstruction* op, float value) { case HloOpcode::kBroadcast: return IsAllFloat(op->operand(0), value); case HloOpcode::kConstant: + if (HasDynamicConstantFrontendAttributes(op)) { + return false; + } return op->literal().IsAllFloat(value); default: return false; @@ -119,6 +141,9 @@ bool IsAll(const HloInstruction* op, const Literal& scalar) { case HloOpcode::kBroadcast: return IsAll(op->operand(0), scalar); case HloOpcode::kConstant: + if (HasDynamicConstantFrontendAttributes(op)) { + return false; + } return op->literal().IsAll(scalar); default: return false; @@ -171,9 +196,18 @@ bool IsPositive(const HloInstruction* hlo, static bool IsScalarConstant(const HloInstruction* hlo) { return hlo->opcode() == HloOpcode::kConstant && + !HasDynamicConstantFrontendAttributes(hlo) && ShapeUtil::IsEffectiveScalar(hlo->shape()); } +static bool HasDynamicConstantFrontendAttributes( + const HloInstruction* instruction) { + const auto& attrs = instruction->frontend_attributes().map(); + return HasDynamicConstantContents(instruction) || + attrs.contains("dynamic_constant_index") || + attrs.contains("dynamic_constant_expr"); +} + std::optional GetConstantValue(const HloInstruction* inst) { if (!IsScalarConstant(inst)) { return std::nullopt; @@ -701,8 +735,9 @@ absl::Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction( HloInstruction* multiplier; // Pattern match a scalar multiply. if (Match(inst, m::MultiplyAnyOrder( - m::Op(&operand), - m::Broadcast(m::ConstantScalar(&multiplier))))) { + m::Op(&operand), m::Broadcast(m::ConstantScalar( + &multiplier)))) && + !HasDynamicConstantFrontendAttributes(multiplier)) { CHECK_LT(index, user->operand_count()); CHECK_EQ(inst, user->operands()[index]); @@ -732,8 +767,9 @@ absl::Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction( HloInstruction* operand; HloInstruction* multiplier; if (Match(inst, m::MultiplyAnyOrder( - m::Op(&operand), - m::Broadcast(m::ConstantScalar(&multiplier))))) { + m::Op(&operand), m::Broadcast(m::ConstantScalar( + &multiplier)))) && + !HasDynamicConstantFrontendAttributes(multiplier)) { values.push_back(*GetConstantValue(multiplier)); TF_RETURN_IF_ERROR(inst->ReplaceAllUsesWith(operand)); @@ -903,6 +939,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { Match(add, m::Add(m::Add(m::NonConstant(&a), m::Broadcast(m::ConstantScalar(&c1))), m::Broadcast(m::ConstantScalar(&c2))))) { + if (HasDynamicConstantFrontendAttributes(c1) || + HasDynamicConstantFrontendAttributes(c2)) { + return absl::OkStatus(); + } TF_ASSIGN_OR_RETURN(auto* sum_of_constants, MakeBinaryHlo(HloOpcode::kAdd, c1, c2)); if (ShapeUtil::IsScalar(sum_of_constants->shape()) && @@ -921,6 +961,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { Match(add, m::Add(m::Subtract(m::Broadcast(m::ConstantScalar(&c1)), m::NonConstant(&a)), m::Broadcast(m::ConstantScalar(&c2))))) { + if (HasDynamicConstantFrontendAttributes(c1) || + HasDynamicConstantFrontendAttributes(c2)) { + return absl::OkStatus(); + } TF_ASSIGN_OR_RETURN(HloInstruction * sum_of_constants, MakeBinaryHlo(HloOpcode::kAdd, c1, c2)); if (ShapeUtil::IsScalar(sum_of_constants->shape()) && @@ -1201,13 +1245,15 @@ absl::StatusOr AlgebraicSimplifierVisitor::TrySimplifyTautologicalCompare( m::Shape().IsEffectiveScalar().WithElementType(PrimitiveType::S32); if (Match(cmp, m::Compare(m::Op(&lhs), m::Constant(&rhs).WithShape(scalar_shape_matcher)) - .WithComparisonDirection(ComparisonDirection::kLt))) { + .WithComparisonDirection(ComparisonDirection::kLt)) && + !HasDynamicConstantFrontendAttributes(rhs)) { return {LessThanCompareInfo{lhs, *rhs->literal().GetFirstInteger()}}; } else if (Match( cmp, m::Compare(m::Constant(&lhs).WithShape(scalar_shape_matcher), m::Op(&rhs)) - .WithComparisonDirection(ComparisonDirection::kGt))) { + .WithComparisonDirection(ComparisonDirection::kGt)) && + !HasDynamicConstantFrontendAttributes(lhs)) { return {LessThanCompareInfo{rhs, *lhs->literal().GetFirstInteger()}}; } return std::nullopt; @@ -1232,6 +1278,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAllGather( if (all_gather->shape().IsArray() && Match(all_gather->mutable_operand(0), m::Broadcast(m::ConstantScalar()))) { + if (HasDynamicConstantFrontendAttributes( + all_gather->mutable_operand(0)->operand(0))) { + return absl::OkStatus(); + } return ReplaceWithNewInstruction( all_gather, all_gather->mutable_operand(0)->CloneWithNewShape(all_gather->shape())); @@ -1244,6 +1294,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAllToAll( if (all_to_all->shape().IsArray() && Match(all_to_all->mutable_operand(0), m::Broadcast(m::ConstantScalar()))) { + if (HasDynamicConstantFrontendAttributes( + all_to_all->mutable_operand(0)->operand(0))) { + return absl::OkStatus(); + } return ReplaceInstruction(all_to_all, all_to_all->mutable_operand(0)); } return absl::OkStatus(); @@ -2185,6 +2239,10 @@ std::unique_ptr TryDivideToShift( return nullptr; } + if (HasDynamicConstantFrontendAttributes(c)) { + return nullptr; + } + if (ShapeUtil::ElementIsSigned(divide->shape())) { int64_t b_value = static_cast(c->literal().GetFirstElement()); if (b_value > 0 && absl::has_single_bit(static_cast(b_value))) { @@ -2234,6 +2292,13 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b)))); // A/1 => A VLOG(10) << "trying transform [A/1 => A]: " << divide->ToString(); + if ((b->opcode() == HloOpcode::kConstant && + HasDynamicConstantFrontendAttributes(b)) || + (b->opcode() == HloOpcode::kBroadcast && + b->operand(0)->opcode() == HloOpcode::kConstant && + HasDynamicConstantFrontendAttributes(b->operand(0)))) { + return absl::OkStatus(); + } if (IsAll(b, 1) && ReplaceInstructionIfCompatible(divide, a)) { return absl::OkStatus(); } @@ -4476,6 +4541,10 @@ absl::StatusOr> MinMaxToClamp( CHECK(Match(clamp_upper_bound_bcast, m::Broadcast(m::ConstantEffectiveScalar(&clamp_upper_bound)))) << clamp_upper_bound_bcast->ToString(); + if (HasDynamicConstantFrontendAttributes(clamp_lower_bound) || + HasDynamicConstantFrontendAttributes(clamp_upper_bound)) { + return nullptr; + } const Literal& lower_bound = Cast(clamp_lower_bound)->literal(); @@ -4720,6 +4789,9 @@ absl::Status AlgebraicSimplifierVisitor::HandleClamp(HloInstruction* clamp) { if ((Match(to_clamp, m::PartitionId()) || Match(to_clamp, m::ReplicaId())) && Match(clamp_lower_bound, m::ConstantScalar(0U)) && Match(clamp_upper_bound, m::ConstantScalar())) { + if (HasDynamicConstantFrontendAttributes(clamp_upper_bound)) { + return absl::OkStatus(); + } int64_t upper_bound = Cast(clamp_upper_bound) ->literal() .GetFirstElement(); @@ -4947,6 +5019,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply( m::Multiply( m::Multiply(m::Op(&a), m::Broadcast(m::ConstantScalar(&c1))), m::Broadcast(m::ConstantScalar(&c2))))) { + if (HasDynamicConstantFrontendAttributes(c1) || + HasDynamicConstantFrontendAttributes(c2)) { + return absl::OkStatus(); + } TF_ASSIGN_OR_RETURN(auto* product_of_constants, MakeBinaryHlo(HloOpcode::kMultiply, c1, c2)); if (ShapeUtil::IsScalar(product_of_constants->shape()) && @@ -5998,6 +6074,9 @@ std::unique_ptr TryRemainderToAnd( !Match(b, m::Broadcast(m::ConstantEffectiveScalar(&c)))) { return nullptr; } + if (HasDynamicConstantFrontendAttributes(c)) { + return nullptr; + } if (ShapeUtil::ElementIsSigned(remainder->shape())) { int64_t b_value = static_cast(c->literal().GetFirstElement()); @@ -6083,7 +6162,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder( HloInstruction* divisor; if (Match(remainder, m::Remainder(m::Iota(&iota), - m::Broadcast(m::ConstantEffectiveScalar(&divisor))))) { + m::Broadcast(m::ConstantEffectiveScalar(&divisor)))) && + !HasDynamicConstantFrontendAttributes(divisor)) { // The iota counts {0, ..., iota_upper_bound - 1}. (Actually this is // conservative; the iota may overflow and count up to a smaller value than // this. But that's OK for our purposes here.) @@ -6111,7 +6191,7 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder( m::AddAnyOrder(m::Iota(&iota), m::Broadcast(m::ConstantEffectiveScalar(&addend))), m::Broadcast(&bcast, m::ConstantEffectiveScalar(&divisor)))) && - addend == divisor) { + addend == divisor && !HasDynamicConstantFrontendAttributes(divisor)) { // The iota counts {0, ...iota_upper_bound - 1}, with the same caveat above // that iota_upper_bound is conservative, and the true upper bound may be // smaller. @@ -7332,6 +7412,9 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice( std::vector slice_strides(rank, 1); for (int64_t i = 0; i < rank; ++i) { + if (HasDynamicConstantFrontendAttributes(dynamic_slice->operand(i + 1))) { + return absl::OkStatus(); + } std::optional offset = dynamic_slice->operand(i + 1)->literal().GetFirstInteger(); if (!offset || *offset < 0) { @@ -7525,6 +7608,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice( compatible = false; break; } + if (HasDynamicConstantFrontendAttributes(slice_dim_start)) { + compatible = false; + break; + } VLOG(2) << "slice: " << slice_dim_start->ToString(); std::optional start = slice_dim_start->literal().GetFirstInteger(); diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 5a9ef91900d850..7ed0be00ab1d44 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" #include "xla/hlo/evaluator/hlo_evaluator.h" @@ -225,6 +226,56 @@ absl::StatusOr HloConstantFolding::Run( continue; } + bool source_has_dynamic_constant_marker = false; + std::vector marked_constant_operands; + if (instruction->has_contents()) { + bool has_dynamic_content = false; + for (const auto& content : instruction->contents()) { + has_dynamic_content = + content.node_type_case() != ExpressionProto::kConstantValue && + content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET; + if (has_dynamic_content) { + break; + } + } + if (has_dynamic_content) { + source_has_dynamic_constant_marker = true; + marked_constant_operands.push_back(absl::StrFormat( + "%s:contents=%d", instruction->name(), + instruction->contents().size())); + } + } + for (const HloInstruction* operand : instruction->operands()) { + if (operand->opcode() != HloOpcode::kConstant || + !operand->has_contents()) { + continue; + } + bool has_dynamic_content = false; + for (const auto& content : operand->contents()) { + has_dynamic_content = + content.node_type_case() != ExpressionProto::kConstantValue && + content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET; + if (has_dynamic_content) { + break; + } + } + if (!has_dynamic_content) { + continue; + } + source_has_dynamic_constant_marker = true; + marked_constant_operands.push_back(absl::StrFormat( + "%s:contents=%d literal=%s", operand->name(), + operand->contents().size(), operand->literal().ToString())); + } + if (source_has_dynamic_constant_marker) { + VLOG(1) << "Skipping HloConstantFolding for " << instruction->name() + << " (" << HloOpcodeString(instruction->opcode()) + << ") because it or its source constant operands carry dynamic contents"; + VLOG(1) << "Marked constant operands: " + << absl::StrJoin(marked_constant_operands, ", "); + continue; + } + // Check for instructions that we can't fold even if they appear inside of // a subcomputation (e.g. a kCall). if (IsOrContainsIllegalInstr(instruction)) { @@ -323,6 +374,17 @@ absl::StatusOr HloConstantFolding::Run( changed = true; HloInstruction* new_constant = instruction->AddInstruction( HloInstruction::CreateConstant(std::move(result))); + VLOG(1) << "HloConstantFolding created constant from " + << instruction->name() << " (" + << HloOpcodeString(instruction->opcode()) + << "), source_has_dynamic_constant_marker=" + << source_has_dynamic_constant_marker; + if (!marked_constant_operands.empty()) { + VLOG(1) << "Marked constant operands: " + << absl::StrJoin(marked_constant_operands, ", "); + } + VLOG(1) << "Folded constant literal -> " + << new_constant->literal().ToString(); if (new_constant->shape().has_layout()) { // Update element_size_in_bits on the new instruction's layout. Literals // always have element_size_in_bits set to 0, and CreateConstant copies diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 7cd03e529ae86a..bc3afa6f73ba2a 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -121,6 +121,26 @@ cc_library( ], ) +cc_library( + name = "dynamic_constant_rewriter", + srcs = ["dynamic_constant_rewriter.cc"], + hdrs = ["dynamic_constant_rewriter.h"], + deps = [ + "//xla:literal_util", + "//xla:shape_util", + "//xla:status_macros", + "//xla/hlo/ir:hlo", + "//xla/hlo/pass:hlo_pass", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@local_tsl//tsl/platform:errors", + "@local_tsl//tsl/platform:protobuf", + ], +) + cc_library( name = "all_reduce_promotion", srcs = ["all_reduce_promotion.cc"], diff --git a/third_party/xla/xla/service/dynamic_constant_rewriter.cc b/third_party/xla/xla/service/dynamic_constant_rewriter.cc new file mode 100644 index 00000000000000..3c734006456b79 --- /dev/null +++ b/third_party/xla/xla/service/dynamic_constant_rewriter.cc @@ -0,0 +1,165 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 "xla/service/dynamic_constant_rewriter.h" + +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/literal_util.h" +#include "xla/shape.h" +#include "xla/shape_util.h" +#include "xla/status_macros.h" +#include "tsl/platform/errors.h" + +namespace xla { +namespace { + +absl::StatusOr BuildDynamicConstantReplacement( + HloInstruction* constant_instr) { + TF_RET_CHECK(constant_instr->opcode() == HloOpcode::kConstant); + TF_RET_CHECK(constant_instr->has_contents()); + + const Shape& shape = constant_instr->shape(); + TF_RET_CHECK(shape.IsArray()); + TF_RET_CHECK(shape.element_type() == S32 || shape.element_type() == S64) + << "Only s32/s64 marked constants are supported"; + TF_RET_CHECK(shape.dimensions_size() <= 1) + << "Only scalar and rank-1 marked constants are supported"; + + int64_t dynamic_index = 0; + DExpr expr; + for (int64_t i = 0; i < constant_instr->contents().size(); ++i) { + DExpr candidate = DExprFromProto(constant_instr->contents()[i]); + if (candidate && candidate->is_dynamic()) { + dynamic_index = i; + expr = std::move(candidate); + break; + } + } + TF_RET_CHECK(expr) << "Marked dynamic constant is missing dynamic contents"; + + int64_t carrier_bound; + if (shape.dimensions_size() == 0) { + dynamic_index = 0; + carrier_bound = shape.element_type() == S32 + ? constant_instr->literal().GetFirstElement() + : constant_instr->literal().GetFirstElement(); + } else { + TF_RET_CHECK(dynamic_index >= 0 && dynamic_index < shape.dimensions(0)) + << "dynamic content index=" << dynamic_index + << " out of bounds for shape " << shape.ToString(); + carrier_bound = shape.element_type() == S32 + ? constant_instr->literal().data()[dynamic_index] + : constant_instr->literal().data()[dynamic_index]; + } + + TF_RET_CHECK(carrier_bound <= std::numeric_limits::max() && + carrier_bound >= std::numeric_limits::min()) + << "GetExpressionValue carriers must fit in s32, got " << carrier_bound; + + HloComputation* computation = constant_instr->parent(); + HloInstruction* carrier = computation->AddInstruction( + HloInstruction::CreateConstant( + LiteralUtil::CreateR1( + {static_cast(carrier_bound)}))); + ExpressionProto expr_proto; + expr.to_proto(&expr_proto); + + HloInstruction* runtime_value = computation->AddInstruction( + HloInstruction::CreateCustomCall( + ShapeUtil::MakeShape(S32, {}), {carrier}, "GetExpressionValue")); + runtime_value->set_contents({std::move(expr_proto)}); + if (shape.element_type() == S64) { + runtime_value = computation->AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::MakeShape(S64, {}), runtime_value)); + } + + if (shape.dimensions_size() == 0) { + return runtime_value; + } + + HloInstruction* base_constant = + computation->AddInstruction(constant_instr->Clone()); + base_constant->set_contents({}); + Shape update_shape = ShapeUtil::MakeShape(shape.element_type(), {1}); + HloInstruction* update = computation->AddInstruction( + HloInstruction::CreateReshape(update_shape, runtime_value)); + HloInstruction* start_index = computation->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0( + static_cast(dynamic_index)))); + return computation->AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + shape, base_constant, update, {start_index})); +} + +bool IsMarkedDynamicConstant(const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kConstant && instr->has_contents(); +} + +} // namespace + +absl::StatusOr DynamicConstantRewriter::Run( + HloModule* module, + const absl::flat_hash_set& execution_threads) { + bool changed = false; + for (HloComputation* computation : + module->MakeNonfusionComputations(execution_threads)) { + std::vector marked_constants; + for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { + if (instruction->opcode() == HloOpcode::kConstant) { + const bool is_marked = IsMarkedDynamicConstant(instruction); + VLOG(1) << "DynamicConstantRewriter sees constant " + << instruction->name() << " shape=" + << instruction->shape().ToString() + << " literal=" << instruction->literal().ToString() + << " marked=" << is_marked; + if (is_marked) { + VLOG(1) << " contents_size=" << instruction->contents().size(); + marked_constants.push_back(instruction); + } + } + } + + for (HloInstruction* constant_instr : marked_constants) { + if (constant_instr->IsDead()) { + continue; + } + TF_ASSIGN_OR_RETURN(HloInstruction * replacement, + BuildDynamicConstantReplacement(constant_instr)); + VLOG(1) << "Rewriting marked constant " << constant_instr->name() + << " into " << replacement->name() << " (" + << HloOpcodeString(replacement->opcode()) << ")"; + TF_RETURN_IF_ERROR( + computation->ReplaceInstruction(constant_instr, replacement)); + replacement->erase_frontend_attribute("dynamic_constant_index"); + replacement->erase_frontend_attribute("dynamic_constant_expr"); + changed = true; + } + } + return changed; +} + +} // namespace xla diff --git a/third_party/xla/xla/service/dynamic_constant_rewriter.h b/third_party/xla/xla/service/dynamic_constant_rewriter.h new file mode 100644 index 00000000000000..e1bd5c57cadf5c --- /dev/null +++ b/third_party/xla/xla/service/dynamic_constant_rewriter.h @@ -0,0 +1,41 @@ +/* Copyright 2026 The OpenXLA Authors. + +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. +==============================================================================*/ + +#ifndef XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_ +#define XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_ + +#include "absl/container/flat_hash_set.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/pass/hlo_pass_interface.h" + +namespace xla { + +class DynamicConstantRewriter : public HloModulePass { + public: + absl::string_view name() const override { + return "dynamic_constant_rewriter"; + } + + using HloPassInterface::Run; + absl::StatusOr Run( + HloModule* module, + const absl::flat_hash_set& execution_threads) override; +}; + +} // namespace xla + +#endif // XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_ From 5b64b8f4a4a32b82cf8a7ea6e431604ca988513c Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:04:08 +0100 Subject: [PATCH 10/21] Cluster graphs using normalized dynamic subexpressions Normalize compatible dynamic expressions around their smallest shared variable-covering subexpression before clustering. Reject cluster merges whose dynamic inputs and outputs cannot use one consistent symbolic core. --- .../jit/encapsulate_subgraphs_pass.cc | 88 +- tensorflow/compiler/jit/encapsulate_util.cc | 5 + tensorflow/compiler/jit/encapsulate_util.h | 13 + tensorflow/compiler/jit/flags.cc | 11 + tensorflow/compiler/jit/flags.h | 6 + .../compiler/jit/mark_for_compilation_pass.cc | 752 +++++++++++++++++- .../compiler/jit/mark_for_compilation_pass.h | 12 + .../jit/mark_for_compilation_pass_test.cc | 69 ++ .../compiler/jit/partially_decluster_pass.cc | 14 + tensorflow/core/graph/subgraph.cc | 14 + 10 files changed, 980 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 3e8a43ce08ed58..080bef205441f2 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/types/optional.h" +#include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/mark_for_compilation_pass.h" #include "tensorflow/compiler/jit/shape_inference_helpers.h" @@ -54,9 +55,14 @@ limitations under the License. #include "tensorflow/core/public/version.h" #include "tensorflow/core/util/device_name_utils.h" #include "tensorflow/core/util/dump_graph.h" - +#include "tensorflow/core/framework/tensor_shape.pb.h" namespace tensorflow { +static const absl::flat_hash_set kFailingOps = { + "Where", + // add more here +}; + const char* const kXlaCompiledKernelAttr = "_XlaCompiledKernel"; const char* const kXlaNumConstantArgsAttr = "_XlaNumConstantArgs"; const char* const kXlaNumResourceArgsAttr = "_XlaNumResourceArgs"; @@ -114,6 +120,41 @@ void MarkGuaranteedConstants( } } +// Helper to convert ExpressionProto to a readable string. +std::string ExprProtoToString(const ExpressionProto& e) { + switch (e.node_type_case()) { + case ExpressionProto::kConstantValue: + return std::to_string(e.constant_value()); + case ExpressionProto::kVariableId: + return absl::StrCat("Var(", e.variable_id(), ")"); + case ExpressionProto::kAddNode: + return absl::StrCat("(", ExprProtoToString(e.add_node().lhs()), " + ", + ExprProtoToString(e.add_node().rhs()), ")"); + case ExpressionProto::kSubNode: + return absl::StrCat("(", ExprProtoToString(e.sub_node().lhs()), " - ", + ExprProtoToString(e.sub_node().rhs()), ")"); + case ExpressionProto::kMulNode: + return absl::StrCat("(", ExprProtoToString(e.mul_node().lhs()), " * ", + ExprProtoToString(e.mul_node().rhs()), ")"); + case ExpressionProto::kDivNode: + return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ", + ExprProtoToString(e.div_node().rhs()), ")"); + case ExpressionProto::kMaxNode: + return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ", + ExprProtoToString(e.max_node().rhs()), ")"); + case ExpressionProto::kGtNode: + return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ", + ExprProtoToString(e.gt_node().rhs()), ")"); + case ExpressionProto::kSelectNode: + return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()), + ", ", ExprProtoToString(e.select_node().on_true()), + ", ", ExprProtoToString(e.select_node().on_false()), + ")"); + default: + return ""; + } +} + struct OutputInputTensorPairHasher { uint64 operator()(std::pair const& s) const { return Hash64Combine(OutputTensor::Hash()(s.first), @@ -369,6 +410,19 @@ class Encapsulator { namespace { +bool BuildOutputShapeProto(const Node& node, int output_slot, + TensorShapeProto* proto) { + AttrSlice attrs = node.attrs(); + auto shape_attr = + attrs.FindByString(kXlaInferredOutputTensorShapesAttrName); + if (shape_attr == nullptr || !shape_attr->has_list() || + shape_attr->list().shape_size() <= output_slot) { + return false; + } + *proto = shape_attr->list().shape(output_slot); + return true; +} + // Return in 'sorted' a topological sort of clusters according to the // dependencies encoded in ancestors. clusters is the list of all clusters // including clusters that are not present in the ancestors map. has_successors @@ -470,6 +524,26 @@ absl::Status Encapsulator::Subgraph::RecordArg( DataType dtype = edge->dst()->input_type(edge->dst_input()); builder.Attr("T", dtype); builder.Attr("index", arg_index); + AttrSlice attrs = src_node->attrs(); + TensorShapeProto output_shape_proto; + if (BuildOutputShapeProto(*src_node, src_slot, &output_shape_proto)) { + VLOG(1) << "Adding following output shapes for node " << src_node->name() + << " : " << output_shape_proto.DebugString(); + builder.Attr("_output_shapes", {output_shape_proto}); + builder.Attr(kXlaInferredOutputShapesAttrName, {output_shape_proto}); + } else { + // if cluster argument is the real argument. + auto build_attr = attrs.FindByString("_dynamic_dim"); + if (build_attr) { + VLOG(1) << "Found Dynamic dimension in " << src_node->name() << ":" + << src_slot; + builder.Attr("_dynamic_dim", *build_attr); + } + } + auto shape_derived_attr = attrs.FindByString(kXlaShapeDerivedAttrName); + if (shape_derived_attr) { + builder.Attr(kXlaShapeDerivedAttrName, *shape_derived_attr); + } absl::Status s = builder.Finalize(&arg_def); if (!s.ok()) return s; @@ -1143,6 +1217,14 @@ static absl::Status RenumberArguments(Graph* graph, return absl::OkStatus(); } +static bool SubgraphHasFailingOps(const Graph& g) { + for (Node* n : g.op_nodes()) { + if (n->IsRetval()) continue; + if (kFailingOps.contains(n->def().op())) return true; + } + return false; +} + absl::Status EncapsulateSubgraphsPass::Run( const GraphOptimizationPassOptions& options) { VLOG(1) << "EncapsulateSubgraphsPass::Run"; @@ -1289,8 +1371,8 @@ absl::Status EncapsulateSubgraphsPass::Run( // TODO(phawkins): add a forward is-constant analysis, similarly split // outputs into host-memory constants and device-memory non-constants. - - AddNodeAttr(kXlaCompiledKernelAttr, true, node); + bool compile_enabled = !SubgraphHasFailingOps(**subgraph); + AddNodeAttr(kXlaCompiledKernelAttr, compile_enabled, node); AddNodeAttr(kXlaNumConstantArgsAttr, num_consts, node); AddNodeAttr(kXlaNumResourceArgsAttr, num_resources, node); return absl::OkStatus(); diff --git a/tensorflow/compiler/jit/encapsulate_util.cc b/tensorflow/compiler/jit/encapsulate_util.cc index fa94a341bbabc6..e680849a0fb8f6 100644 --- a/tensorflow/compiler/jit/encapsulate_util.cc +++ b/tensorflow/compiler/jit/encapsulate_util.cc @@ -303,6 +303,11 @@ absl::Status PostprocessControlEdgesBetweenOutsideCompilations( } // namespace const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes"; +const char kXlaInferredOutputTensorShapesAttrName[] = + "_xla_inferred_output_tensor_shapes"; +const char kXlaInferredOutputShapesAttrName[] = + "_xla_inferred_output_shapes"; +const char kXlaShapeDerivedAttrName[] = "_xla_shape_derived"; const char kXlaConnectedToXlaComputationAttrName[] = "_xla_connected_to_xla_computation"; diff --git a/tensorflow/compiler/jit/encapsulate_util.h b/tensorflow/compiler/jit/encapsulate_util.h index 7c99763c770728..f0f0b0570fd19c 100644 --- a/tensorflow/compiler/jit/encapsulate_util.h +++ b/tensorflow/compiler/jit/encapsulate_util.h @@ -28,6 +28,19 @@ namespace tensorflow { // a list of PartialTensorShape objects. extern const char kXlaInferredShapesAttrName[]; +// Attribute marking Grappler-inferred output TensorShapeProtos. Attribute +// value is a list of TensorShapeProto objects and may include ExpressionProto +// annotations when available. +extern const char kXlaInferredOutputTensorShapesAttrName[]; + +// Attribute carrying inferred output TensorShapeProtos on encapsulated _Arg +// nodes for XlaCompileOp argument reconstruction. +extern const char kXlaInferredOutputShapesAttrName[]; + +// Attribute indicating that a node is shape-derived and worth tracking across +// JIT boundaries, e.g. Shape and shape-fed Cast. +extern const char kXlaShapeDerivedAttrName[]; + // Infers output shapes for all nodes in graph `g`. The output shapes will be // stored in node attribute `kXlaInferredShapesAttrName`. // diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index 622b9a369fcb12..8dd2ac0af31614 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -108,6 +108,15 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { Flag("tf_xla_max_cluster_size", &mark_for_compilation_flags->tf_xla_max_cluster_size, "Maximum number of operators in an XLA compilation."), + Flag("tf_xla_annotate_cluster_id", + &mark_for_compilation_flags->tf_xla_annotate_cluster_id, + "Allow operator names to influence clustering scheme." + "Operators whose name starting with .cluster.{id} will likely" + "to be clustered together if the ids are the same number. " + ".cluster.none will not be clustered with those having numbered id"), + Flag("tf_xla_cluster_parallel", + &mark_for_compilation_flags->tf_xla_cluster_parallel, + "Split parallel compute subgraph info different clusters"), Flag( "tf_xla_ops_to_cluster", &mark_for_compilation_flags->tf_xla_ops_to_cluster, @@ -238,6 +247,8 @@ void AllocateAndParseFlags() { mark_for_compilation_flags->tf_xla_min_cluster_size = 4; mark_for_compilation_flags->tf_xla_max_cluster_size = std::numeric_limits::max(); + mark_for_compilation_flags->tf_xla_annotate_cluster_id = false; + mark_for_compilation_flags->tf_xla_cluster_parallel = false; mark_for_compilation_flags->tf_xla_clustering_debug = false; mark_for_compilation_flags->tf_xla_cpu_global_jit = false; mark_for_compilation_flags->tf_xla_clustering_fuel = diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index fb443e619e278f..c37391d2d1ff56 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -62,6 +62,12 @@ struct MarkForCompilationPassFlags { // Maximum number of operators in an XLA compilation. int32 tf_xla_max_cluster_size; + // Enable operator name to influence clustering decision + bool tf_xla_annotate_cluster_id; + + // Split parallel compute subgraph info different clusters + bool tf_xla_cluster_parallel; + // If non-empty, limit XLA clustering to the following TF operations. string tf_xla_ops_to_cluster; diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index c3a24f3e0f7163..81365b13e5a9d5 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -23,12 +23,14 @@ limitations under the License. #include #include #include +#include #include #include #include #include #include #include +#include #include "absl/base/call_once.h" #include "absl/container/flat_hash_map.h" @@ -39,6 +41,7 @@ limitations under the License. #include "tensorflow/compiler/jit/deadness_analysis.h" #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/device_util.h" +#include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/resource_operation_safety_analysis.h" #include "tensorflow/compiler/jit/xla_cluster_util.h" @@ -68,6 +71,9 @@ limitations under the License. #include "tensorflow/core/platform/types.h" #include "tensorflow/core/public/version.h" #include "tensorflow/core/util/dump_graph.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/grappler/costs/graph_properties.h" +#include "tensorflow/core/grappler/grappler_item.h" namespace tensorflow { @@ -109,6 +115,8 @@ class MarkForCompilationPassImpl { // stable from run to rum. bool deterministic_cluster_names; + bool enable_dynamic_sizes; + int max_cluster_size; int min_cluster_size; @@ -123,6 +131,11 @@ class MarkForCompilationPassImpl { std::atomic* fuel; bool dump_graphs; + + // Enable models to influcence clustering with operator names + int annotate_cluster_id; + + bool enable_cluster_parallel; }; MarkForCompilationPassImpl(DebugOptions debug_options, Graph* graph, @@ -245,7 +258,28 @@ class MarkForCompilationPassImpl { " others #", cycles_graph_node_id(), ">"); } + int annotated_id() const { return annotated_id_; } + void set_annotated_id(int id) { annotated_id_ = id; } + int chain_id() const {return chain_id_;} + void set_chain_id(int id) {chain_id_ = id;} + void add_dim_var(int dim_var) { dim_vars_.insert(dim_var); } + void merge_dim_vars(const std::set& dim_vars) { + 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_; int effective_cluster_size_; @@ -317,6 +351,21 @@ class MarkForCompilationPassImpl { return compilation_candidates_.find(n) != compilation_candidates_.end(); } + 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); + void collectPathNodes(Node* start, std::set &path_nodes, + std::set& merger_nodes); + std::map> collectParallelNode( + const std::vector& nodeSet); + absl::Status AssignParallelChains(); + // Tries to contract the edge from cluster `from` to cluster `to`. Returns // true if successful. absl::StatusOr TryToContractEdge(Cluster* from, Cluster* to); @@ -606,6 +655,11 @@ void MarkForCompilationPassImpl::Cluster::Merge(Cluster* other) { xla_scope_ = std::move(other->xla_scope_); } + 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() + other->resource_var_operation_node_ids_.size()); @@ -651,11 +705,320 @@ absl::Status IgnoreResourceOpForSafetyAnalysis( } return absl::OkStatus(); } +// node mapping to multiple vectors of expressions (one for each output in +// order) +static std::map>>> + expr_map; +// Helper to convert ExpressionProto to a readable string. +std::string ExprProtoToString(const ExpressionProto& e) { + switch (e.node_type_case()) { + case ExpressionProto::kConstantValue: + return std::to_string(e.constant_value()); + case ExpressionProto::kVariableId: + return absl::StrCat("Var(", e.variable_id(), ")"); + case ExpressionProto::kAddNode: + return absl::StrCat("(", ExprProtoToString(e.add_node().lhs()), " + ", + ExprProtoToString(e.add_node().rhs()), ")"); + case ExpressionProto::kSubNode: + return absl::StrCat("(", ExprProtoToString(e.sub_node().lhs()), " - ", + ExprProtoToString(e.sub_node().rhs()), ")"); + case ExpressionProto::kMulNode: + return absl::StrCat("(", ExprProtoToString(e.mul_node().lhs()), " * ", + ExprProtoToString(e.mul_node().rhs()), ")"); + case ExpressionProto::kDivNode: + return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ", + ExprProtoToString(e.div_node().rhs()), ")"); + case ExpressionProto::kMaxNode: + return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ", + ExprProtoToString(e.max_node().rhs()), ")"); + case ExpressionProto::kGtNode: + return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ", + ExprProtoToString(e.gt_node().rhs()), ")"); + case ExpressionProto::kSelectNode: + return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()), + ", ", ExprProtoToString(e.select_node().on_true()), + ", ", ExprProtoToString(e.select_node().on_false()), + ")"); + default: + return ""; + } +} + +bool HasDynamicInputExpression(const Node* node) { + for (const Edge* edge : node->in_edges()) { + if (edge->IsControlEdge()) { + continue; + } + auto it = expr_map.find(edge->src()->name()); + if (it == expr_map.end()) { + continue; + } + const int output_index = edge->src_output(); + if (output_index < 0 || output_index >= it->second.size()) { + continue; + } + for (const auto& expr : it->second[output_index]) { + if (expr == nullptr) { + continue; + } + xla::DExpr dynamic_expr = *expr; + if (dynamic_expr && dynamic_expr->is_dynamic()) { + return true; + } + } + } + return false; +} + +std::string DynamicExpressionToString(const xla::DExpr& expr) { + xla::DExpr simplified = expr.simplify(); + if (!simplified && !simplified.is_unknown()) { + return ""; + } + xla::StringPrinter printer; + simplified->print(&printer); + return std::move(printer).ToString(); +} + +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(); + } + int replacement_id = std::numeric_limits::min(); + while (expected_ids.count(replacement_id) != 0) { + ++replacement_id; + } + const std::set replacement_ids = {replacement_id}; + + 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=", + DynamicExpressionToString(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=", + DynamicExpressionToString(*expected_core), ", expr=", + DynamicExpressionToString(expr), ", core=", + DynamicExpressionToString(core)); + } + const xla::DExpr normalized = + expr.replace_subexpression(core, xla::DExpr::Var(replacement_id)) + .simplify(); + if (normalized->get_all_ids() != replacement_ids) { + return absl::StrCat( + "dynamic expression core does not cover every variable occurrence: " + "expr=", + DynamicExpressionToString(expr), ", core=", + DynamicExpressionToString(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 = *expr_ptr; + 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 = *expr_ptr; + 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) { + using tensorflow::ExpressionProto; + using tensorflow::GraphDef; + using tensorflow::NodeDef; + using tensorflow::TensorShapeProto; + using tensorflow::grappler::GraphProperties; + using tensorflow::grappler::GrapplerItem; + + expr_map.clear(); + + GraphDef graph_def; + graph.ToGraphDef(&graph_def); + auto node_name_index = graph.BuildNodeNameIndex(); + + GrapplerItem item; + item.id = "mark_for_compilation_pass_expr_dump"; + item.graph = graph_def; + + GraphProperties props(item); + + absl::Status st = props.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true); + + if (!st.ok()) { + LOG(ERROR) << "[EXPR][GP] InferStatically failed: " << st.message(); + return; + } + + int found = 0; + VLOG(1) << "[EXPR][GP] === GraphProperties output expr dump ==="; + + auto convert_graph_properties_shape = [](const TensorShapeProto& gp_shape) { + TensorShapeProto out; + out.set_unknown_rank(gp_shape.unknown_rank()); + for (int i = 0; i < gp_shape.dim_size(); ++i) { + const auto& dim = gp_shape.dim(i); + out.add_dim()->set_size(dim.size()); + ExpressionProto* expr = out.add_expressions(); + if (i < gp_shape.expressions_size() && + gp_shape.expressions(i).node_type_case() != + ExpressionProto::NODE_TYPE_NOT_SET) { + *expr = gp_shape.expressions(i); + } else { + expr->set_constant_value(dim.size()); + } + } + return out; + }; + + for (const NodeDef& n : graph_def.node()) { + if (!props.HasOutputProperties(n.name())) continue; + const auto& outs = props.GetOutputProperties(n.name()); + std::vector inferred_output_shapes; + inferred_output_shapes.reserve(outs.size()); + std::vector>> list_exprs(outs.size()); + for (int out_idx = 0; out_idx < static_cast(outs.size()); ++out_idx) { + const auto& tp = outs[out_idx]; + const TensorShapeProto& shp = tp.shape(); + inferred_output_shapes.push_back(convert_graph_properties_shape(shp)); + + std::vector> exprs; + for (int d = 0; d < shp.dim_size(); ++d) { + if (d >= shp.expressions_size()) continue; + const ExpressionProto& expr = shp.expressions(d); + if (expr.node_type_case() == ExpressionProto::NODE_TYPE_NOT_SET) + continue; + + VLOG(1) << "Node " << n.name() << " has expression " + << ExprProtoToString(expr); + + exprs.push_back( + std::make_unique(DimExprFromProto(expr))); + + ++found; + } + if (shp.dim_size() == 0 && shp.unknown_rank()) { + // Add two dummy variables to represent the unknown rank + exprs.push_back(std::make_unique(DimExpr::Var(-888))); + exprs.push_back(std::make_unique(DimExpr::Var(-889))); + } + + list_exprs[out_idx] = std::move(exprs); + } + expr_map[n.name()] = std::move(list_exprs); + auto node_it = node_name_index.find(n.name()); + if (node_it != node_name_index.end()) { + node_it->second->AddAttr(kXlaInferredOutputTensorShapesAttrName, + inferred_output_shapes); + } + + } + + VLOG(1) << "[EXPR][GP] === Found " << found + << " expressions via GraphProperties ==="; +} 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()) { @@ -686,6 +1049,44 @@ absl::StatusOr MarkForCompilationPassImpl::Initialize() { // representative names the node in the 'cycles' graph that represents the // cluster. TF_RETURN_IF_ERROR(BuildInitialClusterSet()); + + // Source model may be annotated with preferred clusters. This function + // just interpreter the annotations and assign preferred IDs + if (debug_options_.annotate_cluster_id) { + TF_RETURN_IF_ERROR(AssignAnnotatedClusterIDs()); + } + if (debug_options_.enable_dynamic_sizes) { + TF_RETURN_IF_ERROR(AssignDimVars()); + for (Node* n : graph_->op_nodes()) { + bool mark_shape_derived = false; + auto is_shape_like = [](const Node* node) { + const string& type = node->type_string(); + return type == "Shape" || type == "ShapeN" || type == "Size"; + }; + if (is_shape_like(n)) { + mark_shape_derived = HasDynamicInputExpression(n); + } else if (n->type_string() == "Cast") { + for (const Edge* edge : n->in_edges()) { + if (edge->IsControlEdge()) { + continue; + } + const Node* src = edge->src(); + if (is_shape_like(src) && HasDynamicInputExpression(src)) { + mark_shape_derived = true; + break; + } + } + } + if (mark_shape_derived) { + n->AddAttr(kXlaShapeDerivedAttrName, true); + VLOG(1) << "MarkForCompilation marked shape-derived node " + << n->name() << " op=" << n->type_string(); + } + } + } + if (debug_options_.enable_cluster_parallel) { + TF_RETURN_IF_ERROR(AssignParallelChains()); + } return true; } @@ -893,7 +1294,12 @@ absl::Status MarkForCompilationPassImpl::RunEdgeContractionLoop() { ForEachEdgeInPostOrder([&](Cluster* from, Cluster* to) { return TryToContractEdge(from, to); })); + /* Clustering conditions for dynamic shapes may break the assumption of + * fixed point at phase 2, so this check may fail. + * Now just disable this check, and we can re-enable it after we have more + * confidence in the code. TF_RET_CHECK(!changed); + */ return absl::OkStatus(); } @@ -1028,6 +1434,7 @@ absl::Status MarkForCompilationPassImpl::CreateClusters() { // trouble. if (cluster->effective_cluster_size() >= debug_options_.min_cluster_size || + cluster->chain_id() != -1 || cluster->has_functional_control_flow() || cluster->is_xla_compile_attr_true()) { string& name = cluster_names[cluster->cycles_graph_node_id()]; @@ -1429,6 +1836,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( @@ -1548,14 +1964,322 @@ bool MarkForCompilationPassImpl::CompilationDisallowedByXlaCompileAttr( return false; } +absl::Status MarkForCompilationPassImpl::AssignAnnotatedClusterIDs(void) { + VLOG(1) << "Run AssignAnnotatedClusterIDs"; + + for (Node* node : graph_->nodes()) { + Cluster * cluster = GetClusterForNode(node); + auto name = node->name(); + if (cluster) { + std::string pat = "^\\.cluster\\.(\\d+|none)"; + std::regex idPattern(pat); + std::smatch matched; + if (std::regex_search(name, matched, idPattern)) { + auto m = matched.str(1); + if (m == "none") { + // Prefer not to cluster + VLOG(1) << name << " : Default annotated cluster id -1"; + cluster->set_annotated_id(-1); + } + else { + try { + int id = std::stoi(m); + cluster->set_annotated_id(id); + VLOG(1) << name << " : Set annotated cluster id " << m; + } + catch (...) { + VLOG(1) << name << " : Invalid cluster id: " << m; + } + } + } + else { + VLOG(1) << "Not matched: " << name << " pattern is " << pat; + } + } + else { + VLOG(1) << name << ": Not initially clustered"; + } + } + return absl::OkStatus(); +} + +absl::Status MarkForCompilationPassImpl::AssignDimVars(void) { + for (Node* node : graph_->nodes()) { + auto node_name = node->name(); + Cluster * cluster = GetClusterForNode(node); + if (!cluster) continue; + for (const tensorflow::Edge* edge : node->in_edges()) { + if (edge->IsControlEdge()) { + // Skip control edges if you are only interested in data edges + continue; + } + + const tensorflow::Node* input = edge->src(); // Source node of the edge + auto it = expr_map.find(input->name()); + if (it == expr_map.end()) { + VLOG(2) << "No expression found for node " << input->name(); + continue; + } + + auto output_index = edge->src_output(); // Output index of the source node + if (output_index >= (it->second).size()) { + LOG(INFO) << "Warning: Output index " << output_index << " is out of bounds for node " << input->name(); + continue; + } + for (auto& pDim: (it->second)[output_index]) { + DimExpr * d= pDim.get(); + xla::DExpr dyn = *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); + VLOG(2) << "Add dim var " << id << " to cluster of node "<< node_name; + } + } + } + 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 = *expr_ptr; + 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()) { + VLOG(2) << "Cluster of node " << node_name << " has no dim vars."; + } + else { + std::string id_str; + for (auto id : cluster->dim_vars()) { + id_str += "Dim var " + std::to_string(id) + ", "; + } + VLOG(2) << "Cluster of node " << node_name << " has dim vars:\n" << id_str; + } + } + } + return absl::OkStatus(); +} + bool MarkForCompilationPassImpl::LogNotContractableAndReturnFalse( Cluster* from, Cluster* to, absl::string_view reason) { VLOG(3) << EdgeContractionFailureMsg(from, to, reason); return false; } +void MarkForCompilationPassImpl::collectInputNodes(std::set &path_nodes) { + std::unordered_map out_degree_count; + + // 4. Initialize the queue and add nodes from path_nodes + std::queue queue; + for (auto node : path_nodes) { + queue.push(node); + } + + // 5. BFS search + while (!queue.empty()) { + auto u = queue.front(); + queue.pop(); + + // Traverse all predecessor nodes + for (const Edge* e : u->in_edges()) { + Node* p = e->src(); + if (path_nodes.find(p) != path_nodes.end() || + !IsCompilationCandidate(p)) { + continue; + } + if (out_degree_count.count(p) == 0) { + // Initialize the out-degree count for the node + out_degree_count[p] = p->out_edges().size(); + } + // Decrease the out-degree count for the predecessor node + out_degree_count[p] -= 1; + + // If the predecessor node's out-degree count is 0 and not in path_nodes, + // add it to path_nodes and queue + if (out_degree_count[p] == 0) { + path_nodes.insert(p); + queue.push(p); + VLOG(3) << p->DebugString(); + } + } + } +} + +void MarkForCompilationPassImpl::collectMergeNodes( + const std::vector& nodeSet, std::set &merger_nodes) { + // 1. Collect the number of nodeSet that can reach each node + std::map> reach_map; + for (Node* start : nodeSet) { + std::set visited; + std::vector stack = {start}; + while (!stack.empty()) { + Node* cur = stack.back(); + stack.pop_back(); + if (visited.count(cur)) continue; + visited.insert(cur); + for (const Edge* e : cur->out_edges()) { + Node* next = e->dst(); + if (!visited.count(next)) + stack.push_back(next); + } + } + reach_map[start] = std::move(visited); + } + + // 2. Determine the merger node + std::map node_reach_count; + std::set all_nodes; + for (const auto& kv : reach_map) { + for (Node* n : kv.second) { + node_reach_count[n]++; + all_nodes.insert(n); + } + } + for (Node* n : all_nodes) { + // Condition 1: Reached by multiple sources + if (node_reach_count[n] >= 2) merger_nodes.insert(n); + // Condition 2: No output edges + if (n->out_edges().empty()) merger_nodes.insert(n); + } +} + +void MarkForCompilationPassImpl::collectPathNodes( + Node* start, std::set &path_nodes, std::set& merger_nodes) { + std::vector stack = {start}; + std::set visited; + + while (!stack.empty()) { + Node* cur = stack.back(); + stack.pop_back(); + if (visited.count(cur) || !IsCompilationCandidate(cur)) { + continue; + } + visited.insert(cur); + + // Stop search met the merger node + if (merger_nodes.count(cur)) { + if (cur->out_edges().empty()) path_nodes.insert(cur); + continue; + } + path_nodes.insert(cur); + + for (const Edge* e : cur->out_edges()) { + Node* next = e->dst(); + if (!visited.count(next)) { + stack.push_back(next); + } + } + } +} + +// collectParallelNode +// Search the serial merger nodes based on the parallel matmul starting points +// Search along the output edge to get the boundary from start to all merger points +// Search along the input edge to get the entire parallel computation graph +std::map> +MarkForCompilationPassImpl::collectParallelNode( + const std::vector& nodeSet) { + std::set merger_nodes; + collectMergeNodes(nodeSet, merger_nodes); + + // Collect path nodes + std::map> result; + for (Node* start : nodeSet) { + VLOG(4) << "Search parallel graph form node: " << start->DebugString(); + std::set path_nodes; + + // Search along output edge form start to merger nodes + collectPathNodes(start, path_nodes, merger_nodes); + + VLOG(4) << "Collect path nodes:"; + for (auto node : path_nodes) { + VLOG(4) << node->type_string(); + } + + VLOG(4) << "Collect input nodes:"; + // search along input edge + collectInputNodes(path_nodes); + + result[start] = std::vector(path_nodes.begin(), path_nodes.end()); + } + return result; +} + +// Collect parallel matmuls as input nodes into nodeSet +// Use collectParallelNode to get the parallel subgraph +// Mark parallel nodes change ID +absl::Status MarkForCompilationPassImpl::AssignParallelChains() { + VLOG(4) << "Run AssignParallelChains"; + // Record the matmuls that can be paralleled + std::vector> parallel_matmuls; + int minParallelMatmulNum = 2; + int next_chain_id = 0; + // Collect matmul nodes with shared input to parallel_matmuls + for (Node* node : graph_->nodes()) { + if (node->out_edges().size() < 2) continue; + std::vector matmul_nodes; + for (const Edge* e : node->out_edges()) { + if (e->IsControlEdge()) continue; + Node* succ = e->dst(); + VLOG(4) << "Find matmul node: " << succ->type_string() << " : " << succ->DebugString(); + if (succ->type_string() == "MatMul") + matmul_nodes.push_back(succ); + } + if (matmul_nodes.size() >= minParallelMatmulNum) + parallel_matmuls.push_back(matmul_nodes); + } + + for (auto matmul_nodes : parallel_matmuls) { + VLOG(4) << "Process matmul nodes: Total " << matmul_nodes.size() + << " sub matmuls"; + bool visited = false; + for (auto matmul : matmul_nodes) { + VLOG(4) << matmul->name(); + Cluster* cluster = GetClusterForNode(matmul); + if (!cluster || cluster->chain_id() != -1) { + visited = true; + break; + } + } + if (visited) { + VLOG(4) << "stop collect: has visited matmul"; + continue; + } + + std::map> subgraphMap = + collectParallelNode(matmul_nodes); + for (auto it : subgraphMap) { + auto nodeSet = it.second; + VLOG(4) << "One of Parallel sub-graph is: "; + for (auto node : nodeSet) { + VLOG(4) << node->DebugString(); + Cluster* cluster = GetClusterForNode(node); + cluster->set_chain_id(next_chain_id); + } + next_chain_id++; + } + } + return absl::OkStatus(); +} + absl::StatusOr MarkForCompilationPassImpl::TryToContractEdge( Cluster* from, Cluster* to) { + if (from->chain_id() != to->chain_id()) { + return LogNotContractableAndReturnFalse( + from, to, "nodes are in different parallel chains"); + } DCHECK(from->deadness_predicate().has_value() == to->deadness_predicate().has_value()); if (from->deadness_predicate() != to->deadness_predicate()) { @@ -1569,6 +2293,22 @@ absl::StatusOr MarkForCompilationPassImpl::TryToContractEdge( return false; } + if (debug_options_.annotate_cluster_id && from->annotated_id() != to->annotated_id()) { + return LogNotContractableAndReturnFalse( + from, to, "the two nodes do not have same annotated ids"); + } + + if (debug_options_.enable_dynamic_sizes) { + 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); + } + } + TF_ASSIGN_OR_RETURN(bool devices_compatible, AreDevicesCompatible(*from, *to)); if (!devices_compatible) { @@ -1639,7 +2379,6 @@ absl::Status MarkForCompilationPassImpl::Run() { // MarkForCompilationPassImpl is not set up to run the subsequent phases. return absl::OkStatus(); } - TF_RETURN_IF_ERROR(RunEdgeContractionLoop()); TF_RETURN_IF_ERROR(DeclusterNodes()); TF_RETURN_IF_ERROR(CreateClusters()); @@ -1958,10 +2697,14 @@ absl::Status MarkForCompilationPass::Run( debug_options.ignore_xla_compile_attr = false; debug_options.deterministic_cluster_names = flags->tf_xla_deterministic_cluster_names; + debug_options.enable_dynamic_sizes = + flags->tf_xla_enable_dynamic_sizes; debug_options.max_cluster_size = flags->tf_xla_max_cluster_size; debug_options.min_cluster_size = flags->tf_xla_min_cluster_size; debug_options.fuel = GetPointerToFuel(flags->tf_xla_clustering_fuel); debug_options.dump_graphs = flags->tf_xla_clustering_debug; + debug_options.annotate_cluster_id = flags->tf_xla_annotate_cluster_id; + debug_options.enable_cluster_parallel = flags->tf_xla_cluster_parallel; return MarkForCompilation(options, debug_options); } @@ -1977,10 +2720,12 @@ absl::Status MarkForCompilationPass::RunForTest( flags->tf_xla_disable_resource_variable_safety_checks_for_debugging; debug_options.ignore_xla_compile_attr = true; debug_options.deterministic_cluster_names = deterministic_cluster_names; + debug_options.enable_dynamic_sizes = false; debug_options.max_cluster_size = flags->tf_xla_max_cluster_size; debug_options.min_cluster_size = flags->tf_xla_min_cluster_size; debug_options.fuel = GetPointerToFuel(flags->tf_xla_clustering_fuel); debug_options.dump_graphs = flags->tf_xla_clustering_debug; + debug_options.annotate_cluster_id = flags->tf_xla_annotate_cluster_id; return MarkForCompilation(options, debug_options); } @@ -2056,6 +2801,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..bdc4344af22547 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,74 @@ 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, + DynamicExpressionCompatibilityUsesCompleteRepeatedVariableCore) { + const xla::DExpr a = xla::DExpr::Var(1); + const xla::DExpr b = xla::DExpr::Var(2); + const xla::DExpr shared_core = (a + b) * (a - b); + const std::vector exprs = {shared_core + 2, shared_core + 3}; + + EXPECT_FALSE( + testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value()); +} + +TEST(XlaCompilationTest, + DynamicExpressionCompatibilityAcceptsSharedCoreWithDifferentScaling) { + const xla::DExpr shared_core = + xla::DExpr::Var(1) + xla::DExpr::Var(2); + const std::vector exprs = { + shared_core / 240, shared_core, xla::DExpr::Const(3) * shared_core}; + + EXPECT_FALSE( + testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value()); +} + +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/partially_decluster_pass.cc b/tensorflow/compiler/jit/partially_decluster_pass.cc index 442635c9a29696..bdc123a9e3d9d3 100644 --- a/tensorflow/compiler/jit/partially_decluster_pass.cc +++ b/tensorflow/compiler/jit/partially_decluster_pass.cc @@ -19,8 +19,10 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/jit/device_util.h" +#include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/framework/function.h" @@ -47,6 +49,10 @@ absl::Status FindNodesToDecluster(const Graph& graph, MemoryTypeVector input_mtypes, output_mtypes; for (Node* n : post_order) { + if (SymbolicContentEnabled() && + n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) { + continue; + } std::optional from_cluster = GetXlaClusterForNode(*n); if (!from_cluster) { continue; @@ -308,6 +314,10 @@ absl::Status PartiallyDeclusterGraph(Graph* graph, if (!compile_time_const_nodes[n->id()]) { continue; } + if (SymbolicContentEnabled() && + n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) { + continue; + } absl::string_view cluster_name = *GetXlaClusterForNode(*n); bool node_on_cluster_edge = @@ -379,6 +389,10 @@ absl::Status PartiallyDeclusterGraph(Graph* graph) { if (!IsShapeConsumerOp(*n)) { continue; } + if (SymbolicContentEnabled() && + n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) { + continue; + } std::optional cluster = GetXlaClusterForNode(*n); if (!cluster.has_value()) { diff --git a/tensorflow/core/graph/subgraph.cc b/tensorflow/core/graph/subgraph.cc index bb47f37ef7fbe3..d1616e2396515e 100644 --- a/tensorflow/core/graph/subgraph.cc +++ b/tensorflow/core/graph/subgraph.cc @@ -79,6 +79,20 @@ absl::Status FeedInputs( TF_RETURN_IF_ERROR( feed_rewrites[i]->AddNode(g, {n, id.second}, &feed_node)); + // Set an attribute in _Arg node to indicate it has a batch dimension + auto node_attrs = n->attrs(); + const AttrValue* shape_attr = node_attrs.FindByString("_output_shapes"); + if (shape_attr && shape_attr->has_list()) { + const TensorShapeProto& shape = shape_attr->list().shape(0); + for (int i = 0; i < shape.dim_size(); ++i) { + if (shape.dim(i).size() == -1) { + feed_node->AddAttr("_dynamic_dim", i); + break; + } + } + // Keep _output_shapes for further runs of shape inference + feed_node->AddAttr("_output_shapes", *shape_attr); + } // Update name_index (*name_index)[feed_node->name()] = feed_node; // Duplicate control edges aren't allowed, but feed_node was *just* created From 80f4a190484a64ac7a15c9c4d9e86e049a21d4b0 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Wed, 19 Aug 2026 15:50:17 +0100 Subject: [PATCH 11/21] Solve dynamic expressions at XLA cluster boundaries Map compiled XLA parameters back to their runtime TensorFlow inputs and solve the shared dynamic value from concrete dimensions. Filter broadcast-singleton evidence, pass the solved value through execution, and substitute it into dynamic output shapes after execution. --- tensorflow/compiler/jit/kernels/BUILD | 21 +- tensorflow/compiler/jit/kernels/xla_ops.cc | 619 +++++++++++++++++- .../compiler/jit/kernels/xla_ops_internal.h | 84 +++ .../compiler/jit/kernels/xla_ops_test.cc | 282 ++++++++ tensorflow/compiler/jit/xla_launch_util.cc | 133 +++- tensorflow/compiler/jit/xla_launch_util.h | 11 +- tensorflow/compiler/tf2xla/xla_compiler.cc | 68 +- tensorflow/core/BUILD | 1 + tensorflow/core/framework/BUILD | 2 + .../core/framework/batch_size_resource.h | 40 ++ tensorflow/core/kernels/function_ops.cc | 48 +- tensorflow/core/kernels/function_ops.h | 2 + third_party/xla/xla/executable_run_options.h | 8 + 13 files changed, 1298 insertions(+), 21 deletions(-) create mode 100644 tensorflow/compiler/jit/kernels/xla_ops_internal.h create mode 100644 tensorflow/compiler/jit/kernels/xla_ops_test.cc create mode 100644 tensorflow/core/framework/batch_size_resource.h diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index dec057ebfba9ab..587d04b30c5274 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -1,3 +1,4 @@ +load("//tensorflow:tensorflow.bzl", "tf_cc_test") load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") package( @@ -48,7 +49,10 @@ XLA_OPS_DEPS = [ cc_library( name = "xla_ops_no_jit_rewrite_registration", srcs = ["xla_ops.cc"], - hdrs = ["xla_ops.h"], + hdrs = [ + "xla_ops.h", + "xla_ops_internal.h", + ], deps = XLA_OPS_DEPS + [ "//tensorflow/compiler/jit:device_compilation_cache", "//tensorflow/compiler/jit:device_compilation_profiler", @@ -60,6 +64,7 @@ cc_library( "//tensorflow/core/platform:refcount", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:node_hash_map", + "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -74,6 +79,20 @@ cc_library( alwayslink = 1, ) +tf_cc_test( + name = "xla_ops_test", + srcs = ["xla_ops_test.cc"], + env = {"TF_XLA_FLAGS": "--tf_xla_enable_dynamic_sizes"}, + deps = [ + ":xla_ops_no_jit_rewrite_registration", + "//tensorflow/compiler/tf2xla:xla_argument", + "//tensorflow/core:framework", + "//tensorflow/core/platform:test", + "@com_google_googletest//:gtest_main", + "@local_xla//xla:shape_util", + ], +) + cc_library( name = "xla_ops", hdrs = ["xla_ops.h"], diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 468b85280e2a47..f18c11c37bdc21 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/jit/kernels/xla_ops.h" +#include "tensorflow/compiler/jit/kernels/xla_ops_internal.h" #include #include @@ -36,6 +37,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" @@ -61,10 +63,13 @@ limitations under the License. #include "xla/client/local_client.h" #include "xla/executable_run_options.h" #include "xla/pjrt/pjrt_client.h" +#include "xla/printer.h" #include "xla/service/gpu/gpu_executable_run_options.h" +#include "xla/shape_expr.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/tsl/protobuf/error_codes.pb.h" #include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/batch_size_resource.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" @@ -199,6 +204,571 @@ using PjRtExecutableClosure = using PjRtExecutableClosureStore = ExecutableClosureStore; +struct DynamicSolveEvidence { + int64_t solved_value; + int64_t observed_value; + std::string expr; +}; + +struct DynamicSolveSelection { + std::set singleton_values; + std::set nonsingleton_values; + std::set expressions; + std::optional chosen_value; + std::string reason; + bool rejected = false; +}; + +template +DynamicSolveSelection SelectDynamicSolveEvidence( + absl::Span candidates) { + DynamicSolveSelection selection; + for (const Candidate& candidate : candidates) { + const DynamicSolveEvidence& evidence = candidate.evidence; + selection.expressions.insert(evidence.expr); + if (evidence.observed_value == 1) { + selection.singleton_values.insert(evidence.solved_value); + } else { + selection.nonsingleton_values.insert(evidence.solved_value); + } + } + + // A runtime singleton may be an implicitly broadcast operand. Prefer + // consistent non-singleton evidence for the same normalized variable. + if (!selection.nonsingleton_values.empty()) { + if (selection.nonsingleton_values.size() != 1) { + selection.rejected = true; + selection.reason = "conflicting non-singleton candidates"; + return selection; + } + selection.chosen_value = *selection.nonsingleton_values.begin(); + selection.reason = + selection.singleton_values.empty() + ? "used non-singleton evidence" + : "preferred non-singleton evidence over singleton-derived " + "candidates"; + return selection; + } + + if (selection.singleton_values.size() == 1) { + selection.chosen_value = *selection.singleton_values.begin(); + selection.reason = "used singleton-derived evidence only"; + } else if (selection.singleton_values.size() > 1) { + selection.rejected = true; + selection.reason = "conflicting singleton-derived candidates"; + } else { + selection.reason = "no dynamic values were solved"; + } + return selection; +} + +struct DynamicSolveCandidate { + DynamicSolveEvidence evidence; + std::string context; +}; + +} // namespace + +namespace xla_ops_internal { + +int GetRuntimeInputIndex(absl::Span input_mapping, + int xla_input_index, int num_constant_args, + bool constants_omitted) { + if (xla_input_index < 0 || xla_input_index >= input_mapping.size()) { + return -1; + } + const int missing_input_prefix = constants_omitted ? num_constant_args : 0; + return input_mapping[xla_input_index] - missing_input_prefix; +} + +std::optional GetConstantArgumentElementValue( + const XlaArgument& 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(XlaArgument* 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 { + DynamicSolveEvidence evidence; + IgnoredDynamicArgumentOccurrence occurrence; + }; + + DynamicSolveFilterDecision result; + std::map, std::vector> variable_ids_to_candidates; + + for (int arg_index = 0; arg_index < args.size(); ++arg_index) { + const XlaArgument& 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); + variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back( + Candidate{ + DynamicSolveEvidence{*solved_value, shape.dim_size(dim), + expr_string}, + 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); + variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back( + Candidate{ + DynamicSolveEvidence{*solved_value, *observed_value, + expr_string}, + IgnoredDynamicArgumentOccurrence{ + IgnoredDynamicArgumentOccurrence::Source:: + kConstantValueElement, + arg_index, + element_index, + *observed_value, + *solved_value, + expr_string}}); + } + } + + std::vector expr_summaries; + for (const auto& [variable_ids, candidates] : variable_ids_to_candidates) { + DynamicSolveSelection selection = SelectDynamicSolveEvidence( + absl::MakeConstSpan(candidates)); + + if (selection.rejected) { + result.can_run = false; + } else if (selection.chosen_value.has_value() && + !selection.nonsingleton_values.empty()) { + for (const auto& candidate : candidates) { + if (candidate.evidence.observed_value == 1 && + candidate.evidence.solved_value != *selection.chosen_value) { + VLOG(1) << "Ignoring singleton-derived dynamic " + << "occurrence during XLA signature filtering: " + << "variable_ids={" << absl::StrJoin(variable_ids, ", ") + << "} expr=" << candidate.occurrence.expr + << " chosen_value=" << *selection.chosen_value + << " ignored_observed_value=" + << candidate.evidence.observed_value + << " ignored_solved_value=" + << candidate.evidence.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); + } + } + } + + expr_summaries.push_back(absl::StrCat( + "variable_ids={", absl::StrJoin(variable_ids, ", "), + "} expressions={", absl::StrJoin(selection.expressions, ", "), + "} chosen=", + selection.chosen_value.has_value() + ? std::to_string(*selection.chosen_value) + : std::string(""), + " singleton_derived_values={", + absl::StrJoin(selection.singleton_values, ", "), + "} nonsingleton_values={", + absl::StrJoin(selection.nonsingleton_values, ", "), "} reason=", + selection.reason)); + } + + 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 singleton-derived dynamic occurrences for XLA " + "signature/HLO: ", + absl::StrJoin(ignored_summaries, "; ")); + } + return result; +} + +std::vector BuildStaticCompilationArguments( + absl::Span args) { + std::vector static_args(args.begin(), args.end()); + auto clear_shape_expressions = [&](auto&& self, xla::Shape* shape) -> void { + if (shape->IsTuple()) { + for (xla::Shape& subshape : *shape->mutable_tuple_shapes()) { + self(self, &subshape); + } + } else if (shape->IsArray()) { + shape->set_expressions({}); + for (int dim = 0; dim < shape->dimensions().size(); ++dim) { + if (shape->dimensions(dim) >= 0) { + shape->set_dynamic_dimension(dim, false); + } + } + } + }; + + // Preserve concrete runtime dimensions and values so CompileIfNeeded uses + // its ordinary static-shape cache key. Only symbolic/dynamic annotations + // are removed. + for (XlaArgument& arg : static_args) { + if (absl::holds_alternative(arg.shape)) { + std::get(arg.shape).set_expressions({}); + } else { + clear_shape_expressions(clear_shape_expressions, + &std::get(arg.shape)); + } + arg.constant_value_expressions.clear(); + } + return static_args; +} + +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; + } + XlaArgument& arg = (*args)[occurrence.arg_index]; + if (occurrence.source == + IgnoredDynamicArgumentOccurrence::Source::kShapeDimension) { + if (!absl::holds_alternative(arg.shape)) { + continue; + } + VLOG(1) << "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; + } + + VLOG(1) << "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 ResolveDynamicBatchSizeFromRuntimeShapes( + absl::Span xla_input_shapes, + int runtime_input_count, + absl::FunctionRef get_runtime_input_index, + absl::FunctionRef get_runtime_input_shape, + absl::FunctionRef get_runtime_input_name, + bool log_solves, absl::string_view cluster_name) { + DynamicBatchResolutionResult result; + std::set dyn_vals; + std::map, std::vector> + variable_ids_to_candidates; + + for (int i = 0; i < xla_input_shapes.size(); ++i) { + const xla::Shape& xla_shape = xla_input_shapes[i]; + const int input_idx = get_runtime_input_index(i); + if (!xla_shape.IsArray() || xla_shape.expressions().empty() || + input_idx < 0 || input_idx >= runtime_input_count) { + continue; + } + + const TensorShape& runtime_shape = get_runtime_input_shape(input_idx); + const absl::string_view input_name = get_runtime_input_name(input_idx); + for (int dim = 0; dim < xla_shape.expressions().size(); ++dim) { + const xla::DExpr& expr = xla_shape.expressions(dim); + if (!(expr && expr->is_dynamic())) { + continue; + } + xla::DExpr simplified_expr = expr.simplify(); + const int64_t size = runtime_shape.dim_size(dim); + std::optional dyn_val = simplified_expr->solve(size); + if (log_solves && dyn_val.has_value()) { + VLOG(1) << "Dynamic solve: cluster=" << cluster_name + << " runtime_input_index=" << input_idx + << " xla_input_index=" << i << " dim=" << dim + << " input_name=" << input_name + << " expr=" << DExprToString(simplified_expr) + << " target_size=" << size << " result=" << *dyn_val; + } + if (!dyn_val.has_value()) { + if (log_solves) { + LOG(WARNING) << "Dynamic solve failed: cluster=" << cluster_name + << " runtime_input_index=" << input_idx + << " xla_input_index=" << i << " dim=" << dim + << " input_name=" << input_name + << " expr=" << DExprToString(simplified_expr) + << " target_size=" << size; + } + continue; + } + const std::string expr_string = DExprToString(simplified_expr); + const std::string context = absl::StrCat( + "xla_input_index=", i, " runtime_input_index=", input_idx, + " input_name=", input_name, " dim=", dim, + " runtime_dim_size=", size, + " solved_dynamic_value=", *dyn_val); + variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back( + DynamicSolveCandidate{ + DynamicSolveEvidence{*dyn_val, size, expr_string}, context}); + } + } + + std::vector expr_summaries; + std::optional> expected_dyn_ids; + bool mismatched_dyn_ids = false; + bool candidate_filter_rejected = false; + for (const auto& [variable_ids, candidates] : variable_ids_to_candidates) { + std::vector singleton_contexts; + std::vector nonsingleton_contexts; + for (const auto& candidate : candidates) { + if (candidate.evidence.observed_value == 1) { + singleton_contexts.push_back(candidate.context); + } else { + nonsingleton_contexts.push_back(candidate.context); + } + } + DynamicSolveSelection selection = + SelectDynamicSolveEvidence(absl::MakeConstSpan(candidates)); + + if (selection.rejected) { + candidate_filter_rejected = true; + LOG(WARNING) << "Dynamic solve rejected candidates for variable_ids={" + << absl::StrJoin(variable_ids, ", ") + << "}: " << selection.reason; + } + if (selection.chosen_value.has_value()) { + dyn_vals.insert(*selection.chosen_value); + } + + expr_summaries.push_back(absl::StrCat( + "variable_ids={", absl::StrJoin(variable_ids, ", "), + "} expressions={", absl::StrJoin(selection.expressions, ", "), + "} chosen=", + selection.chosen_value.has_value() + ? std::to_string(*selection.chosen_value) + : std::string(""), + " singleton_derived_values={", + absl::StrJoin(selection.singleton_values, ", "), + "} nonsingleton_values={", + absl::StrJoin(selection.nonsingleton_values, ", "), + "} reason=", selection.reason, " singleton_derived_contexts=[", + absl::StrJoin(singleton_contexts, "; "), + "] nonsingleton_contexts=[", + absl::StrJoin(nonsingleton_contexts, "; "), "]")); + } + + for (const xla::Shape& xla_shape : xla_input_shapes) { + if (!xla_shape.IsArray() || xla_shape.expressions().empty()) { + continue; + } + for (const xla::DExpr& expr : xla_shape.expressions()) { + 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 (variable_ids_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, " | "), "] all_values={", + absl::StrJoin(dyn_vals, ", "), "}"); + } 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, " | ")); + } 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, " | ")); + } 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, " | "), "] all_values={", + absl::StrJoin(dyn_vals, ", "), "}"); + } + return result; +} + + +} // namespace xla_ops_internal + +namespace { + +using xla_ops_internal::AnalyzeIgnoredDynamicArgumentOccurrences; +using xla_ops_internal::BuildStaticCompilationArguments; +using xla_ops_internal::DynamicBatchResolutionResult; +using xla_ops_internal::DynamicSolveFilterDecision; +using xla_ops_internal::GetRuntimeInputIndex; +using xla_ops_internal::StripIgnoredDynamicArgumentOccurrences; + +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; + if (!GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) { + return result; + } + + auto get_runtime_input_index = [&](int i) { + int input_idx = + GetRuntimeInputIndex(comp_result.input_mapping, i, num_constant_args, + /*constants_omitted=*/op_def.op() == "_XlaRun"); + const bool is_runtime_key_input = + input_idx >= 0 && input_idx < ctx->num_inputs() && + ctx->input_dtype(input_idx) == DT_STRING && + input_idx == op_def.input_size() - 1; + return is_runtime_key_input ? -1 : input_idx; + }; + auto get_runtime_input_shape = [&](int i) -> const TensorShape& { + return ctx->input(i).shape(); + }; + auto get_runtime_input_name = [&](int i) -> absl::string_view { + return i < op_def.input_size() ? op_def.input(i) : ""; + }; + + result = xla_ops_internal::ResolveDynamicBatchSizeFromRuntimeShapes( + comp_result.xla_input_shapes, ctx->num_inputs(), + get_runtime_input_index, get_runtime_input_shape, + get_runtime_input_name, log_solves, cluster_name); + if (result.has_batch_size || !result.can_run) { + 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; @@ -881,6 +1451,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(""); + VLOG(1) << "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 @@ -917,6 +1495,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(""); + VLOG(1) << "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 = @@ -953,6 +1539,36 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { xla::ExecutableRunOptions run_options; + 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; + if (!batch_resolution.can_run) { + LOG(ERROR) << batch_resolution.diagnostic; + ctx->CtxFailure(errors::InvalidArgument(batch_resolution.diagnostic)); + return; + } + if (batch_resolution.has_batch_size) { + VLOG(1) << "Setting run_options.batch_size " + << batch_resolution.diagnostic; + run_options.set_batch_size(batch_resolution.batch_size); + is_set = true; + } + if (!is_set) { + LOG(WARNING) << "Entering XLA cluster without run_options.batch_size " + << "being set because " << batch_resolution.diagnostic + << ". 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; + } + } + // Host callbacks used for HLO send/recv. xla::SendDeviceMemoryFunction send_function = GetSendDeviceMemoryFunction(ctx, key); @@ -981,7 +1597,8 @@ void XlaRunOp::Compute(OpKernelContext* ctx) { launch_context.PopulateOutputs( ctx, closure.compilation_result(), execution_output->ConsumeResult(), /*missing_ctx_input_prefix=*/closure.num_constant_args(), - absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs)); + absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs, + &run_options)); } XlaMergeOp::XlaMergeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} diff --git a/tensorflow/compiler/jit/kernels/xla_ops_internal.h b/tensorflow/compiler/jit/kernels/xla_ops_internal.h new file mode 100644 index 00000000000000..546bc46d93534a --- /dev/null +++ b/tensorflow/compiler/jit/kernels/xla_ops_internal.h @@ -0,0 +1,84 @@ +/* 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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_ +#define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_ + +#include +#include +#include + +#include "absl/functional/function_ref.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/tf2xla/xla_argument.h" + +namespace tensorflow { +namespace xla_ops_internal { + +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; +}; + +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; +}; + +int GetRuntimeInputIndex(absl::Span input_mapping, + int xla_input_index, int num_constant_args, + bool constants_omitted); + +DynamicBatchResolutionResult ResolveDynamicBatchSizeFromRuntimeShapes( + absl::Span xla_input_shapes, + int runtime_input_count, + absl::FunctionRef get_runtime_input_index, + absl::FunctionRef get_runtime_input_shape, + absl::FunctionRef get_runtime_input_name, + bool log_solves = false, absl::string_view cluster_name = {}); + +DynamicSolveFilterDecision AnalyzeIgnoredDynamicArgumentOccurrences( + absl::Span args); + +std::vector BuildStaticCompilationArguments( + absl::Span args); + +void StripIgnoredDynamicArgumentOccurrences( + const std::vector& ignored_occurrences, + std::vector* args); + +} // namespace xla_ops_internal +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_ diff --git a/tensorflow/compiler/jit/kernels/xla_ops_test.cc b/tensorflow/compiler/jit/kernels/xla_ops_test.cc new file mode 100644 index 00000000000000..bcea59771626df --- /dev/null +++ b/tensorflow/compiler/jit/kernels/xla_ops_test.cc @@ -0,0 +1,282 @@ +/* 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/kernels/xla_ops_internal.h" + +#include +#include +#include + +#include "xla/shape_expr.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace xla_ops_internal { +namespace { + +TEST(XlaInputMappingTest, UsesCompilationInputMapping) { + const std::vector input_mapping = {1, 3}; + + EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 0, /*num_constant_args=*/1, + /*constants_omitted=*/false), + 1); + EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 1, /*num_constant_args=*/1, + /*constants_omitted=*/false), + 3); + EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 0, /*num_constant_args=*/1, + /*constants_omitted=*/true), + 0); + EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 1, /*num_constant_args=*/1, + /*constants_omitted=*/true), + 2); + EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 2, /*num_constant_args=*/1, + /*constants_omitted=*/true), + -1); +} + +XlaArgument MakeDynamicArgument(int64_t observed_size, + xla::DExpr expr = xla::DExpr::Var(1)) { + XlaArgument arg; + arg.kind = XlaArgument::kParameter; + arg.type = DT_FLOAT; + TensorShape shape({observed_size, 16}); + shape.set_expression(0, std::move(expr)); + shape.set_expression(1, xla::DExpr::Const(16)); + arg.shape = shape; + return arg; +} + +xla::Shape MakeDynamicXlaInput(xla::DExpr leading_expr) { + xla::Shape shape; + shape.set_element_type(xla::F32); + shape.add_dimensions(1024, /*is_dynamic=*/true, + std::move(leading_expr)); + shape.add_dimensions(16, /*is_dynamic=*/false, xla::DExpr::Const(16)); + return shape; +} + +DynamicBatchResolutionResult ResolveForTest( + absl::Span xla_shapes, + absl::Span runtime_shapes, + absl::Span runtime_input_indices, + absl::Span runtime_input_names) { + return ResolveDynamicBatchSizeFromRuntimeShapes( + xla_shapes, runtime_shapes.size(), + [&](int i) { return runtime_input_indices[i]; }, + [&](int i) -> const TensorShape& { return runtime_shapes[i]; }, + [&](int i) -> absl::string_view { return runtime_input_names[i]; }); +} + +TEST(DynamicRuntimeSolveTest, UsesMappedRuntimeShapesToSolveExpressions) { + const xla::DExpr variable = xla::DExpr::Var(1); + const std::vector xla_shapes = { + MakeDynamicXlaInput(variable), + MakeDynamicXlaInput(variable / 2), + }; + // XLA parameters are deliberately ordered differently from runtime inputs. + const std::vector runtime_shapes = { + TensorShape({120, 16}), + TensorShape({240, 16}), + }; + const std::vector runtime_input_indices = {1, 0}; + const std::vector runtime_input_names = {"half", "full"}; + + DynamicBatchResolutionResult result = ResolveForTest( + xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names); + + EXPECT_TRUE(result.can_run); + ASSERT_TRUE(result.has_batch_size); + EXPECT_EQ(result.batch_size, 240); +} + +TEST(DynamicRuntimeSolveTest, + PrefersNonSingletonEvidenceAcrossEquivalentExpressions) { + const xla::DExpr variable = xla::DExpr::Var(1); + const std::vector xla_shapes = { + MakeDynamicXlaInput(variable), + MakeDynamicXlaInput(variable + 1), + }; + // The first occurrence may be a broadcast singleton. A + 1 = 241 gives + // the stronger runtime evidence A = 240. + const std::vector runtime_shapes = { + TensorShape({1, 16}), + TensorShape({241, 16}), + }; + const std::vector runtime_input_indices = {0, 1}; + const std::vector runtime_input_names = {"singleton", "full"}; + + DynamicBatchResolutionResult result = ResolveForTest( + xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names); + + EXPECT_TRUE(result.can_run); + ASSERT_TRUE(result.has_batch_size); + EXPECT_EQ(result.batch_size, 240); +} + +TEST(DynamicRuntimeSolveTest, KeepsConsistentNonSingletonEvidence) { + const xla::DExpr variable = xla::DExpr::Var(1); + const std::vector xla_shapes = { + MakeDynamicXlaInput(variable), + MakeDynamicXlaInput(variable / 3), + }; + const std::vector runtime_shapes = { + TensorShape({720, 16}), + TensorShape({240, 16}), + }; + const std::vector runtime_input_indices = {0, 1}; + const std::vector runtime_input_names = {"full", "third"}; + + DynamicBatchResolutionResult result = ResolveForTest( + xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names); + + EXPECT_TRUE(result.can_run); + ASSERT_TRUE(result.has_batch_size); + EXPECT_EQ(result.batch_size, 720); +} + +TEST(DynamicRuntimeSolveTest, RejectsConflictingNonSingletonEvidence) { + const std::vector xla_shapes = { + MakeDynamicXlaInput(xla::DExpr::Var(1)), + MakeDynamicXlaInput(xla::DExpr::Var(1)), + }; + const std::vector runtime_shapes = { + TensorShape({240, 16}), + TensorShape({720, 16}), + }; + const std::vector runtime_input_indices = {0, 1}; + const std::vector runtime_input_names = {"first", "second"}; + + DynamicBatchResolutionResult result = ResolveForTest( + xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names); + + EXPECT_FALSE(result.can_run); + EXPECT_FALSE(result.has_batch_size); + EXPECT_NE(result.diagnostic.find("conflicting non-singleton candidates"), + std::string::npos); +} + +TEST(DynamicSolveFilterTest, IgnoresSingletonWhenNonSingletonEvidenceExists) { + std::vector args = { + MakeDynamicArgument(1), + MakeDynamicArgument(240), + }; + + DynamicSolveFilterDecision decision = + AnalyzeIgnoredDynamicArgumentOccurrences(args); + + ASSERT_TRUE(decision.can_run); + ASSERT_EQ(decision.ignored_occurrences.size(), 1); + EXPECT_EQ(decision.ignored_occurrences[0].arg_index, 0); + EXPECT_EQ(decision.ignored_occurrences[0].observed_value, 1); + + StripIgnoredDynamicArgumentOccurrences(decision.ignored_occurrences, &args); + const TensorShape& singleton_shape = std::get(args[0].shape); + const TensorShape& nonsingleton_shape = std::get(args[1].shape); + EXPECT_FALSE(singleton_shape.get_expression(0)); + EXPECT_EQ(singleton_shape.dim_size(0), 1); + EXPECT_TRUE(nonsingleton_shape.get_expression(0)); +} + +TEST(DynamicSolveFilterTest, + IgnoresSingletonAcrossExpressionsOfTheSameVariable) { + const xla::DExpr variable = xla::DExpr::Var(1); + std::vector args = { + MakeDynamicArgument(1, variable), + MakeDynamicArgument(241, variable + 1), + }; + + DynamicSolveFilterDecision decision = + AnalyzeIgnoredDynamicArgumentOccurrences(args); + + ASSERT_TRUE(decision.can_run); + ASSERT_EQ(decision.ignored_occurrences.size(), 1); + EXPECT_EQ(decision.ignored_occurrences[0].arg_index, 0); + EXPECT_EQ(decision.ignored_occurrences[0].observed_value, 1); + EXPECT_EQ(decision.ignored_occurrences[0].solved_value, 1); +} + +TEST(DynamicSolveFilterTest, RejectsConflictingNonSingletonEvidence) { + std::vector args = { + MakeDynamicArgument(240), + MakeDynamicArgument(720), + }; + + DynamicSolveFilterDecision decision = + AnalyzeIgnoredDynamicArgumentOccurrences(args); + + EXPECT_FALSE(decision.can_run); + EXPECT_TRUE(decision.ignored_occurrences.empty()); + EXPECT_NE(decision.diagnostic.find("conflicting non-singleton candidates"), + std::string::npos); +} + +TEST(DynamicSolveFilterTest, KeepsConsistentSingletonOnlyEvidence) { + std::vector args = { + MakeDynamicArgument(1), + MakeDynamicArgument(1), + }; + + DynamicSolveFilterDecision decision = + AnalyzeIgnoredDynamicArgumentOccurrences(args); + + EXPECT_TRUE(decision.can_run); + EXPECT_TRUE(decision.ignored_occurrences.empty()); +} + +TEST(DynamicSolveFilterTest, StaticArgumentsKeepConcreteRuntimeShapes) { + XlaArgument dynamic_arg = MakeDynamicArgument(37); + dynamic_arg.constant_value_expressions.resize(1); + dynamic_arg.constant_value_expressions[0].set_variable_id(1); + + std::vector args = {dynamic_arg}; + std::vector static_args = + BuildStaticCompilationArguments(args); + + ASSERT_EQ(static_args.size(), 1); + const TensorShape& static_shape = + std::get(static_args[0].shape); + EXPECT_EQ(static_shape.dim_size(0), 37); + EXPECT_EQ(static_shape.dim_size(1), 16); + EXPECT_TRUE(static_shape.get_expressions().empty()); + EXPECT_TRUE(static_args[0].constant_value_expressions.empty()); +} + +TEST(DynamicSolveFilterTest, StaticArgumentsClearXlaShapeDynamism) { + xla::Shape dynamic_shape; + dynamic_shape.set_element_type(xla::F32); + dynamic_shape.add_dimensions(37, /*is_dynamic=*/true, + xla::DExpr::Var(1)); + + XlaArgument dynamic_arg; + dynamic_arg.kind = XlaArgument::kParameter; + dynamic_arg.type = DT_FLOAT; + dynamic_arg.shape = dynamic_shape; + std::vector args = {dynamic_arg}; + + std::vector static_args = + BuildStaticCompilationArguments(args); + + const xla::Shape& static_shape = std::get(static_args[0].shape); + EXPECT_EQ(static_shape.dimensions(0), 37); + EXPECT_FALSE(static_shape.is_dynamic_dimension(0)); + ASSERT_EQ(static_shape.expressions().size(), 1); + EXPECT_TRUE(static_shape.expressions(0)->is_constant()); + EXPECT_EQ(static_shape.expressions(0)->get_val(), 37); +} + +} // namespace +} // namespace xla_ops_internal +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 588161df309131..b6b1f08fe73dcb 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/types/span.h" #include "tensorflow/compiler/jit/pjrt_tensor_buffer.h" #include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h" @@ -38,6 +40,7 @@ limitations under the License. #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_future.h" +#include "xla/printer.h" #include "xla/shape_util.h" #include "xla/status_macros.h" #include "xla/stream_executor/platform_manager.h" @@ -47,6 +50,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" #include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/batch_size_resource.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_mgr.h" @@ -63,6 +67,17 @@ limitations under the License. #include "tsl/platform/statusor.h" namespace tensorflow { + +std::string DExprToString(const xla::DExpr& expr) { + xla::DExpr simplified = expr.simplify(); + if (!simplified && !simplified.is_unknown()) { + return ""; + } + xla::StringPrinter printer; + simplified->print(&printer); + return std::move(printer).ToString(); +} + namespace { using xla::ScopedShapedBuffer; using xla::ShapedBuffer; @@ -361,7 +376,8 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs( ScopedShapedBuffer output, int missing_ctx_input_prefix, absl::Span variable_infos, const xla::HloInputOutputAliasConfig& input_output_alias, - const std::map& resource_vars) { + const std::map& resource_vars, + const xla::ExecutableRunOptions* run_options) { se::Stream* stream = ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; Allocator* allocator = ctx->device()->GetAllocator({}); @@ -430,10 +446,123 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs( } } else { for (int i = 0; i < ctx->num_outputs(); ++i) { - output_tensor_shapes.push_back(compilation_result->outputs[i].shape); + xla::Shape output_host_shape = output.on_host_shape(); + const xla::Shape& subshape = + xla::ShapeUtil::GetSubshape(output_host_shape, {i}); + VLOG(2) << "PopulateOutputs: subshape[" << i << "]: "<< subshape; + TensorShape shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(subshape, &shape)); + bool has_dynamic = false; + + for (int dim = 0; dim < subshape.expressions().size(); ++dim) { + const auto& expr = subshape.expressions(dim); + if (expr && expr->is_dynamic()) { + has_dynamic = true; + VLOG(1) << "Current expression is " << expr; + if (run_options) { + const int64_t run_options_batch_size = run_options->batch_size(); + if (run_options_batch_size <= 0) { + return absl::InvalidArgumentError(absl::StrCat( + "Cannot substitute dynamic output shape for output ", i, + ", dimension ", dim, + ": the XLA runtime batch size was not initialized")); + } + VLOG(1) << "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(); + VLOG(1) << "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(); + VLOG(1) << "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 " + "constant for output ", + i, ", dimension ", dim, ": ", DExprToString(subst_expr))); + } + shape.set_dim(dim, subst_expr->get_val()); + } else { + // TODO: Fallback to BatchSizeResource for now. Remove it later. + LOG(WARNING) << "PopulateOutputs did not receive run_options for " + << "output " << i << " dimension " << dim + << "; falling back to BatchSizeResource"; + 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"); + } + core::ScopedUnref bsr_ref(bsr); + const int64_t runtime_batch_size = bsr->GetBatchSize(); + if (runtime_batch_size <= 0) { + return absl::InvalidArgumentError(absl::StrCat( + "Cannot substitute dynamic output shape for output ", i, + ", dimension ", dim, + ": the XLA runtime batch size was not initialized")); + } + xla::DExpr batch_size = xla::DExpr::Const(runtime_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(); + VLOG(1) << "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(); + VLOG(1) << "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 " + "constant for output ", + i, ", dimension ", dim, ": ", DExprToString(subst_expr))); + } + shape.set_dim(dim, subst_expr->get_val()); + } + } + } + if (has_dynamic) { + output_tensor_shapes.push_back(shape); + } + else { + output_tensor_shapes.push_back(compilation_result->outputs[i].shape); + } } } + VLOG(2) << "output_tensor_shapes:"; + for (auto s:output_tensor_shapes) { + VLOG(2) << s; + } + // Copy XLA results to the OpOutputList. int output_num = 0; for (int i = 0, end = ctx->num_outputs(); i < end; ++i) { diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 5e5128d515bf97..e7b4c462b306cd 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -18,6 +18,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ #define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ +#include #include #include #include @@ -29,14 +30,21 @@ limitations under the License. #include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" #include "xla/service/shaped_buffer.h" +#include "xla/executable_run_options.h" #include "xla/stream_executor/device_memory_allocator.h" #include "tensorflow/core/framework/allocation_description.pb.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/thread_annotations.h" +namespace xla { +class DExpr; +} + namespace tensorflow { +std::string DExprToString(const xla::DExpr& expr); + // Creates a list of updated resource variables. absl::StatusOr> GatherVariableInfo( OpKernelContext* ctx, @@ -208,7 +216,8 @@ class XlaComputationLaunchContext { xla::ScopedShapedBuffer output, int missing_ctx_input_prefix, absl::Span variable_infos, const xla::HloInputOutputAliasConfig& input_output_alias, - const std::map& resource_vars); + const std::map& resource_vars, + const xla::ExecutableRunOptions* run_options = nullptr); private: xla::LocalClient* client_; diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index b7cff00c8a0bfe..90faea1745d327 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -32,12 +32,14 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/variant.h" #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/shape_inference.h" +#include "tensorflow/compiler/tf2xla/symbolic_content_util.h" #include "tensorflow/compiler/jit/xla_compile_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h" #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h" @@ -93,7 +95,6 @@ namespace { constexpr char kSingleOpComponent[] = "TF2XLA_XLA_COMPILER_COMPILE_SINGLE_OP"; constexpr char kCompileFunctionComponent[] = "TF2XLA_XLA_COMPILER_COMPILE_FUNCTION"; - // Checks that arguments `args` match types `types`. absl::Status CheckSignature(const DataTypeVector& types, absl::Span args) { @@ -511,6 +512,19 @@ std::vector XlaCompiler::Argument::DimensionSizes() const { } } +std::vector XlaCompiler::Argument::DimensionExpressions() const { + if (absl::holds_alternative(shape)) { + return std::get(shape).get_filled_expressions(); + } else { + std::vector expressions; + expressions.reserve(std::get(shape).expressions().size()); + for (const auto& expr : std::get(shape).expressions()) { + expressions.push_back(expr); + } + return expressions; + } +} + absl::InlinedVector XlaCompiler::Argument::DimensionSizesAsInlinedVector() const { if (absl::holds_alternative(shape)) { @@ -839,7 +853,9 @@ absl::Status XlaCompiler::CompileFunction( std::vector{tensor_shape}); } } else { + auto* val_ptr = std::get_if(&args[i].shape); TensorShape tensor_shape = std::get(args[i].shape); + AttrSlice n_attrs = fbody->arg_nodes[i]->attrs(); fbody->arg_nodes[i]->ClearAttr("_output_shapes"); fbody->arg_nodes[i]->AddAttr("_output_shapes", std::vector{tensor_shape}); @@ -930,14 +946,17 @@ absl::Status XlaCompiler::XLAShapeForArgument( TF_RETURN_IF_ERROR(RewriteLayoutWithShardedShape( arg_sharding, /*use_fast_memory=*/false, options_.shape_determination_fns, xla_shape)); - // If the arg is dynamic then we update the shape to reflect that. The - // layout etc above lose it by forcing a swap to TensorShape. - if (std::holds_alternative(arg.shape) && - std::get(arg.shape).is_dynamic()) { - xla::Shape dynamic_shape = std::get(arg.shape); - for (int i = 0; i < xla_shape->dimensions().size(); ++i) { - xla_shape->set_dynamic_dimension( - i, dynamic_shape.is_dynamic_dimension(i)); + // If the arg carries dynamic metadata or symbolic expressions then we + // update the shape to reflect that. The layout logic above routes + // through TensorShape and can otherwise discard this information. + if (std::holds_alternative(arg.shape)) { + const xla::Shape& original_shape = std::get(arg.shape); + if (original_shape.is_dynamic() || original_shape.has_dynamic_expr()) { + for (int i = 0; i < xla_shape->dimensions().size(); ++i) { + xla_shape->set_dynamic_dimension( + i, original_shape.is_dynamic_dimension(i)); + xla_shape->set_expression(i, original_shape.expressions(i)); + } } } } else { @@ -1083,6 +1102,7 @@ absl::Status XlaCompiler::BuildArguments( TF_RET_CHECK(absl::holds_alternative(arg.shape)); // TODO(phawkins): this code assumes that resource arguments do not // alias. + auto* val_ptr = std::get_if(&arg.shape); XlaResource* resource = context->AddResource(std::make_unique( arg.resource_kind, i, arg.name, arg.type, @@ -1108,6 +1128,25 @@ absl::Status XlaCompiler::BuildArguments( } case XlaCompiler::Argument::kConstant: arg_expression = XlaExpression::Constant(arg.constant_value); + if (!arg.constant_value_expressions.empty()) { + VLOG(1) << "BuildArguments attaching " + << arg.constant_value_expressions.size() + << " constant_value_expressions to constant arg " << i + << " (" << arg.name << ")"; + // Preserve symbolic per-element metadata for shape-like constants so + // later tf2xla consumers can recover dynamic contents from them. + std::vector contents; + contents.reserve(arg.constant_value_expressions.size()); + for (const xla::ExpressionProto& expr : + arg.constant_value_expressions) { + xla::DExpr parsed = xla::DExprFromProto(expr); + contents.push_back(parsed && parsed->is_dynamic() + ? std::move(parsed) + : xla::DExpr::Unknown( + xla::kUnknownContentSentinel)); + } + arg_expression.set_contents(std::move(contents)); + } break; case XlaCompiler::Argument::kInvalid: return errors::Internal( @@ -1203,6 +1242,10 @@ absl::Status XlaCompiler::BuildArguments( xla::XlaScopedShardingAssignment assign_sharding( builder, it == arg_shardings.end() ? std::optional() : it->second); + auto& arg = args[input_to_args->at(i)]; + xla::OpMetadata arg_metadata; + arg_metadata.set_op_name(arg.node_name); + builder->SetOneShotOpMetadata(arg_metadata); if (is_entry_computation) { // Add an entry to is_same_across_replicas for every leaf buffer. std::vector is_same_across_replicas( @@ -1222,10 +1265,9 @@ absl::Status XlaCompiler::BuildArguments( // Fill in the handles in non-constant arguments, and reshape parameters // back to their correct shapes. - VLOG(2) << "XLA computation inputs:"; for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; - VLOG(2) << " XLA arg " << i + VLOG(2) << " XLA arg " << i << " shape: " << xla::ShapeUtil::HumanString(arg_shapes[i]) << " name: " << arg.name << " TF arg " << input_to_args->at(i) << " node name: " << arg.node_name @@ -1251,7 +1293,9 @@ absl::Status XlaCompiler::BuildArguments( // return values of functions, and then reshape unconditionally. if (is_entry_computation) { arg_expression = XlaExpression::XlaOp( - xla::Reshape(arg_handles[i], arg.DimensionSizes()), arg.type); + xla::Reshape(arg_handles[i], arg.DimensionSizes(), + arg.DimensionExpressions()), + arg.type); } else { arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type); if (arg.value_bound) { diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 2f0ff5e91867f1..72bd6f712672e3 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -1740,6 +1740,7 @@ tf_cuda_library( "@local_xla//xla/tsl/framework:cancellation", "@local_xla//xla/tsl/util:command_line_flags", "@local_xla//xla/tsl/util:device_name_utils", + "@local_xla//xla:shape_util", ] + if_cuda([ "@local_config_cuda//cuda:cudnn_header", ]) + if_static( diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD index 8a2c2032f3ab25..79f4bd855fd10b 100644 --- a/tensorflow/core/framework/BUILD +++ b/tensorflow/core/framework/BUILD @@ -62,6 +62,7 @@ package( exports_files( srcs = [ "allocator_registry.h", + "batch_size_resource.h", "collective.h", "control_flow.h", "dataset.h", @@ -195,6 +196,7 @@ filegroup( "allocator.h", "allocator_registry.h", "attr_value_util.h", + "batch_size_resource.h", "bfloat16.h", "bounds_check.h", "cancellation.h", diff --git a/tensorflow/core/framework/batch_size_resource.h b/tensorflow/core/framework/batch_size_resource.h new file mode 100644 index 00000000000000..45bdfad02699bb --- /dev/null +++ b/tensorflow/core/framework/batch_size_resource.h @@ -0,0 +1,40 @@ +/* Copyright 2026 The TensorFlow Authors. + +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. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_ +#define TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_ + +#include + +#include "tensorflow/core/framework/resource_mgr.h" + +namespace tensorflow { +const string BatchSizeResourceName = "BatchSizeResource_"; +class BatchSizeResource : public ResourceBase { + public: + ~BatchSizeResource() override { + VLOG(1) << "BatchSizeResource destroyed with batch size: " << batch_size_; + } + string DebugString() const override { return BatchSizeResourceName; } + void SetBatchSize(size_t s) { batch_size_ = s; } + size_t GetBatchSize() { return batch_size_; } + + private: + size_t batch_size_ = 0; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_ diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc index 864855de1d69f6..8a5936980ad578 100644 --- a/tensorflow/core/kernels/function_ops.cc +++ b/tensorflow/core/kernels/function_ops.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/gradients.h" #include "tensorflow/core/common_runtime/graph_constructor.h" #include "tensorflow/core/common_runtime/memory_types.h" +#include "tensorflow/core/framework/batch_size_resource.h" #include "tensorflow/core/framework/cancellation.h" #include "tensorflow/core/framework/full_type.pb.h" #include "tensorflow/core/framework/full_type_util.h" @@ -42,6 +43,13 @@ static constexpr const char* const kGradientOp = ArgOp::ArgOp(OpKernelConstruction* ctx) : OpKernel(ctx) { OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("index", &index_)); + + Status s = ctx->GetAttr("_dynamic_dim", &dynamic_dim_); + if (IsNotFound(s)) { + dynamic_dim_ = -1; + } else { + OP_REQUIRES_OK(ctx, s); + } } void ArgOp::Compute(OpKernelContext* ctx) { @@ -59,16 +67,48 @@ void ArgOp::Compute(OpKernelContext* ctx) { } }; + Tensor t; + int64_t batch_size = -1; if (frame->CanConsumeArg(index_)) { - Tensor val; - frame->ConsumeArg(index_, &val); - OP_REQUIRES_OK(ctx, validate_type(val)); - ctx->set_output(0, std::move(val)); + frame->ConsumeArg(index_, &t); + OP_REQUIRES_OK(ctx, validate_type(t)); + if (dynamic_dim_ >= 0) { + batch_size = t.dim_size(dynamic_dim_); + } + ctx->set_output(0, std::move(t)); } else { OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val)); OP_REQUIRES_OK(ctx, validate_type(*val)); ctx->set_output(0, *val); } + if (dynamic_dim_ >= 0) { + BatchSizeResource* bsr = nullptr; + ScopedStepContainer* step_container = ctx->step_container(); + + OP_REQUIRES_OK(ctx, step_container->LookupOrCreate( + ctx->resource_manager(), BatchSizeResourceName, &bsr, + [](BatchSizeResource** ret) -> Status { + *ret = new BatchSizeResource(); + return OkStatus(); + })); + + if (batch_size < 0) { + batch_size = val->dim_size(dynamic_dim_); + } + VLOG(1) << "Found batch_size in dimension #" << dynamic_dim_; + if (bsr->GetBatchSize() == 0) { + bsr->SetBatchSize(batch_size); + VLOG(1) << "Set batch_size from 0 to " << batch_size + << ". step_id: " << ctx->step_id(); + } else if (bsr->GetBatchSize() != batch_size) { + VLOG(1) << "Warning: Set batch_size from " << bsr->GetBatchSize() + << ". step_id: " << ctx->step_id(); + bsr->SetBatchSize(batch_size); + } else { + VLOG(1) << "batch_size already set to " << batch_size; + } + bsr->Unref(); + } } RetvalOp::RetvalOp(OpKernelConstruction* ctx) : OpKernel(ctx) { diff --git a/tensorflow/core/kernels/function_ops.h b/tensorflow/core/kernels/function_ops.h index 552e1e6c515e3b..fed05ead2007ad 100644 --- a/tensorflow/core/kernels/function_ops.h +++ b/tensorflow/core/kernels/function_ops.h @@ -38,6 +38,7 @@ class ArgOp : public OpKernel { private: int index_; DataType dtype_; + int dynamic_dim_; ArgOp(const ArgOp&) = delete; void operator=(const ArgOp&) = delete; @@ -54,6 +55,7 @@ class RetvalOp : public OpKernel { private: int index_; DataType dtype_; + int dynamic_dim_; RetvalOp(const RetvalOp&) = delete; void operator=(const RetvalOp&) = delete; diff --git a/third_party/xla/xla/executable_run_options.h b/third_party/xla/xla/executable_run_options.h index b377e91670efb9..211932ffcbba98 100644 --- a/third_party/xla/xla/executable_run_options.h +++ b/third_party/xla/xla/executable_run_options.h @@ -195,6 +195,13 @@ class ExecutableRunOptions { return *this; } + ExecutableRunOptions& set_batch_size(int64_t batch_size) { + batch_size_ = batch_size; + return *this; + } + + int64_t batch_size() const { return batch_size_; } + int32_t launch_id() const { return launch_id_; } ExecutableRunOptions& set_run_id(RunId id); @@ -265,6 +272,7 @@ class ExecutableRunOptions { ExecutionProfile* execution_profile_ = nullptr; int rng_seed_ = 0; int32_t launch_id_ = 0; + int64_t batch_size_ = 0; stream_executor::Stream* device_to_host_stream_ = nullptr; stream_executor::Stream* host_to_device_stream_ = nullptr; ThenExecuteFunction* then_execute_function_ = nullptr; From 98f302fa554aaa7b0257b5f031409fedbb384d00 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Wed, 19 Aug 2026 16:05:01 +0100 Subject: [PATCH 12/21] Bucket dynamic compilation arguments for cache reuse Normalize symbolic input metadata, pad solved dynamic values to configured compilation buckets, and include the resulting constants in executable cache signatures. Retry inconsistent dynamic requests with concrete static arguments instead of reusing an incompatible executable. --- tensorflow/compiler/jit/BUILD | 12 + .../device_compilation_cluster_signature.cc | 52 +- .../device_compilation_cluster_signature.h | 8 +- .../jit/device_compilation_profiler.cc | 24 +- tensorflow/compiler/jit/device_compiler.h | 10 + tensorflow/compiler/jit/flags.cc | 6 + tensorflow/compiler/jit/flags.h | 5 + tensorflow/compiler/jit/kernels/BUILD | 1 + tensorflow/compiler/jit/kernels/xla_ops.cc | 512 +++++++++++++++++- tensorflow/compiler/jit/xla_batch_matcher.cc | 203 +++++++ tensorflow/compiler/jit/xla_batch_matcher.h | 49 ++ tensorflow/core/kernels/padding_fifo_queue.cc | 9 +- third_party/xla/xla/debug_options_flags.cc | 10 + third_party/xla/xla/xla.proto | 4 +- 14 files changed, 877 insertions(+), 28 deletions(-) create mode 100644 tensorflow/compiler/jit/xla_batch_matcher.cc create mode 100644 tensorflow/compiler/jit/xla_batch_matcher.h diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 951a50fd1ecfb6..a8197744359fde 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -780,6 +780,8 @@ cc_library( ":flags_headers", ":tf_graph_to_hlo_compiler", ":xla_compile_util", + ":xla_batch_matcher", + "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:framework_lite", @@ -2006,3 +2008,13 @@ tf_cuda_cc_test( "@local_xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", ], ) + +cc_library( + name = "xla_batch_matcher", + srcs = ["xla_batch_matcher.cc"], + hdrs = ["xla_batch_matcher.h"], + deps = [ + "//tensorflow/core/platform:logging", + "@local_xla//xla:debug_options_flags", + ], +) \ No newline at end of file diff --git a/tensorflow/compiler/jit/device_compilation_cluster_signature.cc b/tensorflow/compiler/jit/device_compilation_cluster_signature.cc index 81902a28532dbf..5996c5f0b3a70c 100644 --- a/tensorflow/compiler/jit/device_compilation_cluster_signature.cc +++ b/tensorflow/compiler/jit/device_compilation_cluster_signature.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" +#include "absl/strings/str_cat.h" #include #include #include @@ -22,14 +23,23 @@ limitations under the License. namespace tensorflow { namespace { using Signature = DeviceCompilationClusterSignature; +using ConstantTensor = Signature::ConstantTensor; using TensorTypeAndShape = Signature::TensorTypeAndShape; // Functor that converts a Signature's arg to a human readable string. struct SignatureHumanStringAppender { explicit SignatureHumanStringAppender(std::string* dest) : dest(dest) {} std::string* dest; - void operator()(const Tensor& arg) { - absl::StrAppend(dest, "; ", arg.DebugString()); + void operator()(const ConstantTensor& arg) { + absl::StrAppend(dest, "; ", arg.value.DebugString()); + if (!arg.contents.empty()) { + absl::StrAppend(dest, " contents=["); + for (int i = 0; i < arg.contents.size(); ++i) { + if (i > 0) absl::StrAppend(dest, ","); + absl::StrAppend(dest, arg.contents[i].DebugString()); + } + absl::StrAppend(dest, "]"); + } } void operator()(const TensorTypeAndShape& arg) { absl::StrAppend(dest, ",", DataTypeString(arg.first)); @@ -40,18 +50,29 @@ struct SignatureHumanStringAppender { // Functor that compares the arg values of two different signatures. Returns // true when the args are not equal. struct SignatureNotEqual { - bool operator()(const Tensor& arg, const Tensor& other) { - return arg.dtype() != other.dtype() || arg.shape() != other.shape() || - arg.tensor_data() != other.tensor_data(); + bool operator()(const ConstantTensor& arg, const ConstantTensor& other) { + if (arg.value.dtype() != other.value.dtype() || + arg.value.shape() != other.value.shape() || + arg.value.tensor_data() != other.value.tensor_data() || + arg.contents.size() != other.contents.size()) { + return true; + } + for (int i = 0; i < arg.contents.size(); ++i) { + if (arg.contents[i].SerializeAsString() != + other.contents[i].SerializeAsString()) { + return true; + } + } + return false; } bool operator()(const TensorTypeAndShape& arg, const TensorTypeAndShape& other) { return arg.first != other.first || arg.second != other.second; } - bool operator()(const Tensor& arg, const TensorTypeAndShape& other) { + bool operator()(const ConstantTensor& arg, const TensorTypeAndShape& other) { return true; } - bool operator()(const TensorTypeAndShape& arg, const Tensor& other) { + bool operator()(const TensorTypeAndShape& arg, const ConstantTensor& other) { return true; } }; @@ -61,12 +82,16 @@ struct SignatureNotEqual { struct SignatureHashCombiner { explicit SignatureHashCombiner(const uint64 h) : h(h) {} uint64 h; - uint64 operator()(const Tensor& arg) { - h = Hash64Combine(h, std::hash()(static_cast(arg.dtype()))); + uint64 operator()(const ConstantTensor& arg) { + h = Hash64Combine(h, std::hash()(static_cast(arg.value.dtype()))); h = Hash64Combine( - h, Hash64(arg.tensor_data().data(), arg.tensor_data().size())); - for (int dim = 0; dim < arg.dims(); ++dim) { - h = Hash64Combine(h, std::hash()(arg.dim_size(dim))); + h, Hash64(arg.value.tensor_data().data(), arg.value.tensor_data().size())); + for (int dim = 0; dim < arg.value.dims(); ++dim) { + h = Hash64Combine(h, std::hash()(arg.value.dim_size(dim))); + } + for (const xla::ExpressionProto& expr : arg.contents) { + std::string serialized = expr.SerializeAsString(); + h = Hash64Combine(h, Hash64(serialized.data(), serialized.size())); } return h; } @@ -120,7 +145,8 @@ absl::StatusOr Signature::Build( switch (arg.kind) { case XlaCompiler::Argument::kConstant: case XlaCompiler::Argument::kConstantResource: - signature.args.push_back(arg.constant_value); + signature.args.push_back( + ConstantTensor{arg.constant_value, arg.constant_value_expressions}); break; case XlaCompiler::Argument::kParameter: case XlaCompiler::Argument::kResource: diff --git a/tensorflow/compiler/jit/device_compilation_cluster_signature.h b/tensorflow/compiler/jit/device_compilation_cluster_signature.h index 4acea2a03c2cb4..da2de8de370842 100644 --- a/tensorflow/compiler/jit/device_compilation_cluster_signature.h +++ b/tensorflow/compiler/jit/device_compilation_cluster_signature.h @@ -18,8 +18,10 @@ limitations under the License. #include #include +#include #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" namespace tensorflow { @@ -34,7 +36,11 @@ struct DeviceCompilationClusterSignature { // argument number. Tensors must be in host memory. using TensorTypeAndShape = std::pair>; - absl::InlinedVector, 8> args; + struct ConstantTensor { + Tensor value; + std::vector contents; + }; + absl::InlinedVector, 8> args; bool operator==(const DeviceCompilationClusterSignature& other) const; diff --git a/tensorflow/compiler/jit/device_compilation_profiler.cc b/tensorflow/compiler/jit/device_compilation_profiler.cc index 5e1b3b26e8ecb5..f8a742ccc01148 100644 --- a/tensorflow/compiler/jit/device_compilation_profiler.cc +++ b/tensorflow/compiler/jit/device_compilation_profiler.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/xla_activity.pb.h" #include "tensorflow/compiler/jit/xla_activity_listener.h" #include "tensorflow/core/framework/attr_value.pb.h" @@ -32,9 +33,30 @@ limitations under the License. namespace tensorflow { namespace { bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) { - const int64_t kCompileThreshold = 10; + int64_t kCompileThreshold = 10; const int64_t kMinExecutionsPerCompile = 50; + int64_t tf_xla_threshold_for_megamorphic = + GetMarkForCompilationPassFlags()->tf_xla_threshold_for_megamorphic; + + // Negative values other that -1 cannot be used + if (tf_xla_threshold_for_megamorphic < -1) { + LOG(FATAL) << "The value for the tf_xla_threshold_for_megamorphic flag " + << "is out of range.\n" + << "Allowed ranges are (-1) to " + << std::numeric_limits::max() + << " got " << tf_xla_threshold_for_megamorphic << "."; + } + + // -1: setting clusters as Megamorphic is disabled + // 0 Default behaviour in Tensorflow + // Any other number sets the compilation threshold + if (tf_xla_threshold_for_megamorphic == -1) { + return false; + } else if (tf_xla_threshold_for_megamorphic > 0) { + kCompileThreshold = tf_xla_threshold_for_megamorphic; + } + // This heuristic is trying to capture the following property: have we sunk a // certain minimum amount of compile time into the cluster that didn't quite // "pay off"? diff --git a/tensorflow/compiler/jit/device_compiler.h b/tensorflow/compiler/jit/device_compiler.h index 34b22033129b96..a9e36dbc437673 100644 --- a/tensorflow/compiler/jit/device_compiler.h +++ b/tensorflow/compiler/jit/device_compiler.h @@ -36,6 +36,7 @@ limitations under the License. #include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h" #include "tensorflow/compiler/jit/xla_compile_util.h" +#include "tensorflow/compiler/jit/xla_batch_matcher.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/core/framework/metrics.h" #include "tensorflow/core/framework/op_kernel.h" @@ -125,6 +126,7 @@ class DeviceCompiler : public ResourceBase { DeviceCompilerClient* compiler_client() { return compiler_client_.get(); } + XlaBatchMatcher* xla_batch_matcher() { return xla_batch_matcher_.get(); } string DebugString() const override; @@ -177,6 +179,9 @@ class DeviceCompiler : public ResourceBase { // Pool of threads for asynchronous compilations. std::unique_ptr async_compiler_threads_; + // Specified dynamic batch padding values. + std::unique_ptr xla_batch_matcher_; + mutex cluster_mutexes_mu_; absl::flat_hash_map, DeviceCompilationClusterSignature::Hash> @@ -225,6 +230,11 @@ DeviceCompiler::DeviceCompiler( async_compiler_threads_ = std::make_unique( tensorflow::Env::Default(), "async_compiler_threads", kNumAsyncDeviceCompilerThreads); + + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + if (flags->tf_xla_enable_dynamic_sizes) { + xla_batch_matcher_ = std::make_unique(); + } } template diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index 8dd2ac0af31614..5a6f741a01e972 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -190,6 +190,11 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { &mark_for_compilation_flags->tf_xla_persistent_cache_prefix, "Specifies the persistance cache prefix. Default is " "\"xla_compile_cache\""), + Flag("tf_xla_threshold_for_megamorphic", + &mark_for_compilation_flags->tf_xla_threshold_for_megamorphic, + "Sets the threshold for marking a cluster megamorphic. " + "Setting it to -1 disables marking clusters megamorphic." + "Setting it to 0 uses the default behaviour of TensorFlow."), Flag("tf_xla_sparse_core_disable_table_stacking", &sparse_core_flags->tf_xla_sparse_core_disable_table_stacking, "Disable table stacking for all the tables passed to the SparseCore" @@ -253,6 +258,7 @@ void AllocateAndParseFlags() { mark_for_compilation_flags->tf_xla_cpu_global_jit = false; mark_for_compilation_flags->tf_xla_clustering_fuel = std::numeric_limits::max(); + mark_for_compilation_flags->tf_xla_threshold_for_megamorphic = 0; mark_for_compilation_flags ->tf_xla_disable_deadness_safety_checks_for_debugging = false; mark_for_compilation_flags diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index c37391d2d1ff56..480e20dc7ec296 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -123,6 +123,11 @@ struct MarkForCompilationPassFlags { // Specifies the persistance cache prefix. Default is "xla_compile_cache" string tf_xla_persistent_cache_prefix; + + // Sets the threshold for marking a cluster megamorphic. + // Setting it to -1 disables marking clusters megamorphic. + // Setting it to 0 uses the default behaviour of TensorFlow. + int64_t tf_xla_threshold_for_megamorphic; }; // Flags associated with XLA Sparse Core. diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 587d04b30c5274..1af416550a43dd 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -60,6 +60,7 @@ cc_library( "//tensorflow/compiler/jit:tf_graph_to_hlo_compiler", "//tensorflow/compiler/jit:tf_to_hlo_compiler", "//tensorflow/compiler/jit:xla_compile_util", + "//tensorflow/compiler/jit:xla_batch_matcher", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:refcount", "@com_google_absl//absl/base:core_headers", diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index f18c11c37bdc21..96c096306a3101 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -44,6 +44,7 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" +#include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/pjrt_compile_util.h" #include "tensorflow/compiler/jit/variable_info.h" @@ -56,6 +57,7 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_host_send_device_context.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" +#include "tensorflow/compiler/jit/xla_batch_matcher.h" #include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" @@ -74,6 +76,7 @@ limitations under the License. #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_util.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/monitoring/counter.h" @@ -100,6 +103,8 @@ limitations under the License. namespace tensorflow { namespace { +constexpr char kUserInferredValueContentsAttrName[] = + "_user_inferred_value_contents"; using XlaDeviceCompiler = DeviceCompiler; using PjRtDeviceCompiler = @@ -954,6 +959,7 @@ absl::Status CompileToLocalExecutable( const XlaPlatformInfo& platform_info, const std::vector& args, DeviceCompileMode compile_mode, bool may_alias_resource_update, + bool force_static_shapes, bool* dynamic_solve_conflict, xla::LocalClient** client, const XlaCompiler::CompilationResult** compilation_result, xla::LocalExecutable** executable) { @@ -997,9 +1003,413 @@ absl::Status CompileToLocalExecutable( XlaCompiler::CompileOptions compile_options = GenerateCompileOptions(has_ref_vars, may_alias_resource_update); - return xla_device_compiler->CompileIfNeeded( - options, function, args, compile_options, compile_mode, profiler, - compilation_result, executable); + if (dynamic_solve_conflict != nullptr) { + *dynamic_solve_conflict = false; + } + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + if (flags->tf_xla_enable_dynamic_sizes && !force_static_shapes) { + // 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. + std::vector norm_args(args.begin(), args.end()); + int64_t filled_batch = 0; + bool saw_dynamic_dim_value = false; + // Only supporting one dynamic dimension. + bool has_multiple_dynamic_dim_values = false; + int64_t dynamic_dim_value = 0; + XlaBatchMatcher* xla_batch_matcher = + xla_device_compiler->xla_batch_matcher(); + 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 = + expr.find_smallest_subexpression_covering_all_variables(); + VLOG(1) << "Using shared dynamic subexpression " + << DExprToString(*shared_dynamic_subexpr) + << " for XLA dynamic input normalization. context=" + << context; + } + xla::DExpr normalized_expr = + expr.replace_subexpression(*shared_dynamic_subexpr, + xla::DExpr::Var(1)) + .simplify(); + VLOG(1) << "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) { + auto& arg = norm_args[arg_index]; + if (arg.kind != XlaCompiler::Argument::kConstant) { + return; + } + + auto inferred_shape_it = attr_map.find("user_inferred_shape"); + auto inferred_contents_it = + attr_map.find(kUserInferredValueContentsAttrName); + bool has_dynamic = false; + auto has_dynamic_it = attr_map.find("has_dynamic"); + if (has_dynamic_it != attr_map.end()) { + has_dynamic = has_dynamic_it->second.b(); + } + + if (inferred_contents_it == attr_map.end() && + (!has_dynamic || inferred_shape_it == attr_map.end())) { + return; + } + + TensorShapeProto inferred_shape_proto; + if (inferred_contents_it != attr_map.end()) { + if (!inferred_shape_proto.ParseFromString( + inferred_contents_it->second.s())) { + return; + } + } else { + inferred_shape_proto = inferred_shape_it->second.shape(); + } + + TensorShape inferred_shape(inferred_shape_proto); + if (!((TensorShapeUtils::IsVector(arg.constant_value.shape()) && + arg.constant_value.NumElements() == inferred_shape.dims()) || + (TensorShapeUtils::IsScalar(arg.constant_value.shape()) && + inferred_shape.dims() == 1))) { + return; + } + + arg.constant_value_expressions.clear(); + arg.constant_value_expressions.reserve(inferred_shape.dims()); + for (int64_t i = 0; i < inferred_shape.dims(); ++i) { + xla::ExpressionProto expr; + const xla::DExpr& dim_expr = inferred_shape.get_expression(i); + if (dim_expr && dim_expr->is_dynamic()) { + 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) { + expr.set_constant_value(arg.constant_value.flat()(i)); + } else { + arg.constant_value_expressions.clear(); + return; + } + arg.constant_value_expressions.push_back(std::move(expr)); + } + }; + auto record_dynamic_dim_value = [&](int64_t dim_size, xla::DExpr expr) { + if (!saw_dynamic_dim_value) { + saw_dynamic_dim_value = true; + dynamic_dim_value = dim_size; + dynamic_dim_expr = std::move(expr); + return; + } + if (dynamic_dim_value != dim_size) { + has_multiple_dynamic_dim_values = true; + } + }; + if (options.flib_def != nullptr) { + const FunctionDef* fdef = options.flib_def->Find(function.name()); + if (fdef != nullptr) { + for (const auto& kv : fdef->arg_attr()) { + int arg_index = kv.first; + const auto& attr_map = kv.second.attr(); + const std::string& node_name = + fdef->signature().input_arg(arg_index).name(); + + auto shape_derived_attr = attr_map.find(kXlaShapeDerivedAttrName); + if (shape_derived_attr != attr_map.end()) { + VLOG(1) << "XlaCompileOp retrieved shape-derived marker for arg " + << arg_index << " node=" << node_name; + } + maybe_attach_shape_contents_from_attrs(arg_index, attr_map, node_name); + + // Special case for _dynamic_dim... + auto dyn_dim_attr = attr_map.find("_dynamic_dim"); + if (dyn_dim_attr != attr_map.end()) { + TensorShape& shp = + 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) { + filled_batch = + xla_batch_matcher->get_xla_compile_batch(shp.dim_size(idx)); + } + + std::vector dyn_exprs; + for (int d : shp.dim_sizes()) { + dyn_exprs.push_back(xla::DExpr::Const(d)); + } + dyn_exprs[idx] = *dynamic_dim_expr; + shp.set_expressions(std::move(dyn_exprs)); + continue; + } + auto it = attr_map.find(kXlaInferredOutputShapesAttrName); + if (it == attr_map.end()) continue; + + 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) { + for (int idx = 0; idx < exp.size(); ++idx) { + // Look for dynamic expression. If found then compute padding + // value and exit loop. + auto e = normalize_dynamic_expr(DimExprFromProto(exp[idx])); + if (e->is_dynamic()) { + VLOG(1) << "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)); + VLOG(1) << "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) + << "Failed to solve dynamic dimension for argument " + << arg_index << " dim " << idx << " with size " + << shp.dim_size(idx) + << "; falling back to original dimension size."; + var_value = shp.dim_size(idx); + } else { + var_value = *solved_value; + VLOG(1) << "Solved dynamic dimension from " + << shp.dim_size(idx) << " to " << var_value; + } + record_dynamic_dim_value(var_value, e); + filled_batch = + xla_batch_matcher->get_xla_compile_batch(var_value); + break; + } + } + } + + std::vector dyn_exprs; + for (int d : shp.dim_sizes()) { + dyn_exprs.push_back(xla::DExpr::Const(d)); + } + for (int j = 0; j < exp.size(); ++j) { + auto e = DimExprFromProto(exp[j]); + if (e->is_dynamic()) { + e = normalize_dynamic_expr( + e, absl::StrCat("arg=", arg_index, " dim=", j, + " input_shape")); + dyn_exprs[j] = e; + } + } + shp.set_expressions(std::move(dyn_exprs)); + } + } + } + + DynamicSolveFilterDecision solve_filter_decision = + AnalyzeIgnoredDynamicArgumentOccurrences(norm_args); + if (!solve_filter_decision.can_run) { + if (dynamic_solve_conflict != nullptr) { + *dynamic_solve_conflict = true; + } + 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(WARNING) << 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 (!filled_batch && xla_batch_matcher) { + filled_batch = + xla_batch_matcher->get_xla_compile_batch(*solved_value); + } + } + } + } + + struct SaveOldVar { + int arg_index; + int64_t dyn_dim; + int64_t old_value; + }; + std::vector old_vars; + auto maybe_rewrite_scalar_constant = [&](int arg_index) { + if (!saw_dynamic_dim_value || has_multiple_dynamic_dim_values) { + return; + } + + auto& arg = norm_args[arg_index]; + if (arg.kind != XlaCompiler::Argument::kConstant) { + return; + } + + const bool is_scalar = TensorShapeUtils::IsScalar(arg.constant_value.shape()); + const bool is_vector = TensorShapeUtils::IsVector(arg.constant_value.shape()) && + arg.constant_value.NumElements() > 0; + if (!is_scalar && !is_vector) { + return; + } + + auto set_constant_contents = [&](int rewrite_index) { + arg.constant_value_expressions.clear(); + const int64_t num_elements = arg.constant_value.NumElements(); + arg.constant_value_expressions.reserve(num_elements); + for (int64_t i = 0; i < num_elements; ++i) { + xla::ExpressionProto expr; + if (i == rewrite_index) { + dynamic_dim_expr->to_proto(&expr); + } else { + expr.set_constant_value(arg.constant_value.flat()(i)); + } + arg.constant_value_expressions.push_back(std::move(expr)); + } + }; + + if (arg.constant_value.dtype() == DT_INT32) { + auto flat = arg.constant_value.flat(); + int rewrite_index = -1; + // Heuristic: rewrite only scalar constants or shape-like int vectors. + // In practice we expect at most one entry to match the observed + // runtime batch size, so rewrite the first matching entry. + for (int i = 0; i < arg.constant_value.NumElements(); ++i) { + if (flat(i) == dynamic_dim_value) { + rewrite_index = i; + break; + } + } + if (rewrite_index >= 0) { + arg.constant_value = tensor::DeepCopy(arg.constant_value); + auto mutable_flat = arg.constant_value.flat(); + VLOG(1) << "XlaCompileOp int32 constant arg " << arg_index + << " index " << rewrite_index + << " matches dynamic_dim_value=" << dynamic_dim_value; + mutable_flat(rewrite_index) = filled_batch; + set_constant_contents.template operator()(rewrite_index); + } + } else if (arg.constant_value.dtype() == DT_INT64) { + auto flat = arg.constant_value.flat(); + int rewrite_index = -1; + // Same heuristic for int64 scalar constants or shape-like vectors. + for (int i = 0; i < arg.constant_value.NumElements(); ++i) { + if (flat(i) == dynamic_dim_value) { + rewrite_index = i; + break; + } + } + if (rewrite_index >= 0) { + arg.constant_value = tensor::DeepCopy(arg.constant_value); + auto mutable_flat = arg.constant_value.flat(); + VLOG(1) << "XlaCompileOp int64 constant arg " << arg_index + << " index " << rewrite_index + << " matches dynamic_dim_value=" << dynamic_dim_value; + mutable_flat(rewrite_index) = filled_batch; + set_constant_contents.template operator()(rewrite_index); + } + } + }; + // We rewrite only dynamic dimensions to the padded compile batch and then + // restore the original runtime sizes after compilation. Some scalar + // constants are actually runtime batch sizes folded by earlier TF passes, + // so rewrite only those that match the detected dynamic runtime value. + // Scalar constants are deep-copied before rewrite so the change stays + // local to norm_args and does not require restoration. + if (filled_batch) { + for (int i = 0; i < norm_args.size(); ++i) { + TensorShape& shp = std::get(norm_args[i].shape); + for (int j = 0; j < shp.get_expressions().size(); ++j) { + auto e = shp.get_expression(j); + if (e && e->is_dynamic()) { + int64_t old = shp.dim_size(j); + old_vars.push_back({i, j, old}); + xla::DExpr padded_expr = xla::DExpr::Const(filled_batch); + 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(); + VLOG(1) << "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(); + VLOG(1) << "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 " + "integer constant for argument ", + i, ", dimension ", j, ": ", DExprToString(subst_expr)); + } + int64_t new_dim = subst_expr->get_val(); + if (new_dim >= 0) { + shp.set_dim(j, new_dim); + // Necessary because set_dim removes the expression: + shp.set_expression(j, e); + } + } + } + maybe_rewrite_scalar_constant(i); + } + } + auto status = xla_device_compiler->CompileIfNeeded( + options, function, norm_args, compile_options, compile_mode, profiler, + compilation_result, executable); + // Restore the original runtime dimensions after compilation. + if (filled_batch) { + for (const auto& old_var : old_vars) { + TensorShape& shp = + std::get(norm_args[old_var.arg_index].shape); + shp.set_dim(old_var.dyn_dim, old_var.old_value); + } + } + return status; + } else { + return xla_device_compiler->CompileIfNeeded( + options, function, args, compile_options, compile_mode, profiler, + compilation_result, executable); + } } absl::Status GetUpdatedVariables( @@ -1152,11 +1562,25 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { return; } + bool dynamic_solve_conflict = false; absl::Status status = CompileToLocalExecutable( ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_, xla_compiler_args, DeviceCompileMode::kStrict, - /*may_alias_resource_update=*/true, &client, &compilation_result, + /*may_alias_resource_update=*/true, /*force_static_shapes=*/false, + &dynamic_solve_conflict, &client, &compilation_result, &executable); + if (dynamic_solve_conflict) { + LOG(WARNING) << "Retrying XLA cluster " << function_.name() + << " with concrete static argument shapes"; + std::vector static_args = + BuildStaticCompilationArguments(xla_compiler_args); + status = CompileToLocalExecutable( + ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_, + static_args, DeviceCompileMode::kStrict, + /*may_alias_resource_update=*/true, /*force_static_shapes=*/true, + /*dynamic_solve_conflict=*/nullptr, &client, &compilation_result, + &executable); + } OP_REQUIRES_OK_ASYNC(ctx, status, done); // Continuation of the execution, may be run in a different thread. @@ -1345,6 +1769,27 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { GetXlaOpsCommonFlags() ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( platform_info_.device_type()); + std::vector compiler_args; + auto compile_static_fallback = [&]() -> absl::Status { + std::vector static_args = + BuildStaticCompilationArguments(compiler_args); + kernel = nullptr; + executable = nullptr; + pjrt_executable = nullptr; + LOG(WARNING) << "Retrying XLA cluster " << function_.name() + << " with concrete static argument shapes"; + if (use_pjrt) { + return CompileToPjRtLoadedExecutable( + *ctx, platform_info_, function_, static_args, compile_mode, + has_ref_vars_, /*may_alias_resource_update=*/false, &kernel, + &pjrt_client, &pjrt_executable); + } + return CompileToLocalExecutable( + ctx, function_, has_ref_vars_, platform_info_, static_args, + compile_mode, /*may_alias_resource_update=*/false, + /*force_static_shapes=*/true, + /*dynamic_solve_conflict=*/nullptr, &client, &kernel, &executable); + }; if (GetXlaOpsCommonFlags()->tf_xla_always_defer_compilation || cannot_compile_cluster) { @@ -1353,13 +1798,14 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { auto args_and_variables_snapshot = GetXlaCompilerArgsAndSnapshotVariables( resources_, constants_, inputs, ctx); OP_REQUIRES_OK(ctx, args_and_variables_snapshot.status()); - const std::vector& args = - args_and_variables_snapshot->first; + compiler_args = std::move(args_and_variables_snapshot->first); + const std::vector& args = compiler_args; variables_snapshot = std::move(args_and_variables_snapshot->second); // Do not alias resource updates as locking variables in XlaCompile and // unlocking them in XlaRun may lead to deadlocks. absl::Status status; + bool dynamic_solve_conflict = false; if (use_pjrt) { VLOG(2) << "Using PJRT for compilation. Function name: " << function_.name(); @@ -1370,16 +1816,30 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { } else { status = CompileToLocalExecutable( ctx, function_, has_ref_vars_, platform_info_, args, compile_mode, - /*may_alias_resource_update=*/false, &client, &kernel, &executable); + /*may_alias_resource_update=*/false, /*force_static_shapes=*/false, + &dynamic_solve_conflict, &client, &kernel, &executable); + } + if (dynamic_solve_conflict) { + status = compile_static_fallback(); } + if (compile_mode != DeviceCompileMode::kLazy || status.code() != error::UNIMPLEMENTED) { - OP_REQUIRES_OK(ctx, status); + if ((status != OkStatus()) && + (status.code() != error::UNIMPLEMENTED) && + (compile_mode == DeviceCompileMode::kLazy)) { + // We set the error to error::UNIMPLEMENTED so it falls in the + // conditions of the if to fall back to TensorFlow function call + status = tensorflow::errors::Unimplemented(status.ToString()); + } else { + OP_REQUIRES_OK(ctx, status); + } } if (status.code() == error::UNIMPLEMENTED) { - LOG(WARNING) << "Compilation failed:" << status - << ". Falling back to TF function call."; + LOG(WARNING) << "[HUAWEI] Compilation of the cluster failed with:"; + LOG(WARNING) << "[HUAWEI] " << status; + LOG(WARNING) << "[HUAWEI] Falling back to TF function call.\n"; BroadcastOptimizationRemark( XlaOptimizationRemark::UNIMPLEMENTED_OPERATION, status.ToString()) @@ -1387,10 +1847,41 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { executable = nullptr; pjrt_executable = nullptr; mutex_lock guard(cannot_compile_cluster_mu_); + // TODO: decide if we want to set this flag to true, as we may want to + // allow the cluster to try to compile again later in time. cannot_compile_cluster_ = true; } } + 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); + LOG(WARNING) << error_message; + absl::Status static_status = compile_static_fallback(); + if (!static_status.ok()) { + if (must_compile_) { + OP_REQUIRES_OK(ctx, static_status); + } + LOG(WARNING) << "Static XLA fallback failed for cluster " + << function_.name() << ": " << static_status; + executable = nullptr; + pjrt_executable = nullptr; + kernel = nullptr; + } else { + LOG(INFO) << "Using static XLA fallback for cluster " + << function_.name(); + } + } + } + AllocatorAttributes host_alloc_attrs; host_alloc_attrs.set_gpu_compatible(true); host_alloc_attrs.set_on_host(true); @@ -1441,7 +1932,6 @@ XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) void XlaRunOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaRunOp " << def().name(); Tensor key_tensor = ctx->input(ctx->num_inputs() - 1); - bool use_pjrt = GetXlaOpsCommonFlags() ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( diff --git a/tensorflow/compiler/jit/xla_batch_matcher.cc b/tensorflow/compiler/jit/xla_batch_matcher.cc new file mode 100644 index 00000000000000..8b991ce6e2b865 --- /dev/null +++ b/tensorflow/compiler/jit/xla_batch_matcher.cc @@ -0,0 +1,203 @@ +/* Copyright 2026 The TensorFlow Authors. + +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 +#include + +#include "xla/debug_options_flags.h" +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { + +namespace { + +// Trims spaces and tabs from both ends of a string. +std::string Trim(const std::string& value) { + size_t start = value.find_first_not_of(" \t"); + size_t end = value.find_last_not_of(" \t"); + return start == std::string::npos ? "" + : value.substr(start, end - start + 1); +} + +} // namespace + +XlaBatchMatcher::XlaBatchMatcher() { + env_str_ = xla::GetDebugOptionsFromFlags().xla_compile_batch_sizes(); + parse_env_config(); +} + +std::vector XlaBatchMatcher::parse_single_item(const std::string& item) { + std::vector batch_list; + if (item.empty()) return batch_list; + + // 1. Parse single value (no colon separator) + if (item.find(':') == std::string::npos) { + // Validate all characters are digits (reject non-numeric chars) + for (char c : item) { + if (!isdigit(c)) { + throw std::invalid_argument("Non-numeric characters: " + item); + } + } + auto val = static_cast(std::stoll(item)); + if (val <= 0 || val > kMaxBatch) { + throw std::invalid_argument("Out of valid range (1-" + std::to_string(kMaxBatch) + "): " + item); + } + batch_list.push_back(val); + return batch_list; + } + + // 2. Parse range format (start:end:step) + std::stringstream ss(item); + std::string part; + std::vector parts; + while (std::getline(ss, part, ':')) { + parts.push_back(Trim(part)); + } + if (parts.size() > 3) { + throw std::invalid_argument("Invalid range format (requires start:end:step): " + item); + } + + // Convert and validate range parameters + int64_t start, end, step; + try { + start = std::stoi(parts[0]); + end = std::stoi(parts[1]); + step = parts.size() == 2 ? 1 : std::stoi(parts[2]); + } catch (...) { + throw std::invalid_argument("Invalid numeric values in range: " + item); + } + if (start <= 0 || end <= 0 || step <= 0) { + throw std::invalid_argument("Range parameters must be positive integers: " + item); + } + if (start > end) { + throw std::invalid_argument("Start value > end value in range: " + item); + } + if (end > kMaxBatch) { + throw std::invalid_argument("Range exceeds max limit (" + std::to_string(kMaxBatch) + "): " + item); + } + + // Generate batch list from range + for (int64_t i = start; i <= end; i += step) { + batch_list.push_back(i); + } + return batch_list; +} + +void XlaBatchMatcher::print_all_batches() { + std::ostringstream oss; + oss << "[XLA_BATCH_INFO] Valid batch list update: "; + for (size_t i = 0; i < all_batches_.size(); ++i) { + if (i > 0) oss << ", "; + oss << all_batches_[i]; + } + LOG(INFO) << oss.str(); +} + +// Parse environment variable config into deduplicated, sorted batch list +// For example, export XLA_COMPILE_BATCH_SIZES="10:100:10, 977" +void XlaBatchMatcher::parse_env_config() { + // If the env var not set or is empty, filled with the nearest power of two by default + if (env_str_.empty()) { + VLOG(2) << "[XLA_BATCH_WARN] Env var " << "--tf_xla_compile_batch_sizes" << + "is empty, filled with the nearest power of two by default"; + return; + } + + // Split config by commas + std::stringstream ss(env_str_); + std::string item; + while (std::getline(ss, item, ',')) { + std::string trimmed_item = Trim(item); + if (trimmed_item.empty()) continue; + + // Parse single item (skip on failure to avoid breaking other items) + try { + std::vector item_batches = parse_single_item(trimmed_item); + all_batches_.insert(all_batches_.end(), item_batches.begin(), item_batches.end()); + } catch (const std::exception& e) { + LOG(INFO) << "[XLA_BATCH_WARN] Failed to parse config item, skipping: " << + trimmed_item << " (" << e.what() << ")"; + } + } + + if (!all_batches_.empty()) { + std::sort(all_batches_.begin(), all_batches_.end()); + auto last = std::unique(all_batches_.begin(), all_batches_.end()); + all_batches_.erase(last, all_batches_.end()); + } + + // Print parsed result + if (!all_batches_.empty()) print_all_batches(); + 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) { + if (real_batch <= 0 || real_batch > kMaxBatch) { + LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch; + return real_batch; + } + + if (all_batches_.empty()) { + // Return the next power of two directly without modifying all_batches_ + return GetNextPowerOfTwo(real_batch); + } + + // Edge case 1: Real value < the smallest batch, use smallest + if (real_batch < all_batches_.front()) { + return all_batches_.front(); + } + // 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); + 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(); +} + +int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) { + // Match target batch size + int64_t selected = find_min_larger_batch(real_batch); + 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; + } + return selected; +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_batch_matcher.h b/tensorflow/compiler/jit/xla_batch_matcher.h new file mode 100644 index 00000000000000..dac8079e5081af --- /dev/null +++ b/tensorflow/compiler/jit/xla_batch_matcher.h @@ -0,0 +1,49 @@ +/* Copyright 2026 The TensorFlow Authors. + +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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_ +#define TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_ + +#include +#include +#include + +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; + +class XlaBatchMatcher { + public: + XlaBatchMatcher(); + virtual ~XlaBatchMatcher() = default; + int64_t get_xla_compile_batch(int64_t real_batch); + 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); + + std::vector all_batches_; + std::string env_str_; + int64_t last_batch_ = -1; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_ diff --git a/tensorflow/core/kernels/padding_fifo_queue.cc b/tensorflow/core/kernels/padding_fifo_queue.cc index 3b50099fb9997c..ed669f60550501 100644 --- a/tensorflow/core/kernels/padding_fifo_queue.cc +++ b/tensorflow/core/kernels/padding_fifo_queue.cc @@ -400,7 +400,14 @@ std::vector PaddingFIFOQueue::ConvertShapesPartialDimensionsToZero( for (size_t i = 0; i < shapes.size(); ++i) { const PartialTensorShape& partial = partial_shapes[i]; TensorShape& shape = shapes[i]; - for (int64_t s : partial.dim_sizes()) shape.AddDim(s < 0 ? 0 : s); + for (int d = 0; d < partial.dims(); ++d) { + shape.AddDim(partial.dim_size(d) < 0 ? 0 : partial.dim_size(d)); + xla::DExpr expr = partial.get_filled_expression(d); + if (expr && expr->is_constant() && expr->get_val() < 0) { + expr = xla::DExpr::Const(0); + } + shape.AddExpression(expr); + } } return shapes; } diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 33fa90f7e35e9e..961273151ed5e0 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -103,6 +103,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_cpu_use_fusion_emitters(true); opts.set_xla_cpu_use_thunk_runtime(true); opts.set_xla_cpu_use_xnnpack(false); + opts.set_xla_compile_batch_sizes(""); opts.set_xla_cpu_experimental_xnn_graph_fusion_mode( DebugOptions::XNN_GRAPH_FUSION_MODE_DISABLED); opts.set_xla_cpu_parallel_codegen_split_count(32); @@ -1003,6 +1004,15 @@ void MakeDebugOptionsFlags(std::vector* flag_list, "`XNN_GRAPH_FUSION_MODE_DISABLED` - default value, " "`XNN_GRAPH_FUSION_MODE_GREEDY` - greedy extraction of " "XNNPACK-compatible subgraphs starting from root instructions.")); + flag_list->push_back(tsl::Flag( + "xla_compile_batch_sizes", + string_setter_for( + &DebugOptions::set_xla_compile_batch_sizes), + debug_options->xla_compile_batch_sizes(), + "Comma-separated list of batch sizes to use for compilation, " + "use single value or start:end:step format. " + "e.g. 32, 64, 128, 10:100:10, " + "empty to use the nearest power of two.")); flag_list->push_back(tsl::Flag( "xla_cpu_parallel_codegen_split_count", int32_setter_for(&DebugOptions::set_xla_cpu_parallel_codegen_split_count), diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index ca8ba0553bd56a..18abf310ac67e4 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -222,6 +222,8 @@ message DebugOptions { // When true, XLA:CPU uses XNNPACK to execute supported operations. bool xla_cpu_use_xnnpack = 359; + string xla_compile_batch_sizes = 399; + // Enabling this will enable optimizations that ignore the possibility of NaN. bool xla_enable_fast_math = 335; @@ -1208,7 +1210,7 @@ message DebugOptions { // Note: when adding a new flag, please add it to one of the hardware-specific // or hardware-agnostic sections at the top of this proto message. - // Next id: 389 + // Next id: 400 // Extra options to pass to the compilation backend (e.g. LLVM); specific // interpretation of these values is left to the backend. From 3e4d7d8cce1a074964ea11eee81c15c006134510 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:04:29 +0100 Subject: [PATCH 13/21] Guard MLIR lowering for dynamic shape expressions Prevent MLIR XLA kernels from silently consuming dynamic expressions they cannot preserve. Keep the existing MLIR path for operations without dynamic expression metadata. --- .../compiler/jit/mark_for_compilation_pass.cc | 9 +++ .../compiler/tf2xla/kernels/batch_norm_op.cc | 10 ++- tensorflow/compiler/tf2xla/kernels/fft_ops.cc | 8 +- .../compiler/tf2xla/kernels/identity_op.cc | 7 +- .../compiler/tf2xla/kernels/unary_ops.cc | 57 +++++++++----- .../compiler/tf2xla/mlir_xla_op_kernel.cc | 76 +++++++++++++++++++ .../compiler/tf2xla/mlir_xla_op_kernel.h | 9 +++ tensorflow/compiler/tf2xla/xla_op_registry.cc | 15 +++- tensorflow/compiler/tf2xla/xla_op_registry.h | 24 +++++- .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc | 18 +++-- .../translate/mhlo_to_hlo/type_to_shape.cc | 3 +- 11 files changed, 199 insertions(+), 37 deletions(-) diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 81365b13e5a9d5..d1343c7235fc22 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -1821,6 +1821,15 @@ absl::Status MarkForCompilationPassImpl::FindCompilationCandidates() { continue; } + if (debug_options_.enable_dynamic_sizes && + XlaOpRegistry::IsMlirXlaOp(node->type_string()) && + HasDynamicInputExpression(node)) { + VLOG(1) << "Rejecting " << node->name() + << " from XLA clustering: MlirXlaOpKernel does not support " + "dynamic input expressions"; + continue; + } + if (node->type_string() == "Const") { // Skip Const op with type DT_STRING, since XLA autoclustering doesn't // support it. diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc index 0dd528e3dea173..1db3ab4c78c461 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc @@ -21,9 +21,9 @@ limitations under the License. #include #include "tensorflow/compiler/tf2xla/kernels/relu_op.h" -#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -241,7 +241,9 @@ class FusedBatchNormOpEx : public FusedBatchNormOp { REGISTER_XLA_OP(Name("FusedBatchNorm"), FusedBatchNormOp); REGISTER_XLA_OP(Name("FusedBatchNormV2"), FusedBatchNormOp); -REGISTER_XLA_OP(Name("FusedBatchNormV3"), MlirXlaOpKernel); +REGISTER_XLA_OP_FACTORY( + Name("FusedBatchNormV3"), + CreateDynamicNativeXlaOpKernel); REGISTER_XLA_OP(Name("_FusedBatchNormEx"), FusedBatchNormOpEx); class FusedBatchNormGradOp : public XlaOpKernel { @@ -358,7 +360,9 @@ class FusedBatchNormGradOp : public XlaOpKernel { REGISTER_XLA_OP(Name("FusedBatchNormGrad"), FusedBatchNormGradOp); REGISTER_XLA_OP(Name("FusedBatchNormGradV2"), FusedBatchNormGradOp); -REGISTER_XLA_OP(Name("FusedBatchNormGradV3"), MlirXlaOpKernel); +REGISTER_XLA_OP_FACTORY( + Name("FusedBatchNormGradV3"), + CreateDynamicNativeXlaOpKernel); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc index 8fb04773aafb49..5d4114d8c88d5d 100644 --- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc @@ -20,8 +20,8 @@ limitations under the License. #include #include "absl/container/inlined_vector.h" -#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_builder.h" @@ -139,9 +139,9 @@ class IFFTOp : public GenericFftOp { explicit IFFTOp(OpKernelConstruction* ctx) : GenericFftOp(ctx, /*fft_type=*/FftType::IFFT, /*fft_rank=*/FFTRank) {} }; -REGISTER_XLA_OP(Name("IFFT").TypeConstraint("Tcomplex", - {DT_COMPLEX64, DT_COMPLEX128}), - MlirXlaOpKernel); +REGISTER_XLA_OP_FACTORY( + Name("IFFT").TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128}), + CreateDynamicNativeXlaOpKernel>); REGISTER_XLA_OP(Name("IFFT2D").TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128}), IFFTOp<2>); diff --git a/tensorflow/compiler/tf2xla/kernels/identity_op.cc b/tensorflow/compiler/tf2xla/kernels/identity_op.cc index 3e765b853e110d..f6abb610c456e4 100644 --- a/tensorflow/compiler/tf2xla/kernels/identity_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/identity_op.cc @@ -32,6 +32,10 @@ class IdentityOp : public XlaOpKernel { for (int i = 0; i < ctx->num_inputs(); ++i) { if (IsTensorListInput(ctx, i)) { ctx->SetTensorListOutput(i, ctx->Input(i)); + } else if (ctx->InputExpression(i).kind() != + XlaExpression::Kind::kResource && + ctx->input_type(i) != DT_VARIANT) { + ctx->SetOutputExpression(i, ctx->InputExpression(i)); } else { DCHECK(ctx->input_type(i) != DT_VARIANT); // Forwards using the underlying op_kernel_context so both tensor and @@ -58,7 +62,8 @@ REGISTER_XLA_OP(Name("IdentityN") .CompilationOnly(), IdentityOp); REGISTER_XLA_OP(Name("PlaceholderWithDefault"), IdentityOp); -REGISTER_XLA_OP(Name("PreventGradient"), MlirXlaOpKernel); +REGISTER_XLA_OP_FACTORY( + Name("PreventGradient"), CreateDynamicNativeXlaOpKernel); REGISTER_XLA_OP(Name("StopGradient").AllowVariantTypes(), IdentityOp); REGISTER_XLA_OP(Name("Snapshot"), IdentityOp); REGISTER_XLA_OP(Name("_EagerConst"), IdentityOp); diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index f35e375356516b..8dc75be96d7958 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -45,6 +45,22 @@ namespace { }; \ REGISTER_XLA_OP(Name(#NAME), NAME##Op); +#define XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(NAME, COMPUTATION) \ + class NAME##NativeOp : public XlaOpKernel { \ + public: \ + explicit NAME##NativeOp(OpKernelConstruction* ctx) \ + : XlaOpKernel(ctx) {} \ + void Compile(XlaOpKernelContext* ctx) override { \ + xla::XlaBuilder* b = ctx->builder(); \ + (void)b; \ + xla::XlaOp x = ctx->Input(0); \ + xla::XlaOp y = COMPUTATION; \ + ctx->SetOutput(0, y); \ + } \ + }; \ + REGISTER_XLA_OP_FACTORY( \ + Name(#NAME), CreateDynamicNativeXlaOpKernel); + XLAJIT_MAKE_UNARY(ComplexAbs, xla::Abs(x)); XLAJIT_MAKE_UNARY(Angle, xla::Atan2(xla::Imag(x), xla::Real(x))); @@ -52,29 +68,30 @@ XLAJIT_MAKE_UNARY(Angle, xla::Atan2(xla::Imag(x), xla::Real(x))); XLAJIT_MAKE_UNARY(Conj, xla::Conj(x)); // Return x if x>0, otherwise -x. -REGISTER_XLA_OP(Name("Abs"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Abs, xla::Abs(x)); XLAJIT_MAKE_UNARY(Acos, xla::Acos(x)); XLAJIT_MAKE_UNARY(Acosh, xla::Acosh(x)); XLAJIT_MAKE_UNARY(Asin, xla::Asin(x)) XLAJIT_MAKE_UNARY(Asinh, xla::Asinh(x)); -REGISTER_XLA_OP(Name("Atan"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Atan, xla::Atan(x)); XLAJIT_MAKE_UNARY(Atanh, xla::Atanh(x)); -REGISTER_XLA_OP(Name("Ceil"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Cos"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Ceil, xla::Ceil(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Cos, xla::Cos(x)); XLAJIT_MAKE_UNARY(Cosh, xla::Cosh(x)); XLAJIT_MAKE_UNARY(Sin, xla::Sin(x)); XLAJIT_MAKE_UNARY(Tan, xla::Tan(x)); -REGISTER_XLA_OP(Name("Exp"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Expm1"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Floor"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("IsFinite"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("IsInf"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("IsNan"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Exp, xla::Exp(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Expm1, xla::Expm1(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Floor, xla::Floor(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsFinite, xla::IsFinite(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsInf, xla::IsInf(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsNan, xla::IsNan(x)); // Return 1/x XLAJIT_MAKE_UNARY(Inv, xla::ScalarLike(x, 1.0) / x); -REGISTER_XLA_OP(Name("Reciprocal"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK( + Reciprocal, xla::ScalarLike(x, 1.0) / x); XLAJIT_MAKE_UNARY(Log, xla::Log(x)); -REGISTER_XLA_OP(Name("Log1p"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Log1p, xla::Log1p(x)); XLAJIT_MAKE_UNARY(Invert, xla::Not(x)); XLAJIT_MAKE_UNARY(LogicalNot, xla::Not(x)); @@ -85,12 +102,12 @@ XLAJIT_MAKE_UNARY(Neg, -x); XLAJIT_MAKE_UNARY(Rint, xla::RoundToEven(x)); XLAJIT_MAKE_UNARY(Round, xla::RoundToEven(x)); -REGISTER_XLA_OP(Name("Rsqrt"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Rsqrt, xla::Rsqrt(x)); -REGISTER_XLA_OP(Name("Sigmoid"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sigmoid, xla::Logistic(x)); // Returns NaN if x is NaN, 0 if x is 0, -1 if x < 0 and 1 if x > 0. -REGISTER_XLA_OP(Name("Sign"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sign, xla::Sign(x)); XLAJIT_MAKE_UNARY(Sinh, xla::Sinh(x)); static xla::XlaOp Softplus(xla::XlaBuilder* b, xla::XlaOp features) { @@ -115,11 +132,11 @@ XLAJIT_MAKE_UNARY(Softplus, Softplus(b, x)); // softsign(x) = x / (abs(x) + 1) XLAJIT_MAKE_UNARY(Softsign, x / (xla::Abs(x) + xla::ScalarLike(x, 1.0))); -REGISTER_XLA_OP(Name("Sqrt"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sqrt, xla::Sqrt(x)); XLAJIT_MAKE_UNARY(Square, x* x); -REGISTER_XLA_OP(Name("Tanh"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Real"), MlirXlaOpKernel); -REGISTER_XLA_OP(Name("Imag"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Tanh, xla::Tanh(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Real, xla::Real(x)); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Imag, xla::Imag(x)); XLAJIT_MAKE_UNARY(Erf, xla::Erf(x)); XLAJIT_MAKE_UNARY(Erfc, xla::Erfc(x)); XLAJIT_MAKE_UNARY(Erfinv, xla::ErfInv(x)); @@ -127,7 +144,7 @@ XLAJIT_MAKE_UNARY(Erfinv, xla::ErfInv(x)); XLAJIT_MAKE_UNARY(Ndtri, xla::ScalarLike(x, std::sqrt(2.0)) * xla::ErfInv(xla::ScalarLike(x, 2.0) * x - xla::ScalarLike(x, 1.0))); -REGISTER_XLA_OP(Name("Lgamma"), MlirXlaOpKernel); +XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Lgamma, xla::Lgamma(x)); XLAJIT_MAKE_UNARY(Digamma, xla::Digamma(x)); XLAJIT_MAKE_UNARY(BesselI0e, xla::BesselI0e(x)); XLAJIT_MAKE_UNARY(BesselI1e, xla::BesselI1e(x)); diff --git a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc index b1a93508d92896..28d3447918a2c7 100644 --- a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h" #include +#include #include "absl/status/status.h" #include "absl/strings/str_cat.h" @@ -35,7 +36,9 @@ limitations under the License. #include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/resource_base.h" #include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/graph/graph.h" #include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/refcount.h" #include "tensorflow/core/platform/status.h" @@ -68,6 +71,77 @@ class MLIRContextResource : public ResourceBase { mlir::MLIRContext mlir_ctx_; }; +bool HasDynamicExpressions(const TensorShape& shape) { + for (const auto& expr : shape.get_expressions()) { + if (expr && expr->is_dynamic()) { + return true; + } + } + return false; +} + +bool HasDynamicExpressions(const PartialTensorShape& shape) { + for (const auto& expr : shape.get_expressions()) { + if (expr && expr->is_dynamic()) { + return true; + } + } + return false; +} + +bool HasDynamicExpressions(const xla::Shape& shape) { + for (const auto& expr : shape.expressions()) { + if (expr && expr->is_dynamic()) { + return true; + } + } + return false; +} + +absl::Status RejectDynamicShapeExpressionsInMlirXlaOpKernel( + llvm::ArrayRef args, const Graph& graph) { + for (int i = 0; i < args.size(); ++i) { + const auto& shape = args[i].shape; + const bool has_dynamic_exprs = + std::holds_alternative(shape) + ? HasDynamicExpressions(std::get(shape)) + : HasDynamicExpressions(std::get(shape)); + if (has_dynamic_exprs) { + return errors::Unimplemented( + "MlirXlaOpKernel does not support dynamic shape expressions. " + "Argument ", + i, " carries a dynamic expression."); + } + } + + for (Node* node : graph.nodes()) { + for (const auto& name_attr_pair : node->attrs()) { + const auto& attr_name = name_attr_pair.first; + const auto& attr_value = name_attr_pair.second; + auto maybe_reject_shape = [&](const TensorShapeProto& shape_proto) { + const PartialTensorShape shape(shape_proto); + if (!HasDynamicExpressions(shape)) { + return absl::OkStatus(); + } + return errors::Unimplemented( + "MlirXlaOpKernel does not support dynamic shape expressions. " + "Node '", + node->name(), "' attribute '", attr_name, + "' carries a dynamic expression."); + }; + + if (attr_value.value_case() == AttrValue::kShape) { + TF_RETURN_IF_ERROR(maybe_reject_shape(attr_value.shape())); + } else if (attr_value.value_case() == AttrValue::kList) { + for (const auto& shape_proto : attr_value.list().shape()) { + TF_RETURN_IF_ERROR(maybe_reject_shape(shape_proto)); + } + } + } + } + return absl::OkStatus(); +} + } // namespace absl::Status MlirXlaOpKernel::ContextToXlaArgs( @@ -143,6 +217,8 @@ absl::Status MlirXlaOpKernel::ConstructXlaOp(XlaOpKernelContext* ctx) { // Create a graph that wraps the kernel. TF_ASSIGN_OR_RETURN(auto graph, CreateSingleOpGraph(def(), xla_args, result_dtypes)); + TF_RETURN_IF_ERROR( + RejectDynamicShapeExpressionsInMlirXlaOpKernel(xla_args, *graph)); ResourceMgr* res_manager = ctx->op_kernel_context()->resource_manager(); MLIRContextResource* ctx_res; diff --git a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h index 6053f5d68635d0..7b70096851b3df 100644 --- a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_ #define TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_ +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/core/framework/op_kernel.h" @@ -36,6 +37,14 @@ class MlirXlaOpKernel : public XlaOpKernel { absl::Status ConstructXlaOp(XlaOpKernelContext* ctx); }; +template +OpKernel* CreateDynamicNativeXlaOpKernel(OpKernelConstruction* ctx) { + if (GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) { + return new NativeOp(ctx); + } + return new MlirXlaOpKernel(ctx); +} + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_ diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 445065971f2a6a..92493928bc2d6f 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -415,6 +415,18 @@ std::vector XlaOpRegistry::DeviceKernels( return ops; } +/* static */ bool XlaOpRegistry::IsMlirXlaOp(absl::string_view op) { + XlaOpRegistry& registry = Instance(); + mutex_lock lock(registry.mutex_); + auto it = registry.ops_.find(string(op)); + if (it == registry.ops_.end() || it->second.empty()) { + return false; + } + return absl::c_all_of(it->second, [](const auto& registration) { + return registration->uses_mlir_kernel; + }); +} + /*static*/ const std::unordered_set* XlaOpRegistry::CompileTimeConstantInputArgNames(const string& op) { XlaOpRegistry& registry = Instance(); @@ -604,8 +616,9 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::Label(std::string label) { } std::unique_ptr XlaOpRegistrationBuilder::Build( - XlaOpRegistry::Factory factory) { + XlaOpRegistry::Factory factory, bool uses_mlir_kernel) { registration_->factory = factory; + registration_->uses_mlir_kernel = uses_mlir_kernel; return std::move(registration_); } diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index 5eaf0fb2d42bfa..b6abc2a39f71f2 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.h +++ b/tensorflow/compiler/tf2xla/xla_op_registry.h @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -44,6 +45,8 @@ limitations under the License. namespace tensorflow { +class MlirXlaOpKernel; + // Names of the XLA compilation devices. These are not user-visible, and are // used internally by the Tensorflow/XLA bridge to perform symbolic execution of // a Tensorflow graph. @@ -233,6 +236,10 @@ class XlaOpRegistry { // Returns all operations for which there are XLA kernels on any device. static std::vector GetAllRegisteredOps(); + // Returns true if the operation is lowered exclusively through + // MlirXlaOpKernel. + static bool IsMlirXlaOp(absl::string_view op); + // Returns (via `result`) the indices of inputs to `node_def` that must be // compile-time constants. Returns an empty vector if the op is not // registered. @@ -339,6 +346,8 @@ class XlaOpRegistry { // operands and not their values. bool is_metadata_op = false; + bool uses_mlir_kernel = false; + std::string label; // Factory used to build OpKernels that perform symbolic execution. @@ -383,6 +392,9 @@ class XlaOpRegistry { #define REGISTER_XLA_OP(NAME, OP) \ REGISTER_XLA_OP_UNIQ_HELPER(__COUNTER__, NAME, OP) +#define REGISTER_XLA_OP_FACTORY(NAME, FACTORY) \ + REGISTER_XLA_OP_FACTORY_UNIQ_HELPER(__COUNTER__, NAME, FACTORY) + #define REGISTER_XLA_CONV_OP(BUILDER, OP) \ REGISTER_XLA_OP(BUILDER.TypeConstraint("T", GetXlaConvTypesForNonGpu()), OP) \ REGISTER_XLA_OP(BUILDER.TypeConstraint("T", GetXlaConvTypesForGpu()) \ @@ -431,7 +443,7 @@ class XlaOpRegistrationBuilder { XlaOpRegistrationBuilder& Label(std::string label); std::unique_ptr Build( - XlaOpRegistry::Factory factory); + XlaOpRegistry::Factory factory, bool uses_mlir_kernel = false); private: XlaOpRegistrationBuilder(absl::string_view name); @@ -454,11 +466,19 @@ class XlaOpRegistrar { #define REGISTER_XLA_OP_UNIQ_HELPER(COUNTER, BUILDER, OP) \ REGISTER_XLA_OP_UNIQ(COUNTER, BUILDER, OP) +#define REGISTER_XLA_OP_FACTORY_UNIQ_HELPER(COUNTER, BUILDER, FACTORY) \ + REGISTER_XLA_OP_FACTORY_UNIQ(COUNTER, BUILDER, FACTORY) + #define REGISTER_XLA_OP_UNIQ(CTR, BUILDER, OP) \ static ::tensorflow::XlaOpRegistrar xla_op_registrar__body__##CTR##__object( \ ::tensorflow::XlaOpRegistrationBuilder::BUILDER.Build( \ [](::tensorflow::OpKernelConstruction* context) \ - -> ::tensorflow::OpKernel* { return new OP(context); })); + -> ::tensorflow::OpKernel* { return new OP(context); }, \ + std::is_same_v)); + +#define REGISTER_XLA_OP_FACTORY_UNIQ(CTR, BUILDER, FACTORY) \ + static ::tensorflow::XlaOpRegistrar xla_op_registrar__body__##CTR##__object( \ + ::tensorflow::XlaOpRegistrationBuilder::BUILDER.Build(FACTORY)); class XlaBackendRegistrar { public: diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc index 8391636e00daac..6e50c950e1d7e4 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc @@ -2673,6 +2673,8 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) { SmallVector dimSizes; SmallVector newSizeBounds; std::vector dimsAreDynamic; + std::vector dimExpressions; + for (auto i = 0; i < resultType.getRank(); ++i) { auto runtimeSizeX1 = xla::Slice(outputShape, {i}, {i + 1}, {1}); dimSizes.push_back(xla::Reshape(runtimeSizeX1, {})); @@ -2683,9 +2685,10 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) { return op->emitOpError() << "unbounded dynamism is not supported"; newSizeBounds.push_back(hlo::isStaticDimSize(dimSize) ? dimSize : dimBound); dimsAreDynamic.push_back(!hlo::isStaticDimSize(dimSize)); + dimExpressions.push_back(xla::DExpr::Unknown(40)); } - value_map[op] = - xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic); + value_map[op] = xla::DynamicReshape(operand, dimSizes, newSizeBounds, + dimsAreDynamic, dimExpressions); return success(); } @@ -2696,7 +2699,8 @@ LogicalResult ExportXlaOp(ReshapeOp op, OpLoweringContext ctx) { return failure(); value_map[op] = - xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions()); + xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions(), + xla::TypeToShape(op.getType()).expressions()); return success(); } @@ -3004,6 +3008,7 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) { SmallVector dimSizes; SmallVector newSizeBounds; std::vector dimsAreDynamic; + std::vector dimExpressions; for (auto i = 0; i < resultType.getRank(); ++i) { auto runtimeSizeX1 = xla::Slice(outputShape, {i}, {i + 1}, {1}); dimSizes.push_back(xla::Reshape(runtimeSizeX1, {})); @@ -3014,9 +3019,11 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) { return op->emitOpError() << "unbounded dynamism is not supported"; newSizeBounds.push_back(hlo::isStaticDimSize(dimSize) ? dimSize : dimBound); dimsAreDynamic.push_back(!hlo::isStaticDimSize(dimSize)); + dimExpressions.push_back(xla::DExpr::Unknown(50)); } value_map[op] = - xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic); + xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic, + dimExpressions); return success(); } @@ -4558,7 +4565,8 @@ LogicalResult ExportXlaOp(ReshapeOp op, OpLoweringContext ctx) { return failure(); value_map[op] = - xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions()); + xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions(), + xla::TypeToShape(op.getType()).expressions()); return success(); } diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc index b057d644df8a9f..a06d517914bdf6 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc @@ -150,6 +150,7 @@ Shape TypeToShape(mlir::Type type) { llvm::SmallVector shape(rank, mlir::ShapedType::kDynamic); std::vector is_dynamic(rank, false); + std::vector expressions(rank, DExpr::Unknown(60)); for (int64_t dim = 0; dim < rank; ++dim) { int64_t size = t.getDimSize(dim); if (size == ShapedType::kDynamic) { @@ -192,7 +193,7 @@ Shape TypeToShape(mlir::Type type) { } return ShapeUtil::MakeShape(primitive_type, shape, is_dynamic, - /*expressions=*/{}); + expressions); } else if (auto tuple_type = mlir::dyn_cast(type)) { llvm::SmallVector shapes; shapes.reserve(tuple_type.size()); From 238cd6d9240dac7789303c5da6ce00d4d3c03e26 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:05:01 +0100 Subject: [PATCH 14/21] Pass dynamic batch dimensions through CPU execution Store the solved dynamic batch value in executable run options and expose it to CPU execution paths. Keep thunk and non-thunk execution consistent at the cluster boundary. --- .../xla/xla/backends/cpu/runtime/kernel.cc | 39 +++++++----- .../xla/xla/backends/cpu/runtime/kernel.h | 35 +++++++++-- .../xla/backends/cpu/runtime/kernel_c_api.h | 1 + .../xla/backends/cpu/runtime/kernel_thunk.cc | 8 ++- .../xla/xla/backends/cpu/runtime/thunk.h | 1 + third_party/xla/xla/executable_run_options.h | 4 +- third_party/xla/xla/service/cpu/BUILD | 9 +++ .../xla/xla/service/cpu/cpu_executable.cc | 60 +++++++++++++++---- .../xla/xla/service/cpu/cpu_runtime_test.cc | 9 +++ .../xla/xla/service/cpu/runtime_batch_size.cc | 32 ++++++++++ .../xla/xla/service/cpu/runtime_batch_size.h | 30 ++++++++++ .../service/cpu/runtime_symbol_generator.cc | 2 + .../xla/xla/stream_executor/tpu/c_api_decl.h | 2 + 13 files changed, 197 insertions(+), 35 deletions(-) create mode 100644 third_party/xla/xla/service/cpu/runtime_batch_size.cc create mode 100644 third_party/xla/xla/service/cpu/runtime_batch_size.h diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel.cc b/third_party/xla/xla/backends/cpu/runtime/kernel.cc index ac1a5d7181ec52..668e10f0a0dd09 100644 --- a/third_party/xla/xla/backends/cpu/runtime/kernel.cc +++ b/third_party/xla/xla/backends/cpu/runtime/kernel.cc @@ -64,7 +64,7 @@ template class Kernel::ParallelTask { public: ParallelTask(XLA_CPU_Kernel* kernel, Kernel::ThreadDim thread_dims, - absl::Span args); + size_t batch_size, absl::Span args); // Invokes a host kernel for a given task index. absl::Status operator()(size_t task_index) const; @@ -76,6 +76,7 @@ class Kernel::ParallelTask { XLA_CPU_Kernel* kernel_; XLA_CPU_KernelThreadDim thread_dims_; + size_t batch_size_; absl::InlinedVector args_; size_t num_tasks_; @@ -88,11 +89,13 @@ class Kernel::ParallelTask { template Kernel::ParallelTask::ParallelTask( XLA_CPU_Kernel* kernel, Kernel::ThreadDim thread_dims, + size_t batch_size, absl::Span args) : kernel_(kernel), thread_dims_({thread_dims.x, thread_dims.y, thread_dims.z}), args_(args.begin(), args.end()), num_tasks_(thread_dims_.x * thread_dims_.y * thread_dims_.z), + batch_size_(batch_size), stride_z_(thread_dims_.y * thread_dims_.x), stride_y_(thread_dims_.x) {} @@ -103,7 +106,7 @@ absl::Status Kernel::ParallelTask::operator()( XLA_CPU_KernelThread kernel_thread = Delinearize(task_index); XLA_CPU_KernelCallFrame call_frame = {&thread_dims_, &kernel_thread, - args_.size(), args_.data()}; + args_.size(), args_.data(), batch_size_}; XLA_CPU_KernelError* error = (*kernel_)(&call_frame); @@ -138,12 +141,12 @@ Kernel::Kernel(unsigned arity, XLA_CPU_Kernel* kernel) kernel_(function_->kernel()), arity_(arity) {} -absl::Status Kernel::Launch(const ThreadDim& thread_dims, +absl::Status Kernel::Launch(const ThreadDim& thread_dims, size_t batch_size, absl::Span buffers) const { - return Launch(thread_dims, ConvertBuffersToKernelArgs(buffers)); + return Launch(thread_dims, batch_size, ConvertBuffersToKernelArgs(buffers)); } -absl::Status Kernel::Launch(const ThreadDim& thread_dims, +absl::Status Kernel::Launch(const ThreadDim& thread_dims, size_t batch_size, absl::Span args) const { XLA_CPU_KernelThreadDim kernel_thread_dims = { thread_dims.x, @@ -156,8 +159,9 @@ absl::Status Kernel::Launch(const ThreadDim& thread_dims, for (uint64_t x = 0; x < thread_dims.x; ++x) { XLA_CPU_KernelThread kernel_thread = {x, y, z}; - XLA_CPU_KernelCallFrame call_frame = { - &kernel_thread_dims, &kernel_thread, args.size(), args.data()}; + XLA_CPU_KernelCallFrame call_frame = {&kernel_thread_dims, + &kernel_thread, args.size(), + args.data(), batch_size}; XLA_CPU_KernelError* error = (*kernel_)(&call_frame); @@ -172,20 +176,23 @@ absl::Status Kernel::Launch(const ThreadDim& thread_dims, } tsl::AsyncValueRef Kernel::Launch( - const ThreadDim& thread_dims, absl::Span buffers, + const ThreadDim& thread_dims, size_t batch_size, + absl::Span buffers, const Eigen::ThreadPoolDevice* device) const { - return Launch(thread_dims, ConvertBuffersToKernelArgs(buffers), device); + return Launch(thread_dims, batch_size, ConvertBuffersToKernelArgs(buffers), + device); } tsl::AsyncValueRef Kernel::Launch( - const ThreadDim& thread_dims, absl::Span args, + const ThreadDim& thread_dims, size_t batch_size, + absl::Span args, const Eigen::ThreadPoolDevice* device) const { size_t num_tasks = thread_dims.x * thread_dims.y * thread_dims.z; CHECK_GT(num_tasks, 0) << "Number of tasks must be positive"; // Crash Ok // Short-circuit launch with a single task and run it in the caller thread. if (ABSL_PREDICT_TRUE(num_tasks == 1)) { - absl::Status launched = Launch(thread_dims, args); + absl::Status launched = Launch(thread_dims, batch_size, args); return ABSL_PREDICT_TRUE(launched.ok()) ? OkLaunchEvent() : tsl::MakeErrorAsyncValueRef(std::move(launched)); @@ -197,11 +204,13 @@ tsl::AsyncValueRef Kernel::Launch( std::numeric_limits::max()); if (ABSL_PREDICT_TRUE(thread_dims.y == 1 && thread_dims.z == 1)) { - return Worker::Parallelize(device, num_workers, num_tasks, - ParallelTask(kernel_, thread_dims, args)); + return Worker::Parallelize( + device, num_workers, num_tasks, + ParallelTask(kernel_, thread_dims, batch_size, args)); } else { - return Worker::Parallelize(device, num_workers, num_tasks, - ParallelTask(kernel_, thread_dims, args)); + return Worker::Parallelize( + device, num_workers, num_tasks, + ParallelTask(kernel_, thread_dims, batch_size, args)); } } diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel.h b/third_party/xla/xla/backends/cpu/runtime/kernel.h index 98010905322f67..d03657356bfc7e 100644 --- a/third_party/xla/xla/backends/cpu/runtime/kernel.h +++ b/third_party/xla/xla/backends/cpu/runtime/kernel.h @@ -73,13 +73,25 @@ class Kernel { // Calls the kernel once in the caller thread for a thread dim (0,0,0). // This is a fast path for small host kernels that have just one thread. - absl::Status CallOnce(absl::Span args) const; + absl::Status CallOnce(absl::Span args) const { + return CallOnce(args, 0); + } + absl::Status CallOnce(absl::Span args, + size_t batch_size) const; // Launches the kernel on the current thread by iterating over all threads in // `thread_dims` and calling the kernel function. absl::Status Launch(const ThreadDim& thread_dims, - absl::Span buffers) const; + absl::Span buffers) const { + return Launch(thread_dims, 0, buffers); + } absl::Status Launch(const ThreadDim& thread_dims, + absl::Span args) const { + return Launch(thread_dims, 0, args); + } + absl::Status Launch(const ThreadDim& thread_dims, size_t batch_size, + absl::Span buffers) const; + absl::Status Launch(const ThreadDim& thread_dims, size_t batch_size, absl::Span args) const; // Launches the kernel by iterating over all threads in `thread_dims` and @@ -90,9 +102,21 @@ class Kernel { // get the number of tasks that are expected to be completed. tsl::AsyncValueRef Launch( const ThreadDim& thread_dims, absl::Span buffers, - const Eigen::ThreadPoolDevice* device) const; + const Eigen::ThreadPoolDevice* device) const { + return Launch(thread_dims, 0, buffers, device); + } tsl::AsyncValueRef Launch( const ThreadDim& thread_dims, absl::Span args, + const Eigen::ThreadPoolDevice* device) const { + return Launch(thread_dims, 0, args, device); + } + tsl::AsyncValueRef Launch( + const ThreadDim& thread_dims, size_t batch_size, + absl::Span buffers, + const Eigen::ThreadPoolDevice* device) const; + tsl::AsyncValueRef Launch( + const ThreadDim& thread_dims, size_t batch_size, + absl::Span args, const Eigen::ThreadPoolDevice* device) const; // For host platform, we assume that a core is a thread, and we can run at @@ -123,12 +147,13 @@ class Kernel { }; inline ABSL_ATTRIBUTE_ALWAYS_INLINE absl::Status Kernel::CallOnce( - absl::Span args) const { + absl::Span args, + size_t batch_size) const { constexpr XLA_CPU_KernelThreadDim kernel_thread_dims = {1, 1, 1}; constexpr XLA_CPU_KernelThread kernel_thread = {1, 1, 1}; XLA_CPU_KernelCallFrame call_frame = {&kernel_thread_dims, &kernel_thread, - args.size(), args.data()}; + args.size(), args.data(), batch_size}; XLA_CPU_KernelError* error = (*kernel_)(&call_frame); diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h b/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h index cbe0568506385d..cf8016759ed217 100644 --- a/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h +++ b/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h @@ -72,6 +72,7 @@ typedef struct XLA_CPU_KernelCallFrame { size_t num_args; const XLA_CPU_KernelArg* args; + size_t batch_size; } XLA_CPU_KernelCallFrame; // Error reporting for host kernels. NULL means success. diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc index 90c15a09bd677d..7a1c0ec70c449b 100644 --- a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc @@ -223,7 +223,7 @@ KernelThunk::ExecuteInternal( // Use a fast path if kernel called just once. if (ABSL_PREDICT_TRUE(call_once_)) { - TF_RETURN_IF_ERROR(kernel->CallOnce(kernel_args)); + TF_RETURN_IF_ERROR(kernel->CallOnce(kernel_args, params.batch_size)); return OkExecuteEvent(); } @@ -231,10 +231,12 @@ KernelThunk::ExecuteInternal( // by scheduling tasks into it. HostKernel launch completion will // automatically signal KernelThunk execute completion. if (ABSL_PREDICT_TRUE(params.intra_op_threadpool)) { - return kernel->Launch(thread_dim_, kernel_args, params.intra_op_threadpool); + return kernel->Launch(thread_dim_, params.batch_size, kernel_args, + params.intra_op_threadpool); } - TF_RETURN_IF_ERROR(kernel->Launch(thread_dim_, kernel_args)); + TF_RETURN_IF_ERROR( + kernel->Launch(thread_dim_, params.batch_size, kernel_args)); return OkExecuteEvent(); } diff --git a/third_party/xla/xla/backends/cpu/runtime/thunk.h b/third_party/xla/xla/backends/cpu/runtime/thunk.h index d4a56ae88fa55b..5b7042ccb2803f 100644 --- a/third_party/xla/xla/backends/cpu/runtime/thunk.h +++ b/third_party/xla/xla/backends/cpu/runtime/thunk.h @@ -259,6 +259,7 @@ class Thunk { TaskRunner* task_runner = nullptr; CollectiveExecuteParams* collective_params = nullptr; CustomCallExecuteParams* custom_call_params = nullptr; + int64_t batch_size = 0; ExecuteSession session = ExecuteSession(ExecuteSession::kMaxWorkers, ExecuteSession::kSplitThreshold); }; diff --git a/third_party/xla/xla/executable_run_options.h b/third_party/xla/xla/executable_run_options.h index 211932ffcbba98..38b13e5ec0726a 100644 --- a/third_party/xla/xla/executable_run_options.h +++ b/third_party/xla/xla/executable_run_options.h @@ -195,6 +195,8 @@ class ExecutableRunOptions { return *this; } + int32_t launch_id() const { return launch_id_; } + ExecutableRunOptions& set_batch_size(int64_t batch_size) { batch_size_ = batch_size; return *this; @@ -202,8 +204,6 @@ class ExecutableRunOptions { int64_t batch_size() const { return batch_size_; } - int32_t launch_id() const { return launch_id_; } - ExecutableRunOptions& set_run_id(RunId id); RunId run_id() const; diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 24247e60993937..eef57b0e88409e 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -597,6 +597,7 @@ cc_library( ":onednn_layer_norm", ":onednn_matmul", ":onednn_softmax", + ":runtime_batch_size_api", ":runtime_conv2d", ":runtime_conv2d_acl", ":runtime_conv2d_mkl", @@ -1094,6 +1095,7 @@ cc_library( name = "cpu_runtime", srcs = [ "cpu_runtime.cc", + "runtime_batch_size.cc", "xfeed_manager.cc", ], hdrs = [ @@ -1103,6 +1105,7 @@ cc_library( copts = runtime_copts(), deps = [ ":cpu_executable_run_options", + ":runtime_batch_size_api", "//xla:executable_run_options", "//xla:shape_util", "//xla:types", @@ -1156,6 +1159,11 @@ cc_library( ], ) +cc_library( + name = "runtime_batch_size_api", + hdrs = ["runtime_batch_size.h"], +) + cc_library( name = "runtime_conv3d", srcs = ["runtime_conv3d.cc"], @@ -1461,6 +1469,7 @@ xla_cc_test( tags = ["optonly"], deps = [ ":cpu_runtime", + ":runtime_batch_size_api", ":runtime_custom_call_status", ":runtime_matmul", ":runtime_matmul_acl", diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc index 158a4000d27722..853aa3519924c1 100644 --- a/third_party/xla/xla/service/cpu/cpu_executable.cc +++ b/third_party/xla/xla/service/cpu/cpu_executable.cc @@ -163,14 +163,34 @@ static absl::StatusOr MemoryForAllocation( se::DeviceMemoryAllocator* memory_allocator, int device_ordinal) { VLOG(3) << allocation.ToString(); if (allocation.is_entry_computation_parameter()) { - se::DeviceMemoryBase out = arguments[allocation.parameter_number()] + se::DeviceMemoryBase param_mem = arguments[allocation.parameter_number()] .Buffer(allocation.param_shape_index()) .AsDeviceMemoryBase(); - CHECK_LE(allocation.size(), out.size()) - << "Size mismatch on param " << allocation.parameter_number() - << " at shape index " << allocation.param_shape_index().ToString(); - VLOG(3) << "allocation is a parameter"; - return MaybeOwningDeviceMemory{out}; + + const int64_t compiled_bytes = allocation.size(); + const int64_t runtime_bytes = param_mem.size(); + + if(runtime_bytes == 0){ + return MaybeOwningDeviceMemory{param_mem}; + } + + if (compiled_bytes <= runtime_bytes) { + return MaybeOwningDeviceMemory{param_mem}; + } + + //padded owns that device memory + TF_ASSIGN_OR_RETURN( + se::OwningDeviceMemory padded, + memory_allocator->Allocate(device_ordinal, compiled_bytes)); + + void* dst = padded->opaque(); + void* src = param_mem.opaque(); + + std::memcpy(dst, src, runtime_bytes); + //Fill rest of them with zeros. + std::memset(static_cast(dst) + runtime_bytes, 0, compiled_bytes - runtime_bytes); + return MaybeOwningDeviceMemory{std::move(padded)}; + } else if (allocation.is_constant()) { VLOG(3) << "allocation is a constant"; if (allocation.index() < constants.size()) { @@ -275,11 +295,28 @@ absl::Status CpuExecutable::ExecuteComputeFunction( return absl::OkStatus(); } +void PrintScalars(absl::Span buffers) { + for (int i = 0; i < buffers.size(); ++i) { + const se::DeviceMemoryBase& dmem = buffers[i].AsDeviceMemoryBase(); + if (dmem.opaque() && dmem.size() >= sizeof(int64_t)) { + int64_t val = 0; + std::memcpy(&val, dmem.opaque(), sizeof(val)); + std::cerr << "Buffer " << i << " scalar: " << val << std::endl; + } else { + std::cerr << "Buffer " << i << " empty or too small" << std::endl; + } + } +} + absl::Status CpuExecutable::ExecuteThunks( const ExecutableRunOptions* run_options, absl::Span buffers) { uint64_t start_ns = tsl::Env::Default()->NowNanos(); + #if defined(PRINT_BATCHSIZE) + PrintScalars(buffers); + #endif + size_t profile_counters_size = 0; int64_t* profile_counters = nullptr; @@ -318,7 +355,8 @@ absl::Status CpuExecutable::ExecuteThunks( intra_op_thread_pool, &task_runner, &collective_execute_params, - &custom_call_execute_params}; + &custom_call_execute_params, + run_options->batch_size()}; auto executed_event = thunks_->Execute(execute_params); tsl::BlockUntilReady(executed_event); @@ -339,13 +377,15 @@ absl::StatusOr CpuExecutable::CreateResultShapedBuffer( absl::Span buffers, absl::Span arguments) { se::Stream* stream = run_options->stream(); - ExecutionOutput result(/*on_device_shape=*/result_shape(), - run_options->allocator(), - stream->parent()->device_ordinal()); const HloInputOutputAliasConfig& input_output_alias = module().input_output_alias_config(); HloInstruction* root = hlo_module_->entry_computation()->root_instruction(); const Shape& root_shape = root->shape(); + // Use root_shape to initialize ExecutionOuput as the batch multiplier info + // is only attached the ROOT + ExecutionOutput result(/*on_device_shape=*/root_shape, //result_shape(), + run_options->allocator(), + stream->parent()->device_ordinal()); // Move se::OwningDeviceMemory values which contain the array(s) of the result // into the respective location in ScopedShapedBuffer which is returned to the diff --git a/third_party/xla/xla/service/cpu/cpu_runtime_test.cc b/third_party/xla/xla/service/cpu/cpu_runtime_test.cc index 79df0e29e8ea91..e14835602224a4 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_runtime_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "xla/array2d.h" #include "xla/client/local_client.h" #include "xla/executable_run_options.h" +#include "xla/service/cpu/runtime_batch_size.h" #include "xla/service/cpu/runtime_custom_call_status.h" #include "xla/service/cpu/runtime_matmul.h" #include "xla/service/cpu/runtime_matmul_acl.h" @@ -41,6 +42,14 @@ namespace { class CpuRuntimeTest : public ::testing::Test {}; +TEST_F(CpuRuntimeTest, GetBatchSize) { + ExecutableRunOptions run_options; + run_options.set_batch_size(42); + + EXPECT_EQ(__xla_cpu_runtime_GetBatchSize(&run_options), 42); + EXPECT_EQ(__xla_cpu_runtime_GetBatchSize(nullptr), 0); +} + template std::unique_ptr> MaybeTransposeArray2D(const Array2D& array, bool transpose) { diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.cc b/third_party/xla/xla/service/cpu/runtime_batch_size.cc new file mode 100644 index 00000000000000..951910f9aee51c --- /dev/null +++ b/third_party/xla/xla/service/cpu/runtime_batch_size.cc @@ -0,0 +1,32 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 "xla/service/cpu/runtime_batch_size.h" + +#include "xla/executable_run_options.h" + +namespace xla::cpu::runtime { + +const char* const kGetBatchSizeSymbolName = + "__xla_cpu_runtime_GetBatchSize"; + +} // namespace xla::cpu::runtime + +extern "C" int64_t __xla_cpu_runtime_GetBatchSize( + const void* run_options_ptr) { + if (run_options_ptr == nullptr) return 0; + return static_cast(run_options_ptr) + ->batch_size(); +} diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.h b/third_party/xla/xla/service/cpu/runtime_batch_size.h new file mode 100644 index 00000000000000..9204e7841c4e70 --- /dev/null +++ b/third_party/xla/xla/service/cpu/runtime_batch_size.h @@ -0,0 +1,30 @@ +/* Copyright 2026 The OpenXLA Authors. + +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. +==============================================================================*/ + +#ifndef XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_ +#define XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_ + +#include + +namespace xla::cpu::runtime { + +extern const char* const kGetBatchSizeSymbolName; + +} // namespace xla::cpu::runtime + +extern "C" int64_t __xla_cpu_runtime_GetBatchSize( + const void* run_options_ptr); + +#endif // XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_ diff --git a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc index 87aca6c386751a..a87b9b3ae986ab 100644 --- a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc +++ b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc @@ -37,6 +37,7 @@ limitations under the License. #include "llvm/Support/Error.h" #include "mlir/ExecutionEngine/CRunnerUtils.h" #include "xla/service/cpu/cpu_runtime.h" +#include "xla/service/cpu/runtime_batch_size.h" #include "xla/service/cpu/runtime_conv2d.h" #include "xla/service/cpu/runtime_conv2d_acl.h" #include "xla/service/cpu/runtime_conv2d_mkl.h" @@ -160,6 +161,7 @@ static bool RegisterKnownJITSymbols() { registry->Register("puts", reinterpret_cast(&puts), "Host"); REGISTER_CPU_RUNTIME_SYMBOL(AcquireInfeedBufferForDequeue); + REGISTER_CPU_RUNTIME_SYMBOL(GetBatchSize); REGISTER_CPU_RUNTIME_SYMBOL(AcquireOutfeedBufferForPopulation); REGISTER_CPU_RUNTIME_SYMBOL(AllReduce); REGISTER_CPU_RUNTIME_SYMBOL(CollectivePermute); diff --git a/third_party/xla/xla/stream_executor/tpu/c_api_decl.h b/third_party/xla/xla/stream_executor/tpu/c_api_decl.h index 0f3ba256573de1..e06db894249632 100644 --- a/third_party/xla/xla/stream_executor/tpu/c_api_decl.h +++ b/third_party/xla/xla/stream_executor/tpu/c_api_decl.h @@ -247,6 +247,8 @@ typedef struct XLA_Shape { int element_type; Int64List dimensions; BoolList dynamic_dimensions; + Int64List batch_multipliers; + Int64List batch_offsets; struct XLA_Shape* tuple_shapes; // owned int ntuple_shapes; bool has_layout; From 6753c39b5c4ef7c56c589cbe347798f0b1ff8ce3 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:05:16 +0100 Subject: [PATCH 15/21] Generate runtime-dependent CPU loop bounds Generate LLVM loop bounds from the runtime batch dimension when an HLO dimension carries a dynamic expression. Add focused coverage for runtime-dependent bounds. --- .../cpu/codegen/kernel_api_ir_builder.cc | 35 +++- .../cpu/codegen/kernel_api_ir_builder.h | 4 + third_party/xla/xla/service/cpu/ir_emitter.cc | 124 ++++++++++--- third_party/xla/xla/service/cpu/ir_emitter.h | 7 +- .../xla/xla/service/cpu/ir_emitter2.cc | 26 +++ third_party/xla/xla/service/cpu/ir_emitter2.h | 3 + .../xla/service/cpu/parallel_loop_emitter.cc | 8 +- .../xla/xla/service/cpu/runtime_batch_size.cc | 7 - .../xla/xla/service/cpu/runtime_batch_size.h | 3 +- .../xla/xla/service/cpu/thunk_emitter.cc | 17 ++ .../xla/xla/service/cpu/thunk_emitter.h | 3 + .../xla/xla/service/elemental_ir_emitter.cc | 101 +++++++--- third_party/xla/xla/service/llvm_ir/BUILD | 18 +- .../xla/xla/service/llvm_ir/ir_array.cc | 52 +++++- .../xla/xla/service/llvm_ir/ir_array.h | 1 + .../xla/xla/service/llvm_ir/llvm_loop.cc | 82 ++++++--- .../xla/xla/service/llvm_ir/llvm_loop.h | 10 +- .../xla/xla/service/llvm_ir/llvm_util.cc | 173 +++++++++++++++++- .../xla/xla/service/llvm_ir/llvm_util.h | 13 ++ .../xla/xla/service/llvm_ir/llvm_util_test.cc | 83 +++++++++ .../xla/xla/service/llvm_ir/loop_emitter.cc | 29 +++ .../xla/xla/service/llvm_ir/tuple_ops.cc | 2 +- 22 files changed, 693 insertions(+), 108 deletions(-) create mode 100644 third_party/xla/xla/service/llvm_ir/llvm_util_test.cc diff --git a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc index 2ce99c4b93a878..ffc2ed3f8baf73 100644 --- a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc +++ b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc @@ -147,7 +147,7 @@ llvm::StructType* KernelCallFrameTy(llvm::LLVMContext& ctx) { llvm::PointerType* ptr = llvm::PointerType::getUnqual(ctx); llvm::IntegerType* i64 = llvm::IntegerType::getInt64Ty(ctx); return llvm::StructType::create("XLA_CPU_KernelCallFrame", ptr, ptr, i64, - ptr); + ptr, i64); } llvm::FunctionType* KernelFunctionTy(llvm::LLVMContext& ctx) { @@ -316,6 +316,34 @@ auto KernelApiIrBuilder::EmitKernelPrototype( return EmitKernelPrototype(module, name, arguments, results); } +llvm::Value* KernelApiIrBuilder::EmitGetBatchDim(llvm::IRBuilderBase& builder, + llvm::Value* call_frame) { + llvm::LLVMContext& ctx = builder.getContext(); + llvm::Type* ptr = llvm::PointerType::get(ctx, 0); + llvm::IntegerType* i64 = llvm::IntegerType::getInt64Ty(ctx); + llvm::Value* bdim_gep = + builder.CreateStructGEP(call_frame_ty_, call_frame, 4, "bdim_gep"); + llvm::Value* bdim_value = builder.CreateLoad(i64, bdim_gep, "bdim_value"); + +#if defined(PRINT_BATCHSIZE) + // Print batch size + llvm::Function* function = builder.GetInsertBlock()->getParent(); + llvm::Module* module = function->getParent(); + llvm::FunctionType* printfType = llvm::FunctionType::get( + builder.getInt32Ty(), llvm::PointerType::get(builder.getInt8Ty(), 0), + true); + llvm::Value* funcNameStr = + builder.CreateGlobalStringPtr(function->getName()); + llvm::FunctionCallee printfFunc = + module->getOrInsertFunction("printf", printfType); + llvm::Value* formatStr = + builder.CreateGlobalStringPtr("Function: %s, Batch size is : %lld!\n"); + builder.CreateCall(printfFunc, {formatStr, funcNameStr, bdim_value}); +#endif + + return bdim_value; +} + auto KernelApiIrBuilder::EmitKernelPrototype( llvm::Module& module, absl::string_view name, absl::Span arguments, @@ -394,6 +422,8 @@ auto KernelApiIrBuilder::EmitKernelPrototype( ir_results.push_back(std::move(ir_result)); } + EmitGetBatchDim(b, call_frame); + // Return null pointer to signal success as we do not support error handling // in the compiled host kernel. llvm::BasicBlock* return_block = @@ -503,7 +533,8 @@ llvm_ir::IrArray KernelApiIrBuilder::EmitKernelArgument( const llvm::DataLayout& data_layout = llvm_module->getDataLayout(); int64_t pointer_size = data_layout.getTypeStoreSize(builder.getPtrTy()); int64_t byte_size = ShapeUtil::ByteSizeOf(shape, pointer_size); - llvm_ir::SetDereferenceableMetadataForLoad(data, byte_size); + if (!shape.has_dynamic_expr()) + llvm_ir::SetDereferenceableMetadataForLoad(data,byte_size); // All buffers pointers passed to host kernels are expected to be invariant // over the whole program. Note the metadata is attached only to loading diff --git a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h index 58ffc5bc8594e3..08af675fcb6c37 100644 --- a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h +++ b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h @@ -147,9 +147,13 @@ class KernelApiIrBuilder { llvm_ir::IrArray EmitKernelArgument(llvm::IRBuilderBase& builder, llvm::Value* call_frame, int64_t index, const Shape& shape); + llvm::Function* EmitKernelFunction(llvm::Module& module, absl::string_view name); + llvm::Value* EmitGetBatchDim(llvm::IRBuilderBase& builder, + llvm::Value* call_frame); + private: llvm::LLVMContext& context_; diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc index feca6552d243f8..ae40e76e8e2ce2 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter.cc +++ b/third_party/xla/xla/service/cpu/ir_emitter.cc @@ -38,6 +38,7 @@ limitations under the License. #include "absl/container/inlined_vector.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "xla/shape_expr.h" #include "absl/meta/type_traits.h" #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -76,6 +77,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/map_util.h" +#include "xla/printer.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" @@ -2055,9 +2057,11 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) { } llvm_ir::IrArray source_array = GetIrArrayFor(operand); + const auto* slice_hlo = Cast(slice); const llvm_ir::IrArray::Index source_index = target_index.SourceIndexOfSlice( - /*operand_shape=*/operand->shape(), /*starts=*/slice->slice_starts(), - /*strides=*/slice->slice_strides(), /*builder=*/b()); + /*operand_shape=*/operand->shape(), /*starts=*/slice_hlo->slice_starts(), + /*start_exprs=*/slice_hlo->slice_start_exprs(), + /*strides=*/slice_hlo->slice_strides(), /*builder=*/b()); llvm::Value* memcpy_dest = target_array.EmitArrayElementAddress(target_index, b(), "slice.dest"); @@ -2067,7 +2071,7 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) { const int64_t memcpy_elements = primitive_elements_per_logical_element * memcpy_logical_elements; - EmitTransferElements(memcpy_dest, memcpy_source, memcpy_elements, + EmitTransferElements(memcpy_dest, memcpy_source, DExpr::Const(memcpy_elements), slice->shape().element_type(), target_array, source_array); @@ -2358,6 +2362,26 @@ absl::Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) { return EmitSliceToDynamic(hlo, source_arrays, target_array); } +absl::Status IrEmitter::HandleShapeExprValue(HloInstruction* hlo) { + TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo)); + + llvm_ir::IrArray out_array = GetIrArrayFor(hlo); + TF_RET_CHECK(hlo->has_contents()); + TF_RET_CHECK(!hlo->contents().empty()); + xla::DExpr expr = xla::DExprFromProto(hlo->contents()[0]); + TF_RET_CHECK(expr); + llvm::Value* expr_value = llvm_ir::EmitExpression(b(), expr); + + auto it = emitted_value_.find(hlo); + if (it == emitted_value_.end()) { + LOG(ERROR) << "No buffer assigned for instruction " << hlo->name(); + } + llvm::Value* dest_ptr = it->second; + b()->CreateStore(expr_value, dest_ptr); + + return absl::OkStatus(); +} + absl::Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) { TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo)); @@ -2806,6 +2830,10 @@ absl::Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) { #endif // INTEL_MKL absl::Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { + if (custom_call->custom_call_target() == "GetExpressionValue") { + return HandleShapeExprValue(custom_call); + } + if (custom_call->custom_call_target() == "PadToStatic") { return HandlePadToStatic(custom_call); } @@ -3125,11 +3153,11 @@ absl::Status EmitFastConcatenate( // contiguous subregion in the target buffer starting at target_region_begin. llvm::Value* target_region_begin = target_array.EmitArrayElementAddress(target_index, &b, "target_region"); - int64_t byte_offset_into_target_region = 0; + llvm::Value* byte_offset_into_target_region = b.getInt64(0); - int64_t inner_dims_product = absl::c_accumulate( - inner_dims, int64_t{1}, [&](int64_t product, int64_t inner_dim) { - return product * output_shape.dimensions(inner_dim); + DExpr inner_exprs_product = absl::c_accumulate( + inner_dims, DExpr::Const(1), [&](DExpr product, int64_t inner_dim) { + return product * output_shape.expressions(inner_dim); }); // For each operand, emit a memcpy from the operand to the target of size @@ -3142,18 +3170,24 @@ absl::Status EmitFastConcatenate( llvm::Value* copy_source_address = source_array.EmitArrayElementAddress(source_index, &b, "src_addr"); - llvm::Value* copy_target_address = - b.CreateGEP(b.getInt8Ty(), target_region_begin, - b.getInt64(byte_offset_into_target_region)); + llvm::Value* copy_target_address = b.CreateGEP( + b.getInt8Ty(), target_region_begin, byte_offset_into_target_region); + + auto cexpr = input_shape.expressions(concat_dim); ::xla::cpu::EmitTransferElements( copy_target_address, copy_source_address, - inner_dims_product * input_shape.dimensions(concat_dim), primitive_type, - target_array, source_array, module, b); + inner_exprs_product * cexpr, primitive_type, target_array, + source_array, module, b); + + llvm::Value* concat_dim_count = xla::llvm_ir::EmitExpression( + &b, inner_exprs_product * input_shape.expressions(concat_dim)); - byte_offset_into_target_region += inner_dims_product * - input_shape.dimensions(concat_dim) * - primitive_type_size; + llvm::Value* concat_dim_size = + b.CreateMul(concat_dim_count, b.getInt64(primitive_type_size)); + byte_offset_into_target_region = + b.CreateAdd(byte_offset_into_target_region, concat_dim_size, + "byte_offset_into_target_region"); } if (!outer_dims.empty()) { @@ -3364,7 +3398,7 @@ llvm::Value* IrEmitter::EmitCallToFfi(HloCustomCallInstruction* custom_call, } void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source, - int64_t element_count, + const xla::DExpr& element_count, PrimitiveType primitive_type, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& source_array) { @@ -3374,7 +3408,8 @@ void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source, } void EmitTransferElements(llvm::Value* target, llvm::Value* source, - int64_t element_count, PrimitiveType primitive_type, + const xla::DExpr& element_count, + PrimitiveType primitive_type, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& source_array, llvm::Module* module, llvm::IRBuilderBase& b) { @@ -3386,7 +3421,7 @@ void EmitTransferElements(llvm::Value* target, llvm::Value* source, llvm::Type* primitive_llvm_type = llvm_ir::PrimitiveTypeToIrType(primitive_type, module->getContext()); - if (element_count == 1) { + if (element_count->is_constant() && element_count->get_val() == 1) { auto* load_instruction = b.CreateAlignedLoad(primitive_llvm_type, source, element_alignment); source_array.AnnotateLoadStoreInstructionWithMetadata(load_instruction); @@ -3394,11 +3429,12 @@ void EmitTransferElements(llvm::Value* target, llvm::Value* source, b.CreateAlignedStore(load_instruction, target, element_alignment); target_array.AnnotateLoadStoreInstructionWithMetadata(store_instruction); } else { + auto element_count_value = xla::llvm_ir::EmitExpression(&b, element_count); + llvm::Value* elements_size = + b.CreateMul(element_count_value, b.getInt64(primitive_type_size)); auto* memcpy_instruction = b.CreateMemCpy( target, /*DstAlign=*/llvm::Align(element_alignment), source, - /*SrcAlign=*/llvm::Align(element_alignment), - element_count * primitive_type_size); - + /*SrcAlign=*/llvm::Align(element_alignment), elements_size); // The memcpy does the load and the store internally. The aliasing related // metadata has to reflect that. std::map merged_metadata = @@ -3911,8 +3947,10 @@ llvm::Value* IrEmitter::EmitThreadLocalBufferPointer( if (!target_shape.IsOpaque()) { AttachAlignmentMetadataForLoad(param_address_untyped, target_shape); - AttachDereferenceableMetadataForLoad(param_address_untyped, - target_shape); + if (!target_shape.has_dynamic_expr()) { + AttachDereferenceableMetadataForLoad(param_address_untyped, + target_shape); + } } return param_address_untyped; } @@ -3958,7 +3996,10 @@ llvm::Value* IrEmitter::EmitGlobalBufferPointer( AttachInvariantLoadMetadataForLoad(tempbuf_address_base); AttachAlignmentMetadataForLoad(tempbuf_address_base, allocation.size()); - AttachDereferenceableMetadataForLoad(tempbuf_address_base, allocation.size()); + + if (!target_shape.has_dynamic_expr()) + AttachDereferenceableMetadataForLoad(tempbuf_address_base, + allocation.size()); llvm::Value* tempbuf_address_untyped = tempbuf_address_base; // Any explicit buffer pointer should point to the start of the slice. @@ -4059,9 +4100,38 @@ absl::Status IrEmitter::EmitMemcpy(const HloInstruction& source, llvm::Value* source_value = GetEmittedValueFor(&source); llvm::Value* destination_value = GetEmittedValueFor(&destination); int64_t source_size = ByteSizeOf(source.shape()); - // TODO(b/63762267): Be more aggressive about specifying alignment. - MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value, - /*SrcAlign=*/llvm::Align(1), source_size); + auto shape = source.shape(); + auto expressions = shape.expressions(); + bool is_dynamic = + std::any_of(expressions.begin(), expressions.end(), + [](const DExpr& e) { return e && e->is_dynamic(); }); + if (is_dynamic) { + llvm::LLVMContext& ctx = b()->getContext(); + llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx); + int64_t dimensions_accu = 1; + DExpr expression_accu = DExpr::Const(1); + for (int i = 0; i < shape.dimensions_size(); i++) { + const auto& expression = shape.expressions(i); + if (expression && expression->is_dynamic()) { + dimensions_accu *= shape.dimensions(i); + expression_accu = expression_accu * expression; + } + } + llvm::Value* expr_value = + xla::llvm_ir::EmitExpression(b(), expression_accu); + // Divide the size in bytes by the size of the dynamic dimension(s). + // TODO: make that less hacky + llvm::ConstantInt* size = + llvm::ConstantInt::get(i64Type, source_size / dimensions_accu, true); + llvm::Value* memcopy_size = + b()->CreateMul(expr_value, size, "memcopy_size"); + MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value, + /*SrcAlign=*/llvm::Align(1), memcopy_size); + } else { + // TODO(b/63762267): Be more aggressive about specifying alignment. + MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value, + /*SrcAlign=*/llvm::Align(1), source_size); + } return absl::OkStatus(); } diff --git a/third_party/xla/xla/service/cpu/ir_emitter.h b/third_party/xla/xla/service/cpu/ir_emitter.h index 40f54d2f4bff97..53583810acd5cb 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter.h +++ b/third_party/xla/xla/service/cpu/ir_emitter.h @@ -332,6 +332,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, private: absl::Status HandleSliceToDynamic(HloInstruction* hlo); + absl::Status HandleShapeExprValue(HloInstruction* hlo); absl::Status HandlePadToStatic(HloInstruction* hlo); absl::Status HandleTopK(HloInstruction* hlo) override; absl::Status HandleAllReduceSingleReplica(HloInstruction* crs); @@ -569,7 +570,8 @@ class IrEmitter : public DfsHloVisitorWithDefault, // Emits LLVM IR to transfer "element_count" elements of type "primitive_type" // from the address "source" to the address "target". void EmitTransferElements(llvm::Value* target, llvm::Value* source, - int64_t element_count, PrimitiveType primitive_type, + const xla::DExpr& element_count, + PrimitiveType primitive_type, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& source_array); @@ -859,7 +861,8 @@ class IrEmitter : public DfsHloVisitorWithDefault, // Decoupled implementation of IrEmitter::EmitTransferElements. void EmitTransferElements(llvm::Value* target, llvm::Value* source, - int64_t element_count, PrimitiveType primitive_type, + const xla::DExpr& element_count, + PrimitiveType primitive_type, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& source_array, llvm::Module* module, llvm::IRBuilderBase& b); diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.cc b/third_party/xla/xla/service/cpu/ir_emitter2.cc index bce2108bb87572..e39fba7a0b9098 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter2.cc +++ b/third_party/xla/xla/service/cpu/ir_emitter2.cc @@ -29,6 +29,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/statusor.h" +#include "xla/shape_expr.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" @@ -74,6 +75,7 @@ limitations under the License. #include "xla/service/llvm_ir/ir_array.h" #include "xla/service/llvm_ir/llvm_util.h" #include "xla/service/llvm_ir/loop_emitter.h" +#include "xla/printer.h" #include "xla/shape.h" #include "xla/shape_partition.h" #include "xla/stream_executor/launch_dim.h" @@ -195,6 +197,30 @@ absl::StatusOr IrEmitter2::EmitPadHostKernel( KernelInfo(std::move(kernel_prototype), se::BlockDim(), se::ThreadDim())); } +absl::StatusOr +IrEmitter2::EmitGetExpressionValueHostKernel(const HloInstruction* getBatch) { + VLOG(2) << "Emit GetExpressionValue host kernel: " << getBatch->name(); + + TF_ASSIGN_OR_RETURN(KernelPrototype kernel_prototype, + EmitKernelPrototype(getBatch)); + llvm_ir::IrArray operand_array = kernel_prototype.arguments[0]; + llvm_ir::IrArray output_array = kernel_prototype.results[0]; + TF_RET_CHECK(getBatch->has_contents()); + TF_RET_CHECK(!getBatch->contents().empty()); + xla::DExpr expr = xla::DExprFromProto(getBatch->contents()[0]); + TF_RET_CHECK(expr); + llvm::IRBuilder<> b(module_->getContext()); + b.SetInsertPoint(kernel_prototype.function->getEntryBlock().getTerminator()); + llvm::Value* bdim_value = llvm_ir::EmitExpression(&b, expr); + llvm_ir::IrArray::Index output_index(/*multidimensional_index=*/{}, + getBatch->shape(), b.getInt32Ty()); + llvm::Value* output_ptr = + output_array.EmitArrayElementAddress(output_index, &b); + b.CreateStore(bdim_value, output_ptr); + return kernels_.emplace_back( + KernelInfo(std::move(kernel_prototype), se::BlockDim(), se::ThreadDim())); +} + absl::StatusOr IrEmitter2::EmitFusionHostKernel( const HloFusionInstruction* fusion) { VLOG(2) << "Emit fusion host kernel: " << fusion->name(); diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.h b/third_party/xla/xla/service/cpu/ir_emitter2.h index e720e06f37642c..8c493a164626cf 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter2.h +++ b/third_party/xla/xla/service/cpu/ir_emitter2.h @@ -113,6 +113,9 @@ class IrEmitter2 { // Emits a host kernel for the pad instruction. absl::StatusOr EmitPadHostKernel(const HloInstruction* pad); + absl::StatusOr EmitGetExpressionValueHostKernel( + const HloInstruction* getBatch); + // Emits a host kernel for the given fusion instruction. absl::StatusOr EmitFusionHostKernel( const HloFusionInstruction* fusion); diff --git a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc b/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc index 2bfffd88df937e..e3af0399381b79 100644 --- a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc +++ b/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc @@ -61,6 +61,7 @@ ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name, // performance with a large improvement in compile time. auto unroll_mode = (i == 0) ? llvm_ir::UnrollMode::kDefaultUnroll : llvm_ir::UnrollMode::kNoUnroll; + if (bounds_index < dynamic_loop_bounds_->size()) { // Emit dynamic loop bounds for this dimension. Dynamic loop bounds // are read from ir function dynamic loop bounds argument. @@ -69,14 +70,17 @@ ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name, std::unique_ptr loop = loop_nest.AddLoop( /*suffix=*/absl::StrFormat("dim.%d", dimension), start_index, - end_index, unroll_mode); + end_index, unroll_mode, /*prevent_vectorization*/ false, + /* expression */ shape_.expressions(dimension)); array_multi_index[dimension] = loop->GetIndVarValue(); } else { // Emit static loop bounds for this dimension. std::unique_ptr loop = loop_nest.AddLoop( /*start_index=*/0, /*end_index=*/shape_.dimensions(dimension), - /*suffix=*/absl::StrFormat("dim.%d", dimension), unroll_mode); + /*suffix=*/absl::StrFormat("dim.%d", dimension), unroll_mode, + /*prevent_vectorization*/ false, + /* expression */ shape_.expressions(dimension)); array_multi_index[dimension] = loop->GetIndVarValue(); } } diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.cc b/third_party/xla/xla/service/cpu/runtime_batch_size.cc index 951910f9aee51c..063cd5945b9a61 100644 --- a/third_party/xla/xla/service/cpu/runtime_batch_size.cc +++ b/third_party/xla/xla/service/cpu/runtime_batch_size.cc @@ -17,13 +17,6 @@ limitations under the License. #include "xla/executable_run_options.h" -namespace xla::cpu::runtime { - -const char* const kGetBatchSizeSymbolName = - "__xla_cpu_runtime_GetBatchSize"; - -} // namespace xla::cpu::runtime - extern "C" int64_t __xla_cpu_runtime_GetBatchSize( const void* run_options_ptr) { if (run_options_ptr == nullptr) return 0; diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.h b/third_party/xla/xla/service/cpu/runtime_batch_size.h index 9204e7841c4e70..fbd36294a4ddaf 100644 --- a/third_party/xla/xla/service/cpu/runtime_batch_size.h +++ b/third_party/xla/xla/service/cpu/runtime_batch_size.h @@ -20,7 +20,8 @@ limitations under the License. namespace xla::cpu::runtime { -extern const char* const kGetBatchSizeSymbolName; +inline constexpr char kGetBatchSizeSymbolName[] = + "__xla_cpu_runtime_GetBatchSize"; } // namespace xla::cpu::runtime diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc index 9079ed8a16d1a8..72fe6b95e25825 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc @@ -1086,6 +1086,8 @@ absl::StatusOr ThunkEmitter::EmitCustomCallThunk( return EmitTopKThunk(custom_call); } else if (custom_call_target == "SliceToDynamic") { return EmitSliceToDynamicThunk(instruction); + } else if (custom_call_target == "GetExpressionValue") { + return EmitGetExpressionValueThunk(instruction); } // Check the API version. @@ -1128,6 +1130,21 @@ absl::StatusOr ThunkEmitter::EmitSliceToDynamicThunk( /*min_alignment=*/cpu_function_runtime::MinAlign()); } +absl::StatusOr ThunkEmitter::EmitGetExpressionValueThunk( + const HloInstruction* instruction) { + VLOG(2) << "Handling GetExpressionValue for instruction: " + << instruction->ToString(); + const HloCustomCallInstruction* custom_call = + Cast(instruction); + TF_ASSIGN_OR_RETURN( + auto kernel, ir_emitter_.EmitGetExpressionValueHostKernel(custom_call)); + TF_ASSIGN_OR_RETURN(auto result_buffer, + GetHostKernelAllocationSlices(instruction)); + return MakeKernelThunkSequence( + instruction, result_buffer, kernel, + /*min_alignment=*/cpu_function_runtime::MinAlign()); +} + absl::StatusOr ThunkEmitter::EmitSliceThunk( const HloInstruction* instruction) { // TODO(ezhulenev): Consider implementing slice operations as separate diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.h b/third_party/xla/xla/service/cpu/thunk_emitter.h index 9cdc8ae3981680..40aad964fb8d24 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.h +++ b/third_party/xla/xla/service/cpu/thunk_emitter.h @@ -183,6 +183,9 @@ class ThunkEmitter { absl::StatusOr EmitSliceToDynamicThunk( const HloInstruction* instruction); + absl::StatusOr EmitGetExpressionValueThunk( + const HloInstruction* instruction); + absl::StatusOr EmitTopKThunk( const HloCustomCallInstruction* custom_call); diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc index 89391a3bdd1dd5..3fe608125e8c3e 100644 --- a/third_party/xla/xla/service/elemental_ir_emitter.cc +++ b/third_party/xla/xla/service/elemental_ir_emitter.cc @@ -3278,51 +3278,57 @@ absl::StatusOr ElementalIrEmitter::EmitElementalConcatenate( } // We use bisection to select the input operand. - int64_t current_offset = 0; + int64_t coffset = 0; + llvm::Value* current_offset = source_index.GetConstantWithIndexType(0); // Offset for every operand. - std::vector> cases; + std::vector> cases; cases.reserve(hlo->operand_count()); for (const HloInstruction* operand : hlo->operands()) { cases.emplace_back(current_offset, operand); - current_offset += operand->shape().dimensions(concat_dim); + llvm::Value* cdim = source_index.GetConstantWithIndexType( + operand->shape().dimensions(concat_dim)); + const auto& concat_expr = operand->shape().expressions(concat_dim); + if (concat_expr->is_dynamic()) { + cdim = llvm_ir::EmitExpression(b_, concat_expr); + } + current_offset = b_->CreateAdd(current_offset, cdim, "current_offset"); + coffset += operand->shape().dimensions(concat_dim); } - CHECK_EQ(current_offset, hlo->shape().dimensions(concat_dim)); + CHECK_EQ(coffset, hlo->shape().dimensions(concat_dim)); std::function> operands)> + absl::Span> operands)> emit_tree = - [&](absl::Span> + [&](absl::Span> operands) { llvm::IRBuilder<>::InsertPointGuard guard(*b_); size_t mid = operands.size() / 2; - const std::pair& pivot = + const std::pair& pivot = operands[mid]; llvm::BasicBlock* block = llvm_ir::CreateBasicBlock( exit_block, - absl::StrCat("concatenate.pivot.", pivot.first, "."), b_); + absl::StrCat("concatenate.pivot."), b_); b_->SetInsertPoint(block); // If there's only one element we're done. The range is contiguous // so we can just jump to the block for it. if (operands.size() == 1) { - const std::pair& operand = + const std::pair& operand = operands.back(); int64_t operand_id = to_unique_operand_id[operand.second]; source_index_phis[operand_id]->addIncoming( - source_index.GetConstantWithIndexType(operand.first), + operand.first, b_->GetInsertBlock()); b_->CreateBr(emit_operand_blocks[operand_id]); return block; } // Take the middle element and recurse. - llvm::Constant* pivot_const = llvm::ConstantInt::get( - source_index[concat_dim]->getType(), pivot.first); llvm::Value* comp = - b_->CreateICmpULT(source_index[concat_dim], pivot_const); + b_->CreateICmpULT(source_index[concat_dim], pivot.first); llvm::BasicBlock* left_block = emit_tree(operands.subspan(0, mid)); llvm::BasicBlock* right_block = emit_tree(operands.subspan(mid)); @@ -3616,11 +3622,16 @@ absl::StatusOr ElementalIrEmitter::EmitElementalPad( "in_bounds"); multi_index[i] = SDiv(multi_index[i], index_typed_const(pad_dim.interior_padding() + 1)); - in_bounds = - And(in_bounds, - ICmpSLT(multi_index[i], - index_typed_const(hlo->operand(0)->shape().dimensions(i))), - "in_bounds"); + + int64_t shape_dim = hlo->operand(0)->shape().dimensions(i); + llvm::Value* bound = index_typed_const(shape_dim); + + const auto& operand_expr = hlo->operand(0)->shape().expressions(i); + if (operand_expr->is_dynamic()) { + bound = llvm_ir::EmitExpression(b_, operand_expr); + } + + in_bounds = And(in_bounds, ICmpSLT(multi_index[i], bound), "in_bounds"); } // if (in_bounds) { @@ -3687,9 +3698,20 @@ absl::StatusOr ElementalIrEmitter::EmitElementalDot( return llvm::ConstantInt::get(index_type, c); }; + llvm::Value* contracted_bound = index_typed_const(contracted_dim_size); + + if (!hlo->operand(0) + ->shape() + .expressions(lhs_contracting_dim) + ->is_constant()) { + llvm::Value* expr_value = llvm_ir::EmitExpression( + b_, hlo->operand(0)->shape().expressions(lhs_contracting_dim)); + contracted_bound = expr_value; + } + std::unique_ptr inner_loop = llvm_ir::ForLoop::EmitForLoop( - IrName(hlo, "inner"), index_typed_const(0), - index_typed_const(contracted_dim_size), index_typed_const(1), b_); + IrName(hlo, "inner"), index_typed_const(0), contracted_bound, + index_typed_const(1), b_); SetToFirstInsertPoint(inner_loop->GetPreheaderBasicBlock(), b_); PrimitiveType primitive_type = hlo->shape().element_type(); @@ -3883,9 +3905,17 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( const HloInstruction* operand = hlo->operand(0); std::vector source_multi_index = target_index.multidim(); for (int64_t dim : hlo->dimensions()) { - source_multi_index[dim] = Sub(target_index.GetConstantWithIndexType( - hlo->shape().dimensions(dim) - 1), - target_index[dim]); + const auto& dim_expr = hlo->shape().expressions(dim); + if (dim_expr->is_dynamic()) { + llvm::Value* one = target_index.GetConstantWithIndexType(1); + llvm::Value* expr_value = llvm_ir::EmitExpression(b_, dim_expr); + source_multi_index[dim] = + Sub(Sub(expr_value, one), target_index[dim]); + } else { + source_multi_index[dim] = Sub(target_index.GetConstantWithIndexType( + hlo->shape().dimensions(dim) - 1), + target_index[dim]); + } } llvm_ir::IrArray::Index source_index( source_multi_index, operand->shape(), target_index.GetType()); @@ -3967,11 +3997,12 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kSlice: return [this, hlo, &operand_to_generator]( const IrArray::Index& index) -> absl::StatusOr { - IrArray::Index sliced_index = index.SourceIndexOfSlice( - /*operand_shape=*/hlo->operand(0)->shape(), - /*starts=*/hlo->slice_starts(), - /*strides=*/hlo->slice_strides(), /*builder=*/b_); - return operand_to_generator.at(hlo->operand(0))(sliced_index); + const auto* slice = Cast(hlo); + return operand_to_generator.at(hlo->operand(0))( + index.SourceIndexOfSlice(hlo->operand(0)->shape(), + hlo->slice_starts(), + slice->slice_start_exprs(), + hlo->slice_strides(), b_)); }; case HloOpcode::kDynamicSlice: return [this, hlo, &operand_to_generator]( @@ -4236,11 +4267,21 @@ absl::StatusOr ElementalIrEmitter::EmitElementalReduceWindow( // comparison is equivalent to the unsigned comparison // input_multi_index[i] < bound, as a negative value wraps to a large // positive value. + + int64_t dim_bound = reduce_window->inputs()[0]->shape().dimensions(i); + llvm::Value* shape_bound = index_typed_const(dim_bound); + + const auto& window_expr = + reduce_window->inputs()[0]->shape().expressions(i); + if (window_expr->is_dynamic()) { + llvm::Value* expr_value = llvm_ir::EmitExpression(b_, window_expr); + shape_bound = expr_value; + } + in_bounds = And(in_bounds, ICmpULT(input_multi_index[i], - index_typed_const( - reduce_window->inputs()[0]->shape().dimensions(i)))); + shape_bound)); } llvm_ir::LlvmIfData if_data = diff --git a/third_party/xla/xla/service/llvm_ir/BUILD b/third_party/xla/xla/service/llvm_ir/BUILD index 599060b88eee81..0f3cd53a985654 100644 --- a/third_party/xla/xla/service/llvm_ir/BUILD +++ b/third_party/xla/xla/service/llvm_ir/BUILD @@ -79,9 +79,10 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/service/cpu:cpu_options", + "//xla/service/cpu:runtime_batch_size_api", "//xla/service:dump", "//xla/service:hlo_module_config", - "//xla/service/cpu:cpu_options", "//xla/tsl/platform:byte_order", "//xla/tsl/platform:logging", "@com_google_absl//absl/algorithm:container", @@ -355,3 +356,18 @@ xla_cc_test( "@llvm-project//llvm:ir_headers", ], ) + +xla_cc_test( + name = "llvm_util_test", + srcs = ["llvm_util_test.cc"], + deps = [ + ":llvm_util", + "//xla/service/cpu:runtime_batch_size_api", + "//xla:shape_util", + "//xla/tests:xla_internal_test_main", + "@com_google_googletest//:gtest", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:ir_headers", + ], +) diff --git a/third_party/xla/xla/service/llvm_ir/ir_array.cc b/third_party/xla/xla/service/llvm_ir/ir_array.cc index 139d537c88778a..8288cba9530899 100644 --- a/third_party/xla/xla/service/llvm_ir/ir_array.cc +++ b/third_party/xla/xla/service/llvm_ir/ir_array.cc @@ -281,8 +281,11 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape( // linear index by each dimension size. for (int64_t i = common_factors[k + 1].first - 1; i >= common_factors[k].first; --i) { + const auto& input_expr = input_shape.expressions(i); + bool is_dynamic = input_expr->is_dynamic(); llvm::Value* divisor = - GetConstantWithIndexType(input_shape.dimensions(i)); + is_dynamic ? llvm_ir::EmitExpression(builder, input_expr) + : GetConstantWithIndexType(input_shape.dimensions(i)); if (input_shape.dimensions(i) == 1) { source_multidim_index[i] = GetConstantWithIndexType(0); } else if (i == common_factors[k].first) { @@ -307,17 +310,27 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape( IrArray::Index IrArray::Index::SourceIndexOfSlice( const Shape& operand_shape, absl::Span starts, - absl::Span strides, llvm::IRBuilderBase* builder) const { + absl::Span start_exprs, absl::Span strides, + llvm::IRBuilderBase* builder) const { + static const DExpr kMissingSliceExpr = + DExpr::Unknown(kMissingExpressionSentinel); std::vector source_multi_index(multidim_.size()); for (int i = 0; i < multidim_.size(); ++i) { + llvm::Value* start_value = GetConstantWithIndexType(starts[i]); + const DExpr& start_expr = + i < start_exprs.size() ? start_exprs[i] : kMissingSliceExpr; + if (start_expr && start_expr->is_dynamic()) { + start_value = builder->CreateIntCast( + llvm_ir::EmitExpression(builder, start_expr), index_type_, + /*isSigned=*/true); + } int64_t stride = strides[i]; if (stride != 1) { source_multi_index[i] = builder->CreateAdd( builder->CreateMul(multidim_[i], GetConstantWithIndexType(stride)), - GetConstantWithIndexType(starts[i])); + start_value); } else { - source_multi_index[i] = - builder->CreateAdd(multidim_[i], GetConstantWithIndexType(starts[i])); + source_multi_index[i] = builder->CreateAdd(multidim_[i], start_value); } } return Index(source_multi_index, operand_shape, index_type_); @@ -559,8 +572,33 @@ llvm::Value* IrArray::EmitArrayElementAddress(const IrArray::Index& index, int64_t dimension = LayoutUtil::Major(shape_.layout(), i); gep_indices.push_back(actual_index[dimension]); } - return b->CreateInBoundsGEP(pointee_type_, base_ptr_, gep_indices, - llvm_ir::AsStringRef(name)); + + // Do not make a dynamic "GEP" if only the first dimension is dynamic since + // it's always indiced with 0 (i.e. the dynamic dimension has no impact on the + // address computation). + std::vector gep_dims; + std::vector gep_expressions; + gep_dims.reserve(shape_.dimensions().size()); + gep_expressions.reserve(shape_.dimensions().size()); + for (int64_t i = 0; i < shape_.dimensions().size(); ++i) { + int64_t dimension = LayoutUtil::Major(shape_.layout(), i); + gep_dims.push_back(shape_.dimensions(dimension)); + gep_expressions.push_back(shape_.expressions(dimension)); + } + bool dynamic_first_dim = + gep_expressions[0]->is_dynamic() && + std::all_of(gep_expressions.begin() + 1, gep_expressions.end(), + [](const DExpr& e) { return e->is_constant(); }); + if (!dynamic_first_dim && shape_.has_dynamic_expr()) { + llvm::Type* element_type = + PrimitiveTypeToIrType(shape_.element_type(), b->getContext()); + return llvm_ir::createDynamicGEP( + b, base_ptr_, gep_indices, gep_dims, gep_expressions, + element_type, llvm_ir::AsStringRef(name)); + } else { + return b->CreateInBoundsGEP(pointee_type_, base_ptr_, gep_indices, + llvm_ir::AsStringRef(name)); + } } llvm::Value* IrArray::EmitLinearArrayElementAddress( diff --git a/third_party/xla/xla/service/llvm_ir/ir_array.h b/third_party/xla/xla/service/llvm_ir/ir_array.h index f5b2b7c4fbd792..eca880b39d29f1 100644 --- a/third_party/xla/xla/service/llvm_ir/ir_array.h +++ b/third_party/xla/xla/service/llvm_ir/ir_array.h @@ -158,6 +158,7 @@ class IrArray { // `operand_shape`. Index SourceIndexOfSlice(const Shape& operand_shape, absl::Span starts, + absl::Span start_exprs, absl::Span strides, llvm::IRBuilderBase* builder) const; diff --git a/third_party/xla/xla/service/llvm_ir/llvm_loop.cc b/third_party/xla/xla/service/llvm_ir/llvm_loop.cc index 43ee77ef0e85be..c60fc1cc7fa8ba 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_loop.cc +++ b/third_party/xla/xla/service/llvm_ir/llvm_loop.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/service/llvm_ir/llvm_util.h" #include "xla/shape.h" #include "xla/tsl/platform/logging.h" +#include "llvm/include/llvm/Support/Debug.h" namespace xla { namespace llvm_ir { @@ -185,24 +186,40 @@ llvm::BasicBlock* ForLoop::CreateLoopBB(absl::string_view name, return CreateBasicBlock(insert_before_bb_, GetQualifiedName(name), b); } -std::unique_ptr ForLoopNest::AddLoop(absl::string_view suffix, - llvm::Value* start_index, - llvm::Value* end_index, - UnrollMode unroll_mode, - bool prevent_vectorization) { +std::unique_ptr ForLoopNest::AddLoop( + absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index, + UnrollMode unroll_mode, bool prevent_vectorization, + DExpr expression) { return AddLoop(suffix, start_index, end_index, GetConstantWithIndexType(1), - unroll_mode, prevent_vectorization); + unroll_mode, prevent_vectorization, expression); } std::unique_ptr ForLoopNest::AddLoop( absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index, - llvm::Value* stride, UnrollMode unroll_mode, bool prevent_vectorization) { - if (inner_loop_body_bb_ != nullptr) { + llvm::Value* stride, UnrollMode unroll_mode, bool prevent_vectorization, + DExpr expression) { + if (inner_loop_body_bb_ != nullptr && + b_->GetInsertBlock() != inner_loop_body_bb_) { // Create this loop inside the previous one. b_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt()); } + llvm::Value* actual_end = end_index; + if (expression && expression->is_dynamic()) { + // Get batch dim and compare with end_index to use minimum value + llvm::Value* expr_value = llvm_ir::EmitExpression(b_, expression); + actual_end = b_->CreateSelect(b_->CreateICmpULT(end_index, expr_value), + end_index, expr_value, "loop_end_min"); + } + // Emitting dynamic expressions may materialize helper instructions into the + // current function entry block using a different builder. Re-seat the + // insertion point at the end of the current block so the loop preheader is + // always in a well-formed state before emitting the loop. + llvm::BasicBlock* insert_block = b_->GetInsertBlock(); + if (insert_block != nullptr && insert_block->getTerminator() == nullptr) { + b_->SetInsertPoint(insert_block); + } std::unique_ptr loop(new ForLoop( - /*prefix=*/name_, suffix, start_index, end_index, stride, unroll_mode, + /*prefix=*/name_, suffix, start_index, actual_end, stride, unroll_mode, prevent_vectorization)); loop->Emit(b_); @@ -219,25 +236,46 @@ std::unique_ptr ForLoopNest::AddLoop( return loop; } -std::unique_ptr ForLoopNest::AddLoop(int64_t start_index, - int64_t end_index, - absl::string_view suffix, - UnrollMode unroll_mode, - bool prevent_vectorization) { +llvm::Value* ForLoopNest::MaterializeLoopEndValue(int64_t end_index, + DExpr expression) { + llvm::Value* end = GetConstantWithIndexType(end_index); + if (expression && expression->is_dynamic()) { + if (inner_loop_body_bb_ != nullptr) { + // Keep dynamic loop bound materialization in the same block that will + // act as the nested loop preheader so the resulting values dominate the + // loop header compare. + b_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt()); + } + end = EmitExpression(b_, expression); + } + llvm::BasicBlock* insert_block = b_->GetInsertBlock(); + if (insert_block != nullptr && insert_block->getTerminator() == nullptr) { + b_->SetInsertPoint(insert_block); + } + return end; +} + +std::unique_ptr ForLoopNest::AddLoop( + int64_t start_index, int64_t end_index, absl::string_view suffix, + UnrollMode unroll_mode, bool prevent_vectorization, + DExpr expression) { CHECK_LE(start_index, end_index); - return AddLoop(suffix, GetConstantWithIndexType(start_index), - GetConstantWithIndexType(end_index), unroll_mode, - prevent_vectorization); + + llvm::Value* end = MaterializeLoopEndValue(end_index, expression); + return AddLoop(suffix, GetConstantWithIndexType(start_index), end, + unroll_mode, prevent_vectorization); } std::unique_ptr ForLoopNest::AddLoop(int64_t start_index, int64_t end_index, int64_t stride, absl::string_view suffix, UnrollMode unroll_mode, - bool prevent_vectorization) { + bool prevent_vectorization, + DExpr expression) { CHECK_LE(start_index, end_index); - return AddLoop(suffix, GetConstantWithIndexType(start_index), - GetConstantWithIndexType(end_index), + + llvm::Value* end = MaterializeLoopEndValue(end_index, expression); + return AddLoop(suffix, GetConstantWithIndexType(start_index), end, GetConstantWithIndexType(stride), unroll_mode, prevent_vectorization); } @@ -259,7 +297,9 @@ std::vector ForLoopNest::AddLoopsForShapeOnDimensions( /*start_index=*/0, /*end_index=*/shape.dimensions(dimension), /*suffix=*/ - llvm_ir::IrName(suffix, absl::StrCat(dimension))); + llvm_ir::IrName(suffix, absl::StrCat(dimension)), + /*unroll_mode=*/llvm_ir::UnrollMode::kDefaultUnroll, + /*prevent_vectorization=*/false, shape.expressions(dimension)); multi_index[dimension] = loop->GetIndVarValue(); } return multi_index; diff --git a/third_party/xla/xla/service/llvm_ir/llvm_loop.h b/third_party/xla/xla/service/llvm_ir/llvm_loop.h index 7aa8ce9e32e950..3a15102bf6dabd 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_loop.h +++ b/third_party/xla/xla/service/llvm_ir/llvm_loop.h @@ -201,14 +201,14 @@ class ForLoopNest { absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index, llvm::Value* stride, UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, - bool prevent_vectorization = false); + bool prevent_vectorization = false, DExpr expression = DExpr()); // Like the above, except that it defaults to a stride of one. std::unique_ptr AddLoop( absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index, UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, - bool prevent_vectorization = false); + bool prevent_vectorization = false, DExpr expression = DExpr()); // A convenient wrapper of the other flavor of AddLoop. The given start and // end index are constant. @@ -216,13 +216,13 @@ class ForLoopNest { int64_t start_index, int64_t end_index, int64_t stride, absl::string_view suffix, UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, - bool prevent_vectorization = false); + bool prevent_vectorization = false, DExpr expression = DExpr()); // Like the above, except that it defaults to a stride of one. std::unique_ptr AddLoop( int64_t start_index, int64_t end_index, absl::string_view suffix, UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll, - bool prevent_vectorization = false); + bool prevent_vectorization = false, DExpr expression = DExpr()); // Add loops to iterate through the indices within the specified // shape. The returned index collects the induction variables of the @@ -268,6 +268,8 @@ class ForLoopNest { llvm::BasicBlock* GetInnerLoopBodyBasicBlock() { return inner_loop_body_bb_; } private: + llvm::Value* MaterializeLoopEndValue(int64_t end_index, DExpr expression); + void SetIndexType(llvm::Type* index_ty) { index_type_ = index_ty == nullptr ? b_->getInt64Ty() : index_ty; } diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util.cc b/third_party/xla/xla/service/llvm_ir/llvm_util.cc index b1db780705b41f..8763a85033624b 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_util.cc +++ b/third_party/xla/xla/service/llvm_ir/llvm_util.cc @@ -68,7 +68,9 @@ limitations under the License. #include "xla/layout_util.h" #include "xla/literal.h" #include "xla/primitive_util.h" +#include "xla/printer.h" #include "xla/service/cpu/cpu_options.h" +#include "xla/service/cpu/runtime_batch_size.h" #include "xla/service/dump.h" #include "xla/service/hlo_module_config.h" #include "xla/service/llvm_ir/llvm_type_conversion_util.h" @@ -85,6 +87,8 @@ namespace llvm_ir { namespace { +constexpr llvm::StringLiteral kBdimValueName("bdim_value"); + // This works for most llvm / mlir types. This also accepts a const pointer to // objects which have a const print() method. template @@ -288,9 +292,11 @@ llvm::Type* ShapeToIrType(const Shape& shape, llvm::LLVMContext& context) { result_type = llvm::ArrayType::get(result_type, shape.tuple_shapes().size()); } else if (shape.IsArray()) { - for (int64_t dimension : LayoutUtil::MinorToMajor(shape)) { - result_type = - llvm::ArrayType::get(result_type, shape.dimensions(dimension)); + auto dimensions = LayoutUtil::MinorToMajor(shape); + for (int i = 0; i < dimensions.size(); i++) { + bool is_dynamic = shape.expressions(dimensions[i])->is_dynamic(); + int64_t dim_val = is_dynamic ? 0 : shape.dimensions(dimensions[i]); + result_type = llvm::ArrayType::get(result_type, dim_val); } } return result_type; @@ -811,5 +817,166 @@ void EmitEarlyReturn(llvm::Value* condition, llvm::IRBuilderBase* b, b->SetInsertPoint(continued, continued->getFirstInsertionPt()); } +llvm::Value* GetBatchDimByName(llvm::IRBuilderBase* b, int64_t multiplier, + int64_t offset) { + llvm::Function* function = b->GetInsertBlock()->getParent(); + llvm::LLVMContext& ctx = b->getContext(); + llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx); + llvm::Value* loadedValue = nullptr; + llvm::Value* bdim_scaled = nullptr; + for (auto& inst : function->getEntryBlock()) { + if (inst.getName() == kBdimValueName) { + loadedValue = &inst; + } + } + if (!loadedValue) { + // If missing, materialize it through the stable CPU runtime interface. + llvm::Value* run_options = nullptr; + for (auto& arg : function->args()) { + if (arg.getName() == "run_options") { + run_options = &arg; + break; + } + } + + if (!run_options) { + return nullptr; + } + llvm::IRBuilder<> entry_builder( + &function->getEntryBlock(), + function->getEntryBlock().getFirstInsertionPt()); + llvm::FunctionType* getter_type = llvm::FunctionType::get( + i64Type, {run_options->getType()}, /*isVarArg=*/false); + llvm::FunctionCallee getter = function->getParent()->getOrInsertFunction( + xla::cpu::runtime::kGetBatchSizeSymbolName, getter_type); + loadedValue = entry_builder.CreateCall(getter, {run_options}, + kBdimValueName); + } + if (multiplier < 1) { + llvm::errs() << "Multiplier is less than 1, this should not happen.\n"; + } else if (multiplier == 1) { + bdim_scaled = loadedValue; + } else { + llvm::ConstantInt* m = llvm::ConstantInt::get(i64Type, multiplier, true); + bdim_scaled = b->CreateMul(loadedValue, m, "bdim_scaled"); + } + if (offset != 0){ + llvm::ConstantInt* offset_value = + llvm::ConstantInt::get(i64Type, offset, true); + bdim_scaled = b->CreateAdd(bdim_scaled, offset_value, "bdim_offset"); + } + return bdim_scaled; +} + +static llvm::Value* EmitExpressionImpl(llvm::IRBuilderBase* b, + const DynExpr& expr) { + llvm::LLVMContext& ctx = b->getContext(); + llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx); + if (expr.is_constant()) return llvm::ConstantInt::get(i64Type, expr.get_val(), true); + if (expr.kind() == DExpr::Kind::kUnknown) { + return nullptr; + } + if (expr.kind() == DExpr::Kind::kVariable) { + // For now we can just use %bdim... + return GetBatchDimByName(b); + } + if (expr.kind() == DExpr::Kind::kMul) { + auto* mul_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *mul_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *mul_node->get_rhs()); + return b->CreateMul(v_lhs, v_rhs, "mul_dims"); + } + // TODO: Check if this should ever happen + if (expr.kind() == DExpr::Kind::kDiv) { + auto* div_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *div_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *div_node->get_rhs()); + return b->CreateSDiv(v_lhs, v_rhs, "div_dims"); + } + if (expr.kind() == DExpr::Kind::kAdd) { + auto* add_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *add_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *add_node->get_rhs()); + return b->CreateAdd(v_lhs, v_rhs, "add_dims"); + } + if (expr.kind() == DExpr::Kind::kSub) { + auto* sub_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *sub_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *sub_node->get_rhs()); + return b->CreateSub(v_lhs, v_rhs, "sub_dims"); + } + if (expr.kind() == DExpr::Kind::kMax) { + auto* max_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *max_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *max_node->get_rhs()); + llvm::Value* lhs_is_greater = + b->CreateICmpSGT(v_lhs, v_rhs, "max_dims_pred"); + return b->CreateSelect(lhs_is_greater, v_lhs, v_rhs, "max_dims"); + } + if (expr.kind() == DExpr::Kind::kGt) { + auto* gt_node = static_cast(&expr); + llvm::Value* v_lhs = EmitExpressionImpl(b, *gt_node->get_lhs()); + llvm::Value* v_rhs = EmitExpressionImpl(b, *gt_node->get_rhs()); + llvm::Value* pred = b->CreateICmpSGT(v_lhs, v_rhs, "gt_dims_pred"); + return b->CreateZExt(pred, i64Type, "gt_dims"); + } + if (expr.kind() == DExpr::Kind::kSelect) { + auto* select_node = static_cast(&expr); + llvm::Value* pred = EmitExpressionImpl(b, *select_node->get_pred()); + llvm::Value* v_true = + EmitExpressionImpl(b, *select_node->get_on_true()); + llvm::Value* v_false = + EmitExpressionImpl(b, *select_node->get_on_false()); + llvm::Value* nonzero = b->CreateICmpNE( + pred, llvm::ConstantInt::get(i64Type, 0, true), "select_dims_pred"); + return b->CreateSelect(nonzero, v_true, v_false, "select_dims"); + } + return nullptr; +} + +llvm::Value* EmitExpression(llvm::IRBuilderBase* b, const DExpr& expr) { + if (!expr) return nullptr; + DExpr simplified = expr.simplify(); + StringPrinter printer; + simplified->print(&printer); + VLOG(2) << "EmitExpression expr=" << std::move(printer).ToString() + << " kind=" << static_cast(simplified.kind()); + llvm::Value* value = EmitExpressionImpl(b, *simplified.get()); + return value; +} + +llvm::Value* createDynamicGEP(llvm::IRBuilderBase* builder, + llvm::Value* base_ptr, + const std::vector& indices, + absl::Span dims, + absl::Span expressions, + llvm::Type* elem_type, + const llvm::Twine& name) { + llvm::Value* total_index = builder->getInt64(0); + llvm::Type* int64_ty = builder->getInt64Ty(); + + for (size_t i = 0; i < indices.size(); ++i) { + // The stride is the product of all dimensions to the right of this index. + llvm::Value* stride = builder->getInt64(1); + for (size_t j = i; j < dims.size(); ++j) { + if (expressions[j] && expressions[j]->is_dynamic()) { + llvm::Value* expr_value = + EmitExpression(builder, expressions[j]); + stride = builder->CreateMul(stride, expr_value, "stride.dyn"); + } else { + stride = builder->CreateMul( + stride, llvm::ConstantInt::get(int64_ty, dims[j]), "stride.static"); + } + } + llvm::Value* scaled_index = + builder->CreateMul(indices[i], stride, "idx.scaled"); + total_index = builder->CreateAdd(total_index, scaled_index, "idx.total"); + } + + // Final GEP: result = base + total_index * sizeof(elem_type) + llvm::Value* gep = builder->CreateGEP(elem_type, base_ptr, total_index, name); + return gep; +} + } // namespace llvm_ir } // namespace xla diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util.h b/third_party/xla/xla/service/llvm_ir/llvm_util.h index 88c1287d2f236d..77bd543c9c0482 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_util.h +++ b/third_party/xla/xla/service/llvm_ir/llvm_util.h @@ -334,6 +334,19 @@ llvm::BasicBlock* EmitReturnBlock(llvm::IRBuilderBase* b); void EmitEarlyReturn(llvm::Value* condition, llvm::IRBuilderBase* b, llvm::BasicBlock* return_block = nullptr); +llvm::Value* GetBatchDimByName(llvm::IRBuilderBase* b, int64_t multiplier = 1, + int64_t offset = 0); + +llvm::Value* createDynamicGEP(llvm::IRBuilderBase* builder, + llvm::Value* base_ptr, + const std::vector& indices, + absl::Span dims, + absl::Span expressions, + llvm::Type* elem_type, + const llvm::Twine& name = ""); + +llvm::Value* EmitExpression(llvm::IRBuilderBase* b, const DExpr& expr); + } // namespace llvm_ir } // namespace xla diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc b/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc new file mode 100644 index 00000000000000..4e740b1566e988 --- /dev/null +++ b/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc @@ -0,0 +1,83 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 "xla/service/llvm_ir/llvm_util.h" + +#include +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Type.h" +#include "llvm/Support/Casting.h" +#include "xla/service/cpu/runtime_batch_size.h" +#include "xla/shape_expr.h" + +namespace xla { +namespace llvm_ir { +namespace { + +TEST(LlvmUtilTest, DynamicExpressionDivisionIsSigned) { + llvm::LLVMContext context; + llvm::Module module("llvm_util_test", context); + llvm::IRBuilder<> builder(context); + llvm::FunctionType* function_type = llvm::FunctionType::get( + llvm::Type::getVoidTy(context), /*isVarArg=*/false); + llvm::Function* function = llvm::Function::Create( + function_type, llvm::Function::ExternalLinkage, "test", module); + llvm::BasicBlock* block = + llvm::BasicBlock::Create(context, "entry", function); + builder.SetInsertPoint(block); + + llvm::Type* i64_type = builder.getInt64Ty(); + llvm::Value* batch_dim_address = builder.CreateAlloca(i64_type); + builder.CreateLoad(i64_type, batch_dim_address, "bdim_value"); + + DExpr expression = (DExpr::Var(1) - 5) / 2; + llvm::Value* value = EmitExpression(&builder, expression); + + auto* division = llvm::dyn_cast(value); + ASSERT_NE(division, nullptr); + EXPECT_EQ(division->getOpcode(), llvm::Instruction::SDiv); +} + +TEST(LlvmUtilTest, DynamicBatchSizeUsesRuntimeAccessor) { + llvm::LLVMContext context; + llvm::Module module("llvm_util_test", context); + llvm::IRBuilder<> builder(context); + llvm::FunctionType* function_type = llvm::FunctionType::get( + llvm::Type::getVoidTy(context), {builder.getPtrTy()}, + /*isVarArg=*/false); + llvm::Function* function = llvm::Function::Create( + function_type, llvm::Function::ExternalLinkage, "test", module); + function->getArg(0)->setName("run_options"); + llvm::BasicBlock* block = + llvm::BasicBlock::Create(context, "entry", function); + builder.SetInsertPoint(block); + + llvm::Value* value = GetBatchDimByName(&builder); + + auto* call = llvm::dyn_cast(value); + ASSERT_NE(call, nullptr); + ASSERT_NE(call->getCalledFunction(), nullptr); + EXPECT_EQ(call->getCalledFunction()->getName(), + xla::cpu::runtime::kGetBatchSizeSymbolName); +} + +} // namespace +} // namespace llvm_ir +} // namespace xla diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc index 13f6a67764b346..20c2b9098cec17 100644 --- a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc +++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc @@ -28,11 +28,13 @@ limitations under the License. #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Constants.h" #include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Module.h" #include "llvm/IR/Type.h" #include "llvm/IR/Value.h" #include "xla/layout_util.h" #include "xla/service/llvm_ir/ir_array.h" #include "xla/service/llvm_ir/llvm_loop.h" +#include "xla/service/llvm_ir/llvm_util.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/platform/errors.h" @@ -183,6 +185,33 @@ std::vector LoopEmitter::EmitIndexAndSetExitBasicBlock( ForLoopNest loop_nest(loop_name, b_); + llvm::LLVMContext& ctx = b_->getContext(); + llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx); + + llvm::PointerType* ptr = llvm::PointerType::getUnqual(ctx); + llvm::StructType* callFrameTy = llvm::StructType::create( + "XLA_CPU_KernelArg", ptr, ptr, i64Type, ptr, i64Type); + + std::vector dynamic_dims; + for (auto dim : shape_.dimensions()) { + dynamic_dims.push_back(llvm::ConstantInt::get(i64Type, dim)); + } + + bool dynamic = false; + for (int i = 0; i < shape_.dimensions_size(); i++) { + auto expr = shape_.expressions(i); + if (expr && expr->is_dynamic()) { + dynamic_dims[i] = xla::llvm_ir::EmitExpression(b_, expr); + shape_.set_dynamic_dimension(i, true); + dynamic = true; + } + } + + if (dynamic) { + // Assign dynamic batch + dynamic_dims_ = dynamic_dims; + } + IrArray::Index array_index = dynamic_dims_.empty() ? EmitStaticIndex(&loop_nest, index_type) : EmitDynamicIndex(&loop_nest, index_type); diff --git a/third_party/xla/xla/service/llvm_ir/tuple_ops.cc b/third_party/xla/xla/service/llvm_ir/tuple_ops.cc index b4b3fa30affbcb..14bf83174ca269 100644 --- a/third_party/xla/xla/service/llvm_ir/tuple_ops.cc +++ b/third_party/xla/xla/service/llvm_ir/tuple_ops.cc @@ -104,7 +104,7 @@ llvm::Value* EmitGetTupleElement(const Shape& target_shape, int64_t index, llvm::LoadInst* src_buffer = b->CreateLoad(element_pointee_type, element_ptr); // Mark the loaded pointer as dereferenceable if we know its shape. - if (!target_shape.IsOpaque()) { + if (!target_shape.IsOpaque() && !target_shape.has_dynamic_expr()) { SetDereferenceableMetadataForLoad( src_buffer, ByteSizeOf(target_shape, src_buffer->getModule()->getDataLayout())); From 4534365df9575f8f7d4512858031501073bab95e Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:05:36 +0100 Subject: [PATCH 16/21] Support dynamic dimensions in CPU operation emitters Teach CPU emitters to use runtime dynamic dimensions for affected operations instead of fixed padded extents. Preserve the padded allocation while limiting computation to the real shape. --- third_party/xla/xla/service/cpu/BUILD | 1 + .../xla/xla/service/cpu/cpu_compiler.cc | 24 +++- .../xla/service/cpu/cpu_instruction_fusion.cc | 21 ++++ .../xla/xla/service/cpu/dot_op_emitter.cc | 101 +++++++++++++-- .../service/cpu/tests/cpu_dyn_shape_test.cc | 115 +++++++++++++++++- 5 files changed, 247 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index eef57b0e88409e..897c3367566135 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -295,6 +295,7 @@ cc_library( "//xla/mlir_hlo:all_passes", "//xla/mlir_hlo:transforms_passes", "//xla/service:all_reduce_promotion", + "//xla/service:dynamic_constant_rewriter", "//xla/service:all_to_all_decomposer", "//xla/service:batched_gather_scatter_normalizer", "//xla/service:batchnorm_expander", diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index bdf2d19d9492e8..b5ea2ac9c54739 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -40,6 +40,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" @@ -152,6 +153,7 @@ limitations under the License. #include "xla/map_util.h" #include "xla/mlir_hlo/transforms/passes.h" #include "xla/service/all_reduce_promotion.h" +#include "xla/service/dynamic_constant_rewriter.h" #include "xla/service/all_to_all_decomposer.h" #include "xla/service/batched_gather_scatter_normalizer.h" #include "xla/service/batchnorm_expander.h" @@ -450,6 +452,23 @@ void AddHloVerifier(HloPassPipeline* pipeline, HloVerifierOpts&& opts = {}, } } +bool DynamicConstantRewriterEnabled(const HloModule* module) { + const auto& extra_options = + module->config().debug_options().xla_backend_extra_options(); + auto it = extra_options.find("xla_cpu_disable_dynamic_constant_rewriter"); + if (it == extra_options.end()) { + return true; + } + bool disabled = false; + if (!absl::SimpleAtob(it->second, &disabled)) { + LOG(WARNING) + << "Ignoring invalid xla_cpu_disable_dynamic_constant_rewriter value: " + << it->second; + return true; + } + return !disabled; +} + std::unique_ptr> CreateSimplificationPipeline( absl::string_view name, HloModule* module, bool is_fusion_emitters) { // Run the following passes to a fixed point. @@ -480,7 +499,7 @@ std::unique_ptr> CreateSimplificationPipeline( } // Needs to happen after algebraic simplifier. - pipeline->AddPass(); + // pipeline->AddPass(); // BatchNormExpander can create zero-sized ops, so zero-sized HLO // elimination has to come after that pass. @@ -500,6 +519,9 @@ std::unique_ptr> CreateSimplificationPipeline( options::FoldAllConstants(module->config()) ? HloConstantFolding::Level::kAggressive : HloConstantFolding::Level::kDefault); + if (DynamicConstantRewriterEnabled(module)) { + pipeline->AddPass(); + } pipeline->AddPass(); return pipeline; diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc index bd26c1c87621cd..54c34104f4b7b0 100644 --- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc +++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc @@ -59,6 +59,23 @@ bool IsNonComplexNonBatchedMatrixVectorDot(const HloInstruction* hlo) { hlo->dot_dimension_numbers().lhs_batch_dimensions_size() == 0; } +bool HasDynamicDimensions(const Shape& shape) { + for (int64_t i = 0; i < shape.dimensions().size(); ++i) { + if (shape.is_dynamic_dimension(i) || + shape.expressions(i)->is_dynamic()) { + return true; + } + } + return false; +} + +bool IsDynamicDot(const HloInstruction* hlo) { + return hlo->opcode() == HloOpcode::kDot && + (HasDynamicDimensions(hlo->shape()) || + HasDynamicDimensions(hlo->operand(0)->shape()) || + HasDynamicDimensions(hlo->operand(1)->shape())); +} + bool HasExactlyOneUse(const HloInstruction& hlo_instr) { return hlo_instr.user_count() == 1 && absl::c_count(hlo_instr.users().front()->operands(), &hlo_instr) == 1; @@ -142,6 +159,10 @@ FusionDecision CpuInstructionFusion::ShouldFuse(HloInstruction* consumer, return FusionDecision::Forbid("Don't fuse large constants."); } + if (IsDynamicDot(producer) || IsDynamicDot(consumer)) { + return FusionDecision::Forbid("Do not fuse dynamic dots on CPU."); + } + if (CanBeOutputFused(producer, consumer)) { VLOG(2) << "Fusion OK: Can create output fusion."; return FusionDecision::Allow(); diff --git a/third_party/xla/xla/service/cpu/dot_op_emitter.cc b/third_party/xla/xla/service/cpu/dot_op_emitter.cc index e2e4f914579dbb..8b877429123cbf 100644 --- a/third_party/xla/xla/service/cpu/dot_op_emitter.cc +++ b/third_party/xla/xla/service/cpu/dot_op_emitter.cc @@ -160,6 +160,24 @@ bool CanEmitTiledLlvmIrGemm( return true; } +bool HasDynamicMatmulDims(const DotInfo& dot_info) { + const Shape& lhs_shape = dot_info.lhs_shape; + const Shape& rhs_shape = dot_info.rhs_shape; + const DotDimensionNumbers& dim_nums = dot_info.dim_nums; + + DExpr m_expr = lhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : lhs_shape.expressions( + 1LL - dim_nums.lhs_contracting_dimensions(0)); + DExpr k_expr = lhs_shape.expressions(dim_nums.lhs_contracting_dimensions(0)); + DExpr n_expr = rhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : rhs_shape.expressions( + 1LL - dim_nums.rhs_contracting_dimensions(0)); + + return m_expr->is_dynamic() || k_expr->is_dynamic() || n_expr->is_dynamic(); +} + // Returns dot implementation strategy for non-batch dot operations. DotImplementationStrategy GetNonBatchDotImplementationStrategy( const HloModuleConfig& config, const DotInfo& dot_info, @@ -173,6 +191,10 @@ DotImplementationStrategy GetNonBatchDotImplementationStrategy( dot_info.dim_nums.rhs_batch_dimensions_size() == 0) << "Dot operations must be non-batch"; + if (HasDynamicMatmulDims(dot_info)) { + return DotImplementationStrategy::kEigen; + } + // Any Matrix-Vector product of floating point or integral type, or // a transpose-dot fusion of the same can be lowered to a tiled LLVM // IR implementation. @@ -253,6 +275,12 @@ class DotOpEmitter { // The number of columns on the RHS. int64_t n; + DExpr m_expr; + + DExpr k_expr; + + DExpr n_expr; + // True if the LHS matrix is column major. bool lhs_column_major; @@ -858,16 +886,23 @@ absl::Status DotOpEmitter::EmitCallToRuntime() { if (!mat_mult_dims.lhs_column_major) { std::swap(mat_mult_dims.m, mat_mult_dims.n); + std::swap(mat_mult_dims.m_expr, mat_mult_dims.n_expr); std::swap(lhs, rhs); std::swap(transpose_lhs, transpose_rhs); } - b_->CreateCall(matmul_func, - {executable_run_options_value_, target_array_.GetBasePointer(), - lhs->GetBasePointer(), rhs->GetBasePointer(), - b_->getInt64(mat_mult_dims.m), b_->getInt64(mat_mult_dims.n), - b_->getInt64(mat_mult_dims.k), b_->getInt32(transpose_lhs), - b_->getInt32(transpose_rhs)}); + llvm::Value* m_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.m_expr); + llvm::Value* n_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.n_expr); + llvm::Value* k_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.k_expr); + + b_->CreateCall( + matmul_func, + {executable_run_options_value_, target_array_.GetBasePointer(), + lhs->GetBasePointer(), rhs->GetBasePointer(), m_val, n_val, k_val, + b_->getInt32(transpose_lhs), b_->getInt32(transpose_rhs)}); return absl::OkStatus(); } @@ -942,18 +977,30 @@ absl::Status DotOpEmitter::EmitCallToBatchRuntime() { if (!mat_mult_dims.lhs_column_major) { std::swap(mat_mult_dims.m, mat_mult_dims.n); + std::swap(mat_mult_dims.m_expr, mat_mult_dims.n_expr); std::swap(lhs, rhs); std::swap(transpose_lhs, transpose_rhs); } + llvm::Value* m_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.m_expr); + llvm::Value* n_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.n_expr); + llvm::Value* k_val = + xla::llvm_ir::EmitExpression(b_, mat_mult_dims.k_expr); + DExpr batch_size_expr = + lhs_shape.expressions(0) ? lhs_shape.expressions(0) + : DExpr::Const(lhs_shape.dimensions(0)); + llvm::Value* batch_size_val = + xla::llvm_ir::EmitExpression(b_, batch_size_expr); + VLOG(1) << "Batch dot emitted with runtime:" << fn_name; b_->CreateCall( matmul_func, {executable_run_options_value_, target_array_.GetBasePointer(), - lhs->GetBasePointer(), rhs->GetBasePointer(), - b_->getInt64(mat_mult_dims.m), b_->getInt64(mat_mult_dims.n), - b_->getInt64(mat_mult_dims.k), b_->getInt64(lhs_shape.dimensions(0)), + lhs->GetBasePointer(), rhs->GetBasePointer(), m_val, n_val, k_val, + batch_size_val, b_->getInt32(static_cast(transpose_lhs)), b_->getInt32(static_cast(transpose_rhs))}); return absl::OkStatus(); @@ -983,6 +1030,14 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const { /*n=*/rhs_shape.dimensions().size() <= 1 ? 1LL : rhs_shape.dimensions(1LL - dim_nums.rhs_contracting_dimensions(0)), + /*m_expr=*/lhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : lhs_shape.expressions(1LL - dim_nums.lhs_contracting_dimensions(0)), + /*k_expr=*/ + lhs_shape.expressions(dim_nums.lhs_contracting_dimensions(0)), + /*n_expr=*/rhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : rhs_shape.expressions(1LL - dim_nums.rhs_contracting_dimensions(0)), /*lhs_column_major=*/is_column_major(lhs_shape), /*lhs_canonical=*/lhs_shape.dimensions().size() <= 1 || dim_nums.lhs_contracting_dimensions(0) == 1, @@ -1014,6 +1069,14 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetBatchMatMultDims() const { /*n=*/rhs_shape.dimensions().size() <= 1 ? 1LL : rhs_shape.dimensions(2LL - dim_nums.rhs_contracting_dimensions(0)), + /*m_expr=*/lhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : lhs_shape.expressions(2LL - dim_nums.lhs_contracting_dimensions(0)), + /*k_expr=*/ + lhs_shape.expressions(1LL + dim_nums.lhs_contracting_dimensions(0)), + /*n_expr=*/rhs_shape.dimensions().size() <= 1 + ? DExpr::Const(1) + : rhs_shape.expressions(2LL - dim_nums.rhs_contracting_dimensions(0)), /*lhs_column_major=*/is_column_major(lhs_shape), /*lhs_canonical=*/lhs_shape.dimensions().size() <= 1 || dim_nums.lhs_contracting_dimensions(0) == 1, @@ -1093,22 +1156,34 @@ absl::Status EmitNonBatchDotOperation( Shape DropFirstDim(const Shape& shape) { absl::Span array_shape_dims(shape.dimensions()); + absl::Span array_shape_exprs(shape.expressions()); array_shape_dims.remove_prefix(1); - return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(), - array_shape_dims); + array_shape_exprs.remove_prefix(1); + return ShapeUtil::MakeShapeWithDescendingLayout( + shape.element_type(), array_shape_dims, array_shape_exprs); } Shape CollapseFirstNDims(const Shape& shape, int64_t n) { absl::Span input_shape_dims(shape.dimensions()); + absl::Span input_expressions(shape.expressions()); int64_t prefix_dim = std::accumulate(input_shape_dims.begin(), input_shape_dims.begin() + n, 1ll, std::multiplies()); + + DExpr prefix_expression = std::accumulate( + input_expressions.begin(), input_expressions.begin() + n, DExpr::Const(1), + [](DExpr acc, const DExpr& v) { return acc * v; }); + DimensionVector result_dims; + std::vector result_expressions; result_dims.push_back(prefix_dim); + result_expressions.push_back(prefix_expression); std::copy(input_shape_dims.begin() + n, input_shape_dims.end(), std::back_inserter(result_dims)); - return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(), - result_dims); + std::copy(input_expressions.begin() + n, input_expressions.end(), + std::back_inserter(result_expressions)); + return ShapeUtil::MakeShapeWithDescendingLayout( + shape.element_type(), result_dims, result_expressions); } llvm_ir::IrArray CollapseFirstNDims(llvm::IRBuilderBase* b, diff --git a/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc index f0148de9932561..145beb3948b22d 100644 --- a/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc +++ b/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc @@ -18,6 +18,8 @@ limitations under the License. #include #include +#include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" @@ -32,7 +34,68 @@ namespace xla { namespace cpu { namespace { -using CpuDynamicShapeTest = CpuCodegenTest; +CpuAotCompilationOptions GetAotOptions() { + return CpuAotCompilationOptions{ + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, + /*features=*/"", + /*entry_point_name=*/"entry", + /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; +} + +Shape MakeDynamicDenseShape(absl::Span dimensions, + absl::Span minor_to_major, + int64_t dynamic_dimension) { + Shape shape = + ShapeUtil::MakeShapeWithDenseLayout(F32, dimensions, minor_to_major); + shape.set_dynamic_dimension(dynamic_dimension, true); + return shape; +} + +class CpuDynamicShapeTest : public CpuCodegenTest { + protected: + std::unique_ptr CreateDynamicFastConcatModule( + absl::string_view name, bool add_consumer, bool three_operands = false) { + HloComputation::Builder builder(name); + + Shape operand_shape = MakeDynamicDenseShape({64, 26, 4}, {2, 1, 0}, + /*dynamic_dimension=*/0); + Shape concat_shape = + MakeDynamicDenseShape({64, 26, three_operands ? 12 : 8}, {2, 1, 0}, + /*dynamic_dimension=*/0); + + HloInstruction* lhs = builder.AddInstruction( + HloInstruction::CreateParameter(0, operand_shape, "lhs")); + HloInstruction* rhs = builder.AddInstruction( + HloInstruction::CreateParameter(1, operand_shape, "rhs")); + + HloInstruction* concat = nullptr; + if (three_operands) { + HloInstruction* third = builder.AddInstruction( + HloInstruction::CreateParameter(2, operand_shape, "third")); + concat = builder.AddInstruction(HloInstruction::CreateConcatenate( + concat_shape, {lhs, rhs, third}, /*dimension=*/2)); + } else { + concat = builder.AddInstruction(HloInstruction::CreateConcatenate( + concat_shape, {lhs, rhs}, /*dimension=*/2)); + } + + if (add_consumer) { + int64_t addend_param_number = three_operands ? 3 : 2; + HloInstruction* addend = + builder.AddInstruction(HloInstruction::CreateParameter( + addend_param_number, concat_shape, "addend")); + builder.AddInstruction(HloInstruction::CreateBinary( + concat_shape, HloOpcode::kAdd, concat, addend)); + } + + auto hlo_module = CreateNewVerifiedModule(); + hlo_module->AddEntryComputation(builder.Build()); + auto& debug_options = hlo_module->mutable_config().mutable_debug_options(); + debug_options.set_xla_cpu_use_thunk_runtime(false); + debug_options.add_xla_disable_hlo_passes("fusion"); + return hlo_module; + } +}; TEST_F(CpuDynamicShapeTest, DynamicShapeR2) { HloComputation::Builder builder(TestName()); @@ -70,6 +133,56 @@ TEST_F(CpuDynamicShapeTest, DynamicShapeR2) { /*match_optimized_ir=*/false); } +TEST_F(CpuDynamicShapeTest, DynamicFastConcatenateCompilesWithoutFusion) { + auto hlo_module = + CreateDynamicFastConcatModule(TestName(), /*add_consumer=*/false); + + std::string filecheck_pattern = R"( +; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr +; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64 +; CHECK: br label %concatenate.loop_header.concat.1 +; CHECK: call void @llvm.memcpy + )"; + + CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(), + filecheck_pattern, + /*match_optimized_ir=*/false); +} + +TEST_F(CpuDynamicShapeTest, + DynamicFastConcatenateWithAddCompilesWithoutFusion) { + auto hlo_module = + CreateDynamicFastConcatModule(TestName(), /*add_consumer=*/true); + + std::string filecheck_pattern = R"( +; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr +; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64 +; CHECK: br label %concatenate.loop_header.concat.1 +; CHECK: call void @llvm.memcpy + )"; + + CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(), + filecheck_pattern, + /*match_optimized_ir=*/false); +} + +TEST_F(CpuDynamicShapeTest, + DynamicFastConcatenateWithThreeOperandsCompilesWithoutFusion) { + auto hlo_module = CreateDynamicFastConcatModule( + TestName(), /*add_consumer=*/true, /*three_operands=*/true); + + std::string filecheck_pattern = R"( +; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr +; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64 +; CHECK: br label %concatenate.loop_header.concat.1 +; CHECK: call void @llvm.memcpy + )"; + + CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(), + filecheck_pattern, + /*match_optimized_ir=*/false); +} + } // namespace } // namespace cpu } // namespace xla From 3fc64732d4c56bb6ef2522f320dd8401b5c9d7f0 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:05:55 +0100 Subject: [PATCH 17/21] Prevent CPU HLO profiling crashes Make CPU HLO profiling tolerate dynamic execution metadata and missing profile counters. Avoid invalid accesses while retaining profiling for supported instructions. --- third_party/xla/xla/service/cpu/BUILD | 2 ++ .../xla/xla/service/cpu/cpu_compiler.cc | 27 +++++++++++--- .../xla/xla/service/cpu/cpu_executable.cc | 34 +++++++++++++++++- .../xla/xla/service/cpu/cpu_runtime.cc | 8 +++++ third_party/xla/xla/service/cpu/cpu_runtime.h | 2 ++ third_party/xla/xla/service/cpu/ir_emitter.cc | 35 +++++++++++++++---- .../service/cpu/runtime_symbol_generator.cc | 1 + 7 files changed, 98 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 897c3367566135..8ea07dccc5a25e 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -691,6 +691,7 @@ cc_library( "//xla/service:custom_call_status_internal", "//xla/service:executable", "//xla/service:hlo_execution_profile", + "//xla/service:hlo_profile_printer", "//xla/service:hlo_profile_printer_data_cc", "//xla/service:hlo_value", "//xla/service:maybe_owning_device_memory", @@ -1126,6 +1127,7 @@ cc_library( "//xla/stream_executor:stream_executor_h", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:errors", + "//xla/tsl/platform:env_time", "//xla/tsl/platform:logging", "//xla/tsl/platform:status", "@com_google_absl//absl/algorithm:container", diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index b5ea2ac9c54739..5d04d299cecfc2 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -1073,6 +1073,25 @@ absl::Status CreateHloProfilingArtifacts( return absl::OkStatus(); } +bool ShouldCreateHloProfilingArtifacts(const HloModule& module) { + if (!module.config().hlo_profiling_enabled()) { + return false; + } + + // The thunk runtime launches host kernels through XLA_CPU_KernelCallFrame, + // which currently has no profile-counters field. Emitting profiling loads and + // stores in this mode can make generated code dereference a null counter + // pointer at run time. + if (module.config().debug_options().xla_cpu_use_thunk_runtime()) { + LOG(WARNING) << "--xla_hlo_profile is not supported by XLA:CPU thunk " + "runtime; disabling HLO profiling for module " + << module.name(); + return false; + } + + return true; +} + } // namespace absl::StatusOr> CpuCompiler::RunHloPasses( @@ -1461,7 +1480,7 @@ CpuCompiler::CompileCpuExecutable(std::unique_ptr module) { computation_to_profile_idx; std::unique_ptr hlo_profile_index_map; std::unique_ptr hlo_profile_printer_data; - if (module->config().hlo_profiling_enabled()) { + if (ShouldCreateHloProfilingArtifacts(*module)) { TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts( *module, &instruction_to_profile_idx, &computation_to_profile_idx, &hlo_profile_index_map, &hlo_profile_printer_data)); @@ -2013,7 +2032,7 @@ CpuCompiler::CompileAheadOfTimeLegacy( std::unique_ptr hlo_profile_index_map; std::unique_ptr hlo_profile_printer_data; - if (module->config().hlo_profiling_enabled()) { + if (ShouldCreateHloProfilingArtifacts(*module)) { TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts( *module, &instruction_to_profile_idx, &computation_to_profile_idx, &hlo_profile_index_map, &hlo_profile_printer_data)); @@ -2171,7 +2190,7 @@ CpuCompiler::CompileAheadOfTimeThunks( computation_to_profile_idx; std::unique_ptr hlo_profile_index_map; std::unique_ptr hlo_profile_printer_data; - if (module->config().hlo_profiling_enabled()) { + if (ShouldCreateHloProfilingArtifacts(*module)) { TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts( *module, &instruction_to_profile_idx, &computation_to_profile_idx, &hlo_profile_index_map, &hlo_profile_printer_data)); @@ -2420,7 +2439,7 @@ CpuCompiler::CompileAheadOfTimeThunks( cpu_executable->thunks().thunk_sequence(); std::unique_ptr executable_hlo_profile_printer_data = - cpu_executable->module().config().hlo_profiling_enabled() + cpu_executable->hlo_profiling_enabled() ? std::make_unique( cpu_executable->hlo_profile_printer_data()) : nullptr; diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc index 853aa3519924c1..0307561329ee27 100644 --- a/third_party/xla/xla/service/cpu/cpu_executable.cc +++ b/third_party/xla/xla/service/cpu/cpu_executable.cc @@ -54,6 +54,7 @@ limitations under the License. #include "xla/service/custom_call_status_internal.h" #include "xla/service/executable.h" #include "xla/service/hlo_execution_profile.h" +#include "xla/service/hlo_profile_printer.h" #include "xla/service/hlo_profile_printer_data.pb.h" #include "xla/service/hlo_value.h" #include "xla/service/maybe_owning_device_memory.h" @@ -76,6 +77,20 @@ limitations under the License. namespace xla { namespace cpu { +namespace { + +float GetClockRateGhz(const ExecutableRunOptions* run_options) { + if (run_options->stream() == nullptr || + run_options->stream()->parent() == nullptr) { + return 1.0f; + } + float clock_rate_ghz = + run_options->stream()->parent()->GetDeviceDescription().clock_rate_ghz(); + return clock_rate_ghz > 0 ? clock_rate_ghz : 1.0f; +} + +} // namespace + absl::StatusOr> CpuExecutable::Create( std::unique_ptr function_library, std::unique_ptr assignment, @@ -246,8 +261,16 @@ absl::Status CpuExecutable::ExecuteComputeFunction( absl::Span buffers) { uint64_t start_micros = tsl::Env::Default()->NowMicros(); - size_t profile_counters_size = 0; + std::vector profile_counter_storage; + if (hlo_profiling_enabled()) { + profile_counter_storage.assign( + hlo_profile_printer_data().profile_counters_size(), 0); + } + size_t profile_counters_size = profile_counter_storage.size(); int64_t* profile_counters = nullptr; + if (!profile_counter_storage.empty()) { + profile_counters = profile_counter_storage.data(); + } // Call the computation function following the calling convention. See the // definition of 'ComputeFunctionType' for the details of the calling @@ -286,6 +309,15 @@ absl::Status CpuExecutable::ExecuteComputeFunction( compute_function_(nullptr, run_options, nullptr, buffer_pointers.data(), &status, profile_counters); record_profile(); + if (profile_counters != nullptr) { + std::string hlo_profile = + PrintHloProfile(hlo_profile_printer_data(), profile_counters, + GetClockRateGhz(run_options)); + if (!hlo_profile.empty()) { + LOG(INFO) << "XLA:CPU HLO profile for " << module().name() << "\n" + << hlo_profile; + } + } std::optional error_message = CustomCallStatusGetMessage(&status); if (error_message) { diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.cc b/third_party/xla/xla/service/cpu/cpu_runtime.cc index 7caf9c43b1119b..20c482f663b6c6 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime.cc +++ b/third_party/xla/xla/service/cpu/cpu_runtime.cc @@ -60,6 +60,7 @@ limitations under the License. #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/env_time.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/status.h" #include "xla/util.h" @@ -166,6 +167,8 @@ extern const char* const kParallelForkJoinSymbolName = "__xla_cpu_runtime_ParallelForkJoin"; extern const char* const kPrintfToStderrSymbolName = "__xla_cpu_runtime_PrintfToStderr"; +extern const char* const kReadCycleCounterSymbolName = + "__xla_cpu_runtime_ReadCycleCounter"; extern const char* const kStatusIsSuccessSymbolName = "__xla_cpu_runtime_StatusIsSuccess"; extern const char* const kKeyValueSortSymbolName = @@ -620,6 +623,11 @@ ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int __xla_cpu_runtime_PrintfToStderr( return result; } +ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY uint64_t +__xla_cpu_runtime_ReadCycleCounter() { + return tsl::EnvTime::NowNanos(); +} + ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int64_t __xla_cpu_runtime_TracingStart( const void* /* ExecutableRunOptions* run_options_ptr*/, const char* name, const char* hlo_module, int64_t program_id) { diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h index 71e27ea600ee28..de1ea93a855a35 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime.h +++ b/third_party/xla/xla/service/cpu/cpu_runtime.h @@ -79,6 +79,7 @@ extern const char* const kAcquireOutfeedBufferForPopulationSymbolName; extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName; extern const char* const kParallelForkJoinSymbolName; extern const char* const kPrintfToStderrSymbolName; +extern const char* const kReadCycleCounterSymbolName; extern const char* const kStatusIsSuccessSymbolName; extern const char* const kKeyValueSortSymbolName; extern const char* const kTopKF32SymbolName; @@ -115,6 +116,7 @@ int GetDeviceOrdinal(const xla::ExecutableRunOptions* run_options); extern "C" { extern int __xla_cpu_runtime_PrintfToStderr(const char* format, ...); +extern uint64_t __xla_cpu_runtime_ReadCycleCounter(); extern int64_t __xla_cpu_runtime_TracingStart( const void* /* xla::ExecutableRunOptions* */ run_options_ptr, diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc index ae40e76e8e2ce2..bdf1bc2e60f4df 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter.cc +++ b/third_party/xla/xla/service/cpu/ir_emitter.cc @@ -3715,10 +3715,26 @@ void IrEmitter::ProfilingState::UpdateProfileCounter(llvm::IRBuilderBase* b, llvm::Value* prof_counter, llvm::Value* cycle_end, llvm::Value* cycle_start) { + llvm::Value* profile_counters = prof_counter; + if (auto* gep = llvm::dyn_cast(prof_counter)) { + profile_counters = gep->getPointerOperand(); + } + + auto* profile_counters_type = + llvm::cast(profile_counters->getType()); + llvm::Value* has_profile_counters = b->CreateICmpNE( + profile_counters, llvm::ConstantPointerNull::get(profile_counters_type), + "has_profile_counters"); + + llvm::AllocaInst* dummy_counter = llvm_ir::EmitAllocaAtFunctionEntry( + b->getInt64Ty(), "dummy_profile_counter", b); + b->CreateStore(b->getInt64(0), dummy_counter); + prof_counter = + b->CreateSelect(has_profile_counters, prof_counter, dummy_counter); + auto* cycle_diff = b->CreateSub(cycle_end, cycle_start); llvm::LoadInst* old_cycle_count = b->CreateLoad( - llvm::cast(prof_counter)->getSourceElementType(), - prof_counter, "old_cycle_count"); + b->getInt64Ty(), prof_counter, "old_cycle_count"); auto* new_cycle_count = b->CreateAdd(cycle_diff, old_cycle_count, "new_cycle_count"); b->CreateStore(new_cycle_count, prof_counter); @@ -3728,10 +3744,17 @@ llvm::Value* IrEmitter::ProfilingState::ReadCycleCounter( llvm::IRBuilderBase* b) { llvm::Module* module = b->GetInsertBlock()->getModule(); if (!use_rdtscp_) { - llvm::Function* func_llvm_readcyclecounter = - llvm::Intrinsic::getOrInsertDeclaration( - module, llvm::Intrinsic::readcyclecounter); - return b->CreateCall(func_llvm_readcyclecounter); + llvm::FunctionType* fn_type = + llvm::FunctionType::get(b->getInt64Ty(), /*isVarArg=*/false); + llvm::FunctionCallee read_cycle_counter_func = + module->getOrInsertFunction(runtime::kReadCycleCounterSymbolName, + fn_type); + if (auto* fn = + llvm::dyn_cast(read_cycle_counter_func.getCallee())) { + fn->setCallingConv(llvm::CallingConv::C); + fn->setDoesNotThrow(); + } + return b->CreateCall(read_cycle_counter_func); } llvm::Function* func_llvm_x86_rdtscp = llvm::Intrinsic::getOrInsertDeclaration(module, diff --git a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc index a87b9b3ae986ab..366655036cc038 100644 --- a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc +++ b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc @@ -203,6 +203,7 @@ static bool RegisterKnownJITSymbols() { REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulU8); REGISTER_CPU_RUNTIME_SYMBOL(ParallelForkJoin); REGISTER_CPU_RUNTIME_SYMBOL(PrintfToStderr); + REGISTER_CPU_RUNTIME_SYMBOL(ReadCycleCounter); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseInfeedBufferAfterDequeue); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseOutfeedBufferAfterPopulation); REGISTER_CPU_RUNTIME_SYMBOL(StatusIsSuccess); From 0871cde560ea3ac94f56bd3645f1a56a35a493ef Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 10:06:13 +0100 Subject: [PATCH 18/21] Add dynamic expression integration tests Add end-to-end coverage for symbolic shape propagation, clustering, compilation, and CPU execution. Exercise representative dynamic models across the complete TensorFlow-to-XLA path. --- .../compiler/tf2xla/xla_compiler_test.cc | 3187 ++++++++++++++++- 1 file changed, 3044 insertions(+), 143 deletions(-) diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 1e598502ad65d0..bf2728b693c9d8 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -16,8 +16,10 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include +#include #include "absl/strings/match.h" #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/const_op.h" #include "tensorflow/cc/ops/data_flow_ops.h" @@ -57,8 +59,11 @@ limitations under the License. #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape_expr.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/grappler/costs/graph_properties.h" +#include "tensorflow/core/grappler/grappler_item.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/kernels/ops_testutil.h" @@ -96,6 +101,42 @@ class XlaCompilerTest : public ::testing::Test { std::unique_ptr flib_def_; }; +class ScopedTfXlaDynamicSizesFlag { + public: + ScopedTfXlaDynamicSizesFlag() { + old_value_ = GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes; + old_symbolic_content_value_ = + GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content; + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes = true; + GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content = true; + SetTensorShapeExpressionsEnabledForTesting(true); + } + + ~ScopedTfXlaDynamicSizesFlag() { + SetTensorShapeExpressionsEnabledForTesting(std::nullopt); + GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes = old_value_; + GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content = + old_symbolic_content_value_; + } + + private: + bool old_value_ = false; + bool old_symbolic_content_value_ = false; +}; + +class XlaCompilerDynamicSizesTest : public XlaCompilerTest { + protected: + void SetUp() override { + dynamic_sizes_flag_ = std::make_unique(); + XlaCompilerTest::SetUp(); + } + + void TearDown() override { dynamic_sizes_flag_.reset(); } + + private: + std::unique_ptr dynamic_sizes_flag_; +}; + namespace { // Helper class to test the ability to pass resources through to XLA @@ -106,194 +147,3000 @@ class DummyResourceForTest : public ResourceBase { void Increment() { ++value_; } int Get() { return value_; } - private: - int value_ = 0; -}; + private: + int value_ = 0; +}; + +class DummyReadResourceOp : public XlaOpKernel { + public: + explicit DummyReadResourceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + void Compile(XlaOpKernelContext* ctx) override { + ResourceMgr* rm = ctx->op_kernel_context()->resource_manager(); + OP_REQUIRES(ctx, rm, errors::Internal("No resource manager.")); + DummyResourceForTest* dummy; + OP_REQUIRES_OK(ctx, rm->Lookup( + rm->default_container(), "dummy", &dummy)); + dummy->Increment(); + dummy->Unref(); + + ctx->SetOutput(0, ctx->Input(0)); + ctx->SetOutput(1, ctx->Input(0)); + } +}; + +class DummyReadResourceCC { + public: + DummyReadResourceCC(const Scope& scope, const Input& value) { + if (!scope.ok()) return; + auto _value = ops::AsNodeOut(scope, value); + if (!scope.ok()) return; + Node* ret; + const auto unique_name = scope.GetUniqueNameForOp("DummyReadResource"); + auto builder = NodeBuilder(unique_name, "DummyReadResource").Input(_value); + scope.UpdateBuilder(&builder); + scope.UpdateStatus(builder.Finalize(scope.graph(), &ret)); + if (!scope.ok()) return; + scope.UpdateStatus(scope.DoShapeInference(ret)); + if (!scope.ok()) return; + this->output1_ = Output(ret, 0); + this->output2_ = Output(ret, 1); + } + + Output output1_; + Output output2_; +}; + +REGISTER_OP("DummyReadResource") + .Input("input: int32") + .Output("output1: int32") + .Output("output2: int32") + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +A dummy Op. + +input: dummy input. +output1: dummy output. +output2: dummy output. +)doc"); + +REGISTER_XLA_OP(Name("DummyReadResource"), DummyReadResourceOp); + +// DummyDuplicateOp is present purely to test multiple REGISTER_XLA_OP calls +// on the same Op name below. +class DummyDuplicateOp : public XlaOpKernel { + public: + explicit DummyDuplicateOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + void Compile(XlaOpKernelContext* ctx) override { + ctx->SetOutput(0, ctx->Input(0)); + } +}; + +REGISTER_OP("DummyDuplicateOp") + .Input("input: int32") + .Output("output: int32") + .Doc(R"doc( +A dummy Op. + +input: dummy input. +output: dummy output. +)doc"); + +REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT), + DummyDuplicateOp); +REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT), + DummyDuplicateOp); + +// Tests compilation and execution of an empty graph. +TEST_F(XlaCompilerTest, EmptyReturnValues) { + XlaCompiler compiler(DefaultOptions()); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), + /*args=*/{}, &result)); + + TF_ASSERT_OK(client_->Execute(*result.computation, {}).status()); +} + +// Tests compilation and execution of a graph that adds two tensors. +TEST_F(XlaCompilerTest, Simple) { + // Builds a graph that adds two Tensors. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1); + auto c = ops::Add(scope.WithOpName("C"), a, b); + auto d = ops::_Retval(scope.WithOpName("D"), c, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = TensorShape({2}); + + // Compiles the graph. + XlaCompiler compiler(DefaultOptions()); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); + + // Tests that the generated computation works. + xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); + xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); + + xla::Literal expected0 = xla::LiteralUtil::CreateR1({4, 143}); + xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); +} + +absl::StatusOr> LoadModuleFromHloProto( + const xla::HloModuleProto& module_proto) { + TF_ASSIGN_OR_RETURN(auto module_config, + xla::HloModule::CreateModuleConfigFromProto( + module_proto, xla::GetDebugOptionsFromFlags())); + return xla::CreateModuleFromProto(module_proto, module_config); +} + +// Tests compilation and execution of a graph that adds two tensors with dynamic +// shape parameters. +TEST_F(XlaCompilerTest, SimpleDynamicShapeParameter) { + // Builds a graph that adds two Tensors. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1); + auto c = ops::Add(scope.WithOpName("C"), a, b); + auto d = ops::_Retval(scope.WithOpName("D"), c, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = + xla::ShapeUtil::MakeShape(/*element_type=*/xla::S32, /*dimensions=*/{2}, + /*dynamic_dimensions=*/std::vector{true}, + /*expressions=*/{}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = TensorShape(/*dimensions=*/{2}); + + // Compiles the graph. + XlaCompiler compiler(DefaultOptions()); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); + + auto hlo = result.computation->proto(); + TF_ASSERT_OK_AND_ASSIGN(auto module, LoadModuleFromHloProto(hlo)); + EXPECT_EQ(module->computation_count(), 1); + EXPECT_TRUE(module->mutable_computation(0) + ->parameter_instruction(0) + ->shape() + .is_dynamic()); +} + +TEST_F(XlaCompilerDynamicSizesTest, DynamicShapeParameterPreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto identity = ops::Identity(scope.WithOpName("identity"), input); + auto retval = ops::_Retval(scope.WithOpName("retval"), identity, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6}, std::vector{xla::DExpr::Var(1)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "identity", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(1))); + + TF_ASSERT_OK_AND_ASSIGN(auto module, + LoadModuleFromHloProto(result.computation->proto())); + const xla::Shape& param_shape = + module->entry_computation()->parameter_instruction(0)->shape(); + EXPECT_TRUE( + xla::DynExpr::equal(param_shape.expressions(0), xla::DExpr::Var(1))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(1))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + TensorFlowAndXlaAgreeOnReshapeExpression) { + Scope inference_scope = Scope::NewRootScope().ExitOnError(); + auto placeholder = ops::Placeholder( + inference_scope.WithOpName("placeholder"), DT_FLOAT, + ops::Placeholder::Shape(PartialTensorShape({-1, 24}))); + auto inference_shape = + ops::Const(inference_scope.WithOpName("shape"), {-1, 12}); + auto inferred_reshape = ops::Reshape( + inference_scope.WithOpName("reshape"), placeholder, inference_shape); + + grappler::GrapplerItem item; + TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph)); + grappler::GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true)); + + const TensorShapeProto& inferred_input_shape = + properties.GetOutputProperties("placeholder").at(0).shape(); + const TensorShapeProto& inferred_output_shape = + properties.GetOutputProperties("reshape").at(0).shape(); + ASSERT_EQ(inferred_input_shape.expressions_size(), 2); + ASSERT_EQ(inferred_output_shape.expressions_size(), 2); + const xla::DExpr inferred_batch_expr = + DimExprFromProto(inferred_input_shape.expressions(0)); + const xla::DExpr inferred_output_expr = + DimExprFromProto(inferred_output_shape.expressions(0)); + ASSERT_TRUE(inferred_batch_expr->is_dynamic()); + EXPECT_TRUE(xla::DynExpr::equal( + inferred_output_expr, inferred_batch_expr * xla::DExpr::Const(2))); + + Scope compilation_scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0); + auto compilation_shape = + ops::Const(compilation_scope.WithOpName("shape"), {-1, 12}); + auto reshape = ops::Reshape(compilation_scope.WithOpName("reshape"), arg, + compilation_shape); + auto retval = + ops::_Retval(compilation_scope.WithOpName("retval"), reshape, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(compilation_scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 24}, + std::vector{inferred_batch_expr, xla::DExpr::Const(24)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reshape", std::move(graph), + args, &result)); + + TF_ASSERT_OK_AND_ASSIGN(auto module, + LoadModuleFromHloProto(result.computation->proto())); + const xla::Shape& parameter_shape = + module->entry_computation()->parameter_instruction(0)->shape(); + EXPECT_TRUE(xla::DynExpr::equal(parameter_shape.expressions(0), + inferred_batch_expr)); + EXPECT_TRUE(xla::DynExpr::equal(parameter_shape.expressions(1), + xla::DExpr::Const(24))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_EQ(result_shape.dimensions(0), 16); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), + inferred_output_expr)); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), + xla::DExpr::Const(12))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + ShapeValueChainPreservesExpressionThroughTile) { + // Shape turns the dynamic input dimension into the symbolic value A. + // Slicing, squeezing, and packing must preserve it as {A, 1}, which Tile + // then uses as the shape of a repeated static row. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("arg"), DT_FLOAT, 0); + auto input_shape = ops::Shape(scope.WithOpName("input_shape"), arg); + auto begin = ops::Const(scope.WithOpName("begin"), {0}, {1}); + auto end = ops::Const(scope.WithOpName("end"), {1}, {1}); + auto strides = ops::Const(scope.WithOpName("strides"), {1}, {1}); + auto first_dim_vector = ops::StridedSlice( + scope.WithOpName("first_dim_vector"), input_shape, begin, end, strides); + auto first_dim = ops::Squeeze(scope.WithOpName("first_dim"), + first_dim_vector, ops::Squeeze::Axis({0})); + auto one = ops::Const(scope.WithOpName("one"), 1, {}); + auto multiples = ops::Stack(scope.WithOpName("multiples"), + std::vector{first_dim, one}); + auto base_row = ops::Const( + scope.WithOpName("base_row"), + {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}, + {1, 24}); + auto tile = ops::Tile(scope.WithOpName("tile"), base_row, multiples); + auto retval = ops::_Retval(scope.WithOpName("retval"), tile, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Compile with a physical bound of 8 while keeping its logical size as A. + const xla::DExpr input_expr = xla::DExpr::Var(1); + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 24}, + std::vector{input_expr, xla::DExpr::Const(24)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "tile", + std::move(graph), args, &result)); + + // The physical output remains bounded by 8, but its first dimension must + // still describe the runtime value A rather than the bound. + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_EQ(result_shape.dimensions(0), 8); + EXPECT_EQ(result_shape.dimensions(1), 24); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), input_expr)); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), + xla::DExpr::Const(24))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + SizeValueChainPreservesExpressionThroughRange) { + // Size returns a scalar, so the relationship to A is carried as symbolic + // value content rather than as a dimension expression. Range must consume + // that content when it infers the length of its one-dimensional output. + Scope inference_scope = Scope::NewRootScope().ExitOnError(); + auto placeholder = ops::Placeholder( + inference_scope.WithOpName("placeholder"), DT_FLOAT, + ops::Placeholder::Shape(PartialTensorShape({-1, 1}))); + auto size = ops::Size(inference_scope.WithOpName("size"), placeholder); + auto zero = ops::Const(inference_scope.WithOpName("zero"), 0, {}); + auto three = ops::Const(inference_scope.WithOpName("three"), 3, {}); + auto inferred_range = + ops::Range(inference_scope.WithOpName("range"), zero, size, three); + + grappler::GrapplerItem item; + TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph)); + grappler::GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true)); + + const TensorShapeProto& inferred_input_shape = + properties.GetOutputProperties("placeholder").at(0).shape(); + const TensorShapeProto& inferred_range_shape = + properties.GetOutputProperties("range").at(0).shape(); + ASSERT_EQ(inferred_input_shape.expressions_size(), 2); + ASSERT_EQ(inferred_range_shape.expressions_size(), 1); + const xla::DExpr input_expr = + DimExprFromProto(inferred_input_shape.expressions(0)); + const xla::DExpr expected_output_expr = + DimExprFromProto(inferred_range_shape.expressions(0)); + // With start 0 and delta 3, the number of elements is ceil(A / 3). + EXPECT_TRUE(xla::DynExpr::equal( + expected_output_expr, + ((input_expr + xla::DExpr::Const(2)) / xla::DExpr::Const(3)) + .simplify())); + + // Compile the equivalent graph and check the same expression reaches HLO. + Scope compilation_scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0); + auto compiled_size = ops::Size(compilation_scope.WithOpName("size"), arg); + auto compiled_zero = + ops::Const(compilation_scope.WithOpName("zero"), 0, {}); + auto compiled_three = + ops::Const(compilation_scope.WithOpName("three"), 3, {}); + auto range = ops::Range(compilation_scope.WithOpName("range"), compiled_zero, + compiled_size, compiled_three); + auto retval = + ops::_Retval(compilation_scope.WithOpName("retval"), range, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(compilation_scope.ToGraph(graph.get())); + + // Reuse the variable inferred by TensorFlow at the tf2xla boundary so both + // layers describe the same logical input dimension. + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 1}, + std::vector{input_expr, xla::DExpr::Const(1)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "range", + std::move(graph), args, &result)); + + // The bound is ceil(8 / 3) = 3, while the logical size remains ceil(A / 3). + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_EQ(result_shape.dimensions(0), 3); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), + expected_output_expr)); +} + +TEST_F(XlaCompilerDynamicSizesTest, + ShapeValueChainPreservesExpressionsThroughFill) { + // Shape converts [A, 24] into the value vector {A, 24}. Fill then consumes + // those values as dimensions, so its output should recover [A, 24]. + Scope inference_scope = Scope::NewRootScope().ExitOnError(); + auto placeholder = ops::Placeholder( + inference_scope.WithOpName("placeholder"), DT_FLOAT, + ops::Placeholder::Shape(PartialTensorShape({-1, 24}))); + auto fill_shape = + ops::Shape(inference_scope.WithOpName("fill_shape"), placeholder); + auto fill_value = + ops::Const(inference_scope.WithOpName("fill_value"), 1.0f, {}); + auto inferred_fill = ops::Fill(inference_scope.WithOpName("fill"), + fill_shape, fill_value); + + grappler::GrapplerItem item; + TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph)); + grappler::GraphProperties properties(item); + TF_ASSERT_OK(properties.InferStatically( + /*assume_valid_feeds=*/false, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/false, + /*include_output_tensor_values=*/false, + /*enable_dynamic_value_inference=*/true)); + + const TensorShapeProto& inferred_input_shape = + properties.GetOutputProperties("placeholder").at(0).shape(); + const TensorShapeProto& inferred_fill_shape = + properties.GetOutputProperties("fill").at(0).shape(); + ASSERT_EQ(inferred_input_shape.expressions_size(), 2); + ASSERT_EQ(inferred_fill_shape.expressions_size(), 2); + const xla::DExpr input_expr = + DimExprFromProto(inferred_input_shape.expressions(0)); + const xla::DExpr expected_output_expr = + DimExprFromProto(inferred_fill_shape.expressions(0)); + EXPECT_TRUE(xla::DynExpr::equal(expected_output_expr, input_expr)); + + // Compile the equivalent graph and check both dimensions reach HLO. + Scope compilation_scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0); + auto compiled_fill_shape = + ops::Shape(compilation_scope.WithOpName("fill_shape"), arg); + auto compiled_fill_value = + ops::Const(compilation_scope.WithOpName("fill_value"), 1.0f, {}); + auto fill = ops::Fill(compilation_scope.WithOpName("fill"), + compiled_fill_shape, compiled_fill_value); + auto retval = + ops::_Retval(compilation_scope.WithOpName("retval"), fill, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(compilation_scope.ToGraph(graph.get())); + + // Use the same inferred variable at the tf2xla boundary to compare the two + // inference paths structurally rather than by variable spelling. + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 24}, + std::vector{input_expr, xla::DExpr::Const(24)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "fill", + std::move(graph), args, &result)); + + // Compilation keeps the physical [8, 24] bounds and the logical [A, 24] + // expressions at the cluster output. + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_EQ(result_shape.dimensions(0), 8); + EXPECT_EQ(result_shape.dimensions(1), 24); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), + expected_output_expr)); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), + xla::DExpr::Const(24))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ReverseSequencePreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto seq_lens = ops::_Arg(scope.WithOpName("seq_lens"), DT_INT32, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("reverse_sequence", "ReverseSequence") + .Input(input.node()->name(), 0, DT_INT32) + .Input(seq_lens.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tlen", DT_INT32) + .Attr("batch_dim", 0) + .Attr("seq_dim", 1) + .Finalize(&def)); + absl::Status status; + Node* reverse_sequence = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(reverse_sequence)); + scope.graph()->AddEdge(input.node(), 0, reverse_sequence, 0); + scope.graph()->AddEdge(seq_lens.node(), 0, reverse_sequence, 1); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(reverse_sequence), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {4, 8}, + std::vector{xla::DExpr::Var(1), xla::DExpr::Var(2)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {4}, std::vector{xla::DExpr::Var(1)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reverse_sequence", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(1))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Var(2))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(1))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Var(2))); +} + +TEST_F(XlaCompilerDynamicSizesTest, UniquePreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("unique", "Unique") + .Input(input.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("out_idx", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* unique = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(unique)); + scope.graph()->AddEdge(input.node(), 0, unique, 0); + + auto retval0 = + ops::_Retval(scope.WithOpName("retval0"), Output(unique, 0), 0); + auto retval1 = + ops::_Retval(scope.WithOpName("retval1"), Output(unique, 1), 1); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7}, std::vector{xla::DExpr::Var(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "unique", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 2); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(3))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression( + 0), + xla::DExpr::Var(3))); + + const xla::Shape& values_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + const xla::Shape& indices_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {1}); + EXPECT_TRUE( + xla::DynExpr::equal(values_shape.expressions(0), xla::DExpr::Var(3))); + EXPECT_TRUE( + xla::DynExpr::equal(indices_shape.expressions(0), xla::DExpr::Var(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, DynamicPartitionPreservesPartitionExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto data = ops::_Arg(scope.WithOpName("data"), DT_INT32, 0); + auto partitions = ops::_Arg(scope.WithOpName("partitions"), DT_INT32, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("dynamic_partition", "DynamicPartition") + .Input(data.node()->name(), 0, DT_INT32) + .Input(partitions.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("num_partitions", 2) + .Finalize(&def)); + absl::Status status; + Node* dynamic_partition = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(dynamic_partition)); + scope.graph()->AddEdge(data.node(), 0, dynamic_partition, 0); + scope.graph()->AddEdge(partitions.node(), 0, dynamic_partition, 1); + + auto retval0 = ops::_Retval(scope.WithOpName("retval0"), + Output(dynamic_partition, 0), 0); + auto retval1 = ops::_Retval(scope.WithOpName("retval1"), + Output(dynamic_partition, 1), 1); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6}, std::vector{xla::DExpr::Var(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6}, std::vector{xla::DExpr::Var(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "dynamic_partition", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 2); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression( + 0), + xla::DExpr::Var(4))); + + const xla::Shape& result0_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + const xla::Shape& result1_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {1}); + EXPECT_TRUE( + xla::DynExpr::equal(result0_shape.expressions(0), xla::DExpr::Var(4))); + EXPECT_TRUE( + xla::DynExpr::equal(result1_shape.expressions(0), xla::DExpr::Var(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + DynamicPartitionBroadcastPreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto data = ops::_Arg(scope.WithOpName("data"), DT_INT32, 0); + auto partitions = ops::_Arg(scope.WithOpName("partitions"), DT_INT32, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("dynamic_partition", "DynamicPartition") + .Input(data.node()->name(), 0, DT_INT32) + .Input(partitions.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("num_partitions", 2) + .Finalize(&def)); + absl::Status status; + Node* dynamic_partition = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(dynamic_partition)); + scope.graph()->AddEdge(data.node(), 0, dynamic_partition, 0); + scope.graph()->AddEdge(partitions.node(), 0, dynamic_partition, 1); + + auto retval0 = ops::_Retval(scope.WithOpName("retval0"), + Output(dynamic_partition, 0), 0); + auto retval1 = ops::_Retval(scope.WithOpName("retval1"), + Output(dynamic_partition, 1), 1); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 3}, + std::vector{xla::DExpr::Var(40), xla::DExpr::Const(3)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8}, std::vector{xla::DExpr::Var(40)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "dynamic_partition_broadcast", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 2); + for (int i = 0; i < 2; ++i) { + EXPECT_TRUE(xla::DynExpr::equal( + result.outputs[i].shape.get_filled_expression(0), xla::DExpr::Var(40))); + EXPECT_TRUE(xla::DynExpr::equal( + result.outputs[i].shape.get_filled_expression(1), xla::DExpr::Const(3))); + const xla::Shape& out_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {i}); + EXPECT_TRUE( + xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Var(40))); + EXPECT_TRUE( + xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Const(3))); + } +} + +TEST_F(XlaCompilerDynamicSizesTest, + DenseBincountMatrixPreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto weights = ops::_Arg(scope.WithOpName("weights"), DT_FLOAT, 1); + auto size = ops::Const(scope.WithOpName("size"), 5); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("dense_bincount", "DenseBincount") + .Input(input.node()->name(), 0, DT_INT32) + .Input(size.node()->name(), 0, DT_INT32) + .Input(weights.node()->name(), 0, DT_FLOAT) + .Attr("Tidx", DT_INT32) + .Attr("T", DT_FLOAT) + .Attr("binary_output", false) + .Finalize(&def)); + absl::Status status; + Node* bincount = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, bincount, 0); + scope.graph()->AddEdge(size.node(), 0, bincount, 1); + scope.graph()->AddEdge(weights.node(), 0, bincount, 2); + TF_ASSERT_OK(scope.DoShapeInference(bincount)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(bincount, 0), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {12, 4}, + std::vector{xla::DExpr::Var(50), xla::DExpr::Const(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_FLOAT; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::F32, {12, 4}, + std::vector{xla::DExpr::Var(50), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "dense_bincount", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(50))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); + + const xla::Shape& out_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE(xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Var(50))); + EXPECT_TRUE( + xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + DynamicStitchEmptyPreservesTrailingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto data0 = ops::_Arg(scope.WithOpName("data0"), DT_INT32, 0); + auto data1 = ops::_Arg(scope.WithOpName("data1"), DT_INT32, 1); + Tensor empty_indices_tensor(DT_INT32, TensorShape({0})); + auto indices0 = ops::Const(scope.WithOpName("indices0"), empty_indices_tensor); + auto indices1 = ops::Const(scope.WithOpName("indices1"), empty_indices_tensor); + + NodeDef def; + std::vector indices_inputs = { + {indices0.node()->name(), 0, DT_INT32}, + {indices1.node()->name(), 0, DT_INT32}, + }; + std::vector data_inputs = { + {data0.node()->name(), 0, DT_INT32}, + {data1.node()->name(), 0, DT_INT32}, + }; + TF_ASSERT_OK(NodeDefBuilder("dynamic_stitch_empty", "DynamicStitch") + .Input(indices_inputs) + .Input(data_inputs) + .Attr("N", 2) + .Attr("T", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* dynamic_stitch = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(indices0.node(), 0, dynamic_stitch, 0); + scope.graph()->AddEdge(indices1.node(), 0, dynamic_stitch, 1); + scope.graph()->AddEdge(data0.node(), 0, dynamic_stitch, 2); + scope.graph()->AddEdge(data1.node(), 0, dynamic_stitch, 3); + + auto retval = ops::_Retval(scope.WithOpName("retval"), + Output(dynamic_stitch, 0), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {0, 7}, + std::vector{xla::DExpr::Const(0), xla::DExpr::Var(61)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {0, 7}, + std::vector{xla::DExpr::Const(0), xla::DExpr::Var(61)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "dynamic_stitch_empty", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Const(0))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Var(61))); + + const xla::Shape& out_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Const(0))); + EXPECT_TRUE( + xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Var(61))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + TensorListPushBackStackPreservesElementExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto element = ops::_Arg(scope.WithOpName("element"), DT_INT32, 0); + auto element_shape = ops::Const(scope.WithOpName("element_shape"), + {-1, 3}, {2}); + auto max_num_elements = ops::Const(scope.WithOpName("max_num_elements"), 4); + + NodeDef empty_def; + TF_ASSERT_OK(NodeDefBuilder("empty_list", "EmptyTensorList") + .Input(element_shape.node()->name(), 0, DT_INT32) + .Input(max_num_elements.node()->name(), 0, DT_INT32) + .Attr("element_dtype", DT_INT32) + .Attr("shape_type", DT_INT32) + .Finalize(&empty_def)); + absl::Status status; + Node* empty_list = scope.graph()->AddNode(empty_def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(element_shape.node(), 0, empty_list, 0); + scope.graph()->AddEdge(max_num_elements.node(), 0, empty_list, 1); + TF_ASSERT_OK(scope.DoShapeInference(empty_list)); + + NodeDef push_def; + TF_ASSERT_OK(NodeDefBuilder("push_back", "TensorListPushBack") + .Input(empty_list->name(), 0, DT_VARIANT) + .Input(element.node()->name(), 0, DT_INT32) + .Attr("element_dtype", DT_INT32) + .Finalize(&push_def)); + Node* push_back = scope.graph()->AddNode(push_def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(empty_list, 0, push_back, 0); + scope.graph()->AddEdge(element.node(), 0, push_back, 1); + TF_ASSERT_OK(scope.DoShapeInference(push_back)); + + NodeDef stack_def; + TF_ASSERT_OK(NodeDefBuilder("stack", "TensorListStack") + .Input(push_back->name(), 0, DT_VARIANT) + .Input(element_shape.node()->name(), 0, DT_INT32) + .Attr("element_dtype", DT_INT32) + .Attr("num_elements", 4) + .Finalize(&stack_def)); + Node* stack = scope.graph()->AddNode(stack_def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(push_back, 0, stack, 0); + scope.graph()->AddEdge(element_shape.node(), 0, stack, 1); + TF_ASSERT_OK(scope.DoShapeInference(stack)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(stack, 0), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5, 3}, + std::vector{xla::DExpr::Var(70), xla::DExpr::Const(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "tensor_list_stack", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Var(70))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ShapeThenReshapePreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto shape_source = ops::_Arg(scope.WithOpName("shape_source"), DT_INT32, 0); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 1); + auto shape = ops::Shape(scope.WithOpName("shape"), shape_source); + auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape); + auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 7}, + std::vector{xla::DExpr::Var(41), xla::DExpr::Const(7)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 7}, + std::vector{xla::DExpr::Var(41), xla::DExpr::Const(7)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "shape_then_reshape", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(41))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(7))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(41))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(7))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ZerosLikePreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto zeros = ops::ZerosLike(scope.WithOpName("zeros_like"), input); + auto retval = ops::_Retval(scope.WithOpName("retval"), zeros, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {9, 4}, + std::vector{xla::DExpr::Var(43), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "zeros_like_exprs", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(43))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, OnesLikePreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto ones = ops::OnesLike(scope.WithOpName("ones_like"), input); + auto retval = ops::_Retval(scope.WithOpName("retval"), ones, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {9, 4}, + std::vector{xla::DExpr::Var(44), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "ones_like_exprs", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(44))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, MatrixDiagPreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("matrix_diag", "MatrixDiag") + .Input(input.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* matrix_diag = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(matrix_diag)); + scope.graph()->AddEdge(input.node(), 0, matrix_diag, 0); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(matrix_diag), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 4}, + std::vector{xla::DExpr::Var(45), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "matrix_diag_exprs", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(45))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, WhereBuildsDynamicIndexMatrixShape) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_BOOL, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("where", "Where") + .Input(input.node()->name(), 0, DT_BOOL) + .Finalize(&def)); + absl::Status status; + Node* where = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(where)); + scope.graph()->AddEdge(input.node(), 0, where, 0); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(where), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_BOOL; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::PRED, {8, 4, 6}, + std::vector{xla::DExpr::Var(46), xla::DExpr::Const(4), + xla::DExpr::Var(47)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "where", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_EQ(result_shape.dimensions_size(), 2); + EXPECT_EQ(result_shape.dimensions(1), 3); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, DiagDuplicatesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("diag", "Diag") + .Input(input.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* diag = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(diag)); + scope.graph()->AddEdge(input.node(), 0, diag, 0); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(diag), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5}, std::vector{xla::DExpr::Var(42)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "diag", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(42))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Var(42))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(42))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Var(42))); +} + +TEST_F(XlaCompilerDynamicSizesTest, InTopKPreservesBatchExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto predictions = ops::_Arg(scope.WithOpName("predictions"), DT_FLOAT, 0); + auto targets = ops::_Arg(scope.WithOpName("targets"), DT_INT32, 1); + auto k = ops::Const(scope.WithOpName("k"), 3, {}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("in_topk", "InTopKV2") + .Input(predictions.node()->name(), 0, DT_FLOAT) + .Input(targets.node()->name(), 0, DT_INT32) + .Input(k.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* in_topk = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(predictions.node(), 0, in_topk, 0); + scope.graph()->AddEdge(targets.node(), 0, in_topk, 1); + scope.graph()->AddEdge(k.node(), 0, in_topk, 2); + TF_ASSERT_OK(scope.DoShapeInference(in_topk)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(in_topk), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {5, 7}, + std::vector{xla::DExpr::Var(43), xla::DExpr::Const(7)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5}, std::vector{xla::DExpr::Var(43)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "in_topk", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(43))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(43))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ReshapeCollapsePreservesSymbolicExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto shape = ops::Const(scope.WithOpName("shape"), {96}, {1}); + auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape); + auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {3, 4, 8}, + std::vector{xla::DExpr::Var(5), xla::DExpr::Const(4), + xla::DExpr::Const(8)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reshape_collapse", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(5) * xla::DExpr::Const(32)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected)); +} + +TEST_F(XlaCompilerDynamicSizesTest, ReshapeSplitPreservesSymbolicExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto shape = ops::Const(scope.WithOpName("shape"), {5, 16}, {2}); + auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape); + auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {10, 8}, + std::vector{xla::DExpr::Var(6), xla::DExpr::Const(8)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reshape_split", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(6) / xla::DExpr::Const(2)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected)); +} + +TEST_F(XlaCompilerDynamicSizesTest, ReshapeSplitAndCollapsePreservesSymbolicExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto shape = ops::Const(scope.WithOpName("shape"), {4, 64}, {2}); + auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape); + auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 8, 4}, + std::vector{xla::DExpr::Var(7), xla::DExpr::Const(8), + xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reshape_split_collapse", + std::move(graph), args, &result)); + + xla::DExpr expected = + (xla::DExpr::Var(7) / xla::DExpr::Const(2)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected)); +} + +TEST_F(XlaCompilerDynamicSizesTest, GatherV2PreservesUngatheredExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto params = ops::_Arg(scope.WithOpName("params"), DT_INT32, 0); + auto indices = ops::Const(scope.WithOpName("indices"), {0, 2, 4}, {3}); + auto axis = ops::Const(scope.WithOpName("axis"), 1, {}); + auto gathered = + ops::GatherV2(scope.WithOpName("gather"), params, indices, axis); + auto retval = ops::_Retval(scope.WithOpName("retval"), gathered, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {9, 7}, + std::vector{xla::DExpr::Var(8), xla::DExpr::Const(7)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "gather_preserve", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(8))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(8))); +} + +TEST_F(XlaCompilerDynamicSizesTest, TransposePermutesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto perm = ops::Const(scope.WithOpName("perm"), {1, 0}, {2}); + auto transposed = ops::Transpose(scope.WithOpName("transpose"), input, perm); + auto retval = ops::_Retval(scope.WithOpName("retval"), transposed, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5, 7}, + std::vector{xla::DExpr::Var(9), xla::DExpr::Var(10)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "transpose_exprs", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(10))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Var(9))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ExpandDimsInsertsUnitExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto dim = ops::Const(scope.WithOpName("dim"), 1, {}); + auto expanded = ops::ExpandDims(scope.WithOpName("expand"), input, dim); + auto retval = ops::_Retval(scope.WithOpName("retval"), expanded, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 4}, + std::vector{xla::DExpr::Var(11), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "expand_dims_exprs", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(11))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(1))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, SqueezeRemovesUnitExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("squeeze", "Squeeze") + .Input(input.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("squeeze_dims", {1}) + .Finalize(&def)); + absl::Status status; + Node* squeeze = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, squeeze, 0); + TF_ASSERT_OK(scope.DoShapeInference(squeeze)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(squeeze), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 1, 4}, + std::vector{xla::DExpr::Var(12), xla::DExpr::Const(1), + xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "squeeze_exprs", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(12))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, SplitPreservesAndDividesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto split_dim = ops::Const(scope.WithOpName("split_dim"), 0, {}); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto split = ops::Split(scope.WithOpName("split"), split_dim, input, 2); + auto retval0 = ops::_Retval(scope.WithOpName("retval0"), split.output[0], 0); + auto retval1 = ops::_Retval(scope.WithOpName("retval1"), split.output[1], 1); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 5}, + std::vector{xla::DExpr::Var(13), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "split_exprs", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(13) / xla::DExpr::Const(2)).simplify(); + ASSERT_EQ(result.outputs.size(), 2); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression( + 0), + expected)); +} + +TEST_F(XlaCompilerDynamicSizesTest, TileScalesExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto multiples = + ops::Const(scope.WithOpName("multiples"), {3, 1}, {2}); + auto tiled = ops::Tile(scope.WithOpName("tile"), input, multiples); + auto retval = ops::_Retval(scope.WithOpName("retval"), tiled, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {4, 5}, + std::vector{xla::DExpr::Var(14), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "tile_exprs", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(14) * xla::DExpr::Const(3)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, PackInsertsAxisAndPreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input0 = ops::_Arg(scope.WithOpName("input0"), DT_INT32, 0); + auto input1 = ops::_Arg(scope.WithOpName("input1"), DT_INT32, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("pack", "Pack") + .Input({NodeDefBuilder::NodeOut(input0.node()->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(input1.node()->name(), 0, + DT_INT32)}) + .Attr("T", DT_INT32) + .Attr("N", 2) + .Attr("axis", 1) + .Finalize(&def)); + absl::Status status; + Node* pack = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input0.node(), 0, pack, 0); + scope.graph()->AddEdge(input1.node(), 0, pack, 1); + TF_ASSERT_OK(scope.DoShapeInference(pack)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(pack), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6}, std::vector{xla::DExpr::Var(15)}); + args[1] = args[0]; + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "pack", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(15))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(2))); +} + +TEST_F(XlaCompilerDynamicSizesTest, UnpackRemovesAxisAndPreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("unpack", "Unpack") + .Input(input.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("num", 3) + .Attr("axis", 2) + .Finalize(&def)); + absl::Status status; + Node* unpack = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, unpack, 0); + TF_ASSERT_OK(scope.DoShapeInference(unpack)); + + auto retval0 = ops::_Retval(scope.WithOpName("retval0"), Output(unpack, 0), 0); + auto retval1 = ops::_Retval(scope.WithOpName("retval1"), Output(unpack, 1), 1); + auto retval2 = ops::_Retval(scope.WithOpName("retval2"), Output(unpack, 2), 2); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 4, 3}, + std::vector{xla::DExpr::Var(16), xla::DExpr::Const(4), + xla::DExpr::Const(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "unpack", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 3); + for (int i = 0; i < 3; ++i) { + EXPECT_TRUE(xla::DynExpr::equal( + result.outputs[i].shape.get_filled_expression(0), xla::DExpr::Var(16))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[i].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + } +} + +TEST_F(XlaCompilerDynamicSizesTest, ConcatV2AddsLeadingExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto axis = ops::Const(scope.WithOpName("axis"), 0, {}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("concat", "ConcatV2") + .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(rhs.node()->name(), 0, + DT_INT32)}) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("N", 2) + .Finalize(&def)); + absl::Status status; + Node* concat = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(lhs.node(), 0, concat, 0); + scope.graph()->AddEdge(rhs.node(), 0, concat, 1); + scope.graph()->AddEdge(axis.node(), 0, concat, 2); + TF_ASSERT_OK(scope.DoShapeInference(concat)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5, 4}, + std::vector{xla::DExpr::Var(17), xla::DExpr::Const(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 4}, + std::vector{xla::DExpr::Var(18), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "concat", + std::move(graph), args, &result)); + + xla::DExpr expected = + (xla::DExpr::Var(17) + xla::DExpr::Var(18)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ConcatAddsLeadingExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto axis = ops::Const(scope.WithOpName("axis"), 0, {}); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("concat", "Concat") + .Input(axis.node()->name(), 0, DT_INT32) + .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(rhs.node()->name(), 0, + DT_INT32)}) + .Attr("T", DT_INT32) + .Attr("N", 2) + .Finalize(&def)); + absl::Status status; + Node* concat = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(axis.node(), 0, concat, 0); + scope.graph()->AddEdge(lhs.node(), 0, concat, 1); + scope.graph()->AddEdge(rhs.node(), 0, concat, 2); + TF_ASSERT_OK(scope.DoShapeInference(concat)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5, 4}, + std::vector{xla::DExpr::Var(22), xla::DExpr::Const(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 4}, + std::vector{xla::DExpr::Var(23), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "concat_legacy", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(22) + xla::DExpr::Var(23)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ConcatAddsThreeLeadingExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto mid = ops::_Arg(scope.WithOpName("mid"), DT_INT32, 1); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 2); + auto axis = ops::Const(scope.WithOpName("axis"), 0, {}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("concat0", "ConcatV2") + .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(mid.node()->name(), 0, + DT_INT32)}) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("N", 2) + .Finalize(&def)); + absl::Status status; + Node* concat0 = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(lhs.node(), 0, concat0, 0); + scope.graph()->AddEdge(mid.node(), 0, concat0, 1); + scope.graph()->AddEdge(axis.node(), 0, concat0, 2); + TF_ASSERT_OK(scope.DoShapeInference(concat0)); + + NodeDef def1; + TF_ASSERT_OK(NodeDefBuilder("concat1", "ConcatV2") + .Input({NodeDefBuilder::NodeOut(concat0->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(rhs.node()->name(), 0, + DT_INT32)}) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("N", 2) + .Finalize(&def1)); + Node* concat1 = scope.graph()->AddNode(def1, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(concat0, 0, concat1, 0); + scope.graph()->AddEdge(rhs.node(), 0, concat1, 1); + scope.graph()->AddEdge(axis.node(), 0, concat1, 2); + TF_ASSERT_OK(scope.DoShapeInference(concat1)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat1), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(3); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {3, 4}, + std::vector{xla::DExpr::Var(31), xla::DExpr::Const(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {5, 4}, + std::vector{xla::DExpr::Var(32), xla::DExpr::Const(4)}); + args[2].kind = XlaCompiler::Argument::kParameter; + args[2].type = DT_INT32; + args[2].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 4}, + std::vector{xla::DExpr::Var(33), xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "concat_three", std::move(graph), args, + &result)); + + xla::DExpr expected = + (xla::DExpr::Var(31) + xla::DExpr::Var(32) + xla::DExpr::Var(33)) + .simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + expected)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected)); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ConcatV2PreservesLeadingExpressionOnInnerAxis) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto axis = ops::Const(scope.WithOpName("axis"), 1, {}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("concat", "ConcatV2") + .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0, + DT_INT32), + NodeDefBuilder::NodeOut(rhs.node()->name(), 0, + DT_INT32)}) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Attr("N", 2) + .Finalize(&def)); + absl::Status status; + Node* concat = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(lhs.node(), 0, concat, 0); + scope.graph()->AddEdge(rhs.node(), 0, concat, 1); + scope.graph()->AddEdge(axis.node(), 0, concat, 2); + TF_ASSERT_OK(scope.DoShapeInference(concat)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 4}, + std::vector{xla::DExpr::Var(24), xla::DExpr::Const(4)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 3}, + std::vector{xla::DExpr::Var(24), xla::DExpr::Const(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "concat_inner", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(24))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(7))); +} + +TEST_F(XlaCompilerDynamicSizesTest, AddPreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs); + auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 5}, + std::vector{xla::DExpr::Var(44), xla::DExpr::Const(5)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 5}, + std::vector{xla::DExpr::Var(44), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(44))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + AddSameRankBroadcastPreservesMappedExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs); + auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 1, 3}, + std::vector{xla::DExpr::Var(45), xla::DExpr::Const(1), + xla::DExpr::Const(3)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 4, 3}, + std::vector{xla::DExpr::Var(45), xla::DExpr::Const(4), + xla::DExpr::Const(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "add_broadcast", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(45))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, AddDegenerateBroadcastPreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs); + auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {1, 5}, + std::vector{xla::DExpr::Const(1), xla::DExpr::Const(5)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 5}, + std::vector{xla::DExpr::Var(46), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "add_degenerate", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(46))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + MulSameRankBroadcastPreservesMappedExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1); + auto product = ops::Mul(scope.WithOpName("mul"), lhs, rhs); + auto retval = ops::_Retval(scope.WithOpName("retval"), product, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 1, 3}, + std::vector{xla::DExpr::Var(47), xla::DExpr::Const(1), + xla::DExpr::Const(3)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 4, 3}, + std::vector{xla::DExpr::Var(47), xla::DExpr::Const(4), + xla::DExpr::Const(3)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "mul_broadcast", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(47))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, ReverseV2PreservesExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto axis = ops::Const(scope.WithOpName("axis"), {1}, {1}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("reverse", "ReverseV2") + .Input(input.node()->name(), 0, DT_INT32) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tidx", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* reverse = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, reverse, 0); + scope.graph()->AddEdge(axis.node(), 0, reverse, 1); + TF_ASSERT_OK(scope.DoShapeInference(reverse)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(reverse), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {8, 5}, + std::vector{xla::DExpr::Var(19), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "reverse_v2", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(19))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, BatchMatMulPreservesBatchExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_FLOAT, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_FLOAT, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("batch_matmul", "BatchMatMul") + .Input(lhs.node()->name(), 0, DT_FLOAT) + .Input(rhs.node()->name(), 0, DT_FLOAT) + .Attr("T", DT_FLOAT) + .Attr("adj_x", false) + .Attr("adj_y", false) + .Finalize(&def)); + absl::Status status; + Node* batch_matmul = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(lhs.node(), 0, batch_matmul, 0); + scope.graph()->AddEdge(rhs.node(), 0, batch_matmul, 1); + TF_ASSERT_OK(scope.DoShapeInference(batch_matmul)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(batch_matmul), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 4, 6}, + std::vector{xla::DExpr::Var(25), xla::DExpr::Const(4), + xla::DExpr::Const(6)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_FLOAT; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 6, 5}, + std::vector{xla::DExpr::Var(25), xla::DExpr::Const(6), + xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "batch_matmul", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(25))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, BatchMatMulV2BroadcastsBatchExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_FLOAT, 0); + auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_FLOAT, 1); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("batch_matmul_v2", "BatchMatMulV2") + .Input(lhs.node()->name(), 0, DT_FLOAT) + .Input(rhs.node()->name(), 0, DT_FLOAT) + .Attr("T", DT_FLOAT) + .Attr("adj_x", false) + .Attr("adj_y", false) + .Finalize(&def)); + absl::Status status; + Node* batch_matmul = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(lhs.node(), 0, batch_matmul, 0); + scope.graph()->AddEdge(rhs.node(), 0, batch_matmul, 1); + TF_ASSERT_OK(scope.DoShapeInference(batch_matmul)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(batch_matmul), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {1, 4, 6}, + std::vector{xla::DExpr::Const(1), xla::DExpr::Const(4), + xla::DExpr::Const(6)}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_FLOAT; + args[1].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 6, 5}, + std::vector{xla::DExpr::Var(26), xla::DExpr::Const(6), + xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "batch_matmul_v2", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(26))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, SlicePreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 2}, {2}); + auto size = ops::Const(scope.WithOpName("size"), {-1, 3}, {2}); + auto sliced = ops::Slice(scope.WithOpName("slice"), input, begin, size); + auto retval = ops::_Retval(scope.WithOpName("retval"), sliced, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 8}, + std::vector{xla::DExpr::Var(20), xla::DExpr::Const(8)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "slice", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(20))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, SliceSubtractsFromLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {2, 1}, {2}); + auto size = ops::Const(scope.WithOpName("size"), {-1, 3}, {2}); + auto sliced = ops::Slice(scope.WithOpName("slice"), input, begin, size); + auto retval = ops::_Retval(scope.WithOpName("retval"), sliced, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {9, 8}, + std::vector{xla::DExpr::Var(27), xla::DExpr::Const(8)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "slice_subtract", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + (xla::DExpr::Var(27) - 2).simplify())); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(3))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + StridedSlicePreservesLeadingExpressionOnInnerAxis) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 1}, {2}); + auto end = ops::Const(scope.WithOpName("end"), {0, 7}, {2}); + auto strides = ops::Const(scope.WithOpName("strides"), {1, 2}, {2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 1) + .Attr("end_mask", 1) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0) + .Attr("shrink_axis_mask", 0) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 8}, + std::vector{xla::DExpr::Var(40), xla::DExpr::Const(8)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_inner", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(40))); + EXPECT_EQ(result.outputs[0].shape.dim_size(1), 3); +} + +TEST_F(XlaCompilerDynamicSizesTest, StridedSliceScalesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 0}, {2}); + auto end = ops::Const(scope.WithOpName("end"), {0, 4}, {2}); + auto strides = ops::Const(scope.WithOpName("strides"), {2, 1}, {2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 1) + .Attr("end_mask", 1) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0) + .Attr("shrink_axis_mask", 0) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 4}, + std::vector{ + ((xla::DExpr::Const(2) * xla::DExpr::Var(41)) - + xla::DExpr::Const(1)) + .simplify(), + xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_leading", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal( + result.outputs[0].shape.get_filled_expression(0), xla::DExpr::Var(41))); + EXPECT_EQ(result.outputs[0].shape.dim_size(0), 4); + EXPECT_EQ(result.outputs[0].shape.dim_size(1), 4); +} + +TEST_F(XlaCompilerDynamicSizesTest, StridedSliceNewAxisInsertsUnitExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 0}, {2}); + auto end = ops::Const(scope.WithOpName("end"), {0, 0}, {2}); + auto strides = ops::Const(scope.WithOpName("strides"), {1, 1}, {2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 0x1) + .Attr("end_mask", 0x1) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0x2) + .Attr("shrink_axis_mask", 0) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 5}, + std::vector{xla::DExpr::Var(42), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_new_axis", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(42))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(1))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + StridedSliceShrinkAxisPreservesRemainingExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 2, 0}, {3}); + auto end = ops::Const(scope.WithOpName("end"), {0, 2, 5}, {3}); + auto strides = ops::Const(scope.WithOpName("strides"), {1, 1, 1}, {3}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 0x3) + .Attr("end_mask", 0x3) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0) + .Attr("shrink_axis_mask", 0x2) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 5, 5}, + std::vector{xla::DExpr::Var(43), xla::DExpr::Const(5), + xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_shrink", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(43))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + StridedSliceNegativeStridePreservesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 0}, {2}); + auto end = ops::Const(scope.WithOpName("end"), {0, 0}, {2}); + auto strides = + ops::Const(scope.WithOpName("strides"), {-1, 1}, {2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 0x3) + .Attr("end_mask", 0x3) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0) + .Attr("shrink_axis_mask", 0) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 5}, + std::vector{xla::DExpr::Var(44), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_negative", + std::move(graph), args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(44))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} + +TEST_F(XlaCompilerDynamicSizesTest, + StridedSliceNegativeStrideTwoScalesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto begin = ops::Const(scope.WithOpName("begin"), {0, 0}, {2}); + auto end = ops::Const(scope.WithOpName("end"), {0, 0}, {2}); + auto strides = + ops::Const(scope.WithOpName("strides"), {-2, 1}, {2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice") + .Input(input.node()->name(), 0, DT_INT32) + .Input(begin.node()->name(), 0, DT_INT32) + .Input(end.node()->name(), 0, DT_INT32) + .Input(strides.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Index", DT_INT32) + .Attr("begin_mask", 0x3) + .Attr("end_mask", 0x3) + .Attr("ellipsis_mask", 0) + .Attr("new_axis_mask", 0) + .Attr("shrink_axis_mask", 0) + .Finalize(&def)); + absl::Status status; + Node* strided_slice = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, strided_slice, 0); + scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1); + scope.graph()->AddEdge(end.node(), 0, strided_slice, 2); + scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3); + TF_ASSERT_OK(scope.DoShapeInference(strided_slice)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 5}, + std::vector{xla::DExpr::Var(45), xla::DExpr::Const(5)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "strided_slice_negative_two", + std::move(graph), args, &result)); -class DummyReadResourceOp : public XlaOpKernel { - public: - explicit DummyReadResourceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} - void Compile(XlaOpKernelContext* ctx) override { - ResourceMgr* rm = ctx->op_kernel_context()->resource_manager(); - OP_REQUIRES(ctx, rm, errors::Internal("No resource manager.")); - DummyResourceForTest* dummy; - OP_REQUIRES_OK(ctx, rm->Lookup( - rm->default_container(), "dummy", &dummy)); - dummy->Increment(); - dummy->Unref(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal( + result.outputs[0].shape.get_filled_expression(0), + ((xla::DExpr::Var(45) + xla::DExpr::Const(1)) / xla::DExpr::Const(2)) + .simplify())); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} - ctx->SetOutput(0, ctx->Input(0)); - ctx->SetOutput(1, ctx->Input(0)); - } -}; +TEST_F(XlaCompilerDynamicSizesTest, PadAddsToLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto paddings = + ops::Const(scope.WithOpName("paddings"), {1, 2, 0, 0}, {2, 2}); + auto padded = ops::Pad(scope.WithOpName("pad"), input, paddings); + auto retval = ops::_Retval(scope.WithOpName("retval"), padded, 0); -class DummyReadResourceCC { - public: - DummyReadResourceCC(const Scope& scope, const Input& value) { - if (!scope.ok()) return; - auto _value = ops::AsNodeOut(scope, value); - if (!scope.ok()) return; - Node* ret; - const auto unique_name = scope.GetUniqueNameForOp("DummyReadResource"); - auto builder = NodeBuilder(unique_name, "DummyReadResource").Input(_value); - scope.UpdateBuilder(&builder); - scope.UpdateStatus(builder.Finalize(scope.graph(), &ret)); - if (!scope.ok()) return; - scope.UpdateStatus(scope.DoShapeInference(ret)); - if (!scope.ok()) return; - this->output1_ = Output(ret, 0); - this->output2_ = Output(ret, 1); - } + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); - Output output1_; - Output output2_; -}; + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {7, 5}, + std::vector{xla::DExpr::Var(28), xla::DExpr::Const(5)}); -REGISTER_OP("DummyReadResource") - .Input("input: int32") - .Output("output1: int32") - .Output("output2: int32") - .SetShapeFn(shape_inference::UnknownShape) - .Doc(R"doc( -A dummy Op. + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "pad", + std::move(graph), args, &result)); -input: dummy input. -output1: dummy output. -output2: dummy output. -)doc"); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + (xla::DExpr::Var(28) + 3).simplify())); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); +} -REGISTER_XLA_OP(Name("DummyReadResource"), DummyReadResourceOp); +TEST_F(XlaCompilerDynamicSizesTest, SpaceToBatchNDScalesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0); + auto block_shape = + ops::Const(scope.WithOpName("block_shape"), {2}, {1}); + auto paddings = + ops::Const(scope.WithOpName("paddings"), {0, 0}, {1, 2}); -// DummyDuplicateOp is present purely to test multiple REGISTER_XLA_OP calls -// on the same Op name below. -class DummyDuplicateOp : public XlaOpKernel { - public: - explicit DummyDuplicateOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} - void Compile(XlaOpKernelContext* ctx) override { - ctx->SetOutput(0, ctx->Input(0)); - } -}; + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("space_to_batch", "SpaceToBatchND") + .Input(input.node()->name(), 0, DT_FLOAT) + .Input(block_shape.node()->name(), 0, DT_INT32) + .Input(paddings.node()->name(), 0, DT_INT32) + .Attr("T", DT_FLOAT) + .Attr("Tblock_shape", DT_INT32) + .Attr("Tpaddings", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* space_to_batch = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, space_to_batch, 0); + scope.graph()->AddEdge(block_shape.node(), 0, space_to_batch, 1); + scope.graph()->AddEdge(paddings.node(), 0, space_to_batch, 2); + TF_ASSERT_OK(scope.DoShapeInference(space_to_batch)); -REGISTER_OP("DummyDuplicateOp") - .Input("input: int32") - .Output("output: int32") - .Doc(R"doc( -A dummy Op. + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(space_to_batch), 0); -input: dummy input. -output: dummy output. -)doc"); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); -REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT), - DummyDuplicateOp); -REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT), - DummyDuplicateOp); + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {4, 8}, + std::vector{xla::DExpr::Var(29), xla::DExpr::Const(8)}); -// Tests compilation and execution of an empty graph. -TEST_F(XlaCompilerTest, EmptyReturnValues) { XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "space_to_batch", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + (xla::DExpr::Const(2) * xla::DExpr::Var(29)) + .simplify())); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); +} + +TEST_F(XlaCompilerDynamicSizesTest, BatchToSpaceNDDividesLeadingExpression) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0); + auto block_shape = + ops::Const(scope.WithOpName("block_shape"), {2}, {1}); + auto crops = ops::Const(scope.WithOpName("crops"), {0, 0}, {1, 2}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("batch_to_space", "BatchToSpaceND") + .Input(input.node()->name(), 0, DT_FLOAT) + .Input(block_shape.node()->name(), 0, DT_INT32) + .Input(crops.node()->name(), 0, DT_INT32) + .Attr("T", DT_FLOAT) + .Attr("Tblock_shape", DT_INT32) + .Attr("Tcrops", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* batch_to_space = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, batch_to_space, 0); + scope.graph()->AddEdge(block_shape.node(), 0, batch_to_space, 1); + scope.graph()->AddEdge(crops.node(), 0, batch_to_space, 2); + TF_ASSERT_OK(scope.DoShapeInference(batch_to_space)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(batch_to_space), 0); std::unique_ptr graph(new Graph(OpRegistry::Global())); - XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), - /*args=*/{}, &result)); + TF_ASSERT_OK(scope.ToGraph(graph.get())); - TF_ASSERT_OK(client_->Execute(*result.computation, {}).status()); + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {8, 4}, + std::vector{(xla::DExpr::Const(2) * xla::DExpr::Var(30)) + .simplify(), + xla::DExpr::Const(4)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "batch_to_space", std::move(graph), args, + &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(30))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(8))); } -// Tests compilation and execution of a graph that adds two tensors. -TEST_F(XlaCompilerTest, Simple) { - // Builds a graph that adds two Tensors. +TEST_F(XlaCompilerDynamicSizesTest, SpaceToDepthScalesDepthExpression) { Scope scope = Scope::NewRootScope().ExitOnError(); - auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); - auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1); - auto c = ops::Add(scope.WithOpName("C"), a, b); - auto d = ops::_Retval(scope.WithOpName("D"), c, 0); + auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("space_to_depth", "SpaceToDepth") + .Input(input.node()->name(), 0, DT_FLOAT) + .Attr("T", DT_FLOAT) + .Attr("block_size", 2) + .Attr("data_format", "NHWC") + .Finalize(&def)); + absl::Status status; + Node* space_to_depth = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, space_to_depth, 0); + TF_ASSERT_OK(scope.DoShapeInference(space_to_depth)); + + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(space_to_depth), 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); TF_ASSERT_OK(scope.ToGraph(graph.get())); - // Builds a description of the arguments. - std::vector args(2); + std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; - args[0].type = DT_INT32; - args[0].shape = TensorShape({2}); - args[1].kind = XlaCompiler::Argument::kParameter; - args[1].type = DT_INT32; - args[1].shape = TensorShape({2}); + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {5, 8, 8, 3}, + std::vector{xla::DExpr::Const(5), xla::DExpr::Const(8), + xla::DExpr::Const(8), xla::DExpr::Var(35)}); - // Compiles the graph. XlaCompiler compiler(DefaultOptions()); XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "space_to_depth", std::move(graph), args, + &result)); + + xla::DExpr expected_depth = + (xla::DExpr::Const(4) * xla::DExpr::Var(35)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Const(5))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 3), + expected_depth)); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Const(5))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(4))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(2), xla::DExpr::Const(4))); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(3), expected_depth)); +} - // Tests that the generated computation works. - xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); - xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); +TEST_F(XlaCompilerDynamicSizesTest, DepthToSpaceScalesSpatialExpressions) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0); - std::unique_ptr actual = - client_ - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("depth_to_space", "DepthToSpace") + .Input(input.node()->name(), 0, DT_FLOAT) + .Attr("T", DT_FLOAT) + .Attr("block_size", 2) + .Attr("data_format", "NHWC") + .Finalize(&def)); + absl::Status status; + Node* depth_to_space = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, depth_to_space, 0); + TF_ASSERT_OK(scope.DoShapeInference(depth_to_space)); - xla::Literal expected0 = xla::LiteralUtil::CreateR1({4, 143}); - xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); -} + auto retval = + ops::_Retval(scope.WithOpName("retval"), Output(depth_to_space), 0); -absl::StatusOr> LoadModuleFromHloProto( - const xla::HloModuleProto& module_proto) { - TF_ASSIGN_OR_RETURN(auto module_config, - xla::HloModule::CreateModuleConfigFromProto( - module_proto, xla::GetDebugOptionsFromFlags())); - return xla::CreateModuleFromProto(module_proto, module_config); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_FLOAT; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::F32, {5, 4, 4, 12}, + std::vector{xla::DExpr::Const(5), xla::DExpr::Var(37), + xla::DExpr::Const(4), xla::DExpr::Const(12)}); + + XlaCompiler compiler(DefaultOptions()); + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "depth_to_space", std::move(graph), args, + &result)); + + xla::DExpr expected_height = + (xla::DExpr::Const(2) * xla::DExpr::Var(37)).simplify(); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Const(5))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + expected_height)); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 2), + xla::DExpr::Const(8))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 3), + xla::DExpr::Const(3))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Const(5))); + EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), expected_height)); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(2), xla::DExpr::Const(8))); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(3), xla::DExpr::Const(3))); } -// Tests compilation and execution of a graph that adds two tensors with dynamic -// shape parameters. -TEST_F(XlaCompilerTest, SimpleDynamicShapeParameter) { - // Builds a graph that adds two Tensors. +TEST_F(XlaCompilerDynamicSizesTest, RollPreservesExpressions) { Scope scope = Scope::NewRootScope().ExitOnError(); - auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); - auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1); - auto c = ops::Add(scope.WithOpName("C"), a, b); - auto d = ops::_Retval(scope.WithOpName("D"), c, 0); + auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0); + auto shift = ops::Const(scope.WithOpName("shift"), 2, {}); + auto axis = ops::Const(scope.WithOpName("axis"), 1, {}); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("roll", "Roll") + .Input(input.node()->name(), 0, DT_INT32) + .Input(shift.node()->name(), 0, DT_INT32) + .Input(axis.node()->name(), 0, DT_INT32) + .Attr("T", DT_INT32) + .Attr("Tshift", DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* roll = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + scope.graph()->AddEdge(input.node(), 0, roll, 0); + scope.graph()->AddEdge(shift.node(), 0, roll, 1); + scope.graph()->AddEdge(axis.node(), 0, roll, 2); + TF_ASSERT_OK(scope.DoShapeInference(roll)); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(roll), 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); TF_ASSERT_OK(scope.ToGraph(graph.get())); - // Builds a description of the arguments. - std::vector args(2); + std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = - xla::ShapeUtil::MakeShape(/*element_type=*/xla::S32, /*dimensions=*/{2}, - /*dynamic_dimensions=*/{true}, - /*expressions=*/{}); - args[1].kind = XlaCompiler::Argument::kParameter; - args[1].type = DT_INT32; - args[1].shape = TensorShape(/*dimensions=*/{2}); + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {6, 5}, + std::vector{xla::DExpr::Var(21), xla::DExpr::Const(5)}); - // Compiles the graph. XlaCompiler compiler(DefaultOptions()); XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "roll", std::move(graph), args, &result)); - auto hlo = result.computation->proto(); - TF_ASSERT_OK_AND_ASSIGN(auto module, LoadModuleFromHloProto(hlo)); - EXPECT_EQ(module->computation_count(), 1); - EXPECT_TRUE(module->mutable_computation(0) - ->parameter_instruction(0) - ->shape() - .is_dynamic()); + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(21))); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 1), + xla::DExpr::Const(5))); } // Tests compilation of a graph where the _Retval node is not necessarily last @@ -1015,6 +3862,13 @@ FunctionDef FillFn() { {{{"y"}, "Fill", {"dims", "x"}, {{"T", "$T"}}}}); } +FunctionDef IdentityFn() { + return FunctionDefHelper::Define( + "IdentityFn", {"x: T"}, {"y: T"}, + {"T: {float, double, int32, int64}"}, + {{{"y"}, "Identity", {"x"}, {{"T", "$T"}}}}); +} + TEST_F(XlaCompilerTest, FunctionCallWithConstants) { // Certain operations in a function, "Fill" for example, requires the // operator's argument to be a compile-time constant instead of a parameter. @@ -1058,6 +3912,53 @@ TEST_F(XlaCompilerTest, FunctionCallWithConstants) { std::move(graph), args, &result)); } +TEST_F(XlaCompilerDynamicSizesTest, FunctionCallPreservesDynamicExpressions) { + XlaCompiler compiler(DefaultOptions()); + + FunctionDefLibrary flib; + *flib.add_function() = IdentityFn(); + TF_ASSERT_OK(flib_def_->AddFunctionDef(IdentityFn())); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("arg"), DT_INT32, 0); + TF_EXPECT_OK(scope.graph()->AddFunctionLibrary(flib)); + + NodeDef def; + TF_ASSERT_OK(NodeDefBuilder("identity_fn", "IdentityFn", flib_def_.get()) + .Input(arg.node()->name(), 0, DT_INT32) + .Finalize(&def)); + absl::Status status; + Node* identity_fn = scope.graph()->AddNode(def, &status); + TF_ASSERT_OK(status); + TF_ASSERT_OK(scope.DoShapeInference(identity_fn)); + scope.graph()->AddEdge(arg.node(), 0, identity_fn, 0); + + auto retval = ops::_Retval(scope.WithOpName("retval"), Output(identity_fn), 0); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = xla::ShapeUtil::MakeShape( + xla::S32, {9}, std::vector{xla::DExpr::Var(2)}); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "identity_function", std::move(graph), + args, &result)); + + ASSERT_EQ(result.outputs.size(), 1); + EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression( + 0), + xla::DExpr::Var(2))); + + const xla::Shape& result_shape = + xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0}); + EXPECT_TRUE( + xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(2))); +} + // Tests CompileFunction with a local function lookup failing, fails with // informative error about both lookups. TEST_F(XlaCompilerTest, LocalFunctionWithWrongArgumentsFail) { From b0c7c6816e3b555b200d2dba639998576f5a54f9 Mon Sep 17 00:00:00 2001 From: fengxiao-cui Date: Sat, 30 May 2026 18:05:18 +0800 Subject: [PATCH 19/21] Change dynamic batch padding policy Replace power-of-two padding with configurable factor-preserving batch buckets. Use separate environment-controlled multipliers for small and large batches while retaining discovered bounds per cluster. --- tensorflow/compiler/jit/kernels/xla_ops.cc | 6 +- tensorflow/compiler/jit/xla_batch_matcher.cc | 155 +++++++++++++------ tensorflow/compiler/jit/xla_batch_matcher.h | 26 +++- 3 files changed, 135 insertions(+), 52 deletions(-) diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 96c096306a3101..c28c8c3400414a 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -1144,7 +1144,7 @@ absl::Status CompileToLocalExecutable( record_dynamic_dim_value(shp.dim_size(idx), xla::DExpr::Var(1)); if (!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(function.name(), shp.dim_size(idx)); } std::vector dyn_exprs; @@ -1192,8 +1192,10 @@ absl::Status CompileToLocalExecutable( << shp.dim_size(idx) << " to " << var_value; } record_dynamic_dim_value(var_value, e); + VLOG(1) << "Filled batch for function " << function.name() + << " with value " << var_value; filled_batch = - xla_batch_matcher->get_xla_compile_batch(var_value); + xla_batch_matcher->get_xla_compile_batch(function.name(), var_value); break; } } diff --git a/tensorflow/compiler/jit/xla_batch_matcher.cc b/tensorflow/compiler/jit/xla_batch_matcher.cc index 8b991ce6e2b865..ba7b539500aa2b 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.cc +++ b/tensorflow/compiler/jit/xla_batch_matcher.cc @@ -17,9 +17,11 @@ limitations under the License. #include #include +#include #include #include +#include "absl/strings/numbers.h" #include "xla/debug_options_flags.h" #include "tensorflow/core/platform/logging.h" @@ -99,52 +101,54 @@ std::vector XlaBatchMatcher::parse_single_item(const std::string& item) return batch_list; } -void XlaBatchMatcher::print_all_batches() { +void XlaBatchMatcher::print_all_batches(const std::string& cluster_key, + const std::vector& batches) { std::ostringstream oss; - oss << "[XLA_BATCH_INFO] Valid batch list update: "; - for (size_t i = 0; i < all_batches_.size(); ++i) { + oss << "[XLA_BATCH_INFO] cluster_key=" << cluster_key + << " valid batch list update: "; + for (size_t i = 0; i < batches.size(); ++i) { if (i > 0) oss << ", "; - oss << all_batches_[i]; + oss << batches[i]; } LOG(INFO) << oss.str(); } -// Parse environment variable config into deduplicated, sorted batch list -// For example, export XLA_COMPILE_BATCH_SIZES="10:100:10, 977" +// Parse environment variable config into a seed batch list used to initialize +// each cluster-key state. void XlaBatchMatcher::parse_env_config() { - // If the env var not set or is empty, filled with the nearest power of two by default + // If the env var not set or is empty, leave seed empty and let per-cluster + // logic generate factor-preserving padded batches. if (env_str_.empty()) { VLOG(2) << "[XLA_BATCH_WARN] Env var " << "--tf_xla_compile_batch_sizes" << - "is empty, filled with the nearest power of two by default"; + "is empty, will use factor-preserving padding by default"; return; } - // Split config by commas + // Parse into the default seed state. + ClusterState& seed = clusters_[""]; + std::stringstream ss(env_str_); std::string item; while (std::getline(ss, item, ',')) { std::string trimmed_item = Trim(item); if (trimmed_item.empty()) continue; - // Parse single item (skip on failure to avoid breaking other items) try { std::vector item_batches = parse_single_item(trimmed_item); - all_batches_.insert(all_batches_.end(), item_batches.begin(), item_batches.end()); + seed.all_batches.insert(seed.all_batches.end(), item_batches.begin(), + item_batches.end()); } catch (const std::exception& e) { - LOG(INFO) << "[XLA_BATCH_WARN] Failed to parse config item, skipping: " << - trimmed_item << " (" << e.what() << ")"; + LOG(INFO) << "[XLA_BATCH_WARN] Failed to parse config item, skipping: " + << trimmed_item << " (" << e.what() << ")"; } } - if (!all_batches_.empty()) { - std::sort(all_batches_.begin(), all_batches_.end()); - auto last = std::unique(all_batches_.begin(), all_batches_.end()); - all_batches_.erase(last, all_batches_.end()); + if (!seed.all_batches.empty()) { + std::sort(seed.all_batches.begin(), seed.all_batches.end()); + auto last = std::unique(seed.all_batches.begin(), seed.all_batches.end()); + seed.all_batches.erase(last, seed.all_batches.end()); + print_all_batches("", seed.all_batches); } - - // Print parsed result - if (!all_batches_.empty()) print_all_batches(); - return; } // Calculate the smallest power of two greater than the real batch @@ -161,43 +165,104 @@ static int64_t GetNextPowerOfTwo(int64_t real_batch) { return power; } -int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) { +// Small/large multipliers for padding are configurable via environment variables: +// TF_XLA_BATCH_SMALL_FACTOR (default 10) and TF_XLA_BATCH_LARGE_FACTOR (default 2). +static int64_t GetEnvFactorOrDefault(const char* name, int64_t def) { + const char* val = std::getenv(name); + if (val == nullptr) return def; + int64_t parsed = 0; + if (!absl::SimpleAtoi(val, &parsed) || parsed <= 0) { + LOG(WARNING) << "[XLA_BATCH_WARN] Failed to parse env var " << name + << "=\"" << val << "\", using default: " << def; + return def; + } + return parsed; +} + +static int64_t GetFactorPreservingBatch(int64_t real_batch) { + // Cache env values in function-local statics (thread-safe since C++11). + static const int64_t small_factor = + GetEnvFactorOrDefault("TF_XLA_BATCH_SMALL_FACTOR", 10); + static const int64_t large_factor = + GetEnvFactorOrDefault("TF_XLA_BATCH_LARGE_FACTOR", 2); + + const int64_t k = (real_batch < 10) ? small_factor : large_factor; + // Guard overflow / max range. + if (real_batch > kMaxBatch / k) { + LOG(WARNING) << "[XLA_BATCH_ERR] Out of valid range: " << real_batch; + return real_batch; + } + return real_batch * k; +} + +int64_t XlaBatchMatcher::find_min_larger_batch(ClusterState* state, + int64_t real_batch) { if (real_batch <= 0 || real_batch > kMaxBatch) { LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch; return real_batch; } - if (all_batches_.empty()) { - // Return the next power of two directly without modifying all_batches_ - return GetNextPowerOfTwo(real_batch); + // 1) If there is no configured candidate list, do not use power-of-two. + if (state->all_batches.empty()) { + const int val = GetFactorPreservingBatch(real_batch); + state->all_batches.push_back(val); + state->updated = true; + return val; } - // Edge case 1: Real value < the smallest batch, use smallest - if (real_batch < all_batches_.front()) { - return all_batches_.front(); - } - // 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); - all_batches_.emplace_back(val); - print_all_batches(); - return val; + // 2) Prefer a configured candidate strictly larger than real_batch. + auto ub = std::upper_bound(state->all_batches.begin(), state->all_batches.end(), + real_batch); + if (ub != state->all_batches.end()) { + return *ub; } - // 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(); + // 3) real_batch > all_batches_.back(): generate a factor-preserving padded + // batch, write it back, keep list sorted/unique. + const int64_t val = GetFactorPreservingBatch(real_batch); + auto insert_pos = + std::lower_bound(state->all_batches.begin(), state->all_batches.end(), val); + if (insert_pos == state->all_batches.end() || *insert_pos != val) { + state->all_batches.insert(insert_pos, val); + state->updated = true; + } + return val; } -int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) { - // Match target batch size - int64_t selected = find_min_larger_batch(real_batch); - 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; +int64_t XlaBatchMatcher::get_xla_compile_batch(const std::string& cluster_key, + int64_t real_batch) { + // Lazily initialize per-key state from the seed ("" key) to preserve previous + // behavior when an env list is provided. + ClusterState& state = clusters_[cluster_key]; + state.updated = false; + if (state.all_batches.empty()) { + auto it = clusters_.find(""); + if (it != clusters_.end()) { + state.all_batches = it->second.all_batches; + } + } + + const int64_t selected = find_min_larger_batch(&state, real_batch); + + if (state.updated) { + VLOG(2) << "[XLA_BATCH_INFO] cluster_key=" << cluster_key + << " real batch: " << real_batch + << " -> selected compile batch: " << selected; + print_all_batches(cluster_key, state.all_batches); } + return selected; } +int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) { + return get_xla_compile_batch("", real_batch); +} + +std::vector XlaBatchMatcher::get_all_batches( + const std::string& cluster_key) { + auto it = clusters_.find(cluster_key); + if (it == clusters_.end()) return {}; + return it->second.all_batches; +} + } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_batch_matcher.h b/tensorflow/compiler/jit/xla_batch_matcher.h index dac8079e5081af..08476116aaf336 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.h +++ b/tensorflow/compiler/jit/xla_batch_matcher.h @@ -20,6 +20,8 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" + namespace tensorflow { // Define the maximum allowed batch size: 2147483648 >> 1 = 1073741824 @@ -30,18 +32,32 @@ class XlaBatchMatcher { public: XlaBatchMatcher(); virtual ~XlaBatchMatcher() = default; + + // Per-cluster-key API (recommended): isolates padding candidates per cluster. + int64_t get_xla_compile_batch(const std::string& cluster_key, + int64_t real_batch); + + // Backward-compatible API: uses a default (global) key. int64_t get_xla_compile_batch(int64_t real_batch); - std::vector get_all_batches() { return all_batches_; } + + // For debugging. + std::vector get_all_batches(const std::string& cluster_key); private: + struct ClusterState { + std::vector all_batches; + bool updated = false; + }; + void parse_env_config(); - void print_all_batches(); + void print_all_batches(const std::string& cluster_key, + const std::vector& batches); std::vector parse_single_item(const std::string& item); - int64_t find_min_larger_batch(int64_t real_batch); - std::vector all_batches_; + int64_t find_min_larger_batch(ClusterState* state, int64_t real_batch); + + absl::flat_hash_map clusters_; std::string env_str_; - int64_t last_batch_ = -1; }; } // namespace tensorflow From 35d71da2e39b4855e053258d283bb0c87850c423 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Thu, 20 Aug 2026 11:31:55 +0100 Subject: [PATCH 20/21] Keep padded reshape expressions conservative Avoid representing ceil-aligned reshape padding with ordinary integer division, which can simplify to an incorrect symbolic bound. --- tensorflow/compiler/tf2xla/kernels/reshape_op.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc index eb861e895ac03b..c019f42927bb9b 100644 --- a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc @@ -164,11 +164,11 @@ class ReshapeOp : public XlaOpKernel { input, xla::Zero(ctx->builder(), input_xla_shape->element_type()), 0, 0, padded_input_num - input_num_elements); input_shape.set_dim(0, padded_input_num); - missing_expr = - (input_num_elements_expr + (product - 1)) / product; - missing_expr = missing_expr.simplify(); + // This expression only approximates the padded size: the true value + // uses ceil(input_num_elements / product) * product, which we do not + // model symbolically here. xla::DExpr padded_input_num_expr = - (missing_expr * xla::DExpr::Const(product)) + ((input_num_elements_expr / product_expr) * product_expr) .simplify(); input_shape.set_expression(0, padded_input_num_expr); } From a2f4dc3ad76da350b704fc505e5a306f163b93ee Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Thu, 20 Aug 2026 11:32:03 +0100 Subject: [PATCH 21/21] Partially solve Max expressions with a constant bound Solve the dynamic branch when the observed result is strictly greater than the constant bound, while keeping equality ambiguous. --- third_party/xla/xla/shape_expr.h | 21 ++++++++++++++++++++- third_party/xla/xla/shape_test.cc | 14 +++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/shape_expr.h b/third_party/xla/xla/shape_expr.h index babd157a56b0b9..49b564fbb92a57 100644 --- a/third_party/xla/xla/shape_expr.h +++ b/third_party/xla/xla/shape_expr.h @@ -586,8 +586,27 @@ class MaxExpr : public DynExpr { ids.merge(rhs->get_all_ids()); return ids; } - // Max is not invertible: either operand may have produced the result. std::optional solve(int64_t x) override { + auto solve_dynamic_branch = [x](DynExpr* dynamic, + DynExpr* constant) + -> std::optional { + if (constant->kind() != DExprKind::kConstant || + dynamic->get_all_ids().size() != 1) { + return std::nullopt; + } + const int64_t bound = constant->get_val(); + // max(expr, bound) selects expr only when the result exceeds bound. + // Equality is ambiguous because either branch may have produced it. + return x > bound ? dynamic->solve(x) : std::nullopt; + }; + + if (auto result = solve_dynamic_branch(lhs.get(), rhs.get())) { + return result; + } + if (auto result = solve_dynamic_branch(rhs.get(), lhs.get())) { + return result; + } + StringPrinter printer; print(&printer); LOG(WARNING) << "Cannot solve Max dynamic shape expression for value " << x diff --git a/third_party/xla/xla/shape_test.cc b/third_party/xla/xla/shape_test.cc index 69fdb816bd2f05..e47aaf0d398fb2 100644 --- a/third_party/xla/xla/shape_test.cc +++ b/third_party/xla/xla/shape_test.cc @@ -136,7 +136,6 @@ TEST_F(ShapeTest, DExprSolveRejectsZeroDivisor) { TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) { DExpr expr = DExpr::Max(DExpr::Var(1), DExpr::Const(4)); EXPECT_EQ("max(A, 4)", DExprToString(expr.simplify())); - EXPECT_FALSE(expr->solve(7).has_value()); DExpr clamped = DExpr::Max(DExpr::Var(1), DExpr::Const(0)); EXPECT_EQ("A", DExprToString(clamped.simplify())); @@ -167,6 +166,19 @@ TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) { EXPECT_TRUE(expr == DExprFromProto(proto)); } +TEST_F(ShapeTest, DExprMaxPartiallySolvesSelectedDynamicBranch) { + DExpr expr = DExpr::Max(DExpr::Var(1) - 2, DExpr::Const(0)); + EXPECT_EQ(100, expr->solve(98)); + + DExpr reversed = DExpr::Max(DExpr::Const(0), DExpr::Var(1) - 2); + EXPECT_EQ(100, reversed->solve(98)); + + // A result equal to the constant bound may come from either branch. + EXPECT_FALSE(expr->solve(0).has_value()); + // A result below the constant bound cannot be produced by Max. + EXPECT_FALSE(expr->solve(-1).has_value()); +} + TEST_F(ShapeTest, DExprSelectUsesDynamicPredicate) { DExpr delta = DExpr::Var(1) - 3; DExpr expr = DExpr::Select(DExpr::Gt(delta, DExpr::Const(0)),