From 9a6127fb085aa19c70915a6ad7701fb5d04856ee Mon Sep 17 00:00:00 2001 From: Raghunandan Kumar <54378462+RaghunandanKumar@users.noreply.github.com> Date: Wed, 10 Jun 2026 23:07:19 -0400 Subject: [PATCH 01/26] Fix dataset save in debug mode --- tensorflow/python/data/kernel_tests/BUILD | 1 + tensorflow/python/data/kernel_tests/io_test.py | 10 ++++++++++ tensorflow/python/data/ops/save_op.py | 2 +- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD index 72987b0ee305d4..0e69b0babf98e0 100644 --- a/tensorflow/python/data/kernel_tests/BUILD +++ b/tensorflow/python/data/kernel_tests/BUILD @@ -1408,6 +1408,7 @@ py_test( deps = [ ":checkpoint_test_base", ":test_base", + "//tensorflow/python/data/ops:debug_mode", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:combinations", diff --git a/tensorflow/python/data/kernel_tests/io_test.py b/tensorflow/python/data/kernel_tests/io_test.py index ad2945dce322d5..fe319901d62be1 100644 --- a/tensorflow/python/data/kernel_tests/io_test.py +++ b/tensorflow/python/data/kernel_tests/io_test.py @@ -23,6 +23,7 @@ import numpy as np from tensorflow.python.data.kernel_tests import checkpoint_test_base from tensorflow.python.data.kernel_tests import test_base +from tensorflow.python.data.ops import debug_mode from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import def_function from tensorflow.python.framework import combinations @@ -70,6 +71,15 @@ def testCardinality(self): dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec) self.assertEqual(self.evaluate(dataset2.cardinality()), 42) + @combinations.generate(test_base.eager_only_combinations()) + def testSaveInDebugModeWithoutShardFunction(self): + debug_mode.toggle_debug_mode(True) + self.addCleanup(debug_mode.toggle_debug_mode, False) + dataset = dataset_ops.Dataset.range(42) + self.evaluate(dataset.save(self._test_dir)) + dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec) + self.assertDatasetProduces(dataset2, range(42)) + @combinations.generate(test_base.default_test_combinations()) def testCustomShardFunction(self): dataset = dataset_ops.Dataset.range(42) diff --git a/tensorflow/python/data/ops/save_op.py b/tensorflow/python/data/ops/save_op.py index c5a63477aeeb3f..5cd78a93b502b4 100644 --- a/tensorflow/python/data/ops/save_op.py +++ b/tensorflow/python/data/ops/save_op.py @@ -97,7 +97,7 @@ def set_save_dataset_attributes(dataset, shard_func, path): """Sets parameters for SaveDatasetOp and SaveDatasetV2Op.""" if shard_func is None: use_shard_func = False - shard_func = lambda *x: None # a dummy function that will not be used + shard_func = lambda *x: 0 # a dummy function that will not be used else: use_shard_func = True wrapped_func = structured_function.StructuredFunctionWrapper( From 82dc2cfcfe2684165264989df05a9a4b380bfdd2 Mon Sep 17 00:00:00 2001 From: anupamme Date: Thu, 30 Jul 2026 00:43:22 +0000 Subject: [PATCH 02/26] fix: V-003 security vulnerability Automated security fix generated by OrbisAI Security --- tensorflow/core/data/service/dispatcher_impl.cc | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/data/service/dispatcher_impl.cc b/tensorflow/core/data/service/dispatcher_impl.cc index cb480c8832a774..5a79e479d00097 100644 --- a/tensorflow/core/data/service/dispatcher_impl.cc +++ b/tensorflow/core/data/service/dispatcher_impl.cc @@ -205,12 +205,16 @@ absl::Status ValidateDatasetId(const std::string& dataset_id) { absl::StrCat("Invalid dataset ID: ", dataset_id, ". Dataset IDs must not contain '/'.")); } + if (absl::StrContains(dataset_id, '\\')) { + return absl::InvalidArgumentError( + absl::StrCat("Invalid dataset ID: ", dataset_id, + ". Dataset IDs must not contain '\\'.")); + } #if defined(_WIN32) - if (absl::StrContains(dataset_id, '\\') || - absl::StrContains(dataset_id, ':')) { + if (absl::StrContains(dataset_id, ':')) { return absl::InvalidArgumentError( absl::StrCat("Invalid dataset ID: ", dataset_id, - ". Dataset IDs must not contain '\\' or ':'.")); + ". Dataset IDs must not contain ':'.")); } #endif return absl::OkStatus(); From 1dab113e08e7efc086aa93ffb01e496c6fb278fb Mon Sep 17 00:00:00 2001 From: Anupam Mediratta Date: Fri, 31 Jul 2026 06:59:46 +0530 Subject: [PATCH 03/26] test: add gRPC integration test for invalid dataset ID rejection Adds GetOrRegisterDatasetInvalidDatasetId to grpc_dispatcher_impl_test.cc to verify that backslash, slash, '.', and '..' dataset IDs are rejected via the public gRPC API rather than calling the internal ValidateDatasetId function directly. --- .../data/service/grpc_dispatcher_impl_test.cc | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc b/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc index e68dc402565c55..1c9c295256669d 100644 --- a/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc +++ b/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc @@ -175,6 +175,28 @@ TEST_F(GrpcDispatcherImplTest, GetSplitInvalidProviderIndex) { } } +TEST_F(GrpcDispatcherImplTest, GetOrRegisterDatasetInvalidDatasetId) { + const std::vector invalid_ids = { + "..\\..\\etc\\passwd", + "a\\b", + "a/b", + ".", + "..", + }; + + for (const auto& id : invalid_ids) { + ClientContext ctx; + GetOrRegisterDatasetRequest req; + *req.mutable_dataset()->mutable_graph() = testing::RangeDataset(10).graph(); + req.set_dataset_id(id); + GetOrRegisterDatasetResponse resp; + ::grpc::Status status = + dispatcher_client_stub_->GetOrRegisterDataset(&ctx, req, &resp); + EXPECT_EQ(status.error_code(), ::grpc::StatusCode::INVALID_ARGUMENT) + << "Dataset ID '" << id << "' should be rejected."; + } +} + } // namespace } // namespace data } // namespace tensorflow From 3b963f43aa7c363a6e2af41a2ca5bf64c4aa8116 Mon Sep 17 00:00:00 2001 From: AnupamKumar-1 Date: Wed, 19 Aug 2026 11:17:13 +0530 Subject: [PATCH 04/26] Fix eager tf.while_loop corrupting shape for single-var loop with bare-tensor body --- .../python/ops/control_flow_ops_test.py | 21 ++++++++++++++++++- tensorflow/python/ops/while_loop.py | 12 +++++++---- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py index 83e9493004b18a..e93119e3db5d05 100644 --- a/tensorflow/python/ops/control_flow_ops_test.py +++ b/tensorflow/python/ops/control_flow_ops_test.py @@ -1648,6 +1648,25 @@ def testWhileLoopSameReturnShape_TrueSingleLoopVar(self): c, b, [i], return_same_structure=True, maximum_iterations=50) self.assertEqual(self.evaluate(r), [10]) + @test_util.run_v2_only + def testEagerWhileLoopSingleLoopVarBareTensorBodyPreservesShape(self): + x = constant_op.constant([[5.0]], shape=[1, 1]) + shapes_seen = [] + + def cond(x): + shapes_seen.append(x.shape.as_list()) + return math_ops.greater(x[0, 0], 3) + + def body(x): + return x - 1 + + r = while_loop.while_loop(cond, body, [x], return_same_structure=True) + + self.assertAllEqual(shapes_seen, [[1, 1]] * len(shapes_seen)) + self.assertLen(shapes_seen, 3) + self.assertIsInstance(r, list) + self.assertAllClose(self.evaluate(r), [[[3.0]]]) + @test_util.enable_control_flow_v2 @test_util.run_in_graph_and_eager_modes def testSkipsUnnecessaryCaptureGradients(self): @@ -1858,4 +1877,4 @@ def f(): if __name__ == "__main__": - googletest.main() + googletest.main() \ No newline at end of file diff --git a/tensorflow/python/ops/while_loop.py b/tensorflow/python/ops/while_loop.py index d1964dcfdf437a..4ebf4525be8e3e 100644 --- a/tensorflow/python/ops/while_loop.py +++ b/tensorflow/python/ops/while_loop.py @@ -482,13 +482,17 @@ def while_loop(cond, if executing_eagerly: packed = False # whether the body result was packed into a 1-item tuple + orig_loop_vars_type = ( + type(loop_vars) if isinstance(loop_vars, (list, tuple)) else list) + loop_var_structure = nest.map_structure(type_spec.type_spec_from_value, list(loop_vars)) while cond(*loop_vars): loop_vars = body(*loop_vars) - if try_to_pack and not isinstance(loop_vars, (list, tuple)): - packed = True - loop_vars = (loop_vars,) + if not isinstance(loop_vars, (list, tuple)): + if try_to_pack: + packed = True + loop_vars = orig_loop_vars_type((loop_vars,)) nest.assert_same_structure(loop_var_structure, list(loop_vars)) def convert(x): @@ -520,4 +524,4 @@ def convert(x): if maximum_iterations is not None: return result[1] else: - return result + return result \ No newline at end of file From 03b36ddaf3c7e73f5fc3aa85803c9a232f91841c Mon Sep 17 00:00:00 2001 From: AnupamKumar-1 Date: Wed, 19 Aug 2026 11:46:23 +0530 Subject: [PATCH 05/26] Restrict orig_loop_vars_type to exact list/tuple types, not isinstance --- tensorflow/python/ops/while_loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/while_loop.py b/tensorflow/python/ops/while_loop.py index 4ebf4525be8e3e..3fe5667855564a 100644 --- a/tensorflow/python/ops/while_loop.py +++ b/tensorflow/python/ops/while_loop.py @@ -483,7 +483,7 @@ def while_loop(cond, packed = False # whether the body result was packed into a 1-item tuple orig_loop_vars_type = ( - type(loop_vars) if isinstance(loop_vars, (list, tuple)) else list) + type(loop_vars) if type(loop_vars) in (list, tuple) else list) loop_var_structure = nest.map_structure(type_spec.type_spec_from_value, list(loop_vars)) From 0e196de886cbee006be654bce2480901a0d47070 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Wed, 19 Aug 2026 04:25:00 -0700 Subject: [PATCH 06/26] fix double unlocking by using scope block --- .../core/util/tensor_slice_reader_cache.cc | 63 ++++++++++--------- 1 file changed, 35 insertions(+), 28 deletions(-) diff --git a/tensorflow/core/util/tensor_slice_reader_cache.cc b/tensorflow/core/util/tensor_slice_reader_cache.cc index ddb3e36d1e6dbe..5d1ff1479d6739 100644 --- a/tensorflow/core/util/tensor_slice_reader_cache.cc +++ b/tensorflow/core/util/tensor_slice_reader_cache.cc @@ -52,8 +52,6 @@ TensorSliceReaderCache::~TensorSliceReaderCache() { const TensorSliceReader* TensorSliceReaderCache::GetReader( const std::string& filepattern, TensorSliceReader::OpenTableFunction open_function, int preferred_shard) { - mutex_lock l(mu_); - #if defined(__GXX_RTTI) || defined(_CPPRTTI) // Get the function pointer from the open_function value. TensorSliceReaderCache::OpenFuncType* func_ptr = @@ -72,22 +70,42 @@ const TensorSliceReader* TensorSliceReaderCache::GetReader( return nullptr; } - // Wait if another thread is already trying to open the same files. - while (still_opening_.find(filepattern) != still_opening_.end()) { - cv_.wait(l); + TensorSliceReader* reader = nullptr; + + // scope block for lock + { + mutex_lock l(mu_); + + // Wait if another thread is already trying to open the same files. + while (still_opening_.find(filepattern) != still_opening_.end()) { + cv_.wait(l); + } + + if (readers_.find(filepattern) != readers_.end()) { + auto cached_val = readers_[filepattern]; + if (cached_val.first == *func_ptr) { + reader = cached_val.second; + VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": " + << reader; + } else { + LOG(WARNING) << "Caching disabled because the checkpoint file " + << "is being opened with two different open functions: " + << filepattern; + } + return reader; + } + + still_opening_.insert(filepattern); } - TensorSliceReader* reader = nullptr; - if (readers_.find(filepattern) == readers_.end()) { + // no lock for expensive constructing TensorSliceReader + TensorSliceReader* tmp_reader( + new TensorSliceReader(filepattern, open_function, preferred_shard)); + + // scope block for lock + { VLOG(1) << "Creating new TensorSliceReader for " << filepattern; - still_opening_.insert(filepattern); - // Release the lock temporary as constructing TensorSliceReader is - // expensive. - mu_.unlock(); - TensorSliceReader* tmp_reader( - new TensorSliceReader(filepattern, open_function, preferred_shard)); - // Acquire the lock again. - mu_.lock(); + mutex_lock l(mu_); if (tmp_reader->status().ok()) { reader = tmp_reader; readers_[filepattern] = std::make_pair(*func_ptr, reader); @@ -96,20 +114,9 @@ const TensorSliceReader* TensorSliceReaderCache::GetReader( } CHECK_EQ(size_t{1}, still_opening_.erase(filepattern)); VLOG(1) << "Cached TensorSliceReader for " << filepattern << ": " << reader; - } else { - auto cached_val = readers_[filepattern]; - if (cached_val.first == *func_ptr) { - reader = cached_val.second; - VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": " - << reader; - } else { - LOG(WARNING) << "Caching disabled because the checkpoint file " - << "is being opened with two different open functions: " - << filepattern; - } - } - cv_.notify_all(); + cv_.notify_all(); + } return reader; } From 788ff3356ec53f83f19dfbcec4d004f0ccbcdea9 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Wed, 19 Aug 2026 04:44:56 -0700 Subject: [PATCH 07/26] optimize two lookup by storing value inside local variable --- tensorflow/core/util/tensor_slice_reader_cache.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/util/tensor_slice_reader_cache.cc b/tensorflow/core/util/tensor_slice_reader_cache.cc index 5d1ff1479d6739..1eb3410ae581c1 100644 --- a/tensorflow/core/util/tensor_slice_reader_cache.cc +++ b/tensorflow/core/util/tensor_slice_reader_cache.cc @@ -81,8 +81,9 @@ const TensorSliceReader* TensorSliceReaderCache::GetReader( cv_.wait(l); } - if (readers_.find(filepattern) != readers_.end()) { - auto cached_val = readers_[filepattern]; + auto it = readers_.find(filepattern); + if (it != readers_.end()) { + auto cached_val = it->second; if (cached_val.first == *func_ptr) { reader = cached_val.second; VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": " From 7a777d59a18799965964466f79f33d10e2fbddb1 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Wed, 19 Aug 2026 11:58:32 -0700 Subject: [PATCH 08/26] Re-trigger CI From 2a6a3316edb4cf2c3fe304b6b1cf80840c50a5f0 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 03:40:25 -0400 Subject: [PATCH 09/26] Fix configure.py crash when clang reports no parseable version retrieve_clang_version returns None when the clang executable does not report a parseable version, which happens with wrappers such as ccache or any tool whose --version banner lacks a usable version token. disable_clang_offsetof_extension then crashed with AttributeError while deciding whether to add -Wno-gnu-offsetof-extensions for clang 16 and 17. Non-numeric version strings would crash the same way with ValueError. Both cases now skip the flag, matching every confirmed version outside 16 and 17; the existing "current clang installation version unknown" warning still prints. Test Plan: python -m py_compile configure.py python -m pylint --rcfile=tensorflow/tools/ci_build/pylintrc \ configure.py # rated 10.00/10, same as master Direct call matrix on disable_clang_offsetof_extension: None, empty string and 'unknown' return without writing; '16.0.0' and '17.1.8' write build --copt=-Wno-gnu-offsetof-extensions; '22.1.8', '15.x' and '18.x' write nothing, byte identical to master bazelrc output. On master, None reproduced "AttributeError: 'NoneType' object has no attribute 'split'" from issue 125939. --- configure.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/configure.py b/configure.py index 2bb9c197acf9f2..b9b6de97aa8f7d 100644 --- a/configure.py +++ b/configure.py @@ -937,7 +937,13 @@ def retrieve_clang_version(clang_executable): # offset of in the current version of ubp. See # https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183. def disable_clang_offsetof_extension(clang_version): - if int(clang_version.split('.')[0]) in (16, 17): + if not clang_version: + return + try: + clang_major_version = int(clang_version.split('.')[0]) + except ValueError: + return + if clang_major_version in (16, 17): write_to_bazelrc('build --copt=-Wno-gnu-offsetof-extensions') From 9e25c9724613a5eb1e581861ae9900139880c9f7 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 03:40:25 -0400 Subject: [PATCH 10/26] Treat unparseable clang output as unknown version in configure.py retrieve_clang_version indexed the first character of whatever the clang executable printed, so a wrapper that exits with nothing on stdout crashed with IndexError. The parsed result is now always handled as a token list: an empty or unrecognized banner reaches the existing unknown-version warning and returns None, and an adjacent or truncated "clang version " marker does the same instead of raising IndexError. Together with the previous commit this makes configure survive any clang whose --version output cannot be parsed. Version-bearing outputs are unchanged. One deliberate difference: banners without a version token no longer print the misleading "not a release version" warning that master produced by accident, because that check now runs over tokens instead of characters. A token-free banner starting with a digit used to be reported as a bogus one digit version; it is now reported as unknown. Test Plan: Monkeypatched run_shell harness over retrieve_clang_version and the retrieve plus disable_clang_offsetof_extension flow on this branch: empty stdout, whitespace only stdout, banner without a version token, digit-initial token-free banner, adjacent "clang version " separators and a truncated marker all return None with the unknown-version warning and no crash; "clang version 22.1.8" returns 22.1.8; "Ubuntu clang version 18.0.0git" returns 18.0.0 with the prerelease warning. On master the empty-output input reproduced "IndexError: string index out of range". pylint --rcfile=tensorflow/tools/ci_build/pylintrc rates configure.py 10.00/10. --- configure.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/configure.py b/configure.py index b9b6de97aa8f7d..76077c3e865082 100644 --- a/configure.py +++ b/configure.py @@ -914,7 +914,13 @@ def retrieve_clang_version(clang_executable): curr_version_split = curr_version.lower().split('clang version ') if len(curr_version_split) > 1: - curr_version = curr_version_split[1].split()[0].split('git') + tokens = curr_version_split[1].split() + if not tokens: + print('WARNING: current clang installation version unknown.\n') + return None + curr_version = tokens[0].split('git') + else: + curr_version = [curr_version] if len(curr_version) > 1: print('WARNING: current clang installation is not a release version.\n') From ec2128ab5e1636db035c17d8ace73cd198c6061f Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 05:18:05 -0400 Subject: [PATCH 11/26] Re-run CI From 4e43c5967d8fce9abf8a7010b587258ed4fce25b Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 06:29:04 -0400 Subject: [PATCH 12/26] Restructure unknown clang version handling per review Split the no "clang version " marker case into its own early return instead of wrapping the raw banner in a single-element list. Same observable behavior on every input, clearer control flow. Test Plan: Monkeypatched run_shell matrix rerun after the change: empty stdout, whitespace only, banner without token, digit-initial token-free banner, adjacent separators and truncated marker all return None with the unknown-version warning; 'clang version 22.1.8' returns 22.1.8; 'Ubuntu clang version 18.0.0git' returns 18.0.0 with the prerelease warning. python -m py_compile exit 0; pylint with tensorflow/tools/ci_build/pylintrc rates configure.py 10.00/10. --- configure.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/configure.py b/configure.py index 76077c3e865082..e176e1f09d564a 100644 --- a/configure.py +++ b/configure.py @@ -913,14 +913,16 @@ def retrieve_clang_version(clang_executable): stderr=stderr) curr_version_split = curr_version.lower().split('clang version ') - if len(curr_version_split) > 1: - tokens = curr_version_split[1].split() - if not tokens: - print('WARNING: current clang installation version unknown.\n') - return None - curr_version = tokens[0].split('git') - else: - curr_version = [curr_version] + if len(curr_version_split) <= 1: + print('WARNING: current clang installation version unknown.\n') + return None + + tokens = curr_version_split[1].split() + if not tokens: + print('WARNING: current clang installation version unknown.\n') + return None + + curr_version = tokens[0].split('git') if len(curr_version) > 1: print('WARNING: current clang installation is not a release version.\n') From e6d3f7819a3f3513c6c5bdc24341179f0af349a7 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 10:33:26 -0400 Subject: [PATCH 13/26] Re-run CI From 0ba6980a2106450fa31b9ff5f36b889bd431e17c Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Wed, 26 Aug 2026 10:06:49 -0700 Subject: [PATCH 14/26] Remove the non-value-preserving 1/y to Reciprocal grappler rewrite Grappler's constant folding rewrote Div(ones, y) into Reciprocal(y), assuming the two are numerically equivalent. They are not on x86: the CPU Reciprocal kernel uses Eigen's fast-math preciprocal for float, which computes rcp plus one Newton-Raphson step under EIGEN_FAST_MATH and is accurate to about 1 ulp rather than exactly IEEE division, while the packet tail and the Div kernel divide exactly. As a result, 1.0 / x returned different values in eager and graph mode on x86, for example 0.99999994 instead of 1.0 for x = 1.0. Remove the rewrite so 1 / y stays a true division in optimized graphs. Users who want the faster approximate reciprocal can still call tf.math.reciprocal explicitly. The ReduceDivToReciprocalMul strength reduction for division by constants is documented as such and is left unchanged. The new Python test asserts both that no Reciprocal appears in the optimized graph and that 1.0 / ones is exactly ones, so it fails on any platform if the rewrite comes back. Fixes #102771 --- .../grappler/optimizers/constant_folding.cc | 24 ++++--------------- .../grappler/optimizers/constant_folding.h | 1 - .../optimizers/constant_folding_test.cc | 8 ++++--- .../python/grappler/constant_folding_test.py | 19 +++++++++++++++ 4 files changed, 28 insertions(+), 24 deletions(-) diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc index 6801d61262ae42..e0be95530d5ab1 100644 --- a/tensorflow/core/grappler/optimizers/constant_folding.cc +++ b/tensorflow/core/grappler/optimizers/constant_folding.cc @@ -2042,17 +2042,6 @@ void ConstantFolding::ReplaceBinaryOperationWithBroadcastTo( graph_modified_ = true; } -void ConstantFolding::ReplaceDivisionOfOnesByReciprocal(NodeDef* node, - GraphDef* graph) { - node->set_op("Reciprocal"); - node->mutable_input()->SwapElements(0, 1); - const std::string ctrl_dep = - AddControlDependency(node->input(1), graph, node_map_.get()); - node_map_->UpdateInput(node->name(), node->input(1), ctrl_dep); - node->set_input(1, ctrl_dep); - graph_modified_ = true; -} - void ConstantFolding::ReplaceSubtractionFromZeroByNegation(NodeDef* node, GraphDef* graph) { node->set_op("Neg"); @@ -3047,15 +3036,10 @@ absl::Status ConstantFolding::SimplifyArithmeticOperations( return absl::OkStatus(); } - // Replace 1 / y with Reciprocal op. - if (y_matches_output_shape && is_any_div && x_is_one) { - TF_RETURN_IF_ERROR(CheckAttrExists(*node, "T")); - DataType type = node->attr().at("T").type(); - if (DataTypeIsFloating(type) || DataTypeIsComplex(type)) { - ReplaceDivisionOfOnesByReciprocal(node, optimized_graph); - return absl::OkStatus(); - } - } + // Note: 1 / y is intentionally not rewritten to Reciprocal(y). The CPU + // Reciprocal kernel uses Eigen's fast-math reciprocal for float, which + // is not exactly IEEE division, so the rewrite silently changed results + // between eager and graph execution on x86 (see issue #102771). const bool y_is_zero = IsZeros(*y); const bool y_is_one = y_is_zero ? false : IsOnes(*y); diff --git a/tensorflow/core/grappler/optimizers/constant_folding.h b/tensorflow/core/grappler/optimizers/constant_folding.h index 2ac90cd681a1b8..5e9113414aa408 100644 --- a/tensorflow/core/grappler/optimizers/constant_folding.h +++ b/tensorflow/core/grappler/optimizers/constant_folding.h @@ -131,7 +131,6 @@ class ConstantFolding : public GraphOptimizer { NodeDef* node, GraphDef* graph); - void ReplaceDivisionOfOnesByReciprocal(NodeDef* node, GraphDef* graph); absl::Status FoldGraph( const GraphProperties& properties, GraphDef* output, absl::flat_hash_set* nodes_to_not_simplify); diff --git a/tensorflow/core/grappler/optimizers/constant_folding_test.cc b/tensorflow/core/grappler/optimizers/constant_folding_test.cc index 49137d15140bcc..8ef638accac8ab 100644 --- a/tensorflow/core/grappler/optimizers/constant_folding_test.cc +++ b/tensorflow/core/grappler/optimizers/constant_folding_test.cc @@ -888,9 +888,11 @@ TEST_F(ConstantFoldingTest, NeutralElement) { EXPECT_EQ("x", node.input(0)); EXPECT_EQ(ctrl_ones_name, node.input(1)); } else if (name == "div2") { - EXPECT_EQ("Reciprocal", node.op()); - EXPECT_EQ("y", node.input(0)); - EXPECT_EQ(ctrl_ones_name, node.input(1)); + // ones / y is not rewritten to Reciprocal(y): the CPU Reciprocal + // kernel is not exactly IEEE division for float (see issue #102771). + EXPECT_EQ("Div", node.op()); + EXPECT_EQ(ones_name, node.input(0)); + EXPECT_EQ("y", node.input(1)); } else if (name == "floordiv") { EXPECT_EQ("FloorDiv", node.op()); EXPECT_EQ("x", node.input(0)); diff --git a/tensorflow/python/grappler/constant_folding_test.py b/tensorflow/python/grappler/constant_folding_test.py index 9a0e9fd73b73e1..71f940fc1360c8 100644 --- a/tensorflow/python/grappler/constant_folding_test.py +++ b/tensorflow/python/grappler/constant_folding_test.py @@ -102,6 +102,25 @@ def f(x, y): self.assertEqual(assign_count, 1) self.assertLen(graphs[0].node, 11) + # See GitHub issue #102771. + def testDivisionOfOnesNotRewrittenToReciprocal(self): + # 1 / x must stay a true division in the optimized graph. The CPU + # Reciprocal kernel uses Eigen's fast-math reciprocal for float, which + # is not exactly IEEE division, so rewriting 1 / x to Reciprocal(x) + # changed results between eager and graph execution on x86. + + @def_function.function + def div_by_ones(x): + return 1.0 / x + + with context.eager_mode(): + ones = array_ops.ones([64], dtype=dtypes.float32) + with context.collect_graphs(optimized=True) as graphs: + result = div_by_ones(ones).numpy() + self.assertLen(graphs, 1) + self.assertNotIn('Reciprocal', [node.op for node in graphs[0].node]) + self.assertAllEqual(result, np.ones([64], dtype=np.float32)) + if __name__ == '__main__': test.main() From 75c2e497e0fe7081bc87cc0e33bf7a3e4c99bfe5 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Thu, 27 Aug 2026 04:08:57 -0400 Subject: [PATCH 15/26] Replace new print calls in configure.py with sys.stdout.write Internal presubmit lint rejects newly introduced print() statements in Python source. The two new warning branches in retrieve_clang_version now use sys.stdout.write instead, matching the requested fix. Test Plan: python -c 'import ast; ast.parse(open("configure.py").read())' -> ok --- configure.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/configure.py b/configure.py index e176e1f09d564a..25d70c1bfe56f7 100644 --- a/configure.py +++ b/configure.py @@ -914,12 +914,12 @@ def retrieve_clang_version(clang_executable): curr_version_split = curr_version.lower().split('clang version ') if len(curr_version_split) <= 1: - print('WARNING: current clang installation version unknown.\n') + sys.stdout.write('WARNING: current clang installation version unknown.\n') return None tokens = curr_version_split[1].split() if not tokens: - print('WARNING: current clang installation version unknown.\n') + sys.stdout.write('WARNING: current clang installation version unknown.\n') return None curr_version = tokens[0].split('git') From 249f61ca7262a6aee02264c7da6fd4d8cf4847c9 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Mon, 31 Aug 2026 05:22:01 -0700 Subject: [PATCH 16/26] Add ConvertMlirBytecode to Windows export symbols Exports tflite::ConvertMlirBytecode on Windows in _pywrap_tensorflow.def and symbols_pybind.txt to resolve undefined symbol errors during linking of __pywrap_tensorflow_0_shared_object.dll in Windows x86 presubmit builds. PiperOrigin-RevId: 973824659 --- tensorflow/python/_pywrap_tensorflow.def | 1 + tensorflow/tools/def_file_filter/symbols_pybind.txt | 1 + third_party/xla/tools/def_file_filter/symbols_pybind.txt | 1 + 3 files changed, 3 insertions(+) diff --git a/tensorflow/python/_pywrap_tensorflow.def b/tensorflow/python/_pywrap_tensorflow.def index 70aca4798f4b53..937fd5cc377552 100644 --- a/tensorflow/python/_pywrap_tensorflow.def +++ b/tensorflow/python/_pywrap_tensorflow.def @@ -256,6 +256,7 @@ EXPORTS ?ComputeGradient@Tape@gradients@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@3@V?$Span@QEAVAbstractTensorHandle@tensorflow@@@56@11V?$Span@PEAVAbstractTensorHandle@tensorflow@@@56@@Z ?Convert@PythonTensorConverter@tensorflow@@QEBA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@AEAW4DataType@2@PEA_N@Z ?Convert@tflite@@YAPEAU_object@@PEAU2@00_N0PEBVPyFunctionLibrary@quantization@tensorflow@@@Z + ?ConvertMlirBytecode@tflite@@YAPEAU_object@@PEAU2@00@Z ?ConvertPyObjectToAttributeType@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@W4AttributeType@1@@Z ?ConvertPythonAPIParameters@tensorflow@@YA_NAEBVPythonAPIInfo@1@AEBVPythonTensorConverter@1@V?$Span@PEAU_object@@@lts_20260526@absl@@PEAUInferredAttributes@21@@Z ?ConvertToEagerTensor@tensorflow@@YAPEAUTFE_TensorHandle@@PEAUTFE_Context@@PEAU_object@@W4DataType@1@PEBD@Z diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt index 003a77aa2965b9..f8fce5e7322775 100644 --- a/tensorflow/tools/def_file_filter/symbols_pybind.txt +++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt @@ -259,6 +259,7 @@ tensorflow::ImportFunction [//tensorflow/compiler/mlir/lite/python:converter_python_api] # converter_python_api tflite::Convert +tflite::ConvertMlirBytecode tflite::MlirQuantizeModel tflite::MlirSparsifyModel tflite::RegisterCustomOpdefs diff --git a/third_party/xla/tools/def_file_filter/symbols_pybind.txt b/third_party/xla/tools/def_file_filter/symbols_pybind.txt index 003a77aa2965b9..f8fce5e7322775 100644 --- a/third_party/xla/tools/def_file_filter/symbols_pybind.txt +++ b/third_party/xla/tools/def_file_filter/symbols_pybind.txt @@ -259,6 +259,7 @@ tensorflow::ImportFunction [//tensorflow/compiler/mlir/lite/python:converter_python_api] # converter_python_api tflite::Convert +tflite::ConvertMlirBytecode tflite::MlirQuantizeModel tflite::MlirSparsifyModel tflite::RegisterCustomOpdefs From 7d3f38c32f018f2ef4c92e0769ca7ece23eb3f89 Mon Sep 17 00:00:00 2001 From: Dragan Mladjenovic Date: Mon, 31 Aug 2026 05:32:04 -0700 Subject: [PATCH 17/26] PR #47502: [ROCm] Separate lit test that use FileCheck only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47502 ๐Ÿ“ Summary of Changes Tag FileCheck only gpu tests with gpu tag only ๐ŸŽฏ Justification This allows them to be run on rocm_cpu step of CI ๐Ÿš€ Kind of Contribution ๐Ÿงช Tests ๐Ÿ“Š Benchmark (for Performance Improvements) N\A ๐Ÿงช Unit Tests: None ๐Ÿงช Execution Tests: None Copybara import of the project: -- 1d1f8ebd6c98d540aedb33625d4e1f0f8af756c2 by Dragan Mladjenovic : [ROCm] Separate lit test that use FileCheck only This allows them to be run on rocm_cpu step of CI Merging this change closes #47502 PiperOrigin-RevId: 973827849 --- .../backends/gpu/codegen/emitters/tests/BUILD | 23 ++++++++++- .../xla/xla/codegen/emitters/tests/BUILD | 41 ++++++++++++++++++- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD index fb27b51af9de9f..4201819a1b030f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD @@ -25,9 +25,18 @@ package( licenses = ["notice"], ) +_FILECHECK_ONLY_TESTS = [ + "reduce_column_small/f32_32_v2.hlo", + "reduce_column_small/s8_f32_32_v4.hlo", + "scatter/sorted_indices_large.hlo", +] + lit_test_suite( name = "tests", - srcs = glob(["**/*.hlo"]), + srcs = glob( + ["**/*.hlo"], + exclude = _FILECHECK_ONLY_TESTS, + ), cfg = "//xla:lit.cfg.py", default_tags = tf_cuda_tests_tags(), exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}), @@ -38,3 +47,15 @@ lit_test_suite( "@llvm-project//llvm:FileCheck", ], ) + +lit_test_suite( + name = "filecheck_only_tests", + srcs = _FILECHECK_ONLY_TESTS, + cfg = "//xla:lit.cfg.py", + default_tags = ["gpu"], + tools = [ + "//xla/backends/gpu/codegen/tools:fusion_to_mlir", + "//xla/codegen/tools:emitters_opt", + "@llvm-project//llvm:FileCheck", + ], +) diff --git a/third_party/xla/xla/codegen/emitters/tests/BUILD b/third_party/xla/xla/codegen/emitters/tests/BUILD index a1f04cceb1584c..7cedec6f546e7c 100644 --- a/third_party/xla/xla/codegen/emitters/tests/BUILD +++ b/third_party/xla/xla/codegen/emitters/tests/BUILD @@ -44,9 +44,35 @@ copy_file( is_executable = True, ) +#Need no gpu to run +_FILECHECK_ONLY_TESTS = [ + #TODO(rocm) Most of thease test cuda lowering only + "loop/acos_f32.hlo", + "loop/acos_f64.hlo", + "loop/acosh_f32.hlo", + "loop/acosh_f64.hlo", + "loop/asin_f32.hlo", + "loop/asin_f64.hlo", + "loop/asinh_f32.hlo", + "loop/asinh_f64.hlo", + "loop/atanh_f32.hlo", + "loop/atanh_f64.hlo", + "loop/broadcast_constant_block_dim_limit.hlo", + "loop/cosh_f32.hlo", + "loop/cosh_f64.hlo", + "loop/dot_fp8.hlo", + "loop/large_loop_slow_compile_time.hlo", + "loop/sinh_f32.hlo", + "loop/sinh_f64.hlo", + "concatenate/test_small_dim.hlo", +] + lit_test_suite( name = "tests", - srcs = glob(["**/*.hlo"]), + srcs = glob( + ["**/*.hlo"], + exclude = _FILECHECK_ONLY_TESTS, + ), cfg = "//xla:lit.cfg.py", default_tags = tf_cuda_tests_tags(), exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}), @@ -59,3 +85,16 @@ lit_test_suite( "@llvm-project//llvm:FileCheck", ], ) + +lit_test_suite( + name = "filecheck_only_tests", + srcs = _FILECHECK_ONLY_TESTS, + cfg = "//xla:lit.cfg.py", + default_tags = ["gpu"], + tools = [ + ":cpu_fusion_to_mlir_copy", + ":gpu_fusion_to_mlir_copy", + "//xla/codegen/tools:emitters_opt", + "@llvm-project//llvm:FileCheck", + ], +) From 7b34ce2de44d6b34e35ae10d2f5dbd7218e173a9 Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Mon, 31 Aug 2026 05:53:16 -0700 Subject: [PATCH 18/26] [XLA:CPU] Relaunch new xtile pipeline. PiperOrigin-RevId: 973835112 --- third_party/xla/xla/debug_options_flags.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 6f8b72260dd115..66a1e7eb3767c5 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -265,7 +265,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_cpu_use_acl(true); #endif opts.set_xla_cpu_use_xnnpack(true); - opts.set_xla_cpu_use_new_xtile_lowering(false); + opts.set_xla_cpu_use_new_xtile_lowering(true); opts.set_xla_cpu_experimental_xnn_graph_fusion_mode( DebugOptions::XNN_GRAPH_FUSION_MODE_DISABLED); opts.add_xla_cpu_experimental_ynn_fusion_type( From c143796caefcac1921791c8c2a2813fa4c58500e Mon Sep 17 00:00:00 2001 From: Pablo Zimmermann Date: Mon, 31 Aug 2026 06:08:20 -0700 Subject: [PATCH 19/26] Add GPU peak bandwidth calculation helper Part of a stacked effort towards adding tooling to measure and generate GPU bandwidth derate tables reproducibly. This CL introduces `GetPeakBandwidthBytesPerSec(int device_id)`, which dynamically queries the active `StreamExecutor` to retrieve the theoretical peak GPU memory bandwidth at runtime. Test: Added unit tests. PiperOrigin-RevId: 973840679 --- third_party/xla/xla/tools/cost_model/BUILD | 21 ++++++++++++++- .../cost_model/gpu_bandwidth_benchmark.cc | 26 +++++++++++++++++++ .../cost_model/gpu_bandwidth_benchmark.h | 5 ++++ .../gpu_bandwidth_benchmark_test.cc | 24 ++++++++++++++++- 4 files changed, 74 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/tools/cost_model/BUILD b/third_party/xla/xla/tools/cost_model/BUILD index 5b3e54eef7c5ad..fe1e2715bbdbff 100644 --- a/third_party/xla/xla/tools/cost_model/BUILD +++ b/third_party/xla/xla/tools/cost_model/BUILD @@ -3,6 +3,10 @@ load( "//xla:xla.default.bzl", "xla_cc_test", ) +load( + "//xla/tsl/platform/default:cuda_build_defs.bzl", + "if_cuda_is_configured", +) package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -15,17 +19,32 @@ cc_library( srcs = ["gpu_bandwidth_benchmark.cc"], hdrs = ["gpu_bandwidth_benchmark.h"], deps = [ + "//xla/stream_executor:device_description", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor/gpu:gpu_init", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:span", - ], + ] + if_cuda_is_configured([ + "//xla/stream_executor:cuda_platform", + ]), ) xla_cc_test( name = "gpu_bandwidth_benchmark_test", srcs = ["gpu_bandwidth_benchmark_test.cc"], + tags = [ + "cuda-only", + "requires-gpu-nvidia", + ], deps = [ ":gpu_bandwidth_benchmark", "//xla/tests:xla_internal_test_main", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc index 9ce9311bd2a4fe..2ef01e08c97b7c 100644 --- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc +++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc @@ -15,13 +15,39 @@ limitations under the License. #include "xla/tools/cost_model/gpu_bandwidth_benchmark.h" +#include +#include #include +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/gpu/gpu_init.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" namespace xla::gpu { +absl::StatusOr GetPeakBandwidthBytesPerSec(int device_id) { + ABSL_ASSIGN_OR_RETURN(stream_executor::Platform * platform, + stream_executor::PlatformManager::PlatformWithName( + stream_executor::GpuPlatformName())); + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr description, + platform->DescriptionForDevice(device_id)); + const int64_t bandwidth = description->memory_bandwidth(); + if (bandwidth <= 0) { + return absl::InternalError(absl::StrFormat( + "Failed to determine peak memory bandwidth for device %d: " + "memory_bandwidth is %v.", + device_id, bandwidth)); + } + return static_cast(bandwidth); +} + std::string FormatBandwidthTable(absl::Span entries) { std::string result = "DMA Size (Bytes) Bandwidth Fraction\n" diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h index b6842b9e50623d..aa4daa0acc05b5 100644 --- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h +++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/status/statusor.h" #include "absl/types/span.h" namespace xla::gpu { @@ -30,6 +31,10 @@ struct BandwidthEntry { float bandwidth_fraction = 0.0f; }; +// Returns theoretical peak GPU memory bandwidth in bytes per second for +// `device_id`. +absl::StatusOr GetPeakBandwidthBytesPerSec(int device_id); + // Formats bandwidth table entries into a human-readable table. std::string FormatBandwidthTable(absl::Span entries); diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc index 4d70629374f67b..eafd6f7d3dc574 100644 --- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc +++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc @@ -17,11 +17,19 @@ limitations under the License. #include +#include #include +#include "absl/status/status.h" +#include "absl/status/status_matchers.h" namespace xla::gpu { namespace { +using ::absl_testing::IsOk; +using ::absl_testing::IsOkAndHolds; +using ::testing::Gt; +using ::testing::Not; + TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableEmpty) { EXPECT_EQ(FormatBandwidthTable({}), "DMA Size (Bytes) Bandwidth Fraction\n" @@ -29,7 +37,7 @@ TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableEmpty) { } TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableMultipleEntries) { - std::vector entries = { + const std::vector entries = { {8192, 0.00043418f}, {16384, 0.00092645f}, {32768, 0.00184066f}, @@ -44,5 +52,19 @@ TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableMultipleEntries) { " 8589934592 1.00000000\n"); } +TEST(GpuBandwidthBenchmarkTest, GetPeakBandwidthValidDevice) { + const absl::StatusOr peak_bw = + GetPeakBandwidthBytesPerSec(/*device_id=*/0); + if (!peak_bw.ok()) { + GTEST_SKIP() << "No GPU device available: " << peak_bw.status(); + } + EXPECT_THAT(peak_bw, IsOkAndHolds(Gt(0.0))); +} + +TEST(GpuBandwidthBenchmarkTest, GetPeakBandwidthInvalidDevice) { + EXPECT_THAT(GetPeakBandwidthBytesPerSec(/*device_id=*/-1), Not(IsOk())); + EXPECT_THAT(GetPeakBandwidthBytesPerSec(/*device_id=*/9999), Not(IsOk())); +} + } // namespace } // namespace xla::gpu From b859e395ce12786b14c85de58f46a9de9efe47f8 Mon Sep 17 00:00:00 2001 From: Alexander Lyashuk Date: Mon, 31 Aug 2026 08:16:31 -0700 Subject: [PATCH 20/26] [XLA:GPU] Sort keys in dicts in a backend config in HLO module. Without sorting, they are dumped in arbitrary order, which makes dumps not bit-identical, which makes it harder to hunt for indeterminism. I've checked a model that produced non-deterministic dumps before, this change fixes it. PiperOrigin-RevId: 973891310 --- .../convert_async_collectives_to_sync_test.cc | 2 +- .../convert_triton_gemm_config_test.cc | 8 ++--- .../windowed_einsum_handler_test.cc | 32 +++++++++---------- third_party/xla/xla/hlo/ir/hlo_module.cc | 1 + third_party/xla/xla/hlo/ir/hlo_module_test.cc | 16 ++++++++++ third_party/xla/xla/sort_json_test.cc | 6 ++++ 6 files changed, 44 insertions(+), 21 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc index a8138a37cadea8..a77354cb03eed2 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc @@ -512,7 +512,7 @@ TEST_F(GpuConvertAsyncCollectivesToSyncTest, CHECK-NOT: all-reduce-start CHECK: %id2 = f32[] bitcast(%id) CHECK: ROOT %{{.*}} = u32[] all-reduce(%id) - CHECK-SAME: "is_sync":true,"is_pipelined":true + CHECK-SAME: "is_pipelined":true{{.*}}"is_sync":true )"), IsOkAndHolds(true)); } diff --git a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc index 8fc0976a7170ca..975f1430fb65e5 100644 --- a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc @@ -100,12 +100,12 @@ ENTRY entry { CHECK: ENTRY CHECK: ROOT{{.*}}fusion( CHECK-SAME: kind=kCustom - CHECK-SAME: "kind":"__triton_nested_gemm_fusion" CHECK-SAME: "block_level_fusion_config" - CHECK-SAME: "num_warps":"4" - CHECK-SAME: "output_tiles":[{"sizes":["64","256"]}] CHECK-SAME: "num_ctas":3 CHECK-SAME: "num_stages":5 + CHECK-SAME: "num_warps":"4" + CHECK-SAME: "output_tiles":[{"sizes":["64","256"]}] + CHECK-SAME: "kind":"__triton_nested_gemm_fusion" )")); const HloInstruction* fusion = nullptr; ASSERT_THAT(module->entry_computation()->root_instruction(), @@ -157,8 +157,8 @@ ENTRY entry { CHECK: ROOT {{.*}} = bf16[4,4]{1,0} scaled-dot({{.*}}backend_config={"sizes":["64"]} CHECK: ENTRY CHECK: ROOT{{.*}}fusion( - CHECK-SAME: "kind":"__triton_nested_gemm_fusion" CHECK-SAME: "output_tiles":[{"sizes":["16","32"]}] + CHECK-SAME: "kind":"__triton_nested_gemm_fusion" )")); } diff --git a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc index e0a50c3c1ed5cd..f397673ea5fa85 100644 --- a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc @@ -345,7 +345,7 @@ CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0) CHECK-DAG: %[[SLICE4:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [6144:8192], [0:32768]} -CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false +CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8" CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [4096:6144]} CHECK: %[[A2A1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE1]]), @@ -354,7 +354,7 @@ CHECK: {0,1,2,3},{4,5,6,7} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE5:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [4096:6144], [0:32768]} -CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false +CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7" CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [2048:4096]} CHECK: %[[A2A2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE2]]), @@ -363,7 +363,7 @@ CHECK: {0,1,2,3},{4,5,6,7} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE6:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [2048:4096], [0:32768]} -CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false +CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6" CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:2048]} CHECK: %[[A2A3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE3]]), @@ -372,7 +372,7 @@ CHECK: {0,1,2,3},{4,5,6,7} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE7:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:2048], [0:32768]} -CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"5","force_earliest_schedule":false +CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"5" CHECK-DAG: %[[CONSTANT:.*]] = bf16[] constant(0) CHECK-DAG: %[[BROADCAST:.*]] = bf16[1,4,2048,32768]{3,2,1,0} broadcast(%[[CONSTANT:.*]]), dimensions={} CHECK-DAG: %[[ADD0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} add(%[[DOT0:.*]], %[[BROADCAST:.*]]) @@ -415,7 +415,7 @@ CHECK-DAG: %[[P1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} parameter(1) CHECK-DAG: %[[SLICE0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [24576:32768]} CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0) CHECK-DAG: %[[SLICE4:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [24576:32768]} -CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false +CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8" CHECK: %[[A2A0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT0:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -424,7 +424,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [16384:24576]} CHECK-DAG: %[[SLICE5:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [16384:24576]} -CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false +CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7" CHECK: %[[A2A1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT1:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -433,7 +433,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [8192:16384]} CHECK-DAG: %[[SLICE6:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [8192:16384]} -CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false +CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6" CHECK: %[[A2A2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT2:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -442,7 +442,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:8192]} CHECK-DAG: %[[SLICE7:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [0:8192]} -CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"5","force_earliest_schedule":false +CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"5" CHECK: %[[A2A3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT3:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -505,7 +505,7 @@ CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0) CHECK-DAG: %[[SLICE4:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [6144:8192], [0:32768]} -CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"9","force_earliest_schedule":false +CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"9" CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [4096:6144]} CHECK: %[[A2A1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE1]]), @@ -514,7 +514,7 @@ CHECK: {0,1,2,3} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE5:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [4096:6144], [0:32768]} -CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false +CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8" CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [2048:4096]} CHECK: %[[A2A2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE2]]), @@ -523,7 +523,7 @@ CHECK: {0,1,2,3} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE6:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [2048:4096], [0:32768]} -CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false +CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7" CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [0:2048]} CHECK: %[[A2A3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE3]]), @@ -532,7 +532,7 @@ CHECK: {0,1,2,3} CHECK: } CHECK: dimensions={1} CHECK-DAG: %[[SLICE7:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:2048], [0:32768]} -CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false +CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6" CHECK-DAG: %[[CONSTANT:.*]] = bf16[] constant(0) CHECK-DAG: %[[BROADCAST:.*]] = bf16[1,4,2048,32768]{3,2,1,0} broadcast(%[[CONSTANT:.*]]), dimensions={} CHECK-DAG: %[[ADD0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} add(%[[DOT0:.*]], %[[BROADCAST:.*]]) @@ -581,7 +581,7 @@ CHECK-DAG: %[[P1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} parameter(0) CHECK-DAG: %[[SLICE0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [24576:32768]} CHECK-DAG: %[[P0:.*]] = bf16[1,32768,8192]{2,1,0} parameter(1) CHECK-DAG: %[[SLICE4:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [24576:32768], [0:8192]} -CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"12","force_earliest_schedule":false +CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"12" CHECK: %[[A2A0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT0:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -590,7 +590,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [16384:24576]} CHECK-DAG: %[[SLICE5:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [16384:24576], [0:8192]} -CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"11","force_earliest_schedule":false +CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"11" CHECK: %[[A2A1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT1:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -599,7 +599,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [8192:16384]} CHECK-DAG: %[[SLICE6:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [8192:16384], [0:8192]} -CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"10","force_earliest_schedule":false +CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"10" CHECK: %[[A2A2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT2:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} @@ -608,7 +608,7 @@ CHECK: dimensions={1} CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:8192]} CHECK-DAG: %[[SLICE7:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [0:8192]} -CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"9","force_earliest_schedule":false +CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"9" CHECK: %[[A2A3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT3:.*]]), CHECK: replica_groups={ CHECK: {0,1,2,3} diff --git a/third_party/xla/xla/hlo/ir/hlo_module.cc b/third_party/xla/xla/hlo/ir/hlo_module.cc index fb20e46c044cfa..c8fa7c566d2db2 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module.cc +++ b/third_party/xla/xla/hlo/ir/hlo_module.cc @@ -580,6 +580,7 @@ std::string HloModule::ToString() const { print_options.set_print_inline_stack_frames( db_options.xla_hlo_print_inline_stack_frames()); print_options.set_compact_gte(db_options.xla_dump_compact_gte()); + print_options.set_sort_backend_config(true); return ToString(print_options); } diff --git a/third_party/xla/xla/hlo/ir/hlo_module_test.cc b/third_party/xla/xla/hlo/ir/hlo_module_test.cc index b5cc565d3d8f93..e747795bfd523a 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_module_test.cc @@ -776,6 +776,22 @@ TEST(HloModuleTest, CheckToStringHonorsDebugOptions) { EXPECT_TRUE(filecheck_matched); } +TEST(HloModuleTest, CheckToStringSortsBackendConfig) { + const char* hlo = R"( + HloModule test + + ENTRY main { + ROOT custom-call = () custom-call(), custom_call_target="test", backend_config={"tuning_knobs":{"3":"0","2":"2"}} + })"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo)); + EXPECT_THAT( + module->ToString(), + ::testing::HasSubstr( + R"json(backend_config={"tuning_knobs":{"2":"2","3":"0"}})json")); +} + TEST(HloModuleTest, TestCallersAndCallees) { const char* hlo = R"( HloModule jit_h diff --git a/third_party/xla/xla/sort_json_test.cc b/third_party/xla/xla/sort_json_test.cc index c24d5ab689b40d..cdbd3567c23325 100644 --- a/third_party/xla/xla/sort_json_test.cc +++ b/third_party/xla/xla/sort_json_test.cc @@ -46,5 +46,11 @@ TEST(SortJsonTest, SortsJson) { absl_testing::IsOkAndHolds(R"({"a":"a","a":"a}"})")); } +TEST(SortJsonTest, SortsTuningKnobs) { + EXPECT_THAT(SortJson(R"({"algorithm":{"tuning_knobs":{"3":"0","2":"2"}}})"), + absl_testing::IsOkAndHolds( + R"({"algorithm":{"tuning_knobs":{"2":"2","3":"0"}}})")); +} + } // namespace } // namespace xla From 7f873622a2d25961fb28722b6bfb68cf53b82c6d Mon Sep 17 00:00:00 2001 From: Brian Patton Date: Mon, 31 Aug 2026 08:42:49 -0700 Subject: [PATCH 21/26] Preserve frontend attributes in gather_scatter_handler. PiperOrigin-RevId: 973902250 --- .../service/spmd/gather_scatter_handler.cc | 5 + .../xla/service/spmd/spmd_partitioner_test.cc | 126 ++++++++++++++++++ 2 files changed, 131 insertions(+) diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc index 64c6834351ccd1..9d22a72e8d29a3 100644 --- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc +++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc @@ -977,6 +977,7 @@ absl::StatusOr PartitionGather( output_shape, operand.Replicate().hlo(), indices.Replicate().hlo(), gather->gather_dimension_numbers(), slice_sizes, gather->indices_are_sorted())); + new_gather->set_frontend_attributes(gather->frontend_attributes()); new_gather->set_sharding(HloSharding::Replicate()); new_gather = PartitionedHlo(new_gather, new_gather->shape(), operand.state()) .Reshard(output_sharding) @@ -1020,6 +1021,7 @@ absl::Status SpmdPartitioningVisitor::HandleGatherWithoutConflicts( builder()->AddInstruction(HloInstruction::CreateGather( pshape, operand.hlo(), indices.hlo(), dnums, pslice_sizes, gather->indices_are_sorted())); + phlo->set_frontend_attributes(gather->frontend_attributes()); SetPartitionedHlo(hlo, phlo); return absl::OkStatus(); @@ -1047,6 +1049,7 @@ absl::Status SpmdPartitioningVisitor::HandleGatherWithoutConflicts( HloInstruction* pgather = b->AddInstruction(HloInstruction::CreateGather( pshape, operand.hlo(), adjusted_indices_hlo, dnums, pslice_sizes, gather->indices_are_sorted())); + pgather->set_frontend_attributes(gather->frontend_attributes()); const Shape filter_shape = ShapeUtil::ChangeElementType(indices.hlo()->shape(), PRED); @@ -2007,6 +2010,7 @@ absl::StatusOr PartitionScatter( scatter->to_apply()->Clone()), scatter->scatter_dimension_numbers(), scatter->indices_are_sorted(), scatter->unique_indices())); + new_scatter->set_frontend_attributes(scatter->frontend_attributes()); new_scatter->set_sharding( HloSharding::Replicate().NormalizeTupleSharding(new_scatter->shape())); new_scatter = @@ -2123,6 +2127,7 @@ absl::Status SpmdPartitioningVisitor::HandleScatterWithoutConflicts( scatter->to_apply()->Clone()), scatter->scatter_dimension_numbers(), scatter->indices_are_sorted(), scatter->unique_indices())); + pscatter->set_frontend_attributes(scatter->frontend_attributes()); pscatter->set_sharding(HloSharding::Single( pscatter->shape(), hlo->sharding().IsTuple() ? hlo->sharding().tuple_elements()[0] diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index 38b62c9357248f..4f264492bc84d8 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -10626,6 +10626,132 @@ ENTRY entry { } } +// `_xla_compute_type` selects SparseCore offload and is read long after SPMD, +// so a partitioned gather/scatter must carry it across. +// `need_resolve_conflicts` picks between two entirely separate implementations, +// and each builds its own replacement instruction, so both are exercised. +namespace { + +void ExpectComputeTypeAttr(const HloInstruction* instr) { + ASSERT_NE(instr, nullptr); + std::optional attr = + instr->get_frontend_attribute("_xla_compute_type"); + ASSERT_TRUE(attr.has_value()); + EXPECT_EQ(*attr, "sparseoffload"); +} + +} // namespace + +TEST_P(SpmdPartitioningTest, PassthroughGatherPreservesFrontendAttributes) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + %input = f32[2,9] parameter(0), sharding={devices=[1,2]<=[2]} + %indices = s32[3] parameter(1), sharding={replicated} + ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1}, + collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, + slice_sizes={1,9}, sharding={devices=[1,2]<=[2]}, + frontend_attributes={_xla_compute_type="sparseoffload"} +})"; + for (bool need_resolve_conflicts : {true, false}) { + SpmdPartitionerOptions options; + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/2, options)); + ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kGather)); + } +} + +TEST_P(SpmdPartitioningTest, + GatherPartitionedOnTrivialSliceDimsPreservesFrontendAttributes) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + %input = f32[17,9] parameter(0), sharding={devices=[2,1]<=[2]} + %indices = s32[2,3] parameter(1), sharding={replicated} + ROOT %gather = f32[2,3,9] gather(%input, %indices), offset_dims={2}, + collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, + slice_sizes={1,9}, sharding={replicated}, + frontend_attributes={_xla_compute_type="sparseoffload"} +})"; + for (bool need_resolve_conflicts : {true, false}) { + SpmdPartitionerOptions options; + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/2, options)); + ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kGather)); + } +} + +TEST_P(SpmdPartitioningTest, PassthroughScatterPreservesFrontendAttributes) { + absl::string_view hlo_string = R"( +HloModule module + +add (lhs: f32[], rhs: f32[]) -> f32[] { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT sum = f32[] add(lhs, rhs) +} + +ENTRY entry { + %input = f32[2,9] parameter(0), sharding={devices=[1,2]<=[2]} + %indices = s32[3] parameter(1), sharding={replicated} + %updates = f32[3,9] parameter(2), sharding={devices=[1,2]<=[2]} + ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates), + to_apply=add, + update_window_dims={1}, + inserted_window_dims={0}, + scatter_dims_to_operand_dims={0}, + index_vector_dim=1, sharding={devices=[1,2]<=[2]}, + frontend_attributes={_xla_compute_type="sparseoffload"} +})"; + for (bool need_resolve_conflicts : {true, false}) { + SpmdPartitionerOptions options; + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/2, options)); + ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kScatter)); + } +} + +TEST_P(SpmdPartitioningTest, + ScatterPartitionedOnTrivialSliceDimsPreservesFrontendAttributes) { + absl::string_view hlo_string = R"( +HloModule module + +add (lhs: f32[], rhs: f32[]) -> f32[] { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT sum = f32[] add(lhs, rhs) +} + +ENTRY entry { + %input = f32[17,9] parameter(0), sharding={devices=[2,1]<=[2]} + %indices = s32[2,3] parameter(1), sharding={replicated} + %updates = f32[2,3,9] parameter(2), sharding={replicated} + ROOT %scatter = f32[17,9] scatter(%input, %indices, %updates), + to_apply=add, + update_window_dims={2}, + inserted_window_dims={0}, + scatter_dims_to_operand_dims={0}, + index_vector_dim=2, sharding={devices=[2,1]<=[2]}, + frontend_attributes={_xla_compute_type="sparseoffload"} +})"; + for (bool need_resolve_conflicts : {true, false}) { + SpmdPartitionerOptions options; + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/2, options)); + ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kScatter)); + } +} + TEST_P(SpmdPartitioningTest, ScatterPartitionedOnTrivialSliceDims_PartialReplicate) { absl::string_view hlo_string = R"( From c8353808d9a6ae08da0a58416dec12aa330665c0 Mon Sep 17 00:00:00 2001 From: Allan Renucci CA Date: Mon, 31 Aug 2026 09:45:03 -0700 Subject: [PATCH 22/26] Update rules_proto version to 7.1.0 and remove the obsolete rules_proto patch. This aligns rules_proto version across Bzlmod and WORKSPACE builds. PiperOrigin-RevId: 973931168 --- tensorflow/workspace0.bzl | 10 +++---- third_party/xla/third_party/rules_proto.patch | 27 ------------------- third_party/xla/workspace3.bzl | 7 +++-- 3 files changed, 6 insertions(+), 38 deletions(-) delete mode 100644 third_party/xla/third_party/rules_proto.patch diff --git a/tensorflow/workspace0.bzl b/tensorflow/workspace0.bzl index dec1ab868435ee..890e8896f0a97f 100644 --- a/tensorflow/workspace0.bzl +++ b/tensorflow/workspace0.bzl @@ -40,16 +40,12 @@ def workspace(): models_repositories() bazel_toolchains_repositories() - # Apple rules for Bazel. https://github.com/bazelbuild/rules_apple. - # Note: We add this to fix Kokoro builds. - # The rules below call into `rules_proto` but the hash has changed and - # Bazel refuses to continue. So, we add our own mirror. tf_http_archive( name = "rules_proto", - sha256 = "20b240eba17a36be4b0b22635aca63053913d5c1ee36e16be36499d167a2f533", - strip_prefix = "rules_proto-11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8", + sha256 = "14a225870ab4e91869652cfd69ef2028277fc1dc4910d65d353b62d6e0ae21f4", + strip_prefix = "rules_proto-7.1.0", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_proto/archive/11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8.tar.gz", + "https://github.com/bazelbuild/rules_proto/archive/refs/tags/7.1.0.tar.gz", ), ) diff --git a/third_party/xla/third_party/rules_proto.patch b/third_party/xla/third_party/rules_proto.patch deleted file mode 100644 index 6868665a922075..00000000000000 --- a/third_party/xla/third_party/rules_proto.patch +++ /dev/null @@ -1,27 +0,0 @@ -# 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. -# ============================================================================== ---- a/proto/private/native.bzl -+++ b/proto/private/native.bzl -@@ -26,7 +26,10 @@ - """Lovely workaround to be able to expose native constants pretending to be Starlark.""" - -+load("@com_google_protobuf//bazel/common:proto_common.bzl", "proto_common") -+load("@com_google_protobuf//bazel/common:proto_info.bzl", "ProtoInfo") -+ - # buildifier: disable=native-proto - NativeProtoInfo = ProtoInfo - - # buildifier: disable=native-proto - native_proto_common = proto_common \ No newline at end of file diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index b326a9bc4a7938..79a81c7a644602 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -22,12 +22,11 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def workspace(): tf_http_archive( name = "rules_proto", - sha256 = "20b240eba17a36be4b0b22635aca63053913d5c1ee36e16be36499d167a2f533", - strip_prefix = "rules_proto-11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8", + sha256 = "14a225870ab4e91869652cfd69ef2028277fc1dc4910d65d353b62d6e0ae21f4", + strip_prefix = "rules_proto-7.1.0", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_proto/archive/11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8.tar.gz", + "https://github.com/bazelbuild/rules_proto/archive/refs/tags/7.1.0.tar.gz", ), - patch_file = ["//third_party:rules_proto.patch"], ) # https://github.com/bazelbuild/bazel-skylib/releases From 58cbeea42137fea615e35e575ab214021dd27e24 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 31 Aug 2026 10:01:27 -0700 Subject: [PATCH 23/26] Introduce CudaGraphTopologyMapper helper library for CUDA Graphs host-side telemetry PiperOrigin-RevId: 973939699 --- .../xla/xla/backends/profiler/gpu/BUILD | 25 ++ .../gpu/cuda_graph_topology_mapper.cc | 154 ++++++++ .../profiler/gpu/cuda_graph_topology_mapper.h | 98 +++++ .../gpu/cuda_graph_topology_mapper_test.cc | 338 ++++++++++++++++++ 4 files changed, 615 insertions(+) create mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc create mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h create mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD index 3bf88b2de50a18..7c7ea026cfe4a6 100644 --- a/third_party/xla/xla/backends/profiler/gpu/BUILD +++ b/third_party/xla/xla/backends/profiler/gpu/BUILD @@ -22,6 +22,10 @@ load( "if_google", "internal_visibility", ) +load( + "//xla/tsl:tsl.default.bzl", + "get_compatible_with_portable", +) load("//xla/tsl/platform:rules_cc.bzl", "cc_library") load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", "if_cuda_newer_than") @@ -998,3 +1002,24 @@ cc_library( "@local_config_cuda//cuda:cuda_headers", ], ) + +cc_library( + name = "cuda_graph_topology_mapper", + srcs = ["cuda_graph_topology_mapper.cc"], + hdrs = ["cuda_graph_topology_mapper.h"], + compatible_with = get_compatible_with_portable(), + deps = [ + "//xla/tsl/platform:logging", + "@com_google_absl//absl/container:flat_hash_map", + ], +) + +xla_cc_test( + name = "cuda_graph_topology_mapper_test", + srcs = ["cuda_graph_topology_mapper_test.cc"], + deps = [ + ":cuda_graph_topology_mapper", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc new file mode 100644 index 00000000000000..5f5628e0b5d2c9 --- /dev/null +++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc @@ -0,0 +1,154 @@ +/* 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/backends/profiler/gpu/cuda_graph_topology_mapper.h" + +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "xla/tsl/platform/logging.h" + +namespace xla { +namespace profiler { + +size_t CudaGraphTopologyMapper::CalculateMergedSize( + uint32_t graph_id, const absl::flat_hash_map& base_sizes, + const absl::flat_hash_map>& + child_graphs, + absl::flat_hash_map* merged_sizes) { + DCHECK(merged_sizes != nullptr); + auto merged_it = merged_sizes->find(graph_id); + if (merged_it != merged_sizes->end()) { + return merged_it->second; + } + + size_t size = 0; + auto size_it = base_sizes.find(graph_id); + if (size_it != base_sizes.end()) { + size = size_it->second; + } + + auto children_it = child_graphs.find(graph_id); + if (children_it != child_graphs.end()) { + for (const auto& child : children_it->second) { + size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes, + child_graphs, merged_sizes); + size += child_size; + if (!child.is_conditional && size > 0) { + size -= 1; // Child graph node itself is replaced for inline children. + } + } + } + + (*merged_sizes)[graph_id] = size; + return size; +} + +std::pair CudaGraphTopologyMapper::ResolveMergedNode( + uint32_t graph_id, uint32_t node_index, + const absl::flat_hash_map& base_sizes, + const absl::flat_hash_map>& + child_graphs, + absl::flat_hash_map* merged_sizes) { + DCHECK(merged_sizes != nullptr); + auto children_it = child_graphs.find(graph_id); + if (children_it == child_graphs.end() || children_it->second.empty()) { + return {graph_id, node_index}; + } + + // Calculate total inline size. + size_t total_inline_size = 0; + auto size_it = base_sizes.find(graph_id); + if (size_it != base_sizes.end()) { + total_inline_size = size_it->second; + } + for (const auto& child : children_it->second) { + if (!child.is_conditional) { + size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes, + child_graphs, merged_sizes); + if (child_size > 0) { + total_inline_size += (child_size - 1); + } else if (total_inline_size > 0) { + total_inline_size -= 1; + } + } + } + + if (node_index < total_inline_size) { + // 1. Resolve inline nodes + int64_t inline_offset = 0; + for (const auto& child : children_it->second) { + if (child.is_conditional) { + continue; + } + size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes, + child_graphs, merged_sizes); + int64_t child_start = + static_cast(child.insertion_point) + inline_offset; + + if (child_size > 0) { + int64_t child_end = child_start + static_cast(child_size) - 1; + if (static_cast(node_index) >= child_start && + static_cast(node_index) <= child_end) { + uint32_t local_index = static_cast( + static_cast(node_index) - child_start); + return ResolveMergedNode(child.child_graph_id, local_index, + base_sizes, child_graphs, merged_sizes); + } + } + + if (static_cast(node_index) < child_start) { + return {graph_id, + static_cast(static_cast(node_index) - + inline_offset)}; + } + + inline_offset += (static_cast(child_size) - 1); + } + return {graph_id, static_cast(static_cast(node_index) - + inline_offset)}; + } + + // 2. Resolve conditional nodes + uint32_t conditional_offset = 0; + for (const auto& child : children_it->second) { + if (!child.is_conditional) { + continue; + } + size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes, + child_graphs, merged_sizes); + if (child_size == 0) { + continue; + } + uint32_t child_start = total_inline_size + conditional_offset; + uint32_t child_end = child_start + child_size - 1; + + if (node_index >= child_start && node_index <= child_end) { + uint32_t local_index = node_index - child_start; + return ResolveMergedNode(child.child_graph_id, local_index, base_sizes, + child_graphs, merged_sizes); + } + + conditional_offset += child_size; + } + + return {graph_id, node_index}; +} + +} // namespace profiler +} // namespace xla diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h new file mode 100644 index 00000000000000..6fab464bd06b89 --- /dev/null +++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h @@ -0,0 +1,98 @@ +/* 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_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_ +#define XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_ + +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" + +namespace xla { +namespace profiler { + +// CudaGraphTopologyMapper is a pure, state-less C++ helper library designed to +// mirror the layout and node-flattening sequences that the NVIDIA CUDA driver +// uses to assign executable flat LocalNodeId sequences to CUPTI activity traces +// during the instantiation of hierarchical and conditional nested CUDA Graphs. +// +// Topology Examples: +// +// 1. Standard Inline Child Graph (`is_conditional = false`) +// In this scenario, the nodes of the nested child graph completely replace the +// placeholder node in the parent graph. +// Layout Example: +// Parent Graph (Size 3): [P0, P_Child, P2] +// Child Graph (Size 2): [C0, C1] +// Expected Flattened Execution sequence: P0 -> C0 -> C1 -> P2 +// Expected Flat indices: P0 (0), C0 (1), C1 (2), P2 (3) +// +// 2. Conditional Nested Graph (`is_conditional = true`) +// In this scenario, the child graph's nodes are executed as a separate entity +// at runtime, appended to the end of the parent graph's sequence. The +// placeholder node in the parent retains its original index in the flattened +// space. +// Layout Example: +// Parent Graph (Size 3): [P0, P1_Cond, P2] +// Child Graph (Size 2): [C0, C1] +// Expected Flattened Sequence (Kernels only): P0 -> C0 -> C1 -> P2 +// Expected Flat indices: P0 (0), P1_Cond (1), P2 (2), C0 (3), C1 (4) +// (Note that the CUPTI trace for `C0` reports local node index `3`, while the +// original placeholder `P2` reports index `2`). +// +// All APIs in this library are pure, static, and have zero CUDA SDK +// dependencies, allowing them to run identically on the host CPU for unit +// testing. +class CudaGraphTopologyMapper { + public: + struct ChildGraphEntry { + uint32_t child_graph_id; + uint32_t insertion_point; + bool is_conditional; + + bool operator<(const ChildGraphEntry& other) const { + return insertion_point < other.insertion_point; + } + }; + + // Recursively calculates the total merged size of a graph, caching results in + // the provided mutable merged_sizes map. + static size_t CalculateMergedSize( + uint32_t graph_id, + const absl::flat_hash_map& base_sizes, + const absl::flat_hash_map>& + child_graphs, + absl::flat_hash_map* merged_sizes); + + // Resolves a flat, merged node index into its target (template_graph_id, + // local_node_index) pair. + static std::pair ResolveMergedNode( + uint32_t graph_id, uint32_t node_index, + const absl::flat_hash_map& base_sizes, + const absl::flat_hash_map>& + child_graphs, + absl::flat_hash_map* merged_sizes); + + private: + CudaGraphTopologyMapper() = delete; +}; + +} // namespace profiler +} // namespace xla + +#endif // XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_ diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc new file mode 100644 index 00000000000000..1a6b506bdcc6a0 --- /dev/null +++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc @@ -0,0 +1,338 @@ +/* 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/backends/profiler/gpu/cuda_graph_topology_mapper.h" + +#include +#include +#include + +#include +#include "absl/container/flat_hash_map.h" + +namespace xla { +namespace profiler { +namespace { + +TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeSimpleInline) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + + // 3 + 2 - 1 = 4 + EXPECT_EQ(size, 4); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeInline) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + + // Expected Merged indices: + // 0 -> Parent node 0 + // 1 -> Child node 0 + // 2 -> Child node 1 + // 3 -> Parent node 2 + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + + auto res1 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res1.first, 2); + EXPECT_EQ(res1.second, 0); + + auto res2 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 2, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res2.first, 2); + EXPECT_EQ(res2.second, 1); + + auto res3 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 3, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res3.first, 1); + EXPECT_EQ(res3.second, 2); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeDeepNestingInline) { + // Parent (1) -> Child (2) -> Grandchild (3) + absl::flat_hash_map base_sizes = {{1, 3}, {2, 3}, {3, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + child_graphs[2].push_back({3, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + + // Expected indices: + // 0 -> Parent 0 + // 1 -> Child 0 + // 2 -> Grandchild 0 + // 3 -> Grandchild 1 + // 4 -> Child 2 + // 5 -> Parent 2 + auto res2 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 2, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res2.first, 3); + EXPECT_EQ(res2.second, 0); + + auto res5 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 5, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res5.first, 1); + EXPECT_EQ(res5.second, 2); +} + +TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeConditional) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + + // 3 + 2 = 5 (Conditional child graph nodes are appended, size is not shifted + // by -1) + EXPECT_EQ(size, 5); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeConditional) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + + // Expected Merged indices for Conditional Nested Graph (Hardware Ground + // Truth): 0 -> Parent node 0 1 -> Parent node 1 2 -> Parent node 2 3 -> Child + // node 0 4 -> Child node 1 + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + + auto res1 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res1.first, 1); + EXPECT_EQ(res1.second, 1); + + auto res2 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 2, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res2.first, 1); + EXPECT_EQ(res2.second, 2); + + auto res3 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 3, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res3.first, 2); + EXPECT_EQ(res3.second, 0); + + auto res4 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 4, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res4.first, 2); + EXPECT_EQ(res4.second, 1); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMixedInlineAndConditional) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}, {3, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + child_graphs[1].push_back({3, 2, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + + auto res = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 4, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res.first, 3); + EXPECT_EQ(res.second, 0); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMultipleConditional) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}, {3, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/true}); + child_graphs[1].push_back({3, 2, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + + auto res = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 5, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res.first, 3); + EXPECT_EQ(res.second, 0); +} + +TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeOutOfBounds) { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 2}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + + auto res = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 4, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res.first, 1); + EXPECT_EQ(res.second, 4); +} + +TEST(CudaGraphTopologyMapperTest, EmptyChildGraphHandling) { + // Test 1: Inline empty child graph (child size == 0). + // Parent (1) size 2, Child (2) size 0. + // Merged size should be: 2 + 0 - 1 = 1. + { + absl::flat_hash_map base_sizes = {{1, 2}, {2, 0}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(size, 1); + + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + } + + // Test 2: Conditional empty child graph (child size == 0). + // Parent (1) size 2, Child (2) size 0. + // Merged size should be: 2 + 0 = 2. + { + absl::flat_hash_map base_sizes = {{1, 2}, {2, 0}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(size, 2); + + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + + auto res1 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res1.first, 1); + EXPECT_EQ(res1.second, 1); + } + + // Test 3: Completely empty parent and child graphs. + { + absl::flat_hash_map base_sizes = {{1, 0}, {2, 0}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 0, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(size, 0); + } +} + +TEST(CudaGraphTopologyMapperTest, EmptyChildGraphResolutionWithTrailingNodes) { + // 1. Parent size 3 (nodes P0, P1, P2) with empty inline child at insertion + // point 1: + { + absl::flat_hash_map base_sizes = {{1, 3}, {2, 0}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(size, 2); + + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + + auto res1 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res1.first, 1); + EXPECT_EQ(res1.second, 2); + } + + // 2. Parent size 3 with empty inline child at insertion point 1 AND + // conditional child (ID 3, size 2): + { + absl::flat_hash_map base_sizes = { + {{1, 3}, {2, 0}, {3, 2}}}; + absl::flat_hash_map> + child_graphs; + child_graphs[1].push_back({2, 1, /*is_conditional=*/false}); + child_graphs[1].push_back({3, 2, /*is_conditional=*/true}); + + absl::flat_hash_map merged_sizes; + size_t size = CudaGraphTopologyMapper::CalculateMergedSize( + 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(size, 4); + + auto res0 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 0, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res0.first, 1); + EXPECT_EQ(res0.second, 0); + + auto res1 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 1, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res1.first, 1); + EXPECT_EQ(res1.second, 2); + + auto res2 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 2, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res2.first, 3); + EXPECT_EQ(res2.second, 0); + + auto res3 = CudaGraphTopologyMapper::ResolveMergedNode( + 1, 3, base_sizes, child_graphs, &merged_sizes); + EXPECT_EQ(res3.first, 3); + EXPECT_EQ(res3.second, 1); + } +} + +} // namespace +} // namespace profiler +} // namespace xla From c73819b7d62c53ba8db671884369d80cf848dd08 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 10:01:41 -0700 Subject: [PATCH 24/26] Automated Code Change PiperOrigin-RevId: 973939820 --- .../compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h | 2 +- .../compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h b/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h index e3c0ee5cc23238..1c25d9c00bc55a 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h +++ b/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h @@ -200,7 +200,7 @@ class ClusteringPolicySet { private: template void AddImpl(Args&&... args) { - static_assert(std::is_base_of::value, + static_assert(std::is_base_of_v, "T must implement ClusteringPolicy"); policies_.emplace_back(std::make_unique(std::forward(args)...)); } diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc index b2ab71fa5129cb..88be1de8c034ad 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc @@ -199,7 +199,7 @@ class FuseContractionWithBiasAdd : public OpRewritePattern { attrs.push_back( NamedAttribute(StringAttr::get(context, "epsilon"), epsilon)); - if (std::is_same::value) { + if (std::is_same_v) { // Here TArgs types do not include types of the first two parameters, // i.e. the convolution input and the filter. TArgs are parameters for // the extras like the bias etc. From 7eb9cc76a74274824016578f0e02b2debaceac22 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 10:03:43 -0700 Subject: [PATCH 25/26] Add setup scripts and documentation for TPU microbenchmarks PiperOrigin-RevId: 973941174 --- third_party/xla/xla/benchmarks/README.md | 21 +++++ .../xla/xla/benchmarks/core/benchmark.py | 12 ++- .../jax_microbenchmarks/matmul_lib.py | 2 +- .../dense_matmul_lib.py | 2 +- .../subchannel_matmul_lib.py | 2 +- .../xla/xla/benchmarks/requirements.txt | 20 +++++ third_party/xla/xla/benchmarks/setup.sh | 77 +++++++++++++++++++ 7 files changed, 131 insertions(+), 5 deletions(-) create mode 100644 third_party/xla/xla/benchmarks/README.md create mode 100644 third_party/xla/xla/benchmarks/requirements.txt create mode 100755 third_party/xla/xla/benchmarks/setup.sh diff --git a/third_party/xla/xla/benchmarks/README.md b/third_party/xla/xla/benchmarks/README.md new file mode 100644 index 00000000000000..64d2b447423de9 --- /dev/null +++ b/third_party/xla/xla/benchmarks/README.md @@ -0,0 +1,21 @@ +# TPU Microbenchmarks + +## Quickstart + +From the root directory of the XLA project, run: +```bash +# Run setup script to create the venv and install dependencies +./xla/benchmarks/setup.sh + +# Activate the venv +source xla/benchmarks/.venv/bin/activate + +# Run an individual benchmark +python3 xla/benchmarks/pallas_microbenchmarks/dense_matmul.py --dim=1,2048,2048,2048 --fmt=f8e4m3fn,f8e4m3fn,f32 + +# Run dense matmul benchmark suite and write results to a CSV file +python3 xla/benchmarks/run_benchmarks.py --benchmarks=dense_matmul --csv_path= + +# Run full benchmark suite and write results to multiple CSV files +python3 xla/benchmarks/run_benchmarks.py --csv_path= +``` \ No newline at end of file diff --git a/third_party/xla/xla/benchmarks/core/benchmark.py b/third_party/xla/xla/benchmarks/core/benchmark.py index e85f992707ab60..7889f2c05fd5eb 100644 --- a/third_party/xla/xla/benchmarks/core/benchmark.py +++ b/third_party/xla/xla/benchmarks/core/benchmark.py @@ -238,14 +238,22 @@ def run( return profiler_results -@dataclasses.dataclass(frozen=True) +@dataclasses.dataclass(frozen=True, repr=False) class BenchmarkConfig(abc.ABC): """Base class for benchmark configs.""" def as_dict(self) -> dict[str, Any]: """Returns a dictionary representation of the config.""" + + def _is_dtype(val): + try: + jnp.dtype(val) + return True + except (TypeError, ValueError): + return False + return { - k: dtype_to_str(v) if isinstance(v, jnp.dtype) else v + k: dtype_to_str(v) if _is_dtype(v) else v for k, v in dataclasses.asdict(self).items() } diff --git a/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py b/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py index f10dd1e4987572..7b736c874a20cf 100644 --- a/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py +++ b/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py @@ -24,7 +24,7 @@ from xla.benchmarks.core import benchmark -@dataclasses.dataclass(frozen=True, kw_only=True) +@dataclasses.dataclass(frozen=True, kw_only=True, repr=False) class JaxMatmulConfig(benchmark.BenchmarkConfig): """Config for JAX matmul benchmark. diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py index fa988e6e4bdd9b..7b4cd08420df62 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py @@ -94,7 +94,7 @@ def select_window( return int(block_m), int(block_k), int(block_n) -@dataclasses.dataclass(frozen=True, kw_only=True) +@dataclasses.dataclass(frozen=True, kw_only=True, repr=False) class DenseMatmulConfig(benchmark.BenchmarkConfig): """Config for Pallas dense matmul benchmark. diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py index 20838543385c2d..8a8f29fbbe6cdb 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py @@ -125,7 +125,7 @@ def select_window( return int(block_m), int(block_k), int(block_n) -@dataclasses.dataclass(frozen=True, kw_only=True) +@dataclasses.dataclass(frozen=True, kw_only=True, repr=False) class SubchannelMatmulConfig(benchmark.BenchmarkConfig): """Config for Pallas subchannel quantized matmul benchmark. diff --git a/third_party/xla/xla/benchmarks/requirements.txt b/third_party/xla/xla/benchmarks/requirements.txt new file mode 100644 index 00000000000000..77971c82529003 --- /dev/null +++ b/third_party/xla/xla/benchmarks/requirements.txt @@ -0,0 +1,20 @@ +# 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. + +-f https://storage.googleapis.com/jax-releases/libtpu_releases.html +absl-py +immutabledict +jax[tpu] +numpy +pandas diff --git a/third_party/xla/xla/benchmarks/setup.sh b/third_party/xla/xla/benchmarks/setup.sh new file mode 100755 index 00000000000000..63fe0b17e8f708 --- /dev/null +++ b/third_party/xla/xla/benchmarks/setup.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# 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. +# ============================================================================== + +set -e +set -u +set -o pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &> /dev/null && pwd)" +VENV_DIR="${1:-${VENV_DIR:-${SCRIPT_DIR}/.venv}}" + +echo "============================================================" +echo "Setting up environment for OpenXLA TPU microbenchmarks" +echo "Repository root: ${REPO_ROOT}" +echo "Virtual environment directory: ${VENV_DIR}" +echo "============================================================" + +# Check for python3 +if ! command -v python3 &> /dev/null; then + echo "Error: python3 is not installed or not found in PATH." >&2 + exit 1 +fi + +# Create virtual environment if it does not already exist +if [[ ! -d "${VENV_DIR}" ]]; then + echo "Creating virtual environment at ${VENV_DIR}..." + python3 -m venv "${VENV_DIR}" +else + echo "Using existing virtual environment at ${VENV_DIR}." +fi + +# Activate virtual environment +# shellcheck source=/dev/null +source "${VENV_DIR}/bin/activate" + +# Upgrade pip +echo "Upgrading pip..." +pip install --upgrade pip + +# Install dependencies +echo "Installing requirements from ${SCRIPT_DIR}/requirements.txt..." +pip install -r "${SCRIPT_DIR}/requirements.txt" + +# Configure site-packages .pth file so xla is importable from anywhere +echo "Configuring Python site-packages path for OpenXLA repository..." +python3 -c " +import site, pathlib, sys +site_dirs = site.getsitepackages() +if site_dirs: + pth_file = pathlib.Path(site_dirs[0]) / 'openxla_benchmarks.pth' + pth_file.write_text(sys.argv[1] + '\n') +" "${REPO_ROOT}" + +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +echo "============================================================" +echo "Setup completed successfully!" +echo "" +echo "To activate this virtual environment in your shell, run:" +echo " source ${VENV_DIR}/bin/activate" +echo "" +echo "To run the benchmarks suite, run:" +echo " python3 ${SCRIPT_DIR}/run_benchmarks.py" +echo "============================================================" From e001ab45149973a84375c1068c81cf7c97fdc2d9 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 10:10:44 -0700 Subject: [PATCH 26/26] Fix unaligned external weight serialization in flatbuffer export Storing weights outside of flatbuffer should also respect alignment requirement. For now this is just matching alignment requirement for tensors stored inside flatbuffer with constants stored outside PiperOrigin-RevId: 973945746 --- .../compiler/mlir/lite/flatbuffer_export.cc | 1 + tensorflow/lite/python/lite_v2_test.py | 80 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc index 7d34e71a38aa15..ddf93115835684 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc @@ -4518,6 +4518,7 @@ absl::Status Translator::AppendBufferData() { for (const auto& [index, buffer] : const_buffer_storage_.buffers()) { uint64_t hash = buffer->hash(); if (hashcode_to_pos.find(hash) == hashcode_to_pos.end()) { + export_stream_.get().write_zeros(kFbAlignment - offset() % kFbAlignment); int64_t size = 0; int64_t buffer_offset = offset(); auto status = buffer->ApplyData([this, &size](absl::string_view data) { diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py index ae73b2046f6731..1afeb4ca0ea32b 100644 --- a/tensorflow/lite/python/lite_v2_test.py +++ b/tensorflow/lite/python/lite_v2_test.py @@ -5592,6 +5592,86 @@ def testCOncreteFunctionFloat(self): actual_value = self._evaluateTFLiteModel(tflite_model, [input_data]) self.assertEqual(expected_value.numpy(), actual_value) + @test_util.run_v2_only + def testUseBufferOffsetAlignment(self): + """Test that all external constant buffers are aligned to 16 bytes.""" + + class MultiConstantModel(tf.Module): + + def __init__(self): + super().__init__() + # Constants with non-multiple-of-16 byte sizes to test padding alignment + self.w1 = tf.Variable( + tf.ones([3, 1], dtype=tf.float32), name='w1' + ) # 3 * 4 = 12 bytes + self.w2 = tf.Variable( + tf.ones([5, 1], dtype=tf.float32), name='w2' + ) # 5 * 4 = 20 bytes + self.w3 = tf.Variable( + tf.ones([1, 1], dtype=tf.float32), name='w3' + ) # 1 * 4 = 4 bytes + self.w4 = tf.Variable( + tf.ones([7, 1], dtype=tf.float32), name='w4' + ) # 7 * 4 = 28 bytes + self.w5 = tf.Variable( + tf.ones([6, 1], dtype=tf.float32), name='w5' + ) # 6 * 4 = 24 bytes + + @tf.function + def __call__(self, x1, x2, x3, x4, x5): + return ( + tf.matmul(x1, self.w1) + + tf.matmul(x2, self.w2) + + tf.matmul(x3, self.w3) + + tf.matmul(x4, self.w4) + + tf.matmul(x5, self.w5) + ) + + root = MultiConstantModel() + inputs = [ + tf.constant([[1.0, 2.0, 3.0]], dtype=tf.float32), + tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=tf.float32), + tf.constant([[1.0]], dtype=tf.float32), + tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=tf.float32), + tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], dtype=tf.float32), + ] + concrete_func = root.__call__.get_concrete_function(*inputs) + + converter = lite.TFLiteConverterV2.from_concrete_functions( + [concrete_func], root + ) + converter._experimental_use_buffer_offset = True + tflite_model = converter.convert() + + # Parse flatbuffer model and check all buffer offsets + model_obj = schema_fb.Model.GetRootAsModel(tflite_model, 0) + external_buffer_count = 0 + for i in range(model_obj.BuffersLength()): + buf = model_obj.Buffers(i) + if buf.Offset() > 1: + external_buffer_count += 1 + self.assertEqual( + buf.Offset() % 16, + 0, + f'Buffer {i} offset {buf.Offset()} is not 16-byte aligned (offset %' + f' 16 = {buf.Offset() % 16})', + ) + self.assertGreaterEqual(external_buffer_count, 5) + + # Evaluate converted model + expected_value = root(*inputs) + interp = interpreter.Interpreter(model_content=tflite_model) + runner = interp.get_signature_runner() + output = runner( + x1=inputs[0], + x2=inputs[1], + x3=inputs[2], + x4=inputs[3], + x5=inputs[4], + ) + actual_value = list(output.values())[0] + self.assertEqual(expected_value.numpy(), actual_value) + @test_util.run_v2_only def testConcreteFunctionStringInput(self): class Model(tf.Module):