From faf57fef4709acb0832ac283bf14840127b3a5bf Mon Sep 17 00:00:00 2001 From: Daksh Prajapati Date: Wed, 22 Jul 2026 17:29:21 -0700 Subject: [PATCH 01/28] lite: fix memory leak in TFLiteSavedModelConverterV2.convert() convert() called _load() (tf.saved_model.load) solely to obtain graph_debug_info, which is the GraphDebugInfo proto stored on disk alongside saved_model.pb. For large models like DenseNet121, _load() allocates ~25 MB of variable tensors and registers the model's function defs in TF's C++ EagerContext. Reference cycles in the resulting trackable object graph caused these allocations to survive the explicit `del trackable_obj; gc.collect()` call, producing ~22 MB of leaked RSS per convert() invocation (issue #122598). Fix: replace _load() with _parse_saved_model_with_debug_info(), which reads the debug info proto directly from disk without loading any model weights or registering EagerContext function defs. This function is already imported in lite.py. Since _convert_debug_info_func() ignores original_nodes and returns the proto unchanged, the function-name adjustment performed by _load()'s adjust_debug_info_func_names() has no effect on the TFLite C++ conversion pipeline, making the raw proto equally correct. Also remove the now-unused `import gc`. Add regression test testConvertDoesNotCallLoad that mocks lite._load and asserts it is never called during convert(), then verifies the converted model produces correct inference results. --- tensorflow/lite/python/lite.py | 27 +++++++++-------- tensorflow/lite/python/lite_v2_test.py | 40 ++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 55afc80c329795..4af0a508bab60c 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -16,7 +16,6 @@ import enum import functools -import gc import pprint import shutil import sys @@ -1574,19 +1573,19 @@ def convert(self): graph_def, input_tensors, output_tensors ) - trackable_obj = _load(self.saved_model_dir, self._saved_model_tags) - if trackable_obj is None: - self._debug_info = _get_debug_info( - _build_debug_info_func(self._funcs[0].graph), graph_def - ) - else: - self._debug_info = _get_debug_info( - _convert_debug_info_func(trackable_obj.graph_debug_info), - graph_def, - ) - - del trackable_obj - gc.collect() + # Read debug info directly from the SavedModel directory instead of loading + # the full model with _load(). Loading the model allocates all variable + # tensors (~25MB for DenseNet121) plus registers function defs in the + # EagerContext, and these can leak across convert() calls due to reference + # cycles that are slow to collect. The debug info proto is already written + # to disk alongside saved_model.pb and can be read cheaply. + _, saved_debug_info = _parse_saved_model_with_debug_info( + self.saved_model_dir + ) + self._debug_info = _get_debug_info( + _convert_debug_info_func(saved_debug_info), + graph_def, + ) return self._convert_from_saved_model(graph_def) diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py index 6d5c361361979f..90b51d141bc9dd 100644 --- a/tensorflow/lite/python/lite_v2_test.py +++ b/tensorflow/lite/python/lite_v2_test.py @@ -3143,6 +3143,46 @@ def testQDQConversionMode(self, mode): model = converter.convert() self.assertIsNotNone(model) + @test_util.run_v2_only + def testConvertDoesNotCallLoad(self): + """Regression test for https://github.com/tensorflow/tensorflow/issues/122598. + + TFLiteSavedModelConverterV2.convert() used to call _load() internally just + to obtain graph_debug_info. For large models like DenseNet121, _load() + allocates ~25 MB of variable tensors and registers function defs in TF's + EagerContext. Reference cycles in the loaded object caused this memory to + leak across convert() calls (~22 MB/iter). The fix reads debug info directly + from the SavedModel directory via _parse_saved_model_with_debug_info() + instead of loading the full model. + """ + from unittest import mock + + root = autotrackable.AutoTrackable() + root.f = tf.function(lambda x: x * 2.0) + root.f.get_concrete_function(tf.TensorSpec([10], tf.float32)) + save_dir = os.path.join(self.get_temp_dir(), 'saved_model_no_load') + save.save(root, save_dir) + + converter = lite.TFLiteConverterV2.from_saved_model(save_dir) + # Patch _load at the lite module level to detect if it's called during + # convert(). The converter should NOT call _load() โ€” it should read debug + # info directly from disk instead. + with mock.patch.object(lite, '_load', wraps=lite._load) as mock_load: + tflite_model = converter.convert() + mock_load.assert_not_called() + + self.assertIsNotNone(tflite_model) + # Verify the converted model is runnable. + interp = interpreter.Interpreter(model_content=tflite_model) + interp.allocate_tensors() + input_details = interp.get_input_details() + output_details = interp.get_output_details() + interp.set_tensor(input_details[0]['index'], + np.ones([10], dtype=np.float32)) + interp.invoke() + output = interp.get_tensor(output_details[0]['index']) + np.testing.assert_array_almost_equal(output, np.full([10], 2.0)) + class FromKerasModelTest(lite_v2_test_util.ModelTest): From 44f8b8cf86f71b130cd8c1ec31827accda7ea6dd Mon Sep 17 00:00:00 2001 From: Daksh Prajapati Date: Mon, 3 Aug 2026 21:42:07 -0700 Subject: [PATCH 02/28] Address reviewer feedback: fix test save signature and clarify debug info Three issues raised in review: 1. The remote branch had a bad merge commit (c74f8150) that accidentally discarded the lite.py fix. This force-push restores the correct state: the _load() block is removed and replaced with _parse_saved_model_with_debug_info() as originally intended. 2. In testConvertDoesNotCallLoad, save.save() was called without the concrete function, which can produce an empty signature set and cause TFLiteConverterV2.from_saved_model() to fail with: ValueError: Only support at least one signature key. Fix: capture the concrete function and pass it explicitly as the signatures argument to save.save(). 3. Regarding adjust_debug_info_func_names: add a comment explaining why it is not needed. _convert_debug_info_func() (util.py:382) does `del original_nodes` and returns saved_debug_info unchanged, so the function-name adjustment that _load() applies via adjust_debug_info_func_names() has no effect on the TFLite pipeline regardless of which path is taken. The on-disk proto contains the names as written at save() time, which is what the TFLite C++ pipeline uses for source-location tracking. --- tensorflow/lite/python/lite.py | 8 ++++++++ tensorflow/lite/python/lite_v2_test.py | 4 ++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 4af0a508bab60c..581220e0003912 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -1579,6 +1579,14 @@ def convert(self): # EagerContext, and these can leak across convert() calls due to reference # cycles that are slow to collect. The debug info proto is already written # to disk alongside saved_model.pb and can be read cheaply. + # + # Note on adjust_debug_info_func_names: _convert_debug_info_func() (in + # util.py) ignores its original_nodes argument entirely and returns the + # saved_debug_info proto unchanged, so the function-name adjustment that + # _load() applies via adjust_debug_info_func_names() has no effect on this + # code path. The on-disk proto already contains the names as written during + # save(), which is what the TFLite C++ pipeline uses for source-location + # tracking. _, saved_debug_info = _parse_saved_model_with_debug_info( self.saved_model_dir ) diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py index 90b51d141bc9dd..11a072ebbb500c 100644 --- a/tensorflow/lite/python/lite_v2_test.py +++ b/tensorflow/lite/python/lite_v2_test.py @@ -3159,9 +3159,9 @@ def testConvertDoesNotCallLoad(self): root = autotrackable.AutoTrackable() root.f = tf.function(lambda x: x * 2.0) - root.f.get_concrete_function(tf.TensorSpec([10], tf.float32)) + to_save = root.f.get_concrete_function(tf.TensorSpec([10], tf.float32)) save_dir = os.path.join(self.get_temp_dir(), 'saved_model_no_load') - save.save(root, save_dir) + save.save(root, save_dir, to_save) converter = lite.TFLiteConverterV2.from_saved_model(save_dir) # Patch _load at the lite module level to detect if it's called during From fc0683d4b206c2fdc451a99270fedd40e4153849 Mon Sep 17 00:00:00 2001 From: Cocoa Date: Fri, 21 Aug 2026 18:17:10 +0900 Subject: [PATCH 03/28] Fix the CMake build of the Metal delegate after its sources moved to .mm --- tensorflow/lite/CMakeLists.txt | 37 +++++++++++----------------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 39599a21cc5681..6bd68f61331ca2 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -424,14 +424,14 @@ if(TFLITE_ENABLE_GPU) enable_language(OBJCXX) list(APPEND TFLITE_DELEGATES_METAL_SRCS ${TFLITE_SOURCE_DIR}/delegates/gpu/metal_delegate.mm - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer.cc + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer.mm ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer_convert.mm ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/common.mm - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/compute_task.cc - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.cc - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_arguments.cc - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_device.cc - ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_spatial_tensor.cc + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/compute_task.mm + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.mm + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_arguments.mm + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_device.mm + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_spatial_tensor.mm ) add_library(metal_delegate STATIC ${TFLITE_DELEGATES_METAL_SRCS} @@ -464,30 +464,17 @@ if(TFLITE_ENABLE_GPU) # # supplementary libraries for libmetal_delegate # - list(APPEND CC_SRCS - buffer - compute_task - inference_context - metal_arguments - metal_device - metal_spatial_tensor - ) SET(METAL_DELEGATE_PATH ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/) - foreach(lib_name ${CC_SRCS}) - set_source_files_properties(${METAL_DELEGATE_PATH}${lib_name}.cc PROPERTIES LANGUAGE OBJCXX) - add_library("${lib_name}" STATIC ${METAL_DELEGATE_PATH}${lib_name}.cc) - target_include_directories("${lib_name}" PUBLIC - ${CMAKE_BINARY_DIR}/abseil-cpp - ${CMAKE_BINARY_DIR}/flatbuffers/include - ) - set_target_properties(${lib_name} PROPERTIES LINKER_LANGUAGE OBJCXX) - target_link_libraries(${lib_name}) - endforeach() - list(APPEND MM_SRCS + buffer buffer_convert common + compute_task + inference_context + metal_arguments + metal_device + metal_spatial_tensor ) foreach(lib_name ${MM_SRCS}) add_library("${lib_name}" STATIC ${METAL_DELEGATE_PATH}${lib_name}.mm) From d51d7e3ea8161c18cf7b6a661b990033f7bc3567 Mon Sep 17 00:00:00 2001 From: Cocoa Date: Fri, 21 Aug 2026 18:17:10 +0900 Subject: [PATCH 04/28] Fix a flatc race and a source-tree write in the Metal delegate's CMake rules --- tensorflow/lite/CMakeLists.txt | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 6bd68f61331ca2..62fd5a33e59203 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -440,6 +440,7 @@ if(TFLITE_ENABLE_GPU) ${CMAKE_BINARY_DIR}/abseil-cpp ${CMAKE_BINARY_DIR}/flatbuffers/include PRIVATE ${TENSORFLOW_SOURCE_DIR} + PRIVATE ${PROJECT_BINARY_DIR} ) # # generate flatbuffers header for inference_context @@ -449,16 +450,20 @@ if(TFLITE_ENABLE_GPU) else() set(FLATC flatc) endif() + set(METAL_GENERATED_DIR ${PROJECT_BINARY_DIR}/tensorflow/lite/delegates/gpu/metal) add_custom_command( - OUTPUT ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context_generated.h + OUTPUT ${METAL_GENERATED_DIR}/inference_context_generated.h + COMMAND ${CMAKE_COMMAND} -E make_directory ${METAL_GENERATED_DIR} COMMAND ${FLATC} --scoped-enums -I ${TENSORFLOW_SOURCE_DIR} - -o ${TFLITE_SOURCE_DIR}/delegates/gpu/metal + -o ${METAL_GENERATED_DIR} -c ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs + DEPENDS ${FLATC_TARGET} + ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs ) add_custom_target( inference_context_cc_fbs - DEPENDS ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context_generated.h + DEPENDS ${METAL_GENERATED_DIR}/inference_context_generated.h ) add_dependencies(metal_delegate inference_context_cc_fbs) # From fbe17ed31e1b70592a15fb32b4653856795c3bc7 Mon Sep 17 00:00:00 2001 From: Cocoa Date: Fri, 21 Aug 2026 18:50:50 +0900 Subject: [PATCH 05/28] Remove the Metal supplementary libraries that never built --- tensorflow/lite/CMakeLists.txt | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 62fd5a33e59203..27e25e83750461 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -466,29 +466,6 @@ if(TFLITE_ENABLE_GPU) DEPENDS ${METAL_GENERATED_DIR}/inference_context_generated.h ) add_dependencies(metal_delegate inference_context_cc_fbs) - # - # supplementary libraries for libmetal_delegate - # - SET(METAL_DELEGATE_PATH ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/) - - list(APPEND MM_SRCS - buffer - buffer_convert - common - compute_task - inference_context - metal_arguments - metal_device - metal_spatial_tensor - ) - foreach(lib_name ${MM_SRCS}) - add_library("${lib_name}" STATIC ${METAL_DELEGATE_PATH}${lib_name}.mm) - target_include_directories("${lib_name}" PUBLIC - ${CMAKE_BINARY_DIR}/abseil-cpp - ${CMAKE_BINARY_DIR}/flatbuffers/include - ) - target_link_libraries(${lib_name}) - endforeach() endif() list(APPEND TFLITE_TARGET_PUBLIC_OPTIONS "-DCL_DELEGATE_NO_GL" "-DEGL_NO_X11") list(APPEND TFLITE_TARGET_DEPENDENCIES From d6abed341ce2eefe70d4f81360940e2630c2f8e2 Mon Sep 17 00:00:00 2001 From: Wei Wang Date: Fri, 26 Jun 2026 10:24:29 +0800 Subject: [PATCH 06/28] Initial upload. --- .../internal/optimized/optimized_ops.h | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h index 14f8ab199f7c55..1adaf84d802ae0 100644 --- a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h +++ b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h @@ -1828,20 +1828,32 @@ inline typename std::enable_if::value, void>::type Add( auto input1_map = MapAsVector(input1_data, input1_shape); auto input2_map = MapAsVector(input2_data, input2_shape); auto output_map = MapAsVector(output_data, output_shape); + // The element-wise sum is performed in the unsigned domain so that any + // overflow wraps in a well-defined manner instead of triggering + // signed-integer-overflow UB inside Eigen's expression evaluators. The + // wrapped result is bit-identical to the previous two's-complement behavior + // and is then clamped back in the signed domain by the activation min/max. + using UnsignedT = typename std::make_unsigned::type; if (input1_shape == input2_shape) { - output_map.array() = (input1_map.array() + input2_map.array()) + output_map.array() = (input1_map.array().template cast() + + input2_map.array().template cast()) + .template cast() .cwiseMax(activation_min) .cwiseMin(activation_max); } else if (input2_shape.FlatSize() == 1) { - auto scalar = input2_data[0]; - output_map.array() = (input1_map.array() + scalar) - .cwiseMax(activation_min) - .cwiseMin(activation_max); + UnsignedT scalar = static_cast(input2_data[0]); + output_map.array() = + (input1_map.array().template cast() + scalar) + .template cast() + .cwiseMax(activation_min) + .cwiseMin(activation_max); } else if (input1_shape.FlatSize() == 1) { - auto scalar = input1_data[0]; - output_map.array() = (scalar + input2_map.array()) - .cwiseMax(activation_min) - .cwiseMin(activation_max); + UnsignedT scalar = static_cast(input1_data[0]); + output_map.array() = + (scalar + input2_map.array().template cast()) + .template cast() + .cwiseMax(activation_min) + .cwiseMin(activation_max); } else { reference_ops::BroadcastAdd6DSlow(params, input1_shape, input1_data, input2_shape, input2_data, From 71f7bc8c6c511d52ff11c8cf338b97854891a4c6 Mon Sep 17 00:00:00 2001 From: Wei Wang Date: Tue, 14 Jul 2026 11:08:42 +0800 Subject: [PATCH 07/28] Add test. --- tensorflow/lite/kernels/add_test.cc | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tensorflow/lite/kernels/add_test.cc b/tensorflow/lite/kernels/add_test.cc index 0469a7d0ec7c92..6e15e0c8ecc724 100644 --- a/tensorflow/lite/kernels/add_test.cc +++ b/tensorflow/lite/kernels/add_test.cc @@ -17,8 +17,10 @@ limitations under the License. #include #include +#include #include #include +#include #include #include @@ -736,6 +738,21 @@ TYPED_TEST(IntegerAddOpTest, Int32MultiDimBroadcast) { EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 6, 7, 9})); } +TYPED_TEST(IntegerAddOpTest, OverflowWrapping) { + if (std::is_same::value || + std::is_same::value) { + IntegerAddOpModel m(GetTensorType(), {1, 2}, {1, 2}, + ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), + {std::numeric_limits::max(), 5}); + m.PopulateTensor(m.input2(), {1, 5}); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray( + {std::numeric_limits::min(), 10})); + } +} + template void TestQuantizedBroadcast(QuantizedAddOpModel& m, const std::vector& input1_shape, From dc437303c19ee1342c63b40db450b17e172e756e Mon Sep 17 00:00:00 2001 From: Wei Wang Date: Mon, 24 Aug 2026 09:18:41 +0800 Subject: [PATCH 08/28] Resolve comments. --- tensorflow/lite/kernels/add_test.cc | 16 ++++++++++++++++ tensorflow/lite/kernels/internal/reference/add.h | 5 +++-- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/kernels/add_test.cc b/tensorflow/lite/kernels/add_test.cc index 6e15e0c8ecc724..92249345046f7b 100644 --- a/tensorflow/lite/kernels/add_test.cc +++ b/tensorflow/lite/kernels/add_test.cc @@ -753,6 +753,22 @@ TYPED_TEST(IntegerAddOpTest, OverflowWrapping) { } } +TYPED_TEST(IntegerAddOpTest, OverflowWrappingBroadcast) { + if (std::is_same::value || + std::is_same::value) { + IntegerAddOpModel m(GetTensorType(), {1, 2}, {2, 1}, + ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), + {std::numeric_limits::max(), 5}); + m.PopulateTensor(m.input2(), {1, 1}); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray( + {std::numeric_limits::min(), 6, + std::numeric_limits::min(), 6})); + } +} + template void TestQuantizedBroadcast(QuantizedAddOpModel& m, const std::vector& input1_shape, diff --git a/tensorflow/lite/kernels/internal/reference/add.h b/tensorflow/lite/kernels/internal/reference/add.h index 198d576c4863c0..fe031644da1403 100644 --- a/tensorflow/lite/kernels/internal/reference/add.h +++ b/tensorflow/lite/kernels/internal/reference/add.h @@ -41,7 +41,8 @@ inline void Add(const ArithmeticParams& params, MatchingElementsSize(input1_shape, input2_shape, output_shape); for (int i = 0; i < flat_size; ++i) { output_data[i] = ActivationFunctionWithMinMax( - input1_data[i] + input2_data[i], activation_min, activation_max); + WrappingAdd(input1_data[i], input2_data[i]), activation_min, + activation_max); } } @@ -280,7 +281,7 @@ BroadcastAdd6DSlow(const ArithmeticParams& params, T activation_min, activation_max; GetActivationParams(params, &activation_min, &activation_max); auto op = [activation_min, activation_max](T a, T b) { - return ActivationFunctionWithMinMax(a + b, activation_min, + return ActivationFunctionWithMinMax(WrappingAdd(a, b), activation_min, activation_max); }; BroadcastBinaryOpSimple(input1_shape, input1_data, input2_shape, input2_data, From 6314d19e5ea1a8c6f1a649d250f3589148c82918 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 09:44:26 +0300 Subject: [PATCH 09/28] Add an OSS-Fuzz target for the TensorFlow Lite interpreter TensorFlow Lite currently has no OSS-Fuzz coverage. The 29 fuzz targets built by projects/tensorflow cover TF core ops, framework types, path/string helpers and the graph/SavedModel loaders, but nothing under tensorflow/lite: the builtin kernels, the arena planner and the shape-propagation paths are unfuzzed. This adds a single end-to-end harness. An arbitrary buffer is verified as a flatbuffer model, an interpreter is built for it, tensors are allocated, inputs are filled deterministically and the graph is invoked. That reaches the builtin kernel implementations through the same path a real caller uses. Notes on the harness: * VerifyAndBuildFromBuffer runs tflite::VerifyModelBuffer first, so structurally invalid buffers are rejected cheaply rather than being fed to the interpreter. * BuiltinOpResolverWithoutDefaultDelegates is used so the fuzzer exercises the reference and optimized CPU kernels rather than a delegate's own implementation. * Model size and total input-arena size are bounded (1 MiB / 64 MiB) to keep runs inside the OSS-Fuzz memory budget; larger models exercise the allocator rather than the kernels. * String, resource and variant inputs are skipped: those tensors own a dynamic buffer with its own layout, and writing raw bytes into it would corrupt interpreter state instead of testing a kernel. --- tensorflow/lite/fuzzing/BUILD | 22 +++++ tensorflow/lite/fuzzing/interpreter_fuzz.cc | 97 +++++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 tensorflow/lite/fuzzing/BUILD create mode 100644 tensorflow/lite/fuzzing/interpreter_fuzz.cc diff --git a/tensorflow/lite/fuzzing/BUILD b/tensorflow/lite/fuzzing/BUILD new file mode 100644 index 00000000000000..2332243c18ef24 --- /dev/null +++ b/tensorflow/lite/fuzzing/BUILD @@ -0,0 +1,22 @@ +# Fuzzing harnesses for the TensorFlow Lite runtime. + +load( + "//tensorflow/security/fuzzing:tf_fuzzing.bzl", + "tf_cc_fuzz_test", +) + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + default_visibility = ["//visibility:private"], + licenses = ["notice"], +) + +tf_cc_fuzz_test( + name = "interpreter_fuzz", + srcs = ["interpreter_fuzz.cc"], + deps = [ + "//tensorflow/lite:framework", + "//tensorflow/lite/core:framework", + "//tensorflow/lite/kernels:builtin_ops", + ], +) diff --git a/tensorflow/lite/fuzzing/interpreter_fuzz.cc b/tensorflow/lite/fuzzing/interpreter_fuzz.cc new file mode 100644 index 00000000000000..b95255c85f9883 --- /dev/null +++ b/tensorflow/lite/fuzzing/interpreter_fuzz.cc @@ -0,0 +1,97 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Fuzzes the TensorFlow Lite interpreter end to end: an arbitrary buffer is +// verified as a flatbuffer model, an interpreter is built for it, tensors are +// allocated, inputs are filled deterministically, and the graph is invoked. +// +// This exercises the builtin kernel implementations, the arena planner and the +// shape-propagation paths, none of which previously had OSS-Fuzz coverage. + +#include +#include +#include +#include +#include + +#include "fuzztest/fuzztest.h" +#include "tensorflow/lite/core/interpreter.h" +#include "tensorflow/lite/core/interpreter_builder.h" +#include "tensorflow/lite/core/model_builder.h" +#include "tensorflow/lite/kernels/register.h" + +namespace tflite { +namespace fuzzing { +namespace { + +// Keep the fuzzer inside the OSS-Fuzz memory budget. Models and arenas larger +// than this are not interesting: they exercise the allocator, not the kernels. +constexpr size_t kMaxModelBytes = 1 << 20; // 1 MiB +constexpr size_t kMaxArenaBytes = 1 << 26; // 64 MiB + +void FuzzInterpreter(const std::string& model_bytes) { + if (model_bytes.size() < 8 || model_bytes.size() > kMaxModelBytes) { + return; + } + + // VerifyAndBuildFromBuffer applies tflite::VerifyModelBuffer first, so + // structurally invalid buffers are rejected cheaply. + std::unique_ptr model = + FlatBufferModel::VerifyAndBuildFromBuffer(model_bytes.data(), + model_bytes.size()); + if (model == nullptr) { + return; + } + + // Delegates are excluded so the fuzzer exercises the reference and optimized + // CPU kernels rather than a delegate's own implementation. + ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver; + std::unique_ptr interpreter; + if (InterpreterBuilder(*model, resolver)(&interpreter) != kTfLiteOk || + interpreter == nullptr) { + return; + } + + if (interpreter->AllocateTensors() != kTfLiteOk) { + return; + } + + size_t total_bytes = 0; + for (const int tensor_index : interpreter->inputs()) { + TfLiteTensor* tensor = interpreter->tensor(tensor_index); + if (tensor == nullptr || tensor->data.raw == nullptr) { + continue; + } + // String tensors own a dynamic buffer with its own layout; writing raw + // bytes into it would corrupt the interpreter rather than the kernel under + // test. + if (tensor->type == kTfLiteString || tensor->type == kTfLiteResource || + tensor->type == kTfLiteVariant) { + return; + } + total_bytes += tensor->bytes; + if (total_bytes > kMaxArenaBytes) { + return; + } + std::memset(tensor->data.raw, 1, tensor->bytes); + } + + interpreter->Invoke(); +} +FUZZ_TEST(TfLiteFuzz, FuzzInterpreter); + +} // namespace +} // namespace fuzzing +} // namespace tflite From 446a09e9f6597b359597fdaa39b14dbf2a0d0589 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 10:06:50 +0300 Subject: [PATCH 10/28] Check the arena budget before accumulating, not after Avoids any possibility of total_bytes wrapping before the limit is compared. --- tensorflow/lite/fuzzing/interpreter_fuzz.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/fuzzing/interpreter_fuzz.cc b/tensorflow/lite/fuzzing/interpreter_fuzz.cc index b95255c85f9883..dd6ed5df5a657a 100644 --- a/tensorflow/lite/fuzzing/interpreter_fuzz.cc +++ b/tensorflow/lite/fuzzing/interpreter_fuzz.cc @@ -81,10 +81,12 @@ void FuzzInterpreter(const std::string& model_bytes) { tensor->type == kTfLiteVariant) { return; } - total_bytes += tensor->bytes; - if (total_bytes > kMaxArenaBytes) { + // Check before accumulating so the sum itself cannot wrap. + if (tensor->bytes > kMaxArenaBytes || + total_bytes > kMaxArenaBytes - tensor->bytes) { return; } + total_bytes += tensor->bytes; std::memset(tensor->data.raw, 1, tensor->bytes); } From 397406ffd85990e83a0bd6ccd735db22aa9d7c9d Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 10:16:58 +0300 Subject: [PATCH 11/28] Require the TFLite TRANSPOSE permutation to be a bijection ResizeOutputTensor checks that every entry of `perm` is in [-dims, dims), but never that the entries form a permutation of [0, dims). A repeated entry passes the range check. Both the output shape and the element offsets are derived from `perm`: output_size->data[idx] = input_size->data[new_perm_data[idx]]; so with `perm = {0, 0}` on an input of shape [3, N] the output is sized [3, 3] while the offsets the kernel computes for it are those of a [3, N] traversal. The two disagree, and the transpose reads outside the input tensor. The values read are copied into the output, so they are observable by the caller. TensorFlow's own Transpose op already enforces this (tensorflow/core/kernels/transpose_op.cc): it records each index in a `bits` array and then requires every position to have been seen, rejecting `{0, 0}` with "0 is missing from {0,0}". The TFLite kernel was missing the equivalent check. Track which normalised indices have been used and reject duplicates. The check runs after negative entries are normalised, so `{0, -2}` on a rank-2 input is rejected as well. Valid permutations are unaffected. Adds two regression tests alongside the existing TestPermOutOfBounds. --- tensorflow/lite/kernels/transpose.cc | 9 +++++++++ tensorflow/lite/kernels/transpose_test.cc | 17 +++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc index 0b1f2b783b05bd..8bb223e8bbfde1 100644 --- a/tensorflow/lite/kernels/transpose.cc +++ b/tensorflow/lite/kernels/transpose.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/kernels/internal/portable_tensor_utils.h" @@ -50,12 +51,20 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, // Ensure validity of the permutations tensor as a 1D tensor. TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1); TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims); + // `perm` must be a permutation of [0, dims), not merely a set of in-range + // values: the output shape and the element offsets are both derived from it, + // and a repeated entry makes them disagree with the input extent. + std::vector seen(dims, false); for (int idx = 0; idx < dims; ++idx) { TF_LITE_ENSURE_MSG(context, (perm_data[idx] >= -dims && perm_data[idx] < dims), "Transpose op permutations array is out of bounds."); new_perm_data[idx] = perm_data[idx]; if (new_perm_data[idx] < 0) new_perm_data[idx] += dims; + TF_LITE_ENSURE_MSG( + context, !seen[new_perm_data[idx]], + "Transpose op permutations array must not contain duplicate values."); + seen[new_perm_data[idx]] = true; } // Determine size of output tensor. diff --git a/tensorflow/lite/kernels/transpose_test.cc b/tensorflow/lite/kernels/transpose_test.cc index 601e8bf4355540..11930f342ce91a 100644 --- a/tensorflow/lite/kernels/transpose_test.cc +++ b/tensorflow/lite/kernels/transpose_test.cc @@ -193,6 +193,23 @@ TEST(TransposeTest, TestPermOutOfBounds) { EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}), "Transpose op permutations array is out of bounds."); } + +// `perm` must be a permutation of [0, dims). A repeated entry passes the range +// check but makes the derived output shape disagree with the input extent, so +// the kernel reads outside the input tensor. +TEST(TransposeTest, TestPermDuplicateValues) { + EXPECT_DEATH( + TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 2}), + "Transpose op permutations array must not contain duplicate values."); +} + +// Duplicates must also be rejected after negative entries are normalised: +// on a rank-2 input {0, -2} normalises to {0, 0}. +TEST(TransposeTest, TestPermDuplicateValuesAfterNegativeNormalization) { + EXPECT_DEATH( + TransposeOpConstModel({2, 3}, {2}, {0, -2}), + "Transpose op permutations array must not contain duplicate values."); +} #endif TEST(TransposeTest, TestInt41DInputConstTensor) { From 61020a09b034666ef1ea7f77992f30f891c09836 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 10:22:15 +0300 Subject: [PATCH 12/28] Validate SLICE bounds when the input extent is dynamic Prepare() returns early when the output shape is fully specified, leaving the output static. Eval() only re-runs ResizeOutputShape() for a dynamic output, so CalculateOutputShapeVector() -- the one place begin and size are checked against the input -- never runs. When the input has an unspecified dimension the declared output shape says nothing about whether the slice fits the actual extent, and the kernel reads past the input. Fall through when the input has an unspecified dimension so the output is marked dynamic and the bounds are validated on every invocation. --- tensorflow/lite/kernels/slice.cc | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc index 86dca38bb7c427..53ea68878535a9 100644 --- a/tensorflow/lite/kernels/slice.cc +++ b/tensorflow/lite/kernels/slice.cc @@ -145,7 +145,16 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size)); // If the shape of output is fully specified then resize even if // the input shape is not staticly defined. - if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims)) { + // + // A fully specified output shape does not imply the slice is in bounds when + // the input extent is only known at run time. Taking this path leaves the + // output static, so Eval() -- which only re-runs ResizeOutputShape() for a + // dynamic output -- never reaches CalculateOutputShapeVector(), the one place + // `begin` and `size` are checked against the input. Fall through when the + // input has an unspecified dimension so the output is marked dynamic and the + // bounds are validated against the actual extent on every invocation. + if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims) && + !HasUnspecifiedDimension(input)) { return kTfLiteOk; } // Postpone allocation of output if any of the indexing tensors is not From 8ac1138dd14d2f6a15cf88a9d9e4c8a91da57b7c Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 10:22:15 +0300 Subject: [PATCH 13/28] Use a bitmask for the permutation check Avoids a heap allocation in the kernel. dims is bounded by kTransposeMaxDimensions (8), enforced in Prepare() before this function is reachable, so 64 bits are sufficient; asserted at compile time. --- tensorflow/lite/kernels/transpose.cc | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc index 8bb223e8bbfde1..94bdac9439a21b 100644 --- a/tensorflow/lite/kernels/transpose.cc +++ b/tensorflow/lite/kernels/transpose.cc @@ -18,7 +18,6 @@ limitations under the License. #include #include -#include #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/kernels/internal/portable_tensor_utils.h" @@ -54,17 +53,23 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, // `perm` must be a permutation of [0, dims), not merely a set of in-range // values: the output shape and the element offsets are both derived from it, // and a repeated entry makes them disagree with the input extent. - std::vector seen(dims, false); + // `dims` is bounded by kTransposeMaxDimensions, which Prepare() enforces + // before this function is reachable, so a 64-bit mask is sufficient and + // avoids a heap allocation in the kernel. + static_assert(kTransposeMaxDimensions <= 64, + "Permutation bitmask assumes at most 64 dimensions."); + uint64_t seen = 0; for (int idx = 0; idx < dims; ++idx) { TF_LITE_ENSURE_MSG(context, (perm_data[idx] >= -dims && perm_data[idx] < dims), "Transpose op permutations array is out of bounds."); new_perm_data[idx] = perm_data[idx]; if (new_perm_data[idx] < 0) new_perm_data[idx] += dims; + const uint64_t bit = uint64_t{1} << new_perm_data[idx]; TF_LITE_ENSURE_MSG( - context, !seen[new_perm_data[idx]], + context, (seen & bit) == 0, "Transpose op permutations array must not contain duplicate values."); - seen[new_perm_data[idx]] = true; + seen |= bit; } // Determine size of output tensor. From e258a88947c8bd8ab302b924056605625d7c0267 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 19:54:02 +0300 Subject: [PATCH 14/28] Add dynamic-input slice tests and order the rank check first Adds the two tests requested in review, covering a dynamic input dimension with a statically declared output shape: an in-bounds slice that must succeed and mark the output dynamic, and an out-of-bounds window that must be rejected. The suite is named SliceOpDynamicInputTest rather than SliceOpTest because the latter is a TEST_P fixture, and gtest rejects a suite that mixes TEST and TEST_P. Also checks ShapeHasRank before HasUnspecifiedDimension for both tensors. An unranked or scalar shape would otherwise take the early return and skip validation, since HasUnspecifiedDimension only inspects dims_signature. --- tensorflow/lite/kernels/slice.cc | 4 +- tensorflow/lite/kernels/slice_test.cc | 62 +++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc index 53ea68878535a9..597c14f585d25c 100644 --- a/tensorflow/lite/kernels/slice.cc +++ b/tensorflow/lite/kernels/slice.cc @@ -153,8 +153,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // `begin` and `size` are checked against the input. Fall through when the // input has an unspecified dimension so the output is marked dynamic and the // bounds are validated against the actual extent on every invocation. - if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims) && - !HasUnspecifiedDimension(input)) { + if (ShapeHasRank(output->dims) && !HasUnspecifiedDimension(output) && + ShapeHasRank(input->dims) && !HasUnspecifiedDimension(input)) { return kTfLiteOk; } // Postpone allocation of output if any of the indexing tensors is not diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc index 09b68642fff848..97bf99efbb5987 100644 --- a/tensorflow/lite/kernels/slice_test.cc +++ b/tensorflow/lite/kernels/slice_test.cc @@ -102,6 +102,68 @@ class SliceOpModel : public SingleOpModel { class SliceOpTest : public ::testing::TestWithParam {}; +// Model with a dynamic input dimension and a statically shaped output. The +// suite name is deliberately distinct from SliceOpTest, which is a TEST_P +// fixture -- gtest rejects a suite that mixes TEST and TEST_P. +class DynamicInputSliceOpModel : public SingleOpModel { + public: + DynamicInputSliceOpModel(TensorData input_data, + std::initializer_list begin_shape, + std::initializer_list begin_data, + std::initializer_list size_shape, + std::initializer_list size_data, + TensorData output_data) { + input_ = AddInput(input_data); + begin_ = AddConstInput(TensorType_INT32, begin_data, begin_shape); + size_ = AddConstInput(TensorType_INT32, size_data, size_shape); + output_ = AddOutput(output_data); + SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions, + CreateSliceOptions(builder_).Union()); + BuildInterpreter({input_data.shape, begin_shape, size_shape}); + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + const TfLiteTensor* GetOutputTensor() { + return interpreter_->tensor(output_); + } + + private: + int input_; + int begin_; + int size_; + int output_; +}; + +// A dynamic input dimension must force the output dynamic so the bounds are +// re-validated in Eval(), even though the declared output shape is static. +TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputValid) { + TensorData input_data(TensorType_FLOAT32, {1, 8}); + input_data.shape_signature = {1, -1}; + TensorData output_data(TensorType_FLOAT32, {1, 4}); + + DynamicInputSliceOpModel m(input_data, {2}, {0, 2}, {2}, {1, 4}, output_data); + m.SetInput({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.0, 4.0, 5.0, 6.0})); + EXPECT_EQ(m.GetOutputTensor()->allocation_type, kTfLiteDynamic); +} + +// The same path must reject a window that does not fit the actual extent. +TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputOutOfBounds) { + TensorData input_data(TensorType_FLOAT32, {1, 4}); + input_data.shape_signature = {1, -1}; + TensorData output_data(TensorType_FLOAT32, {1, 8}); + + DynamicInputSliceOpModel m(input_data, {2}, {0, 0}, {2}, {1, 8}, output_data); + m.SetInput({1.0, 2.0, 3.0, 4.0}); + EXPECT_EQ(m.Invoke(), kTfLiteError); +} + TEST_P(SliceOpTest, In1D) { SliceOpModel m({4}, {1}, {1}, {1}, {2}, TensorType_INT32, TensorType_FLOAT32, GetParam()); From 15acef56646d175dfb3070c797ae1cbc38fea9de Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 20:29:07 +0300 Subject: [PATCH 15/28] Bypass delegates in DynamicInputSliceOpModel A delegate that claims the SLICE node sets the output allocation type itself, so the kTfLiteDynamic assertion and the out-of-bounds error check would no longer describe the built-in CPU kernel. XNNPACK does handle BuiltinOperator_SLICE, so this is reachable whenever a delegate is supplied via the test delegate providers. AllocateTensors() still runs, since allocate_and_delegate defaults to true. --- tensorflow/lite/kernels/slice_test.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc index 97bf99efbb5987..e3b55f6116e32c 100644 --- a/tensorflow/lite/kernels/slice_test.cc +++ b/tensorflow/lite/kernels/slice_test.cc @@ -119,7 +119,12 @@ class DynamicInputSliceOpModel : public SingleOpModel { output_ = AddOutput(output_data); SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions, CreateSliceOptions(builder_).Union()); - BuildInterpreter({input_data.shape, begin_shape, size_shape}); + // Delegates are bypassed: a delegate that claims the SLICE node would set + // the output allocation type itself, so the assertions below would no + // longer describe the built-in CPU kernel. + BuildInterpreter({input_data.shape, begin_shape, size_shape}, + /*num_threads=*/-1, /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false); } void SetInput(std::initializer_list data) { From 68fc8b287c136c948387e6d104276f899cdd554f Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Thu, 27 Aug 2026 12:15:29 -0700 Subject: [PATCH 16/28] Implement CpuPjRtCompiler::DeserializeExecutable and remove PjRtCpuClient::LoadSerializedExecutable. PiperOrigin-RevId: 972088010 --- third_party/xla/xla/pjrt/cpu/BUILD | 5 ++ third_party/xla/xla/pjrt/cpu/cpu_client.cc | 65 ++++--------------- third_party/xla/xla/pjrt/cpu/cpu_client.h | 31 ++------- .../xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc | 13 ++++ .../xla/xla/pjrt/cpu/cpu_pjrt_compiler.h | 9 +++ 5 files changed, 46 insertions(+), 77 deletions(-) diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD index 709693cc0a4c3a..e8170242698e7f 100644 --- a/third_party/xla/xla/pjrt/cpu/BUILD +++ b/third_party/xla/xla/pjrt/cpu/BUILD @@ -299,6 +299,9 @@ cc_library( "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@eigen_archive//:eigen3", + "@riegeli//riegeli/base:any", + "@riegeli//riegeli/bytes:reader", + "@riegeli//riegeli/messages:parse_message", "@tsl//tsl/platform:denormal", "@tsl//tsl/platform:fingerprint", "@tsl//tsl/platform:protobuf", @@ -396,6 +399,8 @@ cc_library( "//xla/pjrt:pjrt_executable", "//xla/pjrt/plugin/xla_cpu:cpu_topology_description", "@com_google_absl//absl/status:statusor", + "@riegeli//riegeli/base:any", + "@riegeli//riegeli/bytes:reader", ], alwayslink = True, ) diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index 64d780de8066ec..9fa31f3a93ae37 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -44,6 +44,9 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "riegeli/base/any.h" +#include "riegeli/bytes/reader.h" +#include "riegeli/messages/parse_message.h" #include "xla/array.h" #include "xla/backends/cpu/collectives/cpu_collectives.h" #include "xla/backends/cpu/constant_allocation.h" @@ -565,15 +568,12 @@ PjRtCpuExecutable::GetOutputMemoryKinds() const { return out; } -absl::StatusOr> -PjRtCpuClient::LoadSerializedExecutableInternal( - google::protobuf::io::ZeroCopyInputStream* stream, - std::optional options, const LoadOptions& load_options) { +/*static*/ absl::StatusOr> +PjRtCpuExecutable::Deserialize(riegeli::Any reader, + const xla::CpuTopologyDescription& topology, + std::optional&& options) { ExecutableAndOptionsProto proto; - if (!proto.ParseFromZeroCopyStream(stream)) { - return Internal( - "PjRtCpuClient::DeserializeExecutable proto deserialization failed"); - } + ABSL_RETURN_IF_ERROR(riegeli::ParseMessage(reader.get(), proto)); CompileOptions compile_options; if (options.has_value()) { compile_options = *std::move(options); @@ -599,10 +599,9 @@ PjRtCpuClient::LoadSerializedExecutableInternal( ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions( compile_options.compile_portable_executable, &compile_options.executable_build_options, - [this](int num_replicas, int num_partitions) { - return topology().GetDefaultDeviceAssignment(process_index(), - num_replicas, std::nullopt, - num_partitions, nullptr); + [&topology](int num_replicas, int num_partitions) { + return topology.GetDefaultDeviceAssignment( + 0, num_replicas, std::nullopt, num_partitions, nullptr); }, &num_replicas, &num_partitions, &device_assignment)); @@ -630,37 +629,11 @@ PjRtCpuClient::LoadSerializedExecutableInternal( : nullptr); } - auto cpu_executable = std::make_shared( + auto cpu_executable = std::make_unique( num_replicas, num_partitions, std::move(input_options), std::move(executable), std::move(result_buffer_indices), nullptr, - topology()); - return LoadInternal(std::move(cpu_executable), std::move(device_assignment)); -} - -absl::StatusOr> -PjRtCpuClient::LoadSerializedExecutable(absl::string_view serialized, - std::optional options, - const LoadOptions& load_options) { - if (serialized.size() > std::numeric_limits::max()) { - return Internal( - "PjRtCpuClient::DeserializeExecutable proto too large (>2GB)"); - } - google::protobuf::io::ArrayInputStream stream(serialized.data(), serialized.size()); - return LoadSerializedExecutableInternal(&stream, std::move(options), - load_options); -} - -absl::StatusOr> -PjRtCpuClient::LoadSerializedExecutable(const absl::Cord& serialized, - std::optional options, - const LoadOptions& load_options) { - if (serialized.size() > std::numeric_limits::max()) { - return Internal( - "PjRtCpuClient::DeserializeExecutable proto too large (>2GB)"); - } - google::protobuf::io::CordInputStream stream(&serialized); - return LoadSerializedExecutableInternal(&stream, std::move(options), - load_options); + topology); + return cpu_executable; } absl::StatusOr> @@ -851,16 +824,6 @@ PjRtCpuRawClient::CompileAheadOfTime(const XlaComputation& computation, &aot_options); } -absl::StatusOr> -PjRtCpuClient::CompileAheadOfTimeAndLoad( - const XlaComputation& computation, CompileOptions options, - const AotCompilationOptions& aot_options) { - ABSL_ASSIGN_OR_RETURN(auto executable, raw_client()->CompileAheadOfTime( - computation, options, topology(), - process_index(), aot_options)); - return Load(std::move(executable), LoadOptions()); -} - struct CpuCompilationParams { LayoutCanonicalizationCallback layout_canonicalization_callback = nullptr; std::optional num_threads = std::nullopt; diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 2276bb60252b63..6f875a51d22a17 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -267,27 +267,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { std::shared_ptr executable, const LoadOptions& load_options) override; - // TODO(b/403584258): PJRT wants to have just one simple Compile API. When the - // CPU runtime stops supporting the legacy runtime we will unify our compile - // paths better and this will be redundant. - absl::StatusOr> - CompileAheadOfTimeAndLoad(const XlaComputation& computation, - CompileOptions options, - const AotCompilationOptions& aot_options); - - // For PjRtCpuClient, `options` is mandatory. - // This function returns an InvalidArgument error if `std::nullopt` is passed. - // TODO(b/237720161): make it actually optional - absl::StatusOr> - LoadSerializedExecutable(absl::string_view serialized, - std::optional options, - const LoadOptions& load_options) override; - - absl::StatusOr> - LoadSerializedExecutable(const absl::Cord& serialized, - std::optional options, - const LoadOptions& load_options) override; - bool IsOnCpu(PjRtMemorySpace* memory_space) override { return true; } const xla::CpuTopologyDescription& topology() const { @@ -319,11 +298,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { absl::StatusOr> LoadInternal( std::shared_ptr cpu_executable, std::shared_ptr device_assignment); - - absl::StatusOr> - LoadSerializedExecutableInternal(google::protobuf::io::ZeroCopyInputStream* stream, - std::optional options, - const LoadOptions& load_options); }; class PjRtCpuLoadedExecutable; @@ -431,6 +405,11 @@ class PjRtCpuExecutable final : public PjRtExecutable { const CompileOptions& compile_options() const { return compile_options_; } + static absl::StatusOr> Deserialize( + riegeli::Any reader, + const xla::CpuTopologyDescription& topology, + std::optional&& options); + private: friend class PjRtCpuClient; friend class CpuPjRtRawLoadedExecutable; diff --git a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc index 910dc9460103fc..39c600b97c3dcb 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/pjrt/cpu/cpu_pjrt_compiler.h" #include +#include #include #include @@ -25,6 +26,8 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "riegeli/base/any.h" +#include "riegeli/bytes/reader.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/pjrt/cpu/cpu_client.h" #include "xla/pjrt/maybe_owning_mlir_module.h" @@ -99,6 +102,16 @@ CpuPjRtCompiler::DeserializePjRtTopologyDescription( return CpuTopologyDescription::FromProto(proto); } +absl::StatusOr> +CpuPjRtCompiler::DeserializeExecutable( + const PjRtTopologyDescription& topology, + riegeli::Any reader, + std::optional&& options) { + ABSL_ASSIGN_OR_RETURN(auto* cpu_topology, GetCpuTopology(topology)); + return PjRtCpuExecutable::Deserialize(std::move(reader), *cpu_topology, + std::move(options)); +} + } // namespace xla::cpu // Doesn't use STREAM_EXECUTOR_REGISTER_MODULE_INITIALIZER to ensure it is diff --git a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h index 28672b3fd0876f..6cb1307d82619c 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h @@ -17,9 +17,12 @@ limitations under the License. #define XLA_PJRT_CPU_CPU_PJRT_COMPILER_H_ #include +#include #include #include "absl/status/statusor.h" +#include "riegeli/base/any.h" +#include "riegeli/bytes/reader.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/pjrt/maybe_owning_mlir_module.h" #include "xla/pjrt/pjrt_compiler.h" @@ -40,6 +43,12 @@ class CpuPjRtCompiler : public PjRtCompiler { CompileOptions options, MaybeOwningMlirModule module, const PjRtTopologyDescription& topology, PjRtClient* client) override; + // Deserializes a serialized executable. + absl::StatusOr> DeserializeExecutable( + const PjRtTopologyDescription& topology, + riegeli::Any reader, + std::optional&& options) override; + // Deserializes a PjRtTopologyDescription from a string. absl::StatusOr> DeserializePjRtTopologyDescription( From 6db7a6f53526ede6927bbd012485778659a5a81b Mon Sep 17 00:00:00 2001 From: Ionel Gog Date: Thu, 27 Aug 2026 12:21:02 -0700 Subject: [PATCH 17/28] [IFRT IR] Catch a new possible error message in invalid programs that donate twice PiperOrigin-RevId: 972090625 --- .../ir/ifrt_ir_loaded_executable_test_lib.cc | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc index 07be8793b322fe..7730983c23967a 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc @@ -1473,13 +1473,19 @@ module { ExecuteOptionsWithFillStatus(), devices)); EXPECT_THAT( result.status.Await(), - StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("Invalid buffer passed to Execute() as argument 0"))); + AnyOf(StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Invalid buffer passed to Execute() as argument 0")), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("donates or deletes twice an array")))); ASSERT_EQ(result.outputs.size(), 1); EXPECT_THAT( result.outputs[0]->GetReadyFuture().Await(), - StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("Invalid buffer passed to Execute() as argument 0"))); + AnyOf(StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Invalid buffer passed to Execute() as argument 0")), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("donates or deletes twice an array")))); } TEST_F(IfrtIrLoadedExecutableTest, DonatingTwiceAliasedBufferThrowsError) { @@ -1533,6 +1539,8 @@ module { StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr( "Attempt to donate the same buffer twice in Execute()")), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("donates or deletes twice an array")), StatusIs( absl::StatusCode::kNotFound, HasSubstr( @@ -1544,6 +1552,8 @@ module { StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr( "Attempt to donate the same buffer twice in Execute()")), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("donates or deletes twice an array")), StatusIs( absl::StatusCode::kNotFound, HasSubstr( From 327ef3d139c4b26e76eefd36bd542240ad8d8107 Mon Sep 17 00:00:00 2001 From: Oleg Shyshkov Date: Thu, 27 Aug 2026 13:30:34 -0700 Subject: [PATCH 18/28] [XLA:GPU] Add DegenerateDimensionRewriter to GPU pipeline. This pass removes unnecessary size-1 from the module. Degenerate dimensions are generally no-op, but it add unnecessary reshapes/bitcasts to the graph that can sometime prevent better fusion and tiling decision or cause problem with emitter pipelines, like Triton. PiperOrigin-RevId: 972126534 --- third_party/xla/xla/service/gpu/BUILD | 1 + third_party/xla/xla/service/gpu/gpu_compiler.cc | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 1357ef9030f976..eb3e044fe4fc25 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1971,6 +1971,7 @@ cc_library( "//xla/hlo/transforms/simplifiers:broadcast_canonicalizer", "//xla/hlo/transforms/simplifiers:conditional_canonicalizer", "//xla/hlo/transforms/simplifiers:convert_mover", + "//xla/hlo/transforms/simplifiers:degenerate_dimension_rewriter", "//xla/hlo/transforms/simplifiers:dot_merger", "//xla/hlo/transforms/simplifiers:dynamic_dimension_simplifier", "//xla/hlo/transforms/simplifiers:flatten_call_graph", diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 97d916a9feb43f..4ca8fef163aebd 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -209,6 +209,7 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/broadcast_canonicalizer.h" #include "xla/hlo/transforms/simplifiers/conditional_canonicalizer.h" #include "xla/hlo/transforms/simplifiers/convert_mover.h" +#include "xla/hlo/transforms/simplifiers/degenerate_dimension_rewriter.h" #include "xla/hlo/transforms/simplifiers/dot_merger.h" #include "xla/hlo/transforms/simplifiers/dynamic_dimension_simplifier.h" #include "xla/hlo/transforms/simplifiers/flatten_call_graph.h" @@ -977,6 +978,14 @@ absl::Status RunOptimizationPasses( pipeline.AddPass(); pipeline.AddPass( gpu_target_config.device_description.gpu_compute_capability()); + + // It's important to run AlgebraicSimplifier after + // DegenerateDimensionRewriter before ReshapeMover. + // DegenerateDimensionRewriter introduces reshape to remove size-1 dims from + // ops like iota and broadcast, and algebraic simplifier has patterns to + // fold reshape(iota) and reshape(broadcast). If we run ReshapeMover first, + // it will move these reshapes down the graph, and prevent the folding. + pipeline.AddPass(); pipeline.AddPass(layout_insensitive_algsimp_opts, gpu_version); pipeline.AddPass(); From a433858c9f8e0d4fda93cd1666808690237f4da5 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Thu, 27 Aug 2026 14:12:34 -0700 Subject: [PATCH 19/28] Add PjRtExecutable::GetHloModule() which returns a single hlo module for backends which are known to only support a single hlo module. PiperOrigin-RevId: 972151130 --- .../pjrt/c_api_client/pjrt_c_api_client.cc | 17 ++------ .../xla/pjrt/c_api_client/pjrt_c_api_client.h | 3 +- third_party/xla/xla/pjrt/cpu/cpu_client.h | 6 +-- third_party/xla/xla/pjrt/gpu/BUILD | 1 + .../xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc | 41 ++++++++++++++----- third_party/xla/xla/pjrt/pjrt_executable.cc | 6 +++ third_party/xla/xla/pjrt/pjrt_executable.h | 7 +++- third_party/xla/xla/pjrt/se/BUILD | 2 + .../xla/pjrt/se/stream_executor_executable.cc | 2 + .../xla/pjrt/se/stream_executor_executable.h | 8 +--- .../se/stream_executor_executable_test.cc | 5 +++ 11 files changed, 61 insertions(+), 37 deletions(-) diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc index 0b78e60c6d85a1..9e08a8071358bb 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc @@ -2675,8 +2675,8 @@ PjRtCApiExecutable::GetOutputMemoryKinds() const { return std::vector>{std::move(out)}; } -absl::StatusOr>> -PjRtCApiExecutable::GetHloModules() const { +absl::StatusOr> PjRtCApiExecutable::GetHloModule() + const { auto* c_api = pjrt_c_api(); auto* executable = c_executable(); PJRT_Executable_OptimizedProgram_Args args; @@ -2721,23 +2721,14 @@ PjRtCApiExecutable::GetHloModules() const { // equivalent) once implemented. mlir::MlirToHloConversionOptions options; options.return_tuple = false; - ABSL_ASSIGN_OR_RETURN(std::unique_ptr hlo_module, - mlir::ConvertMlirHloToHloModule(module.get(), options)); - - std::vector> out; - out.push_back(std::move(hlo_module)); - return out; + return mlir::ConvertMlirHloToHloModule(module.get(), options); } HloModuleProtoWithConfig proto; if (!proto.ParseFromString(code)) { return InvalidArgument("Failed to deserialize HloModuleProtoWithConfig"); } - std::vector> out; - ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, - HloModule::CreateFromProtoWithConfig(proto)); - out.push_back(std::move(module)); - return out; + return HloModule::CreateFromProtoWithConfig(proto); } absl::StatusOr PjRtCApiExecutable::SerializeExecutable() const { diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h index 5c892ae2ecad28..e9fdfd162fc88f 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h @@ -739,8 +739,7 @@ class PjRtCApiExecutable : public PjRtExecutable { absl::StatusOr> GetCostAnalysis() const override; - absl::StatusOr>> GetHloModules() - const override; + absl::StatusOr> GetHloModule() const override; absl::StatusOr GetCompiledMemoryStats() const override { return pjrt::GetCompiledMemoryStats(c_api_, executable_.get()); diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 6f875a51d22a17..01ca97aafd81d3 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -373,10 +373,8 @@ class PjRtCpuExecutable final : public PjRtExecutable { return cpu_executable_->SizeOfGeneratedCodeInBytes(); } - absl::StatusOr>> GetHloModules() - const override { - return std::vector>{ - cpu_executable_->shared_module()}; + absl::StatusOr> GetHloModule() const override { + return cpu_executable_->shared_module(); } absl::StatusOr>> diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD index ec0db00389c5fa..4b331c67bc86ca 100644 --- a/third_party/xla/xla/pjrt/gpu/BUILD +++ b/third_party/xla/xla/pjrt/gpu/BUILD @@ -917,6 +917,7 @@ xla_test( "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:gpu_topology", + "//xla/service:mock_compiled_module", "//xla/service:mock_compiler", "//xla/service:platform_util", "//xla/stream_executor:stream_executor_h", diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc index bd28695ad9d68a..cdb2f122b4b003 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc @@ -55,6 +55,7 @@ limitations under the License. #include "xla/service/compiled_module.h" #include "xla/service/compiler.h" #include "xla/service/gpu_topology.h" +#include "xla/service/mock_compiled_module.h" #include "xla/service/mock_compiler.h" #include "xla/service/platform_util.h" #include "xla/shape_layout.h" @@ -67,7 +68,8 @@ namespace xla { namespace { using ::absl_testing::IsOkAndHolds; using ::testing::_; -using ::testing::IsEmpty; +using ::testing::ByMove; +using ::testing::ElementsAre; using ::testing::IsNull; using ::testing::Optional; using ::testing::Property; @@ -313,6 +315,17 @@ TEST(StreamExecutorGpuCompilerTest, GetTargetRuntimeAbiVersion) { EXPECT_OK(runtime_abi_version->IsCompatibleWith(*executable_abi_version)); } +std::vector> MakeMockCompiledModule( + std::shared_ptr hlo_module = + std::make_shared("name", HloModuleConfig())) { + auto mock_compiled_module = std::make_unique(); + EXPECT_CALL(*mock_compiled_module, shared_optimized_module()) + .WillRepeatedly(Return(hlo_module)); + std::vector> aot_results; + aot_results.push_back(std::move(mock_compiled_module)); + return aot_results; +} + TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) { auto mock_compiler = std::make_unique(); MockCompiler& mock_compiler_ref = *mock_compiler; @@ -329,6 +342,8 @@ TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) { StreamExecutorGpuTopologyDescription topology_description( CudaId(), CudaName(), gpu_topology); + auto hlo_module = std::make_shared("name", HloModuleConfig()); + // We expect that the underlying compiler is called with no executor given. EXPECT_CALL(mock_compiler_ref, PlatformId) .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId)); @@ -336,14 +351,15 @@ TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) { EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime( _, Property(&AotCompilationOptions::executor, IsNull()))) - .WillOnce(Return(std::vector>{})); + .WillOnce(Return(ByMove(MakeMockCompiledModule(hlo_module)))); ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram)); ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, pjrt_compiler.Compile(CompileOptions(), computation, topology_description, /*client=*/nullptr)); - EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(IsEmpty())); + EXPECT_THAT(executable->GetHloModules(), + IsOkAndHolds(ElementsAre(hlo_module))); } TEST(StreamExecutorGpuCompilerTest, CrossCompilation) { @@ -372,6 +388,8 @@ TEST(StreamExecutorGpuCompilerTest, CrossCompilation) { se::StreamExecutor* stream_executor = se_gpu_client->client()->backend().default_stream_executor(); + auto hlo_module = std::make_shared("name", HloModuleConfig()); + // We expect that the underlying compiler is called with the executor from the // PjRt client. EXPECT_CALL(mock_compiler_ref, PlatformId) @@ -390,14 +408,15 @@ TEST(StreamExecutorGpuCompilerTest, CrossCompilation) { Property(&GpuTopology::num_partitions, 1), Property(&GpuTopology::num_hosts_per_partition, 1), Property(&GpuTopology::num_devices_per_host, 1))))))) - .WillOnce(Return(std::vector>{})); + .WillOnce(Return(ByMove(MakeMockCompiledModule(hlo_module)))); ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram)); ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, pjrt_compiler.Compile(CompileOptions(), computation, topology_description, client.get())); - EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(IsEmpty())); + EXPECT_THAT(executable->GetHloModules(), + IsOkAndHolds(ElementsAre(hlo_module))); } absl::StatusOr> GetSampleH100basedGpuTopology() { @@ -436,12 +455,12 @@ TEST(StreamExecutorGpuCompilerTest, AutoLayoutIsPropagatedInCrossCompilation) { .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId)); EXPECT_CALL(mock_compiler_ref, Compile).Times(0); - std::unique_ptr hlo_module; + std::shared_ptr hlo_module; EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime) .WillOnce([&](std::unique_ptr module, const AotCompilationOptions& options) { hlo_module = std::move(module); - return std::vector>{}; + return MakeMockCompiledModule(hlo_module); }); ASSERT_OK_AND_ASSIGN(std::unique_ptr client, @@ -475,12 +494,12 @@ TEST(StreamExecutorGpuCompilerTest, .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId)); EXPECT_CALL(mock_compiler_ref, Compile).Times(0); - std::unique_ptr hlo_module; + std::shared_ptr hlo_module; EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime) .WillOnce([&](std::unique_ptr module, const AotCompilationOptions& options) { hlo_module = std::move(module); - return std::vector>{}; + return MakeMockCompiledModule(hlo_module); }); ASSERT_OK_AND_ASSIGN( @@ -512,12 +531,12 @@ TEST(StreamExecutorGpuCompilerTest, .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId)); EXPECT_CALL(mock_compiler_ref, Compile).Times(0); - std::unique_ptr hlo_module; + std::shared_ptr hlo_module; EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime) .WillOnce([&](std::unique_ptr module, const AotCompilationOptions& options) { hlo_module = std::move(module); - return std::vector>{}; + return MakeMockCompiledModule(hlo_module); }); CompileOptions options{}; diff --git a/third_party/xla/xla/pjrt/pjrt_executable.cc b/third_party/xla/xla/pjrt/pjrt_executable.cc index 50dd2df937407f..80708b7d4614fd 100644 --- a/third_party/xla/xla/pjrt/pjrt_executable.cc +++ b/third_party/xla/xla/pjrt/pjrt_executable.cc @@ -315,6 +315,12 @@ void GetOpSharding(std::vector& out, const OpSharding& sharding) { } } +absl::StatusOr>> +PjRtExecutable::GetHloModules() const { + ABSL_ASSIGN_OR_RETURN(std::shared_ptr hlo_module, GetHloModule()); + return std::vector>{std::move(hlo_module)}; +} + std::optional> PjRtExecutable::GetOutputShardings() const { auto modules = GetHloModules(); diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h index ee23b502ae4aa3..4a7c456c8457a1 100644 --- a/third_party/xla/xla/pjrt/pjrt_executable.h +++ b/third_party/xla/xla/pjrt/pjrt_executable.h @@ -367,9 +367,14 @@ class PjRtExecutable { // Unique name for this executable, e.g., HloModule name. virtual absl::string_view name() const = 0; + // Return HloModule (optimized). + virtual absl::StatusOr> GetHloModule() const { + return absl::UnimplementedError("GetHloModule is not implemented"); + } + // Return an array of HloModule (optimized) per partition. virtual absl::StatusOr>> - GetHloModules() const = 0; + GetHloModules() const; // Returns an output Shape per program, the size should be equal to // `GetHloModules()`. diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD index aaeb4658eabdde..154374dbfde4e2 100644 --- a/third_party/xla/xla/pjrt/se/BUILD +++ b/third_party/xla/xla/pjrt/se/BUILD @@ -178,11 +178,13 @@ xla_cc_test( srcs = ["stream_executor_executable_test.cc"], deps = [ ":stream_executor_executable", + "//xla/hlo/ir:hlo", "//xla/pjrt:pjrt_abi_version", "//xla/pjrt:pjrt_common", "//xla/pjrt:pjrt_executable", "//xla/pjrt/proto:pjrt_abi_version_proto_cc", "//xla/service:compiled_module", + "//xla/service:hlo_module_config", "//xla/service:mock_compiled_module", "//xla/stream_executor/abi:executable_abi_version", "//xla/stream_executor/abi:executable_abi_version_proto_cc", diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc index 61e578cda19a00..94eaddff745c85 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc @@ -169,6 +169,7 @@ StreamExecutorExecutable::StreamExecutorExecutable( if (mod != nullptr) { hlo_module_ = mod->shared_optimized_module(); } + CHECK(hlo_module_ != nullptr); } StreamExecutorExecutable::StreamExecutorExecutable( @@ -193,6 +194,7 @@ StreamExecutorExecutable::StreamExecutorExecutable( if (local_exec != nullptr) { hlo_module_ = local_exec->executable()->shared_module(); } + CHECK(hlo_module_ != nullptr); } absl::StatusOr diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.h b/third_party/xla/xla/pjrt/se/stream_executor_executable.h index 31b14369015a11..cb0175a729b5d2 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.h +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.h @@ -78,12 +78,8 @@ class StreamExecutorExecutable : public PjRtExecutable { absl::StatusOr GetCompileOptions() const override { return compile_options_; } - absl::StatusOr>> GetHloModules() - const override { - if (hlo_module_ == nullptr) { - return std::vector>{}; - } - return std::vector>{hlo_module_}; + absl::StatusOr> GetHloModule() const override { + return hlo_module_; } const std::shared_ptr& hlo_module() const { return hlo_module_; } diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc index c22e5e255d7bf9..2da6cb9a4bab5b 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc @@ -21,11 +21,13 @@ limitations under the License. #include #include +#include "xla/hlo/ir/hlo_module.h" #include "xla/pjrt/pjrt_abi_version.h" #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/pjrt/proto/pjrt_abi_version.pb.h" #include "xla/service/compiled_module.h" +#include "xla/service/hlo_module_config.h" #include "xla/service/mock_compiled_module.h" #include "xla/stream_executor/abi/executable_abi_version.h" #include "xla/stream_executor/abi/executable_abi_version.pb.h" @@ -49,6 +51,9 @@ TEST(StreamExecutorExecutableTest, GetAbiVersion) { constexpr PjRtPlatformId kPlatformId = 42; auto module = std::make_unique(); + auto hlo_module = std::make_shared("name", HloModuleConfig()); + EXPECT_CALL(*module, shared_optimized_module()) + .WillRepeatedly(Return(hlo_module)); EXPECT_CALL(*module, GetExecutableAbiVersion()) .WillOnce(Return(executable_abi_version)); StreamExecutorExecutable executable(kPlatformId, CompileOptions(), From 5b9e491e7c00a24f96d370cece8366544f40d20e Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 27 Aug 2026 14:36:02 -0700 Subject: [PATCH 20/28] Revert matrix input scaling in EighExpander and TpuEighExpander. Reverting due to compilation failure when handling complex types during TPU lowering: passing Zero(builder, type) where Abs(a) produces a real scalar causes type mismatch and HLO verification/compilation failure. Reverts 6ae3d028916d2138b26b155342b24314bacea6f5 PiperOrigin-RevId: 972162896 --- .../hlo/builder/lib/self_adjoint_eig_test.cc | 80 ------------------- .../hlo/transforms/expanders/eigh_expander.cc | 28 ------- 2 files changed, 108 deletions(-) diff --git a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc index 5778da3b926576..467445ea097988 100644 --- a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc +++ b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc @@ -259,86 +259,6 @@ TEST_F(SelfAdjointEigTest, Test_Orthogonality_8x8) { ErrorSpec(1e-3, 1e-3)); } -TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_2x2) { - XlaBuilder builder(TestName()); - float v = 1e20f; - Array2D input{{v, v}, {v, v}}; - std::vector expected{0.0f, 2e20f}; - - XlaOp a; - auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); - auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, - /*tol=*/1e-5, /*sort_eigenvalues=*/true); - Add(result.w, ZerosLike(result.w)); - - ComputeAndCompareR1(&builder, expected, {&a_data}, - ErrorSpec(1e15f, 1e-4f)); -} - -TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_3x3) { - XlaBuilder builder(TestName()); - float v = 1e20f; - Array2D input{{v, v, v}, {v, v, v}, {v, v, v}}; - std::vector expected{0.0f, 0.0f, 3e20f}; - - XlaOp a; - auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); - auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, - /*tol=*/1e-5, /*sort_eigenvalues=*/true); - Add(result.w, ZerosLike(result.w)); - - ComputeAndCompareR1(&builder, expected, {&a_data}, - ErrorSpec(1e15f, 1e-4f)); -} - -TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_Complex_3x3) { - XlaBuilder builder(TestName()); - float v = 1e20f; - Array input = { - {complex64{v, 0.0f}, complex64{v, -v}, complex64{0.0f, 0.0f}}, - {complex64{v, v}, complex64{v, 0.0f}, complex64{0.0f, 0.0f}}, - {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{v, 0.0f}}, - }; - const Literal a_literal = LiteralUtil::CreateFromArray(input); - XlaOp a = Parameter(&builder, 0, a_literal.shape(), "a"); - auto result = SelfAdjointEig(a); - ComputeMatmulVWVt(result, &builder); - - ComputeAndCompareLiteral(&builder, LiteralUtil::CreateFromArray(input), - {&a_literal}, ErrorSpec(1e15f, 1e-4f)); -} - -TEST_F(SelfAdjointEigTest, Test_Small_Magnitude_2x2) { - XlaBuilder builder(TestName()); - float v = 1e-20f; - Array2D input{{v, v}, {v, v}}; - std::vector expected{0.0f, 2e-20f}; - - XlaOp a; - auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); - auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, - /*tol=*/1e-5, /*sort_eigenvalues=*/true); - Add(result.w, ZerosLike(result.w)); - - ComputeAndCompareR1(&builder, expected, {&a_data}, - ErrorSpec(1e-25f, 1e-4f)); -} - -TEST_F(SelfAdjointEigTest, Test_Zero_Matrix_2x2) { - XlaBuilder builder(TestName()); - Array2D input{{0.0f, 0.0f}, {0.0f, 0.0f}}; - std::vector expected{0.0f, 0.0f}; - - XlaOp a; - auto a_data = CreateR2Parameter(input, 0, "a", &builder, &a); - auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15, - /*tol=*/1e-5, /*sort_eigenvalues=*/true); - Add(result.w, ZerosLike(result.w)); - - ComputeAndCompareR1(&builder, expected, {&a_data}, - ErrorSpec(1e-6f, 1e-6f)); -} - TEST_F(SelfAdjointEigTest, Wrong_Type_Int) { XlaBuilder builder(TestName()); diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc index 7400fb932f566f..3b6c0debf5a26e 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc @@ -477,29 +477,6 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, a = Symmetrize(a, lower); - PrimitiveType real_type = primitive_util::IsComplexType(type) - ? primitive_util::ComplexComponentType(type) - : type; - XlaOp zero_real = Zero(builder, real_type); - XlaOp one_real = One(builder, real_type); - XlaOp abs_a = primitive_util::IsComplexType(type) - ? Max(Abs(Real(a)), Abs(Imag(a))) - : Abs(a); - XlaOp a_max = - Reduce(abs_a, zero_real, CreateScalarMaxComputation(real_type, builder), - {num_dims - 2, num_dims - 1}); - XlaOp scale = Select(Eq(a_max, zero_real), one_real, a_max); - - std::vector batch_broadcast_dims(num_batch_dims); - absl::c_iota(batch_broadcast_dims, 0); - - XlaOp scale_a = primitive_util::IsComplexType(type) - ? Complex(scale, ZerosLike(scale)) - : scale; - scale_a = - BroadcastInDim(scale_a, a_shape.dimensions(), batch_broadcast_dims); - a = a / scale_a; - const int64_t k = CeilOfRatio(n, int64_t{2}); // tl = A[:n // 2, :n // 2] // bl = A[n // 2:, :n // 2] @@ -560,11 +537,6 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol, } v = MaybeConjugate(TransposeInMinorDims(v), true); - ABSL_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w)); - XlaOp scale_w = - BroadcastInDim(scale, w_shape.dimensions(), batch_broadcast_dims); - w = w * scale_w; - if (sort_eigenvalues) { ABSL_RETURN_IF_ERROR(SortByEigenvalues(v, w)); } From 69792ea5fd1415f4a1d9b714ca8a10a5d0866f2e Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Thu, 27 Aug 2026 14:43:26 -0700 Subject: [PATCH 21/28] [XLA:CPU] Exclude the patterns that attempt to insert memref.subview. memref.subview expects standard strided or identity layouts and does not support #xtile.layout, which breaks the verification. PiperOrigin-RevId: 972166650 --- .../non_identity_layout_unit_dim.hlo | 27 +++++++++++++++++++ .../transforms/drop_vector_unit_dims_pass.cc | 2 -- .../tests/drop_vector_unit_dims_pass.mlir | 14 ++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo new file mode 100644 index 00000000000000..e93415f2ad4d6b --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo @@ -0,0 +1,27 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +fusion { + param_0 = f32[2,1,2,1]{3,1,0,2} parameter(0) + copy_0 = f32[2,1,2,1]{3,2,1,0} copy(param_0) + transpose = f32[2,2,1,1]{3,0,2,1} transpose(copy_0), dimensions={2,0,1,3} + copy_1 = f32[2,2,1,1]{3,2,1,0} copy(transpose) + ROOT bitcast = f32[2,2]{1,0} bitcast(copy_1) +} + +ENTRY main { + param_0 = f32[2,1,2,1]{3,1,0,2} parameter(0) + ROOT fusion = f32[2,2]{1,0} fusion(param_0), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["2", "2"]}], + "num_ctas":"1" + } + } + } +} +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM-NOT: llvm.alloc diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc index 798e110957cc43..7663a570e87dc5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc @@ -38,8 +38,6 @@ class DropVectorUnitDimsPass mlir::RewritePatternSet patterns(context); mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(patterns); mlir::vector::populateDropUnitDimWithShapeCastPatterns(patterns); - mlir::vector::populateDropInnerMostUnitDimsXferOpPatterns(patterns); - mlir::vector::populateVectorTransferDropUnitDimsPatterns(patterns); if (mlir::failed( mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) { signalPassFailure(); diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir index 569e8ec080d24b..35e4a68845bc14 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir @@ -65,3 +65,17 @@ func.func @no_unit_dims(%arg0: vector<16xf32>, %arg1: vector<16xf32>) -> vector< // CHECK-SAME: %[[ARG0:.*]]: vector<16xf32>, %[[ARG1:.*]]: vector<16xf32>) -> vector<16xf32> { // CHECK-NOT: vector.shape_cast // CHECK: arith.addf %[[ARG0]], %[[ARG1]] : vector<16xf32> + +// ----- + +func.func @transfer_read_non_identity_layout_unit_dim( + %arg0: memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, %c0: index) -> vector<2x1x2x1xf32> { + %pad = arith.constant 0.0 : f32 + %0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true, true]} : memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, vector<2x1x2x1xf32> + return %0 : vector<2x1x2x1xf32> +} +// CHECK-LABEL: func.func @transfer_read_non_identity_layout_unit_dim( +// CHECK-SAME: %[[ARG0:.*]]: memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, %[[C0:.*]]: index) -> vector<2x1x2x1xf32> { +// CHECK-NOT: memref.subview +// CHECK: vector.transfer_read %[[ARG0]] + From ce016ec14060bc6027a206ca191e98ab2f51ca04 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 27 Aug 2026 14:43:55 -0700 Subject: [PATCH 22/28] [XLA:TSL] Add kDimensions and kType stat types to XPlane schema. Adds `StatType::kDimensions` ("dims") and `StatType::kType` ("type") to the TSL profiler XPlane schema to support events with individual dimension and data type attributes. PiperOrigin-RevId: 972166901 --- third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc | 4 +++- third_party/xla/xla/tsl/profiler/utils/xplane_schema.h | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc index 69fbaef0a48818..a36ff25dc0d6e9 100644 --- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc +++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc @@ -426,7 +426,9 @@ const StatTypeMap& GetStatTypeMap() { {"hbm_power_events", kHbmPowerEvents}, {"transaction_with_chip_core_id", kTransactionWithChipCoreId}, {"program_counter", kProgramCounter}, - {"uses_ici", kUsesIci}}); + {"uses_ici", kUsesIci}, + {"dims", kDimensions}, + {"type", kType}}); DCHECK_EQ(stat_type_map->size(), kNumStatTypes); return *stat_type_map; } diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h index 6aa8200ed1a8be..3f776c23721ae0 100644 --- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h +++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h @@ -408,11 +408,13 @@ enum StatType { // Program Counter in Oci Descriptors, etc kProgramCounter, kUsesIci, + kDimensions, + kType, // LINT.ThenChange(:last_stat_type) // LINT.IfChange(last_stat_type) // Change this to point to the last stat type when adding a new one. - kLastStatType = kUsesIci, + kLastStatType = kType, // LINT.ThenChange(:stat_type_enum) }; From 87a19043223261fc3df0d3e9eed8000fa71defff Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 27 Aug 2026 14:59:09 -0700 Subject: [PATCH 23/28] [XLA][Verifier]Remove output-to-operand aliasing check for call instructions. In preparation for supporting thread filtering in buffer assignment, this change removes the restriction in the HLO verifier that prevents call instructions from having output-to-operand aliasing. This allows call instructions to specify input/output aliasing. PiperOrigin-RevId: 972174275 --- .../xla/xla/hlo/analysis/alias_info.cc | 9 +- .../hlo/analysis/hlo_alias_analysis_test.cc | 42 +++++ third_party/xla/xla/service/hlo_verifier.cc | 38 +++-- .../xla/xla/service/hlo_verifier_test.cc | 152 +++++++++++++++++- 4 files changed, 225 insertions(+), 16 deletions(-) diff --git a/third_party/xla/xla/hlo/analysis/alias_info.cc b/third_party/xla/xla/hlo/analysis/alias_info.cc index b0d32815773303..cdf938757c834b 100644 --- a/third_party/xla/xla/hlo/analysis/alias_info.cc +++ b/third_party/xla/xla/hlo/analysis/alias_info.cc @@ -208,11 +208,10 @@ AliasInfo::GetInPlaceInputOutputPairs(const HloInstruction* user) const { return {}; } } - if (user->opcode() == HloOpcode::kCustomCall) { - // Custom Calls previously assumed that aliased operands were - // forwarded, but now supports modification semantics. - const auto& aliasing_pairs = - Cast(user)->output_to_operand_aliasing(); + if (user->opcode() == HloOpcode::kCall || + user->opcode() == HloOpcode::kCustomCall) { + const auto* callable = Cast(user); + const auto& aliasing_pairs = callable->output_to_operand_aliasing(); std::vector> in_place_pairs; in_place_pairs.reserve(aliasing_pairs.size()); for (const auto& pair : aliasing_pairs) { diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc index fee1d7580d3b89..1f87c30b7ba7d0 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc @@ -24,8 +24,10 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/analysis/hlo_ordering.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/test.h" @@ -1376,5 +1378,45 @@ ENTRY main { analysis.GetUniqueBufferAt(tc_operand1)); } +TEST_F(HloAliasAnalysisTest, CallOutputToOperandAliasing) { + absl::string_view hlo_string = R"( +HloModule Module + +callme { + sub_p0 = f32[16] parameter(0) + sub_p1 = f32[16] parameter(1) + add0 = f32[16] add(sub_p0, sub_p1) + add1 = f32[16] add(sub_p0, sub_p0) + ROOT tuple = (f32[16], f32[16]) tuple(add0, add1) +} + +ENTRY main { + entry_p0 = f32[16] parameter(0) + entry_p1 = f32[16] parameter(1) + ROOT call = (f32[16], f32[16]) call(entry_p0, entry_p1), to_apply=callme +} +)"; + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + auto* call = Cast( + module_->entry_computation()->GetInstructionWithName("call")); + call->set_output_to_operand_aliasing({{{0}, {0, {}}}, {{1}, {1, {}}}}); + + HloAliasAnalysis& analysis = RunAnalysis(); + + const HloInstruction* entry_p0 = + module_->entry_computation()->GetInstructionWithName("entry_p0"); + const HloInstruction* entry_p1 = + module_->entry_computation()->GetInstructionWithName("entry_p1"); + + ASSERT_NE(entry_p0, nullptr); + ASSERT_NE(entry_p1, nullptr); + ASSERT_NE(call, nullptr); + + EXPECT_EQ(analysis.GetUniqueBufferAt(entry_p0), + analysis.GetUniqueBufferAt(call, {0})); + EXPECT_EQ(analysis.GetUniqueBufferAt(entry_p1), + analysis.GetUniqueBufferAt(call, {1})); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index 2ef568d5658deb..b6ab4b9cf7a4b5 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -4387,17 +4387,35 @@ absl::Status InstructionVerifier::HandleWhile(HloInstruction* xla_while) { absl::Status InstructionVerifier::HandleCall(HloInstruction* call) { if (opts_.verify_call_nested_computation_thread_name) { - return CheckCallableInstructionThreadName(call); + ABSL_RETURN_IF_ERROR(CheckCallableInstructionThreadName(call)); + } + const auto* callable = Cast(call); + for (const auto& pair : callable->output_to_operand_aliasing()) { + TF_RET_CHECK(pair.second.first < callable->operand_count()) + << "Invalid aliasing operand index."; + TF_RET_CHECK(ShapeUtil::IndexIsValid( + callable->operand(pair.second.first)->shape(), pair.second.second)) + << "Invalid aliasing operand shape index."; + TF_RET_CHECK(ShapeUtil::IndexIsValid(callable->shape(), pair.first)) + << "Invalid aliasing output shape index."; + const Shape& output_subshape = + ShapeUtil::GetSubshape(callable->shape(), pair.first); + const Shape& operand_subshape = ShapeUtil::GetSubshape( + callable->operand(pair.second.first)->shape(), pair.second.second); + if (opts_.layout_sensitive) { + TF_RET_CHECK(Shape::Equal().IgnoreDynamicDimension()(operand_subshape, + output_subshape)) + << "Different aliasing shapes: " + << operand_subshape.ToString(/*print_layout=*/true) << " vs " + << output_subshape.ToString(/*print_layout=*/true); + } else { + TF_RET_CHECK(Shape::Equal().IgnoreDynamicDimension().IgnoreLayout()( + operand_subshape, output_subshape)) + << "Different aliasing shapes: " + << operand_subshape.ToString(/*print_layout=*/false) << " vs " + << output_subshape.ToString(/*print_layout=*/false); + } } - - // As opposed to other callable instructions, nothing respects input/output - // aliasing for call instructions, so make sure it's not set. - const HloCallableInstruction* callable = - DynCast(call); - TF_RET_CHECK(callable != nullptr); - TF_RET_CHECK(callable->output_to_operand_aliasing().empty()) - << "Call instruction " << call->ToString() - << " may not have an output-to-operand aliasing set."; return absl::OkStatus(); } diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index 26f5b36fc21d51..117e53583e19bb 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -326,10 +326,133 @@ TEST_F(HloVerifierTest, CheckCallOperandOutputAliasing) { module->entry_computation()->GetInstructionWithName("mycall")) ->set_output_to_operand_aliasing({{{}, {0, {}}}}); + EXPECT_OK(verifier().Run(module.get()).status()); +} + +TEST_F(HloVerifierTest, CheckCallTupleOperandOutputAliasing) { + constexpr absl::string_view hlo = R"( + HloModule Module + + callme { + p0 = (s32[], f32[4]) parameter(0) + p1 = f32[4] parameter(1) + p0_0 = s32[] get-tuple-element(p0), index=0 + ROOT result = (s32[], f32[4]) tuple(p0_0, p1) + } + + ENTRY entry { + p0 = (s32[], f32[4]) parameter(0) + p1 = f32[4] parameter(1) + ROOT mycall = (s32[], f32[4]) call(p0, p1), to_apply=callme + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + + Cast( + module->entry_computation()->GetInstructionWithName("mycall")) + ->set_output_to_operand_aliasing({{{0}, {0, {0}}}, {{1}, {1, {}}}}); + + EXPECT_OK(verifier().Run(module.get()).status()); +} + +TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingInvalidOperandIndex) { + constexpr absl::string_view hlo = R"( + HloModule Module + + callme { + ROOT param = (s32[], f32[4]) parameter(0) + } + + ENTRY entry { + p0 = (s32[], f32[4]) parameter(0) + ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + + Cast( + module->entry_computation()->GetInstructionWithName("mycall")) + ->set_output_to_operand_aliasing({{{}, {1, {}}}}); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.message(), HasSubstr("Invalid aliasing operand index.")); +} + +TEST_F(HloVerifierTest, + CheckCallOperandOutputAliasingInvalidOperandShapeIndex) { + constexpr absl::string_view hlo = R"( + HloModule Module + + callme { + ROOT param = (s32[], f32[4]) parameter(0) + } + + ENTRY entry { + p0 = (s32[], f32[4]) parameter(0) + ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + + Cast( + module->entry_computation()->GetInstructionWithName("mycall")) + ->set_output_to_operand_aliasing({{{0}, {0, {2}}}}); + auto status = verifier().Run(module.get()).status(); ASSERT_FALSE(status.ok()); EXPECT_THAT(status.message(), - HasSubstr("may not have an output-to-operand aliasing set.")); + HasSubstr("Invalid aliasing operand shape index.")); +} + +TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingInvalidOutputShapeIndex) { + constexpr absl::string_view hlo = R"( + HloModule Module + + callme { + ROOT param = (s32[], f32[4]) parameter(0) + } + + ENTRY entry { + p0 = (s32[], f32[4]) parameter(0) + ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + + Cast( + module->entry_computation()->GetInstructionWithName("mycall")) + ->set_output_to_operand_aliasing({{{2}, {0, {0}}}}); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.message(), + HasSubstr("Invalid aliasing output shape index.")); +} + +TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingShapeMismatch) { + constexpr absl::string_view hlo = R"( + HloModule Module + + callme { + ROOT param = (s32[], f32[4]) parameter(0) + } + + ENTRY entry { + p0 = (s32[], f32[4]) parameter(0) + ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + + Cast( + module->entry_computation()->GetInstructionWithName("mycall")) + ->set_output_to_operand_aliasing({{{0}, {0, {1}}}}); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.message(), + HasSubstr("Different aliasing shapes: f32[4] vs s32[]")); } TEST_F(HloVerifierTest, CheckCustomCallOperandOutputAliasing) { @@ -1238,6 +1361,33 @@ TEST_F(HloVerifierTestLayoutSensitive, "Different aliasing shapes: f32[32,32]{0,1} vs f32[32,32]{1,0}")); } +TEST_F(HloVerifierTestLayoutSensitive, VerifyCallAliasConfigLayoutMismatch) { + const char* const hlo_string = R"( + HloModule module + + callme { + p0 = f32[32,32]{0,1} parameter(0) + p1 = f32[32,32]{1,0} parameter(1) + ROOT tuple = (f32[32,32]{1,0}) tuple(p1) + } + + ENTRY main { + p0 = f32[32,32]{0,1} parameter(0) + p1 = f32[32,32]{1,0} parameter(1) + ROOT call = (f32[32,32]{1,0}) call(p0, p1), to_apply=callme + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + Cast( + module->entry_computation()->GetInstructionWithName("call")) + ->set_output_to_operand_aliasing({{{0}, {0, {}}}}); + auto status = verifier().Run(module.get()).status(); + EXPECT_FALSE(status.ok()); + EXPECT_THAT( + status.message(), + HasSubstr( + "Different aliasing shapes: f32[32,32]{0,1} vs f32[32,32]{1,0}")); +} + TEST_F(HloVerifierTestLayoutSensitive, VerifyAsyncStartAliasConfigOperandMatchesButComputationMismatches) { const char* const hlo_string = R"( From 5c0a1db6c34d64bb88f3dc36365fa0b61020fb27 Mon Sep 17 00:00:00 2001 From: Dirk Hornung Date: Thu, 27 Aug 2026 15:14:34 -0700 Subject: [PATCH 24/28] [XLA:GPU] Support arbitrary broadcast dimensions in cuDNN convolution fusions. PiperOrigin-RevId: 972181976 --- .../gpu/transforms/cudnn_fusion_compiler.cc | 63 ++++++++++---- .../transforms/cudnn_fusion_compiler_test.cc | 82 +++++++++++++++++++ 2 files changed, 131 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc index c9025869505fc8..d20b178c722844 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc @@ -512,26 +512,61 @@ class ConvDimensionAdapter { dnums_for_layout}; } + int64_t HloDimToCudnnDim(int64_t hlo_dim) const { + if (hlo_dim == dums_.output_batch_dimension()) { + return 0; // Batch (N) + } + if (hlo_dim == dums_.output_feature_dimension()) { + return 1; // Feature / Channel (C) + } + int64_t dummy_spatial_dims = + std::max(0, 2 - dums_.output_spatial_dimensions_size()); + for (int i = 0; i < dums_.output_spatial_dimensions_size(); ++i) { + if (hlo_dim == dums_.output_spatial_dimensions(i)) { + return 2 + dummy_spatial_dims + i; // Spatial dimensions (H, W, ...) + } + } + return -1; + } + std::optional DimensionsAndStrides(const HloInstruction& hlo) { + int64_t spatial_dims = + std::max(2, dums_.input_spatial_dimensions_size()); + int64_t cudnn_rank = spatial_dims + 2; + if (ShapeUtil::IsScalar(hlo.shape())) { Result result; - // cuDNN convolution tensors have a batch and a feature dimension in - // addition to spatial dimensions (at least 2 spatial dimensions for - // cuDNN). - int64_t spatial_dims = - std::max(2, dums_.input_spatial_dimensions_size()); - result.sizes = std::vector(spatial_dims + 2, 1); - result.strides = std::vector(spatial_dims + 2, 1); + result.sizes = std::vector(cudnn_rank, 1); + result.strides = std::vector(cudnn_rank, 1); return result; } - if (hlo.shape().dimensions().size() == 1) { + + int64_t conv_hlo_rank = dums_.input_spatial_dimensions_size() + 2; + if (hlo.shape().dimensions().size() < conv_hlo_rank) { Result result; - int64_t spatial_dims = - std::max(2, dums_.input_spatial_dimensions_size()); - result.sizes = std::vector(spatial_dims + 2, 1); - result.strides = std::vector(spatial_dims + 2, 0); - result.sizes[1] = hlo.shape().dimensions(0); - result.strides[1] = 1; + result.sizes = std::vector(cudnn_rank, 1); + result.strides = std::vector(cudnn_rank, 0); + + // If the parameter is consumed by a broadcast, map its dimensions to the + // corresponding cuDNN canonical axes (N, C, spatial...). + if (hlo.user_count() == 1 && + hlo.users()[0]->opcode() == HloOpcode::kBroadcast) { + const auto& bcast_dims = hlo.users()[0]->dimensions(); + for (int i = 0; i < bcast_dims.size(); ++i) { + int64_t cudnn_dim = HloDimToCudnnDim(bcast_dims[i]); + if (cudnn_dim >= 0 && cudnn_dim < cudnn_rank) { + result.sizes[cudnn_dim] = hlo.shape().dimensions(i); + result.strides[cudnn_dim] = 1; + } + } + } else if (hlo.shape().dimensions().size() == 1) { + // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C, + // 1, 1]. + result.sizes[1] = hlo.shape().dimensions(0); + result.strides[1] = 1; + } else { + return std::nullopt; + } return result; } // Placeholder FP32 data type here, it is not used. diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc index eafc490c1b9e06..06293d680bd5f8 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc @@ -207,6 +207,88 @@ TEST_F(CudnnFusionCompilerConstTest, EXPECT_GT(plan_count, 0); } +TEST_F(CudnnFusionCompilerConstTest, + GetAvailablePlanCountFromConvolutionFusionWithBatchBroadcast) { + std::string hlo_text = R"( + fusion_batch_bcast { + p0 = f32[2,16,16,32] parameter(0) + p1 = f32[32,3,3,32] parameter(1) + conv = f32[2,16,16,32] convolution(p0, p1), + window={size=3x3 pad=1_1x1_1}, + dim_labels=b01f_o01i->b01f, + convolution_kind=fprop + batch_vec = f32[2] parameter(2) + bcast = f32[2,16,16,32] broadcast(batch_vec), dimensions={0} + ROOT mul = f32[2,16,16,32] multiply(conv, bcast) + } + + ENTRY e { + p0 = f32[2,16,16,32] parameter(0) + p1 = f32[32,3,3,32] parameter(1) + p2 = f32[2] parameter(2) + ROOT _ = f32[2,16,16,32] fusion(p0, p1, p2), kind=kCustom, calls=fusion_batch_bcast, + backend_config={ + "fusion_backend_config": { + "kind": "__cudnn$$fusion", + } + } + })"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_text)); + + const HloInstruction* root = + hlo_module->entry_computation()->root_instruction(); + auto* fusion = Cast(root); + + ASSERT_OK_AND_ASSIGN(int plan_count, + CuDnnFusionCompiler::GetAvailablePlanCount( + stream_executor(), + stream_executor()->GetDeviceDescription(), *fusion)); + EXPECT_GT(plan_count, 0); +} + +TEST_F(CudnnFusionCompilerConstTest, + GetAvailablePlanCountFrom1DConvolutionFusionWithBatchBroadcast) { + std::string hlo_text = R"( + fusion_1d_bcast { + p0 = f32[2,16,32] parameter(0) + p1 = f32[32,3,32] parameter(1) + conv = f32[2,16,32] convolution(p0, p1), + window={size=3 pad=1_1}, + dim_labels=b0f_o0i->b0f, + convolution_kind=fprop + batch_vec = f32[2] parameter(2) + bcast = f32[2,16,32] broadcast(batch_vec), dimensions={0} + ROOT mul = f32[2,16,32] multiply(conv, bcast) + } + + ENTRY e { + p0 = f32[2,16,32] parameter(0) + p1 = f32[32,3,32] parameter(1) + p2 = f32[2] parameter(2) + ROOT _ = f32[2,16,32] fusion(p0, p1, p2), kind=kCustom, calls=fusion_1d_bcast, + backend_config={ + "fusion_backend_config": { + "kind": "__cudnn$$fusion", + } + } + })"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_text)); + + const HloInstruction* root = + hlo_module->entry_computation()->root_instruction(); + auto* fusion = Cast(root); + + ASSERT_OK_AND_ASSIGN(int plan_count, + CuDnnFusionCompiler::GetAvailablePlanCount( + stream_executor(), + stream_executor()->GetDeviceDescription(), *fusion)); + EXPECT_GT(plan_count, 0); +} + } // namespace } // namespace gpu } // namespace xla From b8c608ba37691cb812ca240f1a671f23508d74c2 Mon Sep 17 00:00:00 2001 From: Amit Sabne Date: Thu, 27 Aug 2026 15:29:49 -0700 Subject: [PATCH 25/28] Scope async pipelined while loop offset colocation to pipelined while loops in Memory Space Assignment. PiperOrigin-RevId: 972189121 --- .../memory_space_assignment/algorithm.cc | 161 +++++++++++------- 1 file changed, 101 insertions(+), 60 deletions(-) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index 447f8f3e8895f2..abf40b148bf0f9 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -199,8 +199,8 @@ bool HasAsyncPipelinedWhileLoops(const HloModule& module) { return false; } -// Returns true if 'buffer' is defined by or used by an async pipelined while -// loop. +// Returns true if 'buffer' is defined by, used by, or contained within an +// async pipelined while loop computation. bool IsBufferAliasedToAsyncPipelinedWhileLoop(const HloValue* buffer) { if (buffer == nullptr) { return false; @@ -209,11 +209,27 @@ bool IsBufferAliasedToAsyncPipelinedWhileLoop(const HloValue* buffer) { if (IsAsyncPipelinedWhileLoop(pos.instruction)) { return true; } + if (pos.instruction->parent() != nullptr) { + for (const HloInstruction* caller : + pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { + if (IsAsyncPipelinedWhileLoop(caller)) { + return true; + } + } + } } for (const HloUse& use : buffer->GetUses()) { if (IsAsyncPipelinedWhileLoop(use.instruction)) { return true; } + if (use.instruction->parent() != nullptr) { + for (const HloInstruction* caller : + use.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { + if (IsAsyncPipelinedWhileLoop(caller)) { + return true; + } + } + } } return false; } @@ -526,6 +542,41 @@ HloInstruction* GetWhileForBodyRoot(HloInstruction* body_root) { return nullptr; } +// Returns true if the position corresponds to loop-carried state (parameter or +// body root tuple) of an async pipelined while loop, or while/DUS/DS in an +// async pipelined while loop. +bool IsAsyncPipelinedWhilePosition(const HloPosition& pos) { + if (pos.instruction->opcode() == HloOpcode::kParameter) { + for (const HloInstruction* caller : + pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { + if (IsAsyncPipelinedWhileLoop(caller)) { + return true; + } + } + return false; + } + if (pos.instruction->opcode() == HloOpcode::kWhile) { + return IsAsyncPipelinedWhileLoop(pos.instruction); + } + if (pos.instruction->opcode() == HloOpcode::kDynamicUpdateSlice || + pos.instruction->opcode() == HloOpcode::kDynamicSlice) { + for (const HloInstruction* caller : + pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { + if (IsAsyncPipelinedWhileLoop(caller)) { + return true; + } + } + return false; + } + HloComputation* comp = pos.instruction->parent(); + if (comp != nullptr && pos.instruction == comp->root_instruction() && + pos.instruction->opcode() == HloOpcode::kTuple) { + HloInstruction* while_caller = GetWhileForBodyRoot(pos.instruction); + return while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller); + } + return false; +} + // Returns true if the computation is the body or condition of a while loop. bool IsWhileBodyOrConditionComputation(const HloComputation* computation) { if (computation == nullptr) { @@ -672,14 +723,6 @@ MsaAlgorithm::MsaAlgorithm(HloModule* module, AllocationSequence* allocations, namespace { -bool IsAsyncDynamicSliceOrDynamicUpdateSlice( - const HloInstruction* instruction) { - return instruction->IsAsynchronous() && - (instruction->async_wrapped_opcode() == HloOpcode::kDynamicSlice || - instruction->async_wrapped_opcode() == - HloOpcode::kDynamicUpdateSlice); -} - // Finds the allocation value for operand 0 of async updates and dones. We make // sure these uses are associated with the AllocationValue whose position // is the direct source of the use. @@ -687,10 +730,10 @@ AllocationValue* FindAllocationValueForAsyncOperationStateUse( const HloUse& use, absl::Span candidate_allocation_values) { CHECK_EQ(use.operand_number, 0); + HloPosition source_position = GetNonTrivialSourcePosition( + HloPosition{use.instruction->mutable_operand(0), use.operand_index}); for (AllocationValue& allocation_value : candidate_allocation_values) { - if (allocation_value.defining_instruction() == - use.instruction->operand(0) && - allocation_value.defining_position().index == use.operand_index) { + if (allocation_value.defining_position() == source_position) { return &allocation_value; } } @@ -755,14 +798,7 @@ AllocationValue* FindLatestAllocationValueForUse( AllocationValue* MsaAlgorithm::FindAllocationValueForUse( const HloUse& use, absl::Span candidate_allocation_values, int64_t use_time) const { - // Asynchronous DynamicSlice and DynamicUpdateSlice instructions have - // specialized liveness and buffer aliasing semantics that differ from - // communication async operations (such as AllReduce or AllGather). Do not - // restrict their uses to strict start/done pairs. - bool is_async_dus_done = - IsAsyncDynamicSliceOrDynamicUpdateSlice(use.instruction); - - if (IsAsyncOperationStateUse(use) && !is_async_dus_done) { + if (IsAsyncOperationStateUse(use)) { // Case A: Async operation state uses. Note, this case covers uses of // bundled async state in traditional serial async chains and in pipelined // while loops. @@ -775,19 +811,19 @@ AllocationValue* MsaAlgorithm::FindAllocationValueForUse( const absl::flat_hash_map& instruction_schedule = hlo_live_range_.instruction_schedule(); - bool is_pipelined_while_async_start_use = - has_async_pipelined_while_loops_ && - (use.instruction->opcode() == HloOpcode::kWhile || - (use.instruction->opcode() == HloOpcode::kTuple && - GetWhileForBodyRoot(use.instruction) != nullptr)); + HloInstruction* while_for_body_root = GetWhileForBodyRoot(use.instruction); - bool allow_async_state_definitions = - is_async_dus_done || is_pipelined_while_async_start_use; + bool is_pipelined_while_async_start_use = + (use.instruction->opcode() == HloOpcode::kWhile && + IsAsyncPipelinedWhileLoop(use.instruction)) || + (while_for_body_root != nullptr && + IsAsyncPipelinedWhileLoop(while_for_body_root) && + use.instruction->opcode() == HloOpcode::kTuple); return FindLatestAllocationValueForUse( candidate_allocation_values, use_time, use.instruction->parent(), instruction_schedule, - /*skip_async_state_definitions=*/!allow_async_state_definitions); + /*skip_async_state_definitions=*/!is_pipelined_while_async_start_use); } void MsaAlgorithm::CreateAllocationValues( @@ -877,10 +913,13 @@ void MsaAlgorithm::CreateAllocationValues( bool is_async_operation_state = IsAsyncOperationStateDefinition( allocation_value.defining_instruction()) || - absl::c_any_of(allocation_value.uses(), - [](const AllocationValue::Use& use) { - return IsAsyncOperationStateUse(use.hlo_use); - }); + (has_async_pipelined_while_loops_ && + IsBufferAliasedToAsyncPipelinedWhileLoop(allocation_value.value()) && + (IsAsyncPipelinedWhilePosition(allocation_value.defining_position()) || + absl::c_any_of(allocation_value.uses(), + [](const AllocationValue::Use& use) { + return IsAsyncOperationStateUse(use.hlo_use); + }))); // Requiring contiguous allocation ensures that buffers representing inputs // and outputs to the async computation maintain temporal contiguity // (preventing MSA from evicting them to default memory). @@ -1031,6 +1070,7 @@ void MsaAlgorithm::FindAliases( HloInstruction* while_instruction = GetWhileForBodyRoot(use.hlo_use.instruction); if (while_instruction != nullptr && + IsAsyncPipelinedWhileLoop(while_instruction) && use.hlo_use.instruction->opcode() == HloOpcode::kTuple && while_instruction->while_body()->num_parameters() > 0) { ShapeIndex index = use.hlo_use.operand_index; @@ -2418,18 +2458,28 @@ MsaAlgorithm::GetContiguousLiveRangesForBuffer(const HloBuffer* buffer) const { ShapeIndex operand_index = use.operand_index; HloPosition source_position = GetNonTrivialSourcePosition(HloPosition{operand, operand_index}); + bool is_pipelined_async = + (has_async_pipelined_while_loops_ && + IsBufferAliasedToAsyncPipelinedWhileLoop(value) && + (IsAsyncOperationStateDefinition(source_position.instruction) || + IsAsyncPipelinedWhilePosition(source_position) || + IsAsyncOperationStateUse(use))); if (options_.position_requires_contiguous_allocation_fn( source_position) || - IsAsyncOperationStateDefinition(source_position.instruction) || - IsAsyncOperationStateUse(use)) { + is_pipelined_async) { VLOG(3) << "Adding use " << use.ToString() << " to contiguous position " << source_position.ToString(); contiguous_positions_to_uses[source_position].push_back(use); } } for (const HloPosition& position : value->positions()) { + bool is_pipelined_async_pos = + (has_async_pipelined_while_loops_ && + IsBufferAliasedToAsyncPipelinedWhileLoop(value) && + (IsAsyncOperationStateDefinition(position.instruction) || + IsAsyncPipelinedWhilePosition(position))); if (options_.position_requires_contiguous_allocation_fn(position) || - IsAsyncOperationStateDefinition(position.instruction)) { + is_pipelined_async_pos) { if (!(contiguous_positions_to_uses.contains(position))) { LOG(WARNING) << "Position " << position.ToString() << " is required to be contiguous but has no uses, " @@ -2450,9 +2500,13 @@ MsaAlgorithm::GetContiguousLiveRangesForBuffer(const HloBuffer* buffer) const { position.instruction->parent() != nullptr && position.instruction->parent()->root_instruction() != nullptr && IsWhileBodyOrConditionComputation(position.instruction->parent())) { - end_time = std::max( - end_time, hlo_live_range_.instruction_schedule().at( - position.instruction->parent()->root_instruction())); + HloInstruction* while_caller = GetWhileForBodyRoot( + position.instruction->parent()->root_instruction()); + if (while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller)) { + end_time = std::max( + end_time, hlo_live_range_.instruction_schedule().at( + position.instruction->parent()->root_instruction())); + } } for (const HloUse& use : uses) { end_time = std::max(end_time, GetCorrectedUseTime(use)); @@ -5683,18 +5737,6 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( // (including loop input operands, loop parameters, and loop body root // tuple elements) adopt identical starting offsets in alternate memory. AliasedOffset* preferred_offset = nullptr; - auto is_while_tuple_position = [](const HloPosition& pos) { - if (pos.instruction->opcode() == HloOpcode::kParameter || - pos.instruction->opcode() == HloOpcode::kWhile) { - return true; - } - HloComputation* comp = pos.instruction->parent(); - if (comp != nullptr && pos.instruction == comp->root_instruction() && - pos.instruction->opcode() == HloOpcode::kTuple) { - return true; - } - return false; - }; if (has_async_pipelined_while_loops_) { const HloBuffer& hlo_buffer = alias_analysis_.GetUniqueBufferAt( allocation_value_to_update.defining_position().instruction, @@ -5705,7 +5747,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( preferred_offset = buf_it->second; } else { for (const HloPosition& position : hlo_buffer.ComputePositions()) { - if (!is_while_tuple_position(position)) { + if (!IsAsyncPipelinedWhilePosition(position)) { continue; } HloComputation* comp = position.instruction->parent(); @@ -5723,13 +5765,12 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( } } if (preferred_offset == nullptr) { - if (!has_async_pipelined_while_loops_) { - auto comp_it = preferred_offset_for_computation.find( - allocation_value_to_update.computation()); - if (comp_it != preferred_offset_for_computation.end()) { - preferred_offset = comp_it->second; - } - } else if (is_while_tuple_position( + auto comp_it = preferred_offset_for_computation.find( + allocation_value_to_update.computation()); + if (comp_it != preferred_offset_for_computation.end()) { + preferred_offset = comp_it->second; + } else if (has_async_pipelined_while_loops_ && + IsAsyncPipelinedWhilePosition( allocation_value_to_update.defining_position())) { auto comp_it = pipelined_while_preferred_offset_for_computation_.find( @@ -6434,7 +6475,7 @@ void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse( // Special case for while loops since the root offset must agree with // other offsets: remember the preferred offset for the while loop body. AliasedOffset* offset = GetAliasedOffset(*aliased_allocation); - if (has_async_pipelined_while_loops_) { + if (IsAsyncPipelinedWhileLoop(hlo_use.instruction)) { SynchronizeAliasedWhileLoopOffsets(hlo_use, *aliased_allocation, offset); } else { preferred_offset_for_computation[hlo_use.instruction->while_body()] = From 7711f82e0e4f834f6fccd34e1044d56a212e8f62 Mon Sep 17 00:00:00 2001 From: Dragan Mladjenovic Date: Thu, 27 Aug 2026 15:43:37 -0700 Subject: [PATCH 26/28] PR #47522: [ROCm] Move CI to rocm 10 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/47522 ๐Ÿ“ Summary of Changes Move ROCm CI to use upcoming ROCm 10 ๐ŸŽฏ Justification In order to match JAX CI ๐Ÿš€ Kind of Contribution โœจ New Feature ๐Ÿ“Š Benchmark (for Performance Improvements) N\A ๐Ÿงช Unit Tests: None ๐Ÿงช Execution Tests: None Copybara import of the project: -- 0555bf88e02ff894acb63b22251ec5decf762aec by Dragan Mladjenovic : [ROCm] Move CI to rocm 7.14 take two -- 95b0f3e4d9e79c3b17c425608af672bedf65da3b by Dragan Mladjenovic : Fix LLVM symbol clash w/o --dynamic-mode=off -- 1319f537b000c46e3a4f474535e7454edb5c1362 by Dragan Mladjenovic : Remove dup --local_test_jobs=1 -- be17fac4960b2349db3e7b71bc1e4d5ed81da0af by Dragan Mladjenovic : Use rocm-dev-infra -- 5a27cc86870c7ad05a2d47bfbabf1fee347db311 by Dragan Mladjenovic : Move to ROCm 10 -- 24b4136130b5ee7bb648a182423ea1a0dab60c5a by Dragan Mladjenovic : Avoid bzlmod for hermetic rocm path -- 1b5690ac3d9afe5cc3f2c6b96398fe8cd982b9b3 by Dragan Mladjenovic : Update rocm-distro-url to stable Merging this change closes #47522 PiperOrigin-RevId: 972195867 --- third_party/xla/.github/workflows/rocm_ci.yml | 50 ++++++++++++++++--- third_party/xla/tensorflow.bazelrc | 6 +-- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml index 1e08a8013c9b0e..a722371d30ade1 100644 --- a/third_party/xla/.github/workflows/rocm_ci.yml +++ b/third_party/xla/.github/workflows/rocm_ci.yml @@ -42,12 +42,17 @@ jobs: runs-on: ubuntu-latest outputs: docker-image: ${{ steps.out.outputs.docker-image }} + rocm-distro-url: ${{ steps.out.outputs.rocm-distro-url }} + rocm-distro-hash: ${{ steps.out.outputs.rocm-distro-hash }} steps: - id: out shell: bash run: | - # hermetic llvm based on ghcr.io/rocm/jax-ubu22.rocm7.2.4:latest - echo "docker-image=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT" + { + echo "docker-image=ghcr.io/rocm/jax-build-ubu24.rocmless@sha256:68854ed30b800d3e8c390f6e651721e2bead58830da1c59e1b828354eb8f077f" + echo "rocm-distro-url=https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz" + echo "rocm-distro-hash=1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6" + } >> "$GITHUB_OUTPUT" jax: needs: rocm-config @@ -59,6 +64,8 @@ jobs: continue-on-error: true env: DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }} + ROCM_DISTRO_URL: ${{ needs.rocm-config.outputs.rocm-distro-url }} + ROCM_DISTRO_HASH: ${{ needs.rocm-config.outputs.rocm-distro-hash }} # Unique per run so a crashed run cannot collide with this one on a pooled node. CONTAINER_NAME: jax-${{ github.run_id }}-${{ github.run_attempt }} # Stable label so leftover containers from crashed runs can be swept on a pooled node. @@ -150,6 +157,10 @@ jobs: --repo_env=HERMETIC_PYTHON_VERSION=3.14 \ --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \ --jobs=20 \ + --repo_env=ROCM_PATH="" \ + --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \ + --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \ + --repo_env=SYSROOT_DIST=linux_glibc_2_31 \ --bes_keywords=jax \ --bes_keywords=upstream \ --bes_keywords=gpu \ @@ -166,7 +177,10 @@ jobs: timeout-minutes: 220 env: DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }} + ROCM_DISTRO_URL: ${{ needs.rocm-config.outputs.rocm-distro-url }} + ROCM_DISTRO_HASH: ${{ needs.rocm-config.outputs.rocm-distro-hash }} EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh + GLOBAL_SYMBOL_VERSION_LDS_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/global_symbol_version.lds # Unique per run so a crashed run cannot collide with this one on a pooled node. CONTAINER_NAME: xla-${{ github.run_id }}-${{ github.run_attempt }} # Stable label so leftover containers from crashed runs can be swept on a pooled node. @@ -181,10 +195,21 @@ jobs: repository: openxla/xla persist-credentials: false - - name: Fetch ROCm XLA Upstream Script from ROCm/xla + - name: Fetch ROCm XLA Upstream Scripts from ROCm/xla run: | wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}" chmod +x build_tools/rocm/execute_ci_build_upstream.sh + wget -O build_tools/rocm/global_symbol_version.lds "${GLOBAL_SYMBOL_VERSION_LDS_URL}" + LDS_SHA=$(sha256sum build_tools/rocm/global_symbol_version.lds | cut -d' ' -f1) + # Workaround for the symbol clash with ROCm's libLLVM.so + { + echo "LLVM_SYMBOL_CLASH_WAR<> "$GITHUB_ENV" - *start_container @@ -194,16 +219,20 @@ jobs: - name: Test XLA [single_gpu] timeout-minutes: 120 run: | + # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line. + # shellcheck disable=SC2086 docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \ - --config=rocm_ci \ + --config=rocm_ci_hermetic \ --config=rocm_rbe \ --config=ci_single_gpu \ - --local_test_jobs=1 \ --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \ - --repo_env=ROCM_PATH="/opt/rocm" \ + --repo_env=ROCM_PATH="" \ + --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \ + --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \ --local_test_jobs=1 \ --internal_spawn_scheduler \ --strategy=TestRunner=dynamic \ + ${LLVM_SYMBOL_CLASH_WAR} \ --bes_keywords=xla \ --bes_keywords=upstream \ --bes_keywords=gpu \ @@ -212,13 +241,18 @@ jobs: - name: Test XLA [rocm_cpu] timeout-minutes: 80 run: | + # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line. + # shellcheck disable=SC2086 docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \ - --config=rocm_ci \ + --config=rocm_ci_hermetic \ --config=rocm_rbe \ --config=ci_rocm_cpu \ --local_test_jobs=200 \ --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \ - --repo_env=ROCM_PATH="/opt/rocm" \ + --repo_env=ROCM_PATH="" \ + --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \ + --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \ + ${LLVM_SYMBOL_CLASH_WAR} \ --bes_keywords=xla \ --bes_keywords=upstream \ --bes_keywords=cpu \ diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index d6ceb7b9829bc5..054ade6bdf69c5 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -292,20 +292,20 @@ common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_rocm=True common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_cuda=False common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_sycl=False common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_hermetic_cc=True -common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic +common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace common:rocm_clang_hermetic --strategy=CppLink=local common:rocm_clang_hermetic --@rules_ml_toolchain//common:static_libcxx=False common:rocm --config=rocm_clang_hermetic common:rocm_ci --config=rocm -common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic +common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace common:rocm_ci_hermetic --dynamic_mode=off common:rocm_ci_hermetic --config=rocm_clang_hermetic common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a" common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a" common:rocm_ci_hermetic --repo_env=SYSROOT_DIST=linux_glibc_2_31 -common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic +common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace # This config option is used for SYCL as GPU backend. # SYCL Configuration (non-hermetic) From c998cc954fde30b05703da6825741854c99069b3 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 27 Aug 2026 16:29:33 -0700 Subject: [PATCH 27/28] Fix synchronization problems in RamFileBlockCache. This change fixes two problems in RamFileBlockCache. First, RamFileBlockCache::Flush() would remove blocks from the cache and its associated lists without setting the blocks' timestamps to 0. If a block was in the process of being fetched at that time, the "reconcile_state" cleanup callback in MaybeFetch() would treat the lra_iterator as valid, when it in fact it may have been made invalid. Second, if RamFileBlockCache::RemoveFile() were called and removed a block that was in the process of being fetched, RemoveBlock() would access the block's data field (by using data.capacity()), racing with MaybeFetch(), which sets that field. Even if the race were otherwise considered harmless, this might cause the cache_size_ field no longer to reflect the true size of the cache. This change addresses the first problem by making Flush() use code similar to RemoveFile(), thus using RemoveBlock() on every block. RemoveBlock() resets the block's timestamp field to zero. It addresses the second problem by checking the block's state field before accessing the data field, performing the access only if the state is FINISHED. The block's size is not counted in cache_size_ unless the state is FINISHED, and once the state is FINISHED, the data field is immutable, and can be read at will. I added/modified comments that would have helped me understand the code's synchronization invariants, in the hope that they will help future maintainers. PiperOrigin-RevId: 972217969 --- .../plugins/gcs/ram_file_block_cache.cc | 34 +++++++++++++++---- .../plugins/gcs/ram_file_block_cache.h | 10 +++--- .../plugins/gcs/ram_file_block_cache_test.cc | 29 ++++++++++++++++ 3 files changed, 63 insertions(+), 10 deletions(-) diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc index c92a6dfbd3890f..14f3dd69c0ed37 100644 --- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc +++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc @@ -91,6 +91,7 @@ void RamFileBlockCache::UpdateLRU(const Key& key, // in the cache, and our current block is not block size, this likely means // we have inconsistent state within the cache. Note: it's possible some // incomplete reads may still go undetected. + // Read of block->data allowed because block_state==FINISHED here. if (block->data.size() < block_size_) { Key fmax = std::make_pair(key.first, std::numeric_limits::max()); auto fcmp = block_map_.upper_bound(fmax); @@ -105,6 +106,8 @@ void RamFileBlockCache::UpdateLRU(const Key& key, return TF_SetStatus(status, TF_OK, ""); } +// Attempt to fetch data with the given key into block. +// Return with *status TF_OK only if block->state==FINISHED. void RamFileBlockCache::MaybeFetch(const Key& key, const std::shared_ptr& block, TF_Status* status) { @@ -113,10 +116,12 @@ void RamFileBlockCache::MaybeFetch(const Key& key, // Perform this action in a cleanup callback to avoid locking mu_ after // locking block->mu. if (downloaded_block) { + // downloaded_block == (block->state==FINISHED), so reads of block->data + // are legal here absl::MutexLock l(mu_); // Do not update state if the block is already to be evicted. if (block->timestamp != 0) { - // Use capacity() instead of size() to account for all memory + // Use capacity() instead of size() to account for all memory // used by the cache. cache_size_ += block->data.capacity(); // Put to beginning of LRA list. @@ -137,6 +142,7 @@ void RamFileBlockCache::MaybeFetch(const Key& key, // TF_FALLTHROUGH_INTENDED case FetchState::CREATED: block->state = FetchState::FETCHING; + // Thread may modify block->data block->state==FETCHING. block->mu.unlock(); // Release the lock while making the API call. block->data.clear(); block->data.resize(block_size_, 0); @@ -202,6 +208,8 @@ int64_t RamFileBlockCache::Read(const std::string& filename, size_t offset, } MaybeFetch(key, block, status); if (TF_GetCode(status) != TF_OK) return -1; + // At this point, block->state==FINISHED, since MaybeFetch() yielded TF_OK. + // Therefore, it is legal to access block->data. UpdateLRU(key, block, status); if (TF_GetCode(status) != TF_OK) return -1; // Copy the relevant portion of the block into the result buffer. @@ -282,10 +290,14 @@ void RamFileBlockCache::Prune() { void RamFileBlockCache::Flush() { absl::MutexLock lock(mu_); - block_map_.clear(); - lru_list_.clear(); - lra_list_.clear(); - cache_size_ = 0; + // This code mirrors that in RemoveFile_Locked(), + // but iterates over the entire cache. + auto it = block_map_.begin(); + while (it != block_map_.end()) { + auto next = std::next(it); + RemoveBlock(it); + it = next; + } } void RamFileBlockCache::RemoveFile(const std::string& filename) { @@ -309,7 +321,17 @@ void RamFileBlockCache::RemoveBlock(BlockMap::iterator entry) { entry->second->timestamp = 0; lru_list_.erase(entry->second->lru_iterator); lra_list_.erase(entry->second->lra_iterator); - cache_size_ -= entry->second->data.capacity(); + + // Adjust the cache_size_ by the size of the block. + // RemoveBlock() can be called by Flush() on blocks that are not yet FINISHED. + // Only finished blocks are counted in cache_size_, and it would be a race + // to read the data member of a block that is not yet FINISHED. + entry->second->mu.lock(); + if (entry->second->state == FetchState::FINISHED) { + cache_size_ -= entry->second->data.capacity(); + } + entry->second->mu.unlock(); + block_map_.erase(entry); } diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h index 3e972fa6292995..86558a295fbd64 100644 --- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h +++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h @@ -180,8 +180,10 @@ class RamFileBlockCache { /// Thread safety: /// The iterator and timestamp fields should only be accessed while holding /// the block-cache-wide mu_ instance variable. The state variable should only - /// be accessed while holding the Block's mu lock. The data vector should only - /// be accessed after state == FINISHED, and it should never be modified. + /// be accessed while holding the Block's mu lock. Once state==FINISHED, + /// the data vector may be read and may not be written; before + /// state==FINISHED, the data vector may be accessed only by the thread that + /// set state==FETCHING. /// /// In order to prevent deadlocks, never grab the block-cache-wide mu_ lock /// AFTER grabbing any block's mu lock. It is safe to grab mu without locking @@ -212,13 +214,13 @@ class RamFileBlockCache { void Prune() ABSL_LOCKS_EXCLUDED(mu_); bool BlockNotStale(const std::shared_ptr& block) - ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_); + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) ABSL_LOCKS_EXCLUDED(block->mu); /// Look up a Key in the block cache. std::shared_ptr Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_); void MaybeFetch(const Key& key, const std::shared_ptr& block, - TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_); + TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_, block->mu); /// Trim the block cache to make room for another entry. void Trim() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_); diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc index 18438ef0580f72..929fc25a6372bb 100644 --- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc +++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc @@ -610,5 +610,34 @@ TEST(RamFileBlockCacheTest, Flush) { EXPECT_EQ(calls, 2); } +// This test case reproduces a bug, now fixed, in which +// RamFileBlockCache::Flush() used to clear the cache, invalidating the +// lra_iterator field, without resetting the timestamp fields of the blocks. +// This was a problem for the "reconcile_state" cleanup callback in +// MaybeFetch(), which assumes that if the timestamp field is non-zero, the +// lra_iterator is valid. Later versions of Flush() should handle this case +// correctly. +TEST(RamFileBlockCacheTest, FlushDuringFetch) { + // The delays are sized to exceed the delay in RamFileBlockCache::Prune(). + auto delayed_fetcher = [](const std::string& filename, size_t offset, + size_t n, char* buffer, + TF_Status* status) -> int64_t { + memset(buffer, 'x', n); + TF_SetStatus(status, TF_OK, ""); + Env::Default()->SleepForMicroseconds(10 * 1000 * 1000); + return n; + }; + tf_gcs_filesystem::RamFileBlockCache cache(16, 32, /*max_staleness=*/20, + delayed_fetcher); + std::vector out; + std::unique_ptr flush_thread( + Env::Default()->StartThread({}, "delayed_flush", [&cache] { + Env::Default()->SleepForMicroseconds(5 * 1000 * 1000); + cache.Flush(); + })); + TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out)); + Env::Default()->SleepForMicroseconds(5 * 1000 * 1000); +} + } // namespace } // namespace tensorflow From 9c87c11c51498653123c17615079ece32da90123 Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Thu, 27 Aug 2026 16:47:57 -0700 Subject: [PATCH 28/28] Do not allow batch dimensions to be transposed in a way that interleaves them with other dimension types. If they are transposed within themselves, then this is okay as it won't affect the stores within a kernel. PiperOrigin-RevId: 972226337 --- .../backends/gpu/transforms/gemm_fusion.cc | 4 ++ .../gpu/transforms/gemm_fusion_test.cc | 56 +++++++++++++++++++ .../xla/xla/hlo/analysis/shape_tracker.cc | 7 ++- .../xla/xla/hlo/analysis/shape_tracker.h | 9 ++- .../xla/hlo/analysis/shape_tracker_test.cc | 8 +++ 5 files changed, 80 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc index 7715ce308eed56..a61accc39ba232 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc @@ -1306,10 +1306,14 @@ FusionDecision ShouldFuseTranspose(const HloInstruction& transpose, dims->Indices(DotOperandDims::kContracting); absl::Span non_contracting = dims->Indices(DotOperandDims::kNonContracting); + absl::Span batch = dims->Indices(DotOperandDims::kBatch); if (!inverted_tracker->MapsToOneStride(contracting)) { return FusionDecision::Forbid( "Contracting dimension has non-contiguous section."); } + if (!inverted_tracker->MapsToOneStride(batch, /*allow_swaps=*/true)) { + return FusionDecision::Forbid("Batch dimension splits other dimensions."); + } if (operand_index == 1 && !inverted_tracker->MapsToOneStride(non_contracting)) { return FusionDecision::Forbid( diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc index 3e9cd4b3d8144f..5113eefcd9eb41 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc @@ -2394,6 +2394,62 @@ ENTRY main { GmockMatch(TransposeOrBitcastTranspose()))); } +TEST_P(GemmFusionTestV2, AllowTransposeSwappingBatch) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( +HloModule module + +ENTRY main { + p_lhs = s8[8,128,64]{2,1,0} parameter(0) + cvt_lhs = bf16[8,128,64]{2,1,0} convert(p_lhs) + p_rhs = bf16[2,4,1,16,128]{4,3,2,1,0} parameter(1) + + trans = bf16[4,2,1,16,128]{4,3,2,1,0} transpose(p_rhs), dimensions={1,0,2,3,4} + bitcast = bf16[8,16,128]{2,1,0} bitcast(trans) + + ROOT dot = bf16[8,64,16]{2,1,0} dot(cvt_lhs, bitcast), + lhs_batch_dims={0}, lhs_contracting_dims={1}, + rhs_batch_dims={0}, rhs_contracting_dims={2} +} +)")); + + ASSERT_OK_AND_ASSIGN(bool changed, + GemmFusion(gpu_version_).Run(module.get())); + EXPECT_TRUE(changed); + auto* fusion = module->entry_computation()->root_instruction(); + EXPECT_THAT(fusion, GmockMatch(m::Fusion(m::Parameter(), m::Parameter()))); +} + +TEST_P(GemmFusionProfitabilityTest, DisallowTransposeInterleavingBatch) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( +HloModule module + +ENTRY main { + p_lhs = s8[16,64,1024]{2,1,0} parameter(0) + cvt_lhs = bf16[16,64,1024]{2,1,0} convert(p_lhs) + p_rhs = bf16[4,16,4,64,32]{4,3,2,1,0} parameter(1) + + trans = bf16[4,4,16,64,32]{4,3,2,1,0} transpose(p_rhs), dimensions={0,2,1,3,4} + bitcast = bf16[16,1024,32]{2,1,0} bitcast(trans) + + ROOT dot = bf16[16,64,32]{2,1,0} dot(cvt_lhs, bitcast), + lhs_batch_dims={0}, lhs_contracting_dims={2}, + rhs_batch_dims={0}, rhs_contracting_dims={1} +} +)")); + + ASSERT_OK_AND_ASSIGN(bool changed, + GemmFusion(gpu_version_).Run(module.get())); + EXPECT_TRUE(changed); + auto* fusion = module->entry_computation()->root_instruction(); + EXPECT_THAT(fusion, GmockMatch(m::Fusion())); + EXPECT_THAT(fusion->operands(), + ::testing::UnorderedElementsAre( + GmockMatch(m::Parameter()), + GmockMatch(TransposeOrBitcastTranspose()))); +} + TEST_P(GemmFusionTestV2, ConcatResetTrackerCrash) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( HloModule m diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.cc b/third_party/xla/xla/hlo/analysis/shape_tracker.cc index a060b290381b76..e183e2d7a7e398 100644 --- a/third_party/xla/xla/hlo/analysis/shape_tracker.cc +++ b/third_party/xla/xla/hlo/analysis/shape_tracker.cc @@ -1255,7 +1255,8 @@ ShapeTracker::MapInputDimensionsToOutputUnordered( return kept_output_dims; } -bool ShapeTracker::MapsToOneStride(absl::Span input_dims) const { +bool ShapeTracker::MapsToOneStride(absl::Span input_dims, + bool allow_swaps) const { std::optional> mapped = MapInputDimensionsToOutputUnordered(input_dims); if (!mapped.has_value()) { @@ -1282,6 +1283,10 @@ bool ShapeTracker::MapsToOneStride(absl::Span input_dims) const { return false; } + if (allow_swaps) { + return true; + } + // Check for swaps. absl::StatusOr narrowed = Narrow(input_dims); if (!narrowed.ok()) { diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.h b/third_party/xla/xla/hlo/analysis/shape_tracker.h index 2bf34cd3bbf039..f99794f4a0206f 100644 --- a/third_party/xla/xla/hlo/analysis/shape_tracker.h +++ b/third_party/xla/xla/hlo/analysis/shape_tracker.h @@ -143,9 +143,12 @@ class ShapeTracker { absl::Span input_dims) const; // Returns true if the specified input dimensions map to a single contiguous - // stride in the output shape (i.e. they are not swapped, and there are no - // other non-degenerate dimensions between them in the output layout). - bool MapsToOneStride(absl::Span input_dims) const; + // stride in the output shape. If `allow_swaps` is false (default), the + // dimensions must not be swapped with each other. In either case, there must + // be no other non-degenerate dimensions interleaved between them in the + // output layout. + bool MapsToOneStride(absl::Span input_dims, + bool allow_swaps = false) const; // Zips multiple ShapeTrackers into a single one. // For example, suppose we have two trackers: diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc index 59febec3ccc9d7..bba5b2cfec9a62 100644 --- a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc +++ b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc @@ -1904,5 +1904,13 @@ TEST(ShapeTrackerMapsToOneStrideTest, ContiguousTransposeReshapeIsContiguous) { EXPECT_TRUE(tracker.MapsToOneStride({1, 2})); } +TEST(ShapeTrackerMapsToOneStrideTest, AllowSwapsSwappedIsContiguous) { + Shape shape = ShapeUtil::MakeShape(F32, {2, 3, 4}); + ShapeTracker tracker(shape); + ASSERT_TRUE(tracker.AppendTranspose({0, 2, 1}).ok()); + EXPECT_FALSE(tracker.MapsToOneStride({1, 2}, /*allow_swaps=*/false)); + EXPECT_TRUE(tracker.MapsToOneStride({1, 2}, /*allow_swaps=*/true)); +} + } // namespace } // namespace xla