From d6a8e510ac15464656d8aa6af1e5b63b5161ebd3 Mon Sep 17 00:00:00 2001 From: Steven Varoumas Date: Mon, 17 Aug 2026 16:21:24 +0100 Subject: [PATCH] Fix dynamic expression edge cases and batch bucketing --- tensorflow/compiler/jit/BUILD | 12 +++++- tensorflow/compiler/jit/xla_batch_matcher.cc | 27 +++++++----- tensorflow/compiler/jit/xla_batch_matcher.h | 8 ++-- .../compiler/jit/xla_batch_matcher_test.cc | 41 +++++++++++++++++++ tensorflow/core/util/strided_slice_op.cc | 5 ++- .../xla/xla/service/cpu/cpu_executable.cc | 17 -------- .../xla/xla/service/llvm_ir/llvm_util.cc | 4 +- .../xla/xla/service/llvm_ir/llvm_util_test.cc | 8 ++++ third_party/xla/xla/shape.cc | 7 ---- third_party/xla/xla/shape_expr.cc | 3 ++ third_party/xla/xla/shape_expr.h | 6 ++- third_party/xla/xla/shape_test.cc | 15 +++++++ 12 files changed, 108 insertions(+), 45 deletions(-) create mode 100644 tensorflow/compiler/jit/xla_batch_matcher_test.cc diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a8197744359fde..8eb3c620266184 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -2017,4 +2017,14 @@ cc_library( "//tensorflow/core/platform:logging", "@local_xla//xla:debug_options_flags", ], -) \ No newline at end of file +) + +tf_cc_test( + name = "xla_batch_matcher_test", + srcs = ["xla_batch_matcher_test.cc"], + deps = [ + ":xla_batch_matcher", + "//tensorflow/core:test", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tensorflow/compiler/jit/xla_batch_matcher.cc b/tensorflow/compiler/jit/xla_batch_matcher.cc index c6f2f50fecf3e1..ef3eb319e70df1 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.cc +++ b/tensorflow/compiler/jit/xla_batch_matcher.cc @@ -1,4 +1,7 @@ #include "tensorflow/compiler/jit/xla_batch_matcher.h" + +#include + #include "xla/debug_options_flags.h" #include "tensorflow/core/platform/logging.h" @@ -9,6 +12,14 @@ XlaBatchMatcher::XlaBatchMatcher() { parse_env_config(); } +XlaBatchMatcher::XlaBatchMatcher(std::vector batches) + : all_batches_(std::move(batches)) { + std::sort(all_batches_.begin(), all_batches_.end()); + all_batches_.erase( + std::unique(all_batches_.begin(), all_batches_.end()), + all_batches_.end()); +} + // Trim whitespace (spaces/tabs) from both ends of a string std::string trim(const std::string& s) { size_t start = s.find_first_not_of(" \t"); @@ -135,7 +146,7 @@ static int64_t GetNextPowerOfTwo(int64_t real_batch) { return power; } -int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) { +int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) const { if (real_batch <= 0 || real_batch > kMaxBatch) { LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch; return real_batch; @@ -152,10 +163,7 @@ int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) { } // 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; + return GetNextPowerOfTwo(real_batch); } // Find first batch larger than real value (binary search via lower_bound) @@ -163,14 +171,11 @@ int64_t XlaBatchMatcher::find_min_larger_batch(int64_t real_batch) { return (it != all_batches_.end()) ? *it : all_batches_.back(); } -int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) { +int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) const { // 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; - } + VLOG(2) << "[XLA_BATCH_INFO] Real batch: " << real_batch + << " -> Selected compile batch: " << selected; return selected; } diff --git a/tensorflow/compiler/jit/xla_batch_matcher.h b/tensorflow/compiler/jit/xla_batch_matcher.h index 7a294f5430221a..917ecd94d14d2c 100644 --- a/tensorflow/compiler/jit/xla_batch_matcher.h +++ b/tensorflow/compiler/jit/xla_batch_matcher.h @@ -20,19 +20,19 @@ constexpr int kMaxBatch = 2147483648ULL >> 1; class XlaBatchMatcher { public: XlaBatchMatcher(); + explicit XlaBatchMatcher(std::vector batches); virtual ~XlaBatchMatcher() = default; - int64_t get_xla_compile_batch(int64_t real_batch); - std::vector get_all_batches() { return all_batches_; } + int64_t get_xla_compile_batch(int64_t real_batch) const; + std::vector get_all_batches() const { return all_batches_; } private: void parse_env_config(); void print_all_batches(); std::vector parse_single_item(const std::string& item); - int64_t find_min_larger_batch(int64_t real_batch); + int64_t find_min_larger_batch(int64_t real_batch) const; std::vector all_batches_; std::string env_str_; - int64_t last_batch_ = -1; }; } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_batch_matcher_test.cc b/tensorflow/compiler/jit/xla_batch_matcher_test.cc new file mode 100644 index 00000000000000..2175bbf217063f --- /dev/null +++ b/tensorflow/compiler/jit/xla_batch_matcher_test.cc @@ -0,0 +1,41 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/jit/xla_batch_matcher.h" + +#include +#include + +#include + +namespace tensorflow { +namespace { + +TEST(XlaBatchMatcherTest, GeneratedBucketsAreRequestOrderIndependent) { + XlaBatchMatcher matcher({100}); + + EXPECT_EQ(matcher.get_xla_compile_batch(300), 512); + EXPECT_EQ(matcher.get_xla_compile_batch(150), 256); +} + +TEST(XlaBatchMatcherTest, GeneratedBucketsDoNotChangeConfiguredBuckets) { + XlaBatchMatcher matcher({200, 100, 200}); + + EXPECT_EQ(matcher.get_xla_compile_batch(300), 512); + EXPECT_EQ(matcher.get_all_batches(), (std::vector{100, 200})); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc index a2bb25f9c2e923..8d3612309c02a8 100644 --- a/tensorflow/core/util/strided_slice_op.cc +++ b/tensorflow/core/util/strided_slice_op.cc @@ -211,7 +211,8 @@ absl::Status ValidateStridedSliceOp( absl::InlinedVector b; absl::InlinedVector e; - // HACK + // Callers that only request concrete bounds still need matching expressions + // for the shared validation logic below. if (begin_expr == nullptr) { for (int i : *begin) { b.push_back(xla::DExpr::Const(i)); @@ -227,7 +228,7 @@ absl::Status ValidateStridedSliceOp( 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 = diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc index 6ae5465b5b79bb..832940544f8558 100644 --- a/third_party/xla/xla/service/cpu/cpu_executable.cc +++ b/third_party/xla/xla/service/cpu/cpu_executable.cc @@ -327,28 +327,11 @@ 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; 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 4d6804a506a6b5..3ce7dafe89fe63 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_util.cc +++ b/third_party/xla/xla/service/llvm_ir/llvm_util.cc @@ -886,10 +886,12 @@ 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.is_constant()) { + return llvm::ConstantInt::get(i64Type, expr.get_val(), true); + } if (expr.kind() == DExpr::Kind::kVariable) { // For now we can just use %bdim... return GetBatchDimByName(b); 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 index fb3a839085a41b..0d8d3294378cb2 100644 --- a/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc +++ b/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc @@ -54,6 +54,14 @@ TEST(LlvmUtilTest, DynamicExpressionDivisionIsSigned) { EXPECT_EQ(division->getOpcode(), llvm::Instruction::SDiv); } +TEST(LlvmUtilTest, UnknownExpressionDoesNotEmitAConstant) { + llvm::LLVMContext context; + llvm::IRBuilder<> builder(context); + + DExpr expression = DExpr::Unknown(kMissingExpressionSentinel); + EXPECT_EQ(EmitExpression(&builder, expression), nullptr); +} + } // namespace } // namespace llvm_ir } // namespace xla diff --git a/third_party/xla/xla/shape.cc b/third_party/xla/xla/shape.cc index e12d07a25ecfe5..4847125eb14ffa 100644 --- a/third_party/xla/xla/shape.cc +++ b/third_party/xla/xla/shape.cc @@ -99,9 +99,6 @@ 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()) { @@ -154,7 +151,6 @@ 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; } @@ -162,8 +158,6 @@ 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) { @@ -189,7 +183,6 @@ 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; } diff --git a/third_party/xla/xla/shape_expr.cc b/third_party/xla/xla/shape_expr.cc index d74a42175b3380..a1eaf552d54c41 100644 --- a/third_party/xla/xla/shape_expr.cc +++ b/third_party/xla/xla/shape_expr.cc @@ -468,6 +468,9 @@ std::unique_ptr SimplifyFallback(const DynExpr* expr) { } Constant* l = AsConstant(lhs.get()); Constant* r = AsConstant(rhs.get()); + if (r && r->get_val() == 0) { + return std::make_unique
(lhs.release(), rhs.release()); + } if (*lhs == *rhs) { return std::make_unique(1); } diff --git a/third_party/xla/xla/shape_expr.h b/third_party/xla/xla/shape_expr.h index babd157a56b0b9..08d421d7914c49 100644 --- a/third_party/xla/xla/shape_expr.h +++ b/third_party/xla/xla/shape_expr.h @@ -66,7 +66,9 @@ class DynExpr { virtual std::set get_all_ids() = 0; virtual std::optional solve(int64_t x) = 0; - bool is_dynamic() { return !is_constant(); } + bool is_dynamic() const { + return kind() != DExprKind::kUnknown && !is_constant(); + } static DynExpr* zero; static DynExpr* one; @@ -193,7 +195,7 @@ class UnknownExpr : public DynExpr { void to_proto(xla::ExpressionProto* proto) const override { (void)proto; } - bool is_constant() const override { return true; } + bool is_constant() const override { return false; } int get_id() const { return id_; } DynExpr* substitute(int id, DynExpr* v) override { (void)id; diff --git a/third_party/xla/xla/shape_test.cc b/third_party/xla/xla/shape_test.cc index 69fdb816bd2f05..55407b5cdd3060 100644 --- a/third_party/xla/xla/shape_test.cc +++ b/third_party/xla/xla/shape_test.cc @@ -133,6 +133,21 @@ TEST_F(ShapeTest, DExprSolveRejectsZeroDivisor) { EXPECT_FALSE(expr->solve(7).has_value()); } +TEST_F(ShapeTest, DExprSimplifyPreservesZeroDivisor) { + DExpr expr = DExpr::Const(0) / DExpr::Const(0); + DExpr simplified = expr.simplify(); + + EXPECT_EQ(DExpr::Kind::kDiv, simplified.kind()); + EXPECT_FALSE(simplified->is_constant()); +} + +TEST_F(ShapeTest, UnknownExpressionIsNeitherConstantNorDynamic) { + DExpr expr = DExpr::Unknown(kMissingExpressionSentinel); + + EXPECT_FALSE(expr->is_constant()); + EXPECT_FALSE(expr->is_dynamic()); +} + TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) { DExpr expr = DExpr::Max(DExpr::Var(1), DExpr::Const(4)); EXPECT_EQ("max(A, 4)", DExprToString(expr.simplify()));