From ff2452dede7e39d82bc33ad851594b176bb3fba7 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 31 Aug 2026 23:51:21 -0700 Subject: [PATCH 1/6] Validate input rank in the tf2xla QR and SVD kernels The tf2xla kernels for Qr, Svd, and XlaSvd passed their input straight to the XLA builder libraries without checking its rank. The graph-level shape functions reject inputs of known rank below 2, but an input of unknown static rank, for example one produced by StackPopV2, reaches the kernels unchecked at compile time. QrExplicit then reads dimension rank - 2 = -1 and dies on a fatal ShapeUtil check, and SvdOp dies on the dim_size range check in TensorShape, aborting the process instead of failing compilation. Validate rank >= 2 in all three kernels before any trailing dimensions are read, mirroring the validation and message style of the neighboring MatrixSolve kernel and the errors already produced by the Cholesky and SelfAdjointEig builder libraries. The regression tests feed a rank-1 value through an unknown-rank placeholder, which reproduces the abort on unfixed builds, and assert that a regular InvalidArgumentError is raised instead. Fixes #110798 --- tensorflow/compiler/tests/BUILD | 2 ++ tensorflow/compiler/tests/qr_op_test.py | 13 +++++++++++++ tensorflow/compiler/tests/svd_op_test.py | 13 +++++++++++++ tensorflow/compiler/tf2xla/kernels/BUILD | 4 ++++ tensorflow/compiler/tf2xla/kernels/qr_op.cc | 11 +++++++++++ tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc | 14 ++++++++++++++ 6 files changed, 57 insertions(+) diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 5251a2953e7092..00cb78d3a71774 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -594,6 +594,7 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/framework:errors", "//tensorflow/python/framework:tensor_shape", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:linalg_ops", @@ -1371,6 +1372,7 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/framework:errors", "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:linalg_ops", diff --git a/tensorflow/compiler/tests/qr_op_test.py b/tensorflow/compiler/tests/qr_op_test.py index 4eac4a970bf199..b77bc9c46aac28 100644 --- a/tensorflow/compiler/tests/qr_op_test.py +++ b/tensorflow/compiler/tests/qr_op_test.py @@ -21,6 +21,7 @@ import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import linalg_ops @@ -143,6 +144,18 @@ def testRepeatedColumn(self, rows, cols): x_np[:, 1] = x_np[:, 2] self._test(x_np, full_matrices=True, full_rank=False) + def testVectorInputRaisesError(self): + # Regression test for GitHub issue 110798. The graph-level shape check + # only runs for inputs of known rank, so a rank-1 input reaching the + # compiler through an unknown-rank placeholder must be rejected at + # compile time instead of reaching a fatal check inside QrExplicit. + with self.session() as sess: + x_tf = array_ops.placeholder(np.float32) + with self.device_scope(): + q_tf, r_tf = linalg_ops.qr(x_tf, full_matrices=True) + with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"): + sess.run([q_tf, r_tf], feed_dict={x_tf: np.zeros([8], np.float32)}) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/svd_op_test.py b/tensorflow/compiler/tests/svd_op_test.py index 356a4bc7715b0a..7214231579b439 100644 --- a/tensorflow/compiler/tests/svd_op_test.py +++ b/tensorflow/compiler/tests/svd_op_test.py @@ -20,6 +20,7 @@ import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.python.framework import errors from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_linalg_ops @@ -86,6 +87,18 @@ def testSvd(self, n, dtype): self._testSvdCorrectness(dtype, batch_dims + (2 * n, n)) self._testSvdCorrectness(dtype, batch_dims + (n, 2 * n)) + def testVectorInputRaisesError(self): + # Like GitHub issue 110798 for QR: the graph-level shape check only runs + # for inputs of known rank, so a rank-1 input reaching the compiler + # through an unknown-rank placeholder must be rejected at compile time + # instead of reading a negative trailing dimension. + with self.session() as sess: + x_tf = array_ops.placeholder(np.float32) + with self.test_scope(): + s, u, v = linalg_ops.svd(x_tf, full_matrices=True) + with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"): + sess.run([s, u, v], feed_dict={x_tf: np.zeros([8], np.float32)}) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index df541ca511272c..f1b6ab5f79ecb4 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -850,6 +850,8 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", "//tensorflow/core:lib", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder/lib:constants", @@ -1563,6 +1565,8 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder/lib:qr", ], diff --git a/tensorflow/compiler/tf2xla/kernels/qr_op.cc b/tensorflow/compiler/tf2xla/kernels/qr_op.cc index 6120903fe9c991..4f8646be8683dc 100644 --- a/tensorflow/compiler/tf2xla/kernels/qr_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/qr_op.cc @@ -13,12 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/qr.h" #include "xla/hlo/builder/xla_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" +#include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { namespace { @@ -29,6 +32,14 @@ class QROp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_)); } void Compile(XlaOpKernelContext* ctx) override { + // The rank is only known at compile time when the graph-level shape + // inference saw an unknown rank, so it must be validated here before + // QrExplicit reads the trailing two dimensions. + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); xla::XlaOp q, r; xla::QrExplicit(ctx->Input(0), full_matrices_, q, r); ctx->SetOutput(0, q); diff --git a/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc index 9583341992ce7d..941b6785dc3150 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -51,6 +53,11 @@ class XlaSvdOp : public XlaOpKernel { } } void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); auto result = xla::SVD(ctx->Input(0), max_iter_, epsilon_, precision_config_.operand_precision(0)); ctx->SetOutput(0, result.d); @@ -71,7 +78,14 @@ class SvdOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_)); } void Compile(XlaOpKernelContext* ctx) override { + // The rank is only known at compile time when the graph-level shape + // inference saw an unknown rank, so it must be validated here before + // the trailing two dimensions are read. const TensorShape input_shape = ctx->InputShape("input"); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); int m = input_shape.dim_size(input_shape.dims() - 2); int n = input_shape.dim_size(input_shape.dims() - 1); // This is based on heuristics that approx log(n) sweep updates are needed. From c0af952d92bca3122c4a2663293510c98ee7a875 Mon Sep 17 00:00:00 2001 From: Kuldeeep18 Date: Tue, 1 Sep 2026 12:35:44 +0530 Subject: [PATCH 2/6] Fix type promotion and input handling in tnp.logaddexp and tnp.logaddexp2 --- .../python/ops/numpy_ops/np_math_ops.py | 42 ++++++++++++------- .../np_math_ops_no_numpy_methods_test.py | 22 ++++++++++ .../python/ops/numpy_ops/np_math_ops_test.py | 24 +++++++++++ 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index 18e885f5dbdc0d..aef91a5c8ea27f 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -529,25 +529,39 @@ def f(a, b): @tf_export.tf_export('experimental.numpy.logaddexp', v1=[]) @np_utils.np_doc('logaddexp') def logaddexp(x1, x2): - amax = maximum(x1, x2) - delta = x1 - x2 - return np_array_ops.where( - isnan(delta), - x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp(-abs(delta))), - ) + def f(x1, x2): + if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact): + float_dtype = np_utils.result_type(float) + x1 = math_ops.cast(x1, float_dtype) + x2 = math_ops.cast(x2, float_dtype) + amax = maximum(x1, x2) + delta = x1 - x2 + return np_array_ops.where( + isnan(delta), + x1 + x2, # NaNs or infinities of the same sign. + amax + log1p(exp(-abs(delta))), + ) + + return _bin_op(f, x1, x2) @tf_export.tf_export('experimental.numpy.logaddexp2', v1=[]) @np_utils.np_doc('logaddexp2') def logaddexp2(x1, x2): - amax = maximum(x1, x2) - delta = x1 - x2 - return np_array_ops.where( - isnan(delta), - x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp2(-abs(delta))) / np.log(2), - ) + def f(x1, x2): + if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact): + float_dtype = np_utils.result_type(float) + x1 = math_ops.cast(x1, float_dtype) + x2 = math_ops.cast(x2, float_dtype) + amax = maximum(x1, x2) + delta = x1 - x2 + return np_array_ops.where( + isnan(delta), + x1 + x2, # NaNs or infinities of the same sign. + amax + log1p(exp2(-abs(delta))) / np.log(2), + ) + + return _bin_op(f, x1, x2) @tf_export.tf_export('experimental.numpy.polyval', v1=[]) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py index 18c50d31fbdbce..d963f676d09cae 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py @@ -126,6 +126,28 @@ def testAverageAcceptsIntegerInputs(self): err_msg='average({})'.format(arg), ) + def testLogaddexpAcceptsIntegerInputs(self): + # `logaddexp` and `logaddexp2` promote integer arguments to a float dtype, + # which must not depend on `enable_numpy_methods_on_tensor()`. + x1 = np.array([1, 2, 3], dtype=np.int32) + for x2 in (np.int32(1), np.array([4, 5, 6], dtype=np.int32)): + actual = np_math_ops.logaddexp(x1, x2) + np.testing.assert_allclose( + np.asarray(actual), + np.logaddexp(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='logaddexp({}, {})'.format(x1, x2), + ) + actual2 = np_math_ops.logaddexp2(x1, x2) + np.testing.assert_allclose( + np.asarray(actual2), + np.logaddexp2(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='logaddexp2({}, {})'.format(x1, x2), + ) + def testFloatInputsAreUnchanged(self): arg = np.array([1.5, 2.5, 3.5], dtype=np.float32) for name, tf_fun, np_fun in _PROMOTING_UNARY_OPS: diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index cf7e185cfeaf4e..9d68df03e209b3 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -207,6 +207,30 @@ def testHypotSpecialCases(self): expected = np.hypot(x, y) np.testing.assert_equal(actual.tolist(), expected.tolist()) + def testLogaddexp(self): + self._testBinaryOp(np_math_ops.logaddexp, np.logaddexp, 'logaddexp') + self._testBinaryOp(np_math_ops.logaddexp2, np.logaddexp2, 'logaddexp2') + + def testLogaddexpNonFloatInputs(self): + int_args = [ + ([1, 2, 3], [4, 5, 6]), + (np.array([1, 2], dtype=np.int32), np.array([3, 4], dtype=np.int32)), + (np.array([1, 2], dtype=np.int64), np.array([3, 4], dtype=np.int64)), + (1, 2), + ([1.0, 2.0], [3.0, 4.0]), + ] + for x1, x2 in int_args: + self.match( + np_math_ops.logaddexp(x1, x2), + np.logaddexp(np.asarray(x1), np.asarray(x2)), + msg='logaddexp({}, {})'.format(x1, x2), + ) + self.match( + np_math_ops.logaddexp2(x1, x2), + np.logaddexp2(np.asarray(x1), np.asarray(x2)), + msg='logaddexp2({}, {})'.format(x1, x2), + ) + def match(self, actual, expected, msg='', check_dtype=True): self.assertIsInstance(actual, np_arrays.ndarray) if check_dtype: From a5aab0a478b6f1d8568307828f53e81e7b776230 Mon Sep 17 00:00:00 2001 From: Kuldeeep18 Date: Tue, 1 Sep 2026 13:56:40 +0530 Subject: [PATCH 3/6] Cast np.log(2) in logaddexp2 to operand dtype to preserve precision --- tensorflow/python/ops/numpy_ops/np_math_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index aef91a5c8ea27f..054c92606d1506 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -558,7 +558,7 @@ def f(x1, x2): return np_array_ops.where( isnan(delta), x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp2(-abs(delta))) / np.log(2), + amax + log1p(exp2(-abs(delta))) / math_ops.cast(np.log(2), x1.dtype), ) return _bin_op(f, x1, x2) From f83dc0dd70aeb7822237ec15761d4a02929cf9b6 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Thu, 3 Sep 2026 00:02:01 -0700 Subject: [PATCH 4/6] Rollback of PR #126162 Revert #126162 It needs to be aligned with TFLite's security threat model and only do minimal required changes. Reverts 8aaa9f530d689868d48acef0b4a3cf6f445ddf8f PiperOrigin-RevId: 975544861 --- tensorflow/lite/kernels/transpose.cc | 14 -------------- tensorflow/lite/kernels/transpose_test.cc | 17 ----------------- 2 files changed, 31 deletions(-) diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc index 94bdac9439a21b..0b1f2b783b05bd 100644 --- a/tensorflow/lite/kernels/transpose.cc +++ b/tensorflow/lite/kernels/transpose.cc @@ -50,26 +50,12 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, // Ensure validity of the permutations tensor as a 1D tensor. TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1); TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims); - // `perm` must be a permutation of [0, dims), not merely a set of in-range - // values: the output shape and the element offsets are both derived from it, - // and a repeated entry makes them disagree with the input extent. - // `dims` is bounded by kTransposeMaxDimensions, which Prepare() enforces - // before this function is reachable, so a 64-bit mask is sufficient and - // avoids a heap allocation in the kernel. - static_assert(kTransposeMaxDimensions <= 64, - "Permutation bitmask assumes at most 64 dimensions."); - uint64_t seen = 0; for (int idx = 0; idx < dims; ++idx) { TF_LITE_ENSURE_MSG(context, (perm_data[idx] >= -dims && perm_data[idx] < dims), "Transpose op permutations array is out of bounds."); new_perm_data[idx] = perm_data[idx]; if (new_perm_data[idx] < 0) new_perm_data[idx] += dims; - const uint64_t bit = uint64_t{1} << new_perm_data[idx]; - TF_LITE_ENSURE_MSG( - context, (seen & bit) == 0, - "Transpose op permutations array must not contain duplicate values."); - seen |= bit; } // Determine size of output tensor. diff --git a/tensorflow/lite/kernels/transpose_test.cc b/tensorflow/lite/kernels/transpose_test.cc index 11930f342ce91a..601e8bf4355540 100644 --- a/tensorflow/lite/kernels/transpose_test.cc +++ b/tensorflow/lite/kernels/transpose_test.cc @@ -193,23 +193,6 @@ TEST(TransposeTest, TestPermOutOfBounds) { EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}), "Transpose op permutations array is out of bounds."); } - -// `perm` must be a permutation of [0, dims). A repeated entry passes the range -// check but makes the derived output shape disagree with the input extent, so -// the kernel reads outside the input tensor. -TEST(TransposeTest, TestPermDuplicateValues) { - EXPECT_DEATH( - TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 2}), - "Transpose op permutations array must not contain duplicate values."); -} - -// Duplicates must also be rejected after negative entries are normalised: -// on a rank-2 input {0, -2} normalises to {0, 0}. -TEST(TransposeTest, TestPermDuplicateValuesAfterNegativeNormalization) { - EXPECT_DEATH( - TransposeOpConstModel({2, 3}, {2}, {0, -2}), - "Transpose op permutations array must not contain duplicate values."); -} #endif TEST(TransposeTest, TestInt41DInputConstTensor) { From cf38352c52c1885222bf643a2ea49dfb319bf1d1 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 3 Sep 2026 00:24:07 -0700 Subject: [PATCH 5/6] [XLA][HLO][Benchmark] Add microbenchmarks for HloSchedule::Update PiperOrigin-RevId: 975555261 --- third_party/xla/xla/service/BUILD | 3 + .../xla/xla/service/hlo_schedule_test.cc | 156 ++++++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 9e402b8148fed0..39c41bbce90a62 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2311,6 +2311,7 @@ xla_cc_test( srcs = ["hlo_schedule_test.cc"], deps = [ ":buffer_value", + ":hlo_module_config", "//xla:literal_util", "//xla:shape_util", "//xla:xla_data_proto_cc", @@ -2325,6 +2326,8 @@ xla_cc_test( "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_benchmark//:benchmark", "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/service/hlo_schedule_test.cc b/third_party/xla/xla/service/hlo_schedule_test.cc index 212c7f6556cf62..c6567096e7eda9 100644 --- a/third_party/xla/xla/service/hlo_schedule_test.cc +++ b/third_party/xla/xla/service/hlo_schedule_test.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_schedule.h" +#include +#include #include #include #include @@ -22,7 +24,9 @@ limitations under the License. #include #include #include "absl/algorithm/container.h" +#include "absl/log/check.h" #include "absl/log/log.h" +#include "benchmark/benchmark.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -33,6 +37,7 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h" #include "xla/literal_util.h" #include "xla/service/buffer_value.h" +#include "xla/service/hlo_module_config.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/lib/core/status_test_util.h" @@ -660,5 +665,156 @@ ENTRY %test (arg.0: (f32[], f32[])) -> f32[] { ASSERT_OK(module->schedule().Update()); ASSERT_OK(module->schedule().Verify()); } + +std::unique_ptr BuildBenchmarkModule(int64_t num_instructions) { + HloModuleConfig config; + auto module = std::make_unique("bm_module", config); + HloComputation::Builder builder("entry"); + Shape shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* c0 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction* c1 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0f))); + std::vector sequence; + sequence.reserve(num_instructions + 100); + sequence.push_back(c0); + sequence.push_back(c1); + for (int64_t i = 2; i < num_instructions; ++i) { + HloInstruction* op0 = sequence[i - 1]; + HloInstruction* op1 = sequence[i / 2]; + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op0, op1)); + sequence.push_back(add); + } + HloComputation* entry = module->AddEntryComputation(builder.Build()); + CHECK_OK(module->set_schedule(HloSchedule(module.get()))); + module->schedule().set_sequence(entry, sequence); + return module; +} + +// Benchmarks HloSchedule::Update() when a small number of new instructions (5 +// copies) are inserted into an existing scheduled computation. +void BM_HloScheduleUpdate_InsertFewCopies(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + constexpr int64_t kNumCopies = 5; + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + HloComputation* entry = module->entry_computation(); + const auto& seq = module->schedule().sequence(entry).instructions(); + for (int k = 1; k <= kNumCopies; ++k) { + int64_t target_idx = (num_instructions * k) / (kNumCopies + 1); + HloInstruction* new_copy = + entry->AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kCopy, seq[target_idx - 1])); + CHECK_OK(seq[target_idx]->ReplaceOperandWith(0, new_copy)); + } + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_InsertFewCopies) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when a proportional number of instructions +// (~1%) are inserted into an existing scheduled computation. +void BM_HloScheduleUpdate_Insert1PercentCopies(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + const int64_t num_copies = std::max(1, num_instructions / 100); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + HloComputation* entry = module->entry_computation(); + const auto& seq = module->schedule().sequence(entry).instructions(); + for (int k = 1; k <= num_copies; ++k) { + int64_t target_idx = (num_instructions * k) / (num_copies + 1); + HloInstruction* new_copy = + entry->AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kCopy, seq[target_idx - 1])); + CHECK_OK(seq[target_idx]->ReplaceOperandWith(0, new_copy)); + } + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_Insert1PercentCopies) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when instructions are removed from the +// computation. +void BM_HloScheduleUpdate_RemoveFewInstructions(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + constexpr int64_t kNumRemovals = 5; + HloModuleConfig config; + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = std::make_unique("bm_module", config); + HloComputation::Builder builder("entry"); + HloInstruction* c0 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + std::vector sequence; + sequence.push_back(c0); + HloInstruction* prev = c0; + std::vector dead_instructions; + for (int64_t i = 1; i < num_instructions; ++i) { + if (dead_instructions.size() < kNumRemovals && + i % (num_instructions / (kNumRemovals + 1)) == 0) { + HloInstruction* dead = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, c0)); + sequence.push_back(dead); + dead_instructions.push_back(dead); + } else { + prev = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, prev)); + sequence.push_back(prev); + } + } + HloComputation* entry = module->AddEntryComputation(builder.Build()); + CHECK_OK(module->set_schedule(HloSchedule(module.get()))); + module->schedule().set_sequence(entry, sequence); + + for (HloInstruction* dead : dead_instructions) { + CHECK_OK(entry->RemoveInstruction(dead)); + } + + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_RemoveFewInstructions) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when no modifications were made. +void BM_HloScheduleUpdate_NoModifications(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_NoModifications) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + } // namespace } // namespace xla From 0d8b7bfaf147b0d3d1b357bce7802ec7e67a22df Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Thu, 3 Sep 2026 00:34:09 -0700 Subject: [PATCH 6/6] GemmFusion test does not belong here. I could add it to GemmFusion, but I don't think it is a very useful test. Broadcasts would normally get rejected as a user due to IsOutputWorthFusing where output bytes is increased. However, this one doesn't do that because the broadcast is the ROOT and there's nothing else to fuse it to so it says that it's still worth it. V1 actually only rejects it because it doesn't propagate tiling through the broadcast user, but V2 has no issue with this. PiperOrigin-RevId: 975559504 --- .../gpu/triton_fusion_analysis_test.cc | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc index 74325b23ede671..947054a2c0d509 100644 --- a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc +++ b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc @@ -683,27 +683,6 @@ e { ::testing::HasSubstr("Unsupported broadcast"))); } -TEST_F(TritonDotAnalysisTest, OutputBroadcastIsNotAccepted) { - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( -HloModule t - -ENTRY e { - p0 = f16[2,35] parameter(0) - p0c = bf16[2,35] convert(p0) - p1 = bf16[35,2] parameter(1) - dot = bf16[2,2] dot(p0c, p1), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT bc = bf16[2,2,100] broadcast(dot), dimensions={0,1} -})")); - EXPECT_TRUE(GemmFusion(se::GpuComputeCapability{se::CudaComputeCapability{ - se::CudaComputeCapability::kAmpere, 0}}) - .Run(module.get()) - .value()); - EXPECT_EQ(module->entry_computation()->root_instruction()->opcode(), - HloOpcode::kBroadcast); -} - TEST_F(TritonDotAnalysisTest, DegenerateSplitFragmentIsHandled) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"(