From 6739b67a920f24638cadd79fc4880b9915d1adee Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Sat, 5 Sep 2026 09:46:05 -0400 Subject: [PATCH 01/17] Validate FusedPadConv2D input rank --- .../kernels/conv_ops_fused_image_transform.cc | 8 +++---- tensorflow/core/kernels/conv_ops_test.cc | 22 +++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/kernels/conv_ops_fused_image_transform.cc b/tensorflow/core/kernels/conv_ops_fused_image_transform.cc index 1df29d70b320d6..a1cfe77b7c9150 100644 --- a/tensorflow/core/kernels/conv_ops_fused_image_transform.cc +++ b/tensorflow/core/kernels/conv_ops_fused_image_transform.cc @@ -649,6 +649,10 @@ class FusedResizeConv2DUsingGemmOp : public OpKernel { // Input tensor is of the following dimensions: // [ batch, in_rows, in_cols, in_depth ] const Tensor& input = context->input(0); + OP_REQUIRES( + context, input.dims() == 4, + absl::InvalidArgumentError(absl::StrCat( + "input must be 4-dimensional: ", input.shape().DebugString()))); OP_REQUIRES(context, (input.shape().num_elements() > 0), absl::InvalidArgumentError("Input tensor can't be empty")); @@ -764,10 +768,6 @@ class FusedResizeConv2DUsingGemmOp : public OpKernel { const Tensor& filter = context->input(filter_index); // For 2D convolution, there should be 4 dimensions. - OP_REQUIRES( - context, padded_shape.dims() == 4, - absl::InvalidArgumentError(absl::StrCat("input must be 4-dimensional", - padded_shape.DebugString()))); OP_REQUIRES( context, filter.dims() == 4, absl::InvalidArgumentError(absl::StrCat( diff --git a/tensorflow/core/kernels/conv_ops_test.cc b/tensorflow/core/kernels/conv_ops_test.cc index caff583b570092..e29edc7005e2b8 100644 --- a/tensorflow/core/kernels/conv_ops_test.cc +++ b/tensorflow/core/kernels/conv_ops_test.cc @@ -367,6 +367,28 @@ TEST_F(FusedResizePadConvOpTest, NoResizePadOnlySymmetricComparative) { "SAME", DT_FLOAT); } +TEST_F(FusedResizePadConvOpTest, InvalidInputRank) { + TF_ASSERT_OK(NodeDefBuilder("fused_pad_conv_op", "FusedPadConv2D") + .Input(FakeInput(DT_FLOAT)) + .Input(FakeInput(DT_INT32)) + .Input(FakeInput(DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Attr("mode", "REFLECT") + .Attr("strides", {1, 1, 1, 1}) + .Attr("padding", "VALID") + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + AddInputFromArray(TensorShape({3, 2}), {0, 0, 0, 0, 0, 0}); + AddInputFromArray(TensorShape({4, 2}), + {0, 0, 1, 1, 1, 1, 0, 0}); + AddInputFromArray(TensorShape({2, 2, 1, 1}), {0, 0, 0, 0}); + + const absl::Status status = RunOpKernel(); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_NE(status.message().find("input must be 4-dimensional"), + std::string::npos); +} + class ConvOpTest : public OpsTestBase { protected: void HandwrittenConv() { From 12b204fd5edf693f63aa63885314dfc00e863c00 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Fri, 25 Sep 2026 16:48:37 -0700 Subject: [PATCH 02/17] Adjust tolerances for TensorFloat-32 in FP32 TensorFlow GPU tests PiperOrigin-RevId: 988569222 --- .../kernel_tests/linalg/linalg_ops_test.py | 17 ++++++++++++++--- .../kernel_tests/math_ops/approx_topk_test.py | 9 ++++++++- .../math_ops/batch_matmul_op_test.py | 9 ++++++--- .../conv3d_backprop_filter_v2_grad_test.py | 2 +- .../python/kernel_tests/nn_ops/conv_ops_test.py | 5 +++-- tensorflow/python/ops/nn_grad_test.py | 3 ++- tensorflow/python/ops/parallel_for/math_test.py | 8 ++++++-- 7 files changed, 40 insertions(+), 13 deletions(-) diff --git a/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py b/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py index 755ece0f3e1f84..fcb0b0d667207f 100644 --- a/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py +++ b/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py @@ -60,7 +60,12 @@ def test_works_with_five_different_random_pos_def_matrices(self): with self.subTest(n=n, np_type=np_type, atol=atol, k=k): rhs = self.rng.randn(2, n, k).astype(np_type) x = linalg_ops.cholesky_solve(chol, rhs) - self.assertAllClose(rhs, math_ops.matmul(array, x), atol=atol) + rhs_pred = ( + test_util.matmul_without_tf32(array, x) + if test.is_gpu_available() + else math_ops.matmul(array, x) + ) + self.assertAllClose(rhs, rhs_pred, atol=atol) class LogdetTest(test.TestCase): @@ -435,7 +440,10 @@ def test_symmetric(self): expected_a_pinv_ = self.expected_pinv(a_, rcond) a_pinv = linalg.pinv(a, rcond, validate_args=True) a_pinv_ = self.evaluate(a_pinv) - self.assertAllClose(expected_a_pinv_, a_pinv_, atol=2e-5, rtol=2e-5) + use_tf32_tol = self.dtype == np.float32 and test.is_gpu_available() + atol = 3e-3 if use_tf32_tol else 2e-5 + rtol = 1e-3 if use_tf32_tol else 2e-5 + self.assertAllClose(expected_a_pinv_, a_pinv_, atol=atol, rtol=rtol) if not self.use_static_shape: return self.assertAllEqual(expected_a_pinv_.shape, a_pinv.shape) @@ -453,7 +461,10 @@ def test_nonsquare(self): expected_a_pinv_ = self.expected_pinv(a_, rcond) a_pinv = linalg.pinv(a, rcond, validate_args=True) a_pinv_ = self.evaluate(a_pinv) - self.assertAllClose(expected_a_pinv_, a_pinv_, atol=1e-5, rtol=1e-4) + use_tf32_tol = self.dtype == np.float32 and test.is_gpu_available() + atol = 1e-3 if use_tf32_tol else 1e-5 + rtol = 1e-3 if use_tf32_tol else 1e-4 + self.assertAllClose(expected_a_pinv_, a_pinv_, atol=atol, rtol=rtol) if not self.use_static_shape: return self.assertAllEqual(expected_a_pinv_.shape, a_pinv.shape) diff --git a/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py b/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py index a182f442df0690..c9b2aeca400620 100644 --- a/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py +++ b/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py @@ -163,7 +163,14 @@ def ann(qy, db, db_half_norm_sq, k): math_ops.matmul(qy_op, db_op, transpose_b=True)) gt = np.argsort(scores)[:, :k] ann_recall = self.compute_recall(idx, gt) - self.assertGreaterEqual(ann_recall, 0.95) + use_tf32_tol = ( + test_util.is_gpu_available() + and dtype == dtypes.float32 + and k == 1 + and feature_dim == 2 + ) + min_recall = 0.80 if use_tf32_tol else 0.95 + self.assertGreaterEqual(ann_recall, min_recall) def test_highdim(self): db = self._rng.random([2, 10, 200, 3], dtype=np.float32) diff --git a/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py b/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py index c3182d259b1d59..7f1812e71a3674 100644 --- a/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py @@ -168,7 +168,8 @@ def testRank4xRank2(self): w = np.random.normal(size=[4, 16]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) def testRank4xRank2BatchDim(self): np.random.seed(42) @@ -176,7 +177,8 @@ def testRank4xRank2BatchDim(self): w = np.random.normal(size=[5, 7]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) def testRank3xRank2(self): np.random.seed(42) @@ -184,7 +186,8 @@ def testRank3xRank2(self): w = np.random.normal(size=[4, 6]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) class BatchMatmulGradientTest(test.TestCase): diff --git a/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py b/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py index 9d0387da61b5bd..fe06f42e532577 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py @@ -56,7 +56,7 @@ def testGradient(self): [in_val, out_backprop_val], [in_shape, out_backprop_shape], output, filter_shape) print("conv3d_backprop_filter gradient err = %g " % err) - err_tolerance = 1e-3 + err_tolerance = 0.25 if test.is_gpu_available() else 1e-3 self.assertLess(err, err_tolerance) def testBadFilterShape(self): diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py index 7d3aa8fdd59c88..57cdfff5889c0e 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py @@ -1705,8 +1705,9 @@ def _GetVal(data_format, use_gpu): values = [] for (data_format, use_gpu) in GetTestConfigs(): values.append(_GetVal(data_format, use_gpu)) + tol = 1e-3 if test.is_gpu_available() else 2e-4 for i in range(1, len(values)): - self.assertAllClose(values[0], values[i], rtol=2e-4, atol=2e-4) + self.assertAllClose(values[0], values[i], rtol=tol, atol=tol) @test_util.run_in_graph_and_eager_modes def testConv2D2x2Depth1ValidBackpropFilter(self): @@ -3795,7 +3796,7 @@ def Test(self): strides=[stride, stride], dilations=[2, 2], padding=padding, - rtol=5e-4) + rtol=1e-3 if test.is_gpu_available() else 5e-4) return Test diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index 1a2fc541c02ce1..36d6b66f3e79df 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -81,7 +81,8 @@ def run_test(self, x, y): x.get_shape().as_list(), y, y.get_shape().as_list()) - self.assertLess(error, 2e-3) + tol = 0.06 if test.is_gpu_available() else 2e-3 + self.assertLess(error, tol) @test_util.run_deprecated_v1 def testConv2dGradWRTInput(self): diff --git a/tensorflow/python/ops/parallel_for/math_test.py b/tensorflow/python/ops/parallel_for/math_test.py index 2dc84ba913c8ee..133c374c6e7db0 100644 --- a/tensorflow/python/ops/parallel_for/math_test.py +++ b/tensorflow/python/ops/parallel_for/math_test.py @@ -329,7 +329,9 @@ def loop_fn(i): # pylint: enable=cell-var-from-loop - self._test_loop_fn(loop_fn, 2) + rtol = 1e-3 if test.is_gpu_available() else 1e-4 + atol = 1e-3 if test.is_gpu_available() else 1e-5 + self._test_loop_fn(loop_fn, 2, rtol=rtol, atol=atol) def test_batch_matmul_broadcast(self): for broadcast_a in (True, False): @@ -350,7 +352,9 @@ def loop_fn(i): return math_ops.matmul(a, b) # pylint: enable=cell-var-from-loop - self._test_loop_fn(loop_fn, 2) + rtol = 1e-3 if test.is_gpu_available() else 1e-4 + atol = 1e-3 if test.is_gpu_available() else 1e-5 + self._test_loop_fn(loop_fn, 2, rtol=rtol, atol=atol) def test_reduction(self): x = random_ops.random_uniform([2, 3, 4, 5]) From 85baa2150c296ee160543044312aef44178347aa Mon Sep 17 00:00:00 2001 From: Terry Heo Date: Fri, 25 Sep 2026 16:50:45 -0700 Subject: [PATCH 03/17] Stop downloading TensorFlow in LiteRT CMake build The TFLite and LiteRT CMake builds cloned the whole TensorFlow repository at configure time, but only needed a few files from it. - Built converter sources from tflite/converter instead of the TF checkout. - Added Findtensorflow_headers.cmake to download the 6 header-only files (xla/tsl, third_party/fft2d) that the TFLite runtime includes. - Fetched TensorFlow only for benchmark_model and the kernel tests. - Removed TENSORFLOW_SOURCE_DIR usages from LiteRT CMake files. - Removed unused TensorFlow logging include from nnapi_handler, which broke the Android CMake build without the TF checkout. PiperOrigin-RevId: 988570167 --- tensorflow/lite/CMakeLists.txt | 144 +++++++++--------- tensorflow/lite/nnapi/BUILD | 1 - tensorflow/lite/nnapi/nnapi_handler.h | 1 - .../modules/Findtensorflow_headers.cmake | 67 ++++++++ 4 files changed, 142 insertions(+), 71 deletions(-) create mode 100644 tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 27e25e83750461..d6a7e460a32ef5 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -40,32 +40,6 @@ cmake_policy(SET CMP0042 NEW) project(tensorflow-lite C CXX) set(CMAKE_POSITION_INDEPENDENT_CODE ON) -# Use FetchContent to download TensorFlow if not provided -include(FetchContent) - -set(TENSORFLOW_SOURCE_DIR "" CACHE PATH - "Directory that contains the TensorFlow project" -) -if(NOT TENSORFLOW_SOURCE_DIR) - message(STATUS "Downloading TensorFlow repository...") - FetchContent_Declare( - tensorflow - GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git - # v2.21.0-rc0 has updated schema_generated.h file. - GIT_TAG v2.21.0-rc0 - GIT_SHALLOW TRUE - GIT_PROGRESS TRUE - SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/tensorflow-src - ) - FetchContent_GetProperties(tensorflow) - if(NOT tensorflow_POPULATED) - FetchContent_Populate(tensorflow) - endif() - set(TENSORFLOW_SOURCE_DIR "${tensorflow_SOURCE_DIR}") -endif() -set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") -set(TSL_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/third_party/tsl") -set(XLA_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/") set(TFLITE_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}") set(CMAKE_MODULE_PATH "${TFLITE_SOURCE_DIR}/tools/cmake/modules" @@ -93,6 +67,36 @@ option(TFLITE_ENABLE_EXTERNAL_DELEGATE "Enable External Delegate backend" ON) option(TFLITE_KERNEL_TEST "Enable tflite kernel unit test" OFF) +# The TFLite runtime builds from this repository plus a few header-only TensorFlow +# files (see tools/cmake/modules/Findtensorflow_headers.cmake). A full TensorFlow +# source tree is only needed by the benchmark_model tool and the kernel tests. +set(TENSORFLOW_SOURCE_DIR "" CACHE PATH + "Directory that contains the TensorFlow project. Only needed when \ +TFLITE_ENABLE_BENCHMARK_MODEL or TFLITE_KERNEL_TEST is ON." +) +if(TFLITE_ENABLE_BENCHMARK_MODEL OR TFLITE_KERNEL_TEST) + if(NOT TENSORFLOW_SOURCE_DIR) + message(STATUS "Downloading TensorFlow repository...") + include(FetchContent) + FetchContent_Declare( + tensorflow + GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git + GIT_TAG v2.21.0-rc0 + GIT_SHALLOW TRUE + GIT_PROGRESS TRUE + SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/tensorflow-src + ) + FetchContent_GetProperties(tensorflow) + if(NOT tensorflow_POPULATED) + FetchContent_Populate(tensorflow) + endif() + set(TENSORFLOW_SOURCE_DIR "${tensorflow_SOURCE_DIR}") + endif() + set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") + set(TSL_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/third_party/tsl") + set(XLA_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/") +endif() + if(${CMAKE_CROSSCOMPILING}) set(TFLITE_HOST_TOOLS_DIR "" CACHE PATH "Host tools directory") if(TFLITE_HOST_TOOLS_DIR STREQUAL "") @@ -190,11 +194,17 @@ endif() find_package(cpuinfo REQUIRED) #CPUINFO is used by XNNPACK and RUY library find_package(ml_dtypes REQUIRED) find_package(ruy REQUIRED) -# Include TSL, which is in tensorflow/third_party -include_directories( - ${TSL_SOURCE_DIR} - ${XLA_SOURCE_DIR} -) +# Header-only TensorFlow dependencies (xla/tsl/..., third_party/fft2d/...). +find_package(tensorflow_headers REQUIRED) +include_directories(${TENSORFLOW_HEADERS_INCLUDE_DIR}) +if(TENSORFLOW_SOURCE_DIR) + include_directories( + AFTER + ${TENSORFLOW_SOURCE_DIR} + ${TSL_SOURCE_DIR} + ${XLA_SOURCE_DIR} + ) +endif() # Download necessary dependencies. if(TFLITE_ENABLE_XNNPACK) # pthreadpool is used by XNNPACK. @@ -410,7 +420,7 @@ if(TFLITE_ENABLE_GPU) ${TFLITE_DELEGATES_GPU_COMMON_TRANSFORMATIONS_SRCS} ${TFLITE_SOURCE_DIR}/tools/versioning/gpu_compatibility.cc ${TFLITE_SOURCE_DIR}/tools/versioning/op_signature.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/tools/versioning/op_signature.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_signature.cc ) include_directories( AFTER @@ -439,7 +449,7 @@ if(TFLITE_ENABLE_GPU) target_include_directories(metal_delegate PUBLIC ${CMAKE_BINARY_DIR}/abseil-cpp ${CMAKE_BINARY_DIR}/flatbuffers/include - PRIVATE ${TENSORFLOW_SOURCE_DIR} + PRIVATE ${TFLITE_SOURCE_DIR}/.. PRIVATE ${PROJECT_BINARY_DIR} ) # @@ -455,7 +465,7 @@ if(TFLITE_ENABLE_GPU) 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} + -I ${TFLITE_SOURCE_DIR}/.. -o ${METAL_GENERATED_DIR} -c ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs DEPENDS ${FLATC_TARGET} @@ -533,7 +543,7 @@ if(TFLITE_ENABLE_XNNPACK) "${PROJECT_BINARY_DIR}/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h" ) target_include_directories(xnnpack-delegate - PUBLIC $ $ + PUBLIC $ $ PRIVATE "${PROJECT_BINARY_DIR}" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/.." PRIVATE "${PROJECT_BINARY_DIR}/tensorflow" @@ -683,26 +693,26 @@ set(_ALL_TFLITE_SRCS ${TFLITE_SOURCE_DIR}/schema/conversion_metadata_generated.h ${TFLITE_SOURCE_DIR}/schema/schema_generated.h ${TFLITE_SRCS} - ${TF_SOURCE_DIR}/compiler/mlir/lite/allocation.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/allocation.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/error_reporter.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/error_reporter.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/flatbuffer_conversions.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/verifier.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/c/tflite_types.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/macros.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/model_builder_base.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/model_builder_base.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/experimental/remat/metadata_util.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/experimental/remat/metadata_util.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/mmap_allocation.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/mmap_allocation_disabled.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/conversion_metadata_generated.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/schema_generated.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/schema_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/control_edges.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/string_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/string_utils.h + ${TFLITE_SOURCE_DIR}/converter/allocation.cc + ${TFLITE_SOURCE_DIR}/converter/allocation.h + ${TFLITE_SOURCE_DIR}/converter/core/api/error_reporter.cc + ${TFLITE_SOURCE_DIR}/converter/core/api/error_reporter.h + ${TFLITE_SOURCE_DIR}/converter/core/api/flatbuffer_conversions.cc + ${TFLITE_SOURCE_DIR}/converter/core/api/verifier.h + ${TFLITE_SOURCE_DIR}/converter/core/c/tflite_types.h + ${TFLITE_SOURCE_DIR}/converter/core/macros.h + ${TFLITE_SOURCE_DIR}/converter/core/model_builder_base.cc + ${TFLITE_SOURCE_DIR}/converter/core/model_builder_base.h + ${TFLITE_SOURCE_DIR}/converter/experimental/remat/metadata_util.cc + ${TFLITE_SOURCE_DIR}/converter/experimental/remat/metadata_util.h + ${TFLITE_SOURCE_DIR}/converter/mmap_allocation.cc + ${TFLITE_SOURCE_DIR}/converter/mmap_allocation_disabled.cc + ${TFLITE_SOURCE_DIR}/converter/schema/conversion_metadata_generated.h + ${TFLITE_SOURCE_DIR}/converter/schema/schema_generated.h + ${TFLITE_SOURCE_DIR}/converter/schema/schema_utils.cc + ${TFLITE_SOURCE_DIR}/converter/utils/control_edges.h + ${TFLITE_SOURCE_DIR}/converter/utils/string_utils.cc + ${TFLITE_SOURCE_DIR}/converter/utils/string_utils.h # go/keep-sorted end ) @@ -719,7 +729,6 @@ set(_ALL_TFLITE_HDRS ${_ALL_TFLITE_SRCS}) list(FILTER _ALL_TFLITE_HDRS INCLUDE REGEX ".*\\.h$") target_include_directories(tensorflow-lite PUBLIC - $ $ $ $ @@ -811,17 +820,19 @@ endif() # tensorflow/lite directory structure for generated headers. set(TFLITE_GENERATED_HEADERS_DIR ${CMAKE_BINARY_DIR}/tensorflow/lite) -# Add the profiling proto directory. -add_subdirectory(${TFLITE_SOURCE_DIR}/profiling/proto) +if(TFLITE_ENABLE_BENCHMARK_MODEL) + # Add the profiling proto directory. + add_subdirectory(${TFLITE_SOURCE_DIR}/profiling/proto) -# Add the benchmark result proto directory. -add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark/proto) + # Add the benchmark result proto directory. + add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark/proto) -# Add the tf example directory. -add_subdirectory(${TF_SOURCE_DIR}/core/example ${CMAKE_BINARY_DIR}/example_proto_generated) + # Add the tf example directory. + add_subdirectory(${TF_SOURCE_DIR}/core/example ${CMAKE_BINARY_DIR}/example_proto_generated) -# The benchmark tool. -add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark) + # The benchmark tool. + add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark) +endif() # The label_image example. add_subdirectory(${TFLITE_SOURCE_DIR}/examples/label_image) @@ -838,11 +849,6 @@ add_library(_pywrap_tensorflow_interpreter_wrapper SHARED EXCLUDE_FROM_ALL # To remove "lib" prefix. set_target_properties(_pywrap_tensorflow_interpreter_wrapper PROPERTIES PREFIX "") -target_include_directories(_pywrap_tensorflow_interpreter_wrapper - PUBLIC - ${TENSORFLOW_SOURCE_DIR} -) - target_link_libraries(_pywrap_tensorflow_interpreter_wrapper tensorflow-lite ${CMAKE_DL_LIBS} diff --git a/tensorflow/lite/nnapi/BUILD b/tensorflow/lite/nnapi/BUILD index 4fff5610d9ab79..58b6dabe2fc25e 100644 --- a/tensorflow/lite/nnapi/BUILD +++ b/tensorflow/lite/nnapi/BUILD @@ -101,7 +101,6 @@ cc_library( deps = [ ":nnapi_implementation", ":nnapi_lib", - "//tensorflow/core/platform:logging", "//tensorflow/lite:framework", "@com_google_absl//absl/log:check", ], diff --git a/tensorflow/lite/nnapi/nnapi_handler.h b/tensorflow/lite/nnapi/nnapi_handler.h index 57ab9d7dfae339..31e6916dd65e25 100644 --- a/tensorflow/lite/nnapi/nnapi_handler.h +++ b/tensorflow/lite/nnapi/nnapi_handler.h @@ -20,7 +20,6 @@ limitations under the License. #include #include "absl/log/check.h" -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/nnapi/NeuralNetworksTypes.h" #include "tensorflow/lite/nnapi/nnapi_implementation.h" diff --git a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake new file mode 100644 index 00000000000000..8a0514eaab93e5 --- /dev/null +++ b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake @@ -0,0 +1,67 @@ +# +# 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 +# +# https://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. + +# Downloads the few header-only files from the TensorFlow repository that the +# TFLite runtime includes (xla/tsl/... and third_party/fft2d/...), instead of +# cloning the whole TensorFlow repository. +# +# Sets TENSORFLOW_HEADERS_INCLUDE_DIR. + +if(TENSORFLOW_HEADERS_INCLUDE_DIR) + set(tensorflow_headers_FOUND TRUE) + return() +endif() + +set(_tensorflow_headers_base_url + "https://raw.githubusercontent.com/tensorflow/tensorflow/v2.21.0-rc0" +) +set(_tensorflow_headers_dir "${CMAKE_BINARY_DIR}/tensorflow_headers") +# || +set(_tensorflow_headers_files + "third_party/fft2d/fft.h|third_party/fft2d/fft.h|2db045d17dfd4b4fa5201e86a1653f0c0b7741e14c927f0c426007127109e825" + "third_party/fft2d/fft2d.h|third_party/fft2d/fft2d.h|b24c63e77d5daf3affd7386085c41c046f3ad0c8dcab56a621fd8e84a5af5c1a" + "third_party/xla/xla/tsl/framework/convolution/eigen_convolution_helpers.h|xla/tsl/framework/convolution/eigen_convolution_helpers.h|3fd52ebb0f14b9c4f3b3e225d1e7b29c712c3046d34e07936bfec0e0c1f41152" + "third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|c0bec189723d52b4495dec62e197e2a4ffe7725bc313676984384e8b8fbb7a13" + "third_party/xla/xla/tsl/lib/random/philox_random.h|xla/tsl/lib/random/philox_random.h|7a7659f95c59419373261af305736311bfd839b281259703b23f4df9b391a43e" + "third_party/xla/xla/tsl/lib/random/random_distributions_utils.h|xla/tsl/lib/random/random_distributions_utils.h|1d65158a878510a1ec5835c5f26962b0ebe45bd56ecaec7e38314a8ac5e7a517" +) + +foreach(_entry IN LISTS _tensorflow_headers_files) + string(REPLACE "|" ";" _entry "${_entry}") + list(GET _entry 0 _src) + list(GET _entry 1 _path) + list(GET _entry 2 _sha256) + set(_dst "${_tensorflow_headers_dir}/${_path}") + if(EXISTS "${_dst}") + file(SHA256 "${_dst}" _actual) + if(_actual STREQUAL _sha256) + continue() + endif() + endif() + message(STATUS "Downloading ${_src}") + file(DOWNLOAD "${_tensorflow_headers_base_url}/${_src}" "${_dst}" + EXPECTED_HASH SHA256=${_sha256} + STATUS _status + ) + list(GET _status 0 _code) + if(NOT _code EQUAL 0) + message(FATAL_ERROR "Failed to download ${_src}: ${_status}") + endif() +endforeach() + +set(TENSORFLOW_HEADERS_INCLUDE_DIR "${_tensorflow_headers_dir}" CACHE PATH + "Include directory with the TensorFlow headers used by TFLite" +) +set(tensorflow_headers_FOUND TRUE) From 112b42f55dc8423258728c13cc13fa6af501cb5f Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 25 Sep 2026 17:00:43 -0700 Subject: [PATCH 04/17] Manually set scoped VMEM limit for VMEM matmuls. Also set large 2nd minor for all matmuls so user doesn't have to. PiperOrigin-RevId: 988574639 --- .../dense_matmul_lib.py | 15 +++++++++++- .../subchannel_matmul_lib.py | 24 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py index 44a34f79ad6d6b..a7066167f4f1e6 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py @@ -241,10 +241,23 @@ def _target_fn(lhs, rhs): rhs = memory_utils.with_large_2nd_minor_layout(rhs) return pallas_func(lhs, rhs) + compiler_args: dict[str, Any] = { + "xla_detailed_logging": True, + "xla_tpu_control_large_2nd_minor_layout_for_x16": True, + } + if ( + cfg.lhs_mem == pltpu.VMEM + and cfg.rhs_mem == pltpu.VMEM + and cfg.out_mem == pltpu.VMEM + ): + # For VMEM matmuls, if scoped VMEM is too high then we may hit OOM since we + # need space for the input operands outside of scoped VMEM. So set it to a + # relatively low value. + compiler_args["xla_tpu_scoped_vmem_limit_kib"] = 8 * 1024 return _target_fn.lower( jax.ShapeDtypeStruct((m, k), lhs_dtype), jax.ShapeDtypeStruct((k, n), rhs_dtype), - ).compile({"xla_detailed_logging": True}) + ).compile(compiler_args) def dense_matmul_reference(cfg: DenseMatmulConfig) -> Callable[..., Any]: diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py index 9bd532696b8e69..d8ce0734a4a17f 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py @@ -322,7 +322,29 @@ def _target_fn(lhs, rhs, lhs_scales, rhs_scales): rhs = memory_utils.with_large_2nd_minor_layout(rhs) return pallas_func(lhs, rhs, lhs_scales, rhs_scales) - return _target_fn + compiler_args: dict[str, Any] = { + "xla_detailed_logging": True, + "xla_tpu_control_large_2nd_minor_layout_for_x16": True, + } + if ( + cfg.lhs_mem == pltpu.VMEM + and cfg.rhs_mem == pltpu.VMEM + and cfg.out_mem == pltpu.VMEM + ): + # For VMEM matmuls, if scoped VMEM is too high then we may hit OOM since we + # need space for the input operands outside of scoped VMEM. So set it to a + # relatively low value. + compiler_args["xla_tpu_scoped_vmem_limit_kib"] = 8 * 1024 + lower_args = [ + jax.ShapeDtypeStruct((m, k), cfg.lhs_dtype), + jax.ShapeDtypeStruct((k, n), cfg.rhs_quantized_dtype), + None, + jax.ShapeDtypeStruct((subblocks_per_tile, n), cfg.rhs_dtype), + ] + if cfg.pre_quantize_lhs: + lower_args[0] = jax.ShapeDtypeStruct((m, k), cfg.lhs_quantized_dtype) + lower_args[2] = jax.ShapeDtypeStruct((m, subblocks_per_tile), cfg.lhs_dtype) + return _target_fn.lower(*lower_args).compile(compiler_args) def subchannel_matmul_jax( From 02d7ccc3bd780085f74d299a00ed48884bfe84f1 Mon Sep 17 00:00:00 2001 From: Badr Badawi Date: Fri, 25 Sep 2026 17:54:55 -0700 Subject: [PATCH 05/17] Expand the external MegaScaleRuntimeError proto to cover more fields. Breaking change: these messages move to the `xla.megascale.runtime.external` package and drop their `Overlay` suffix, e.g. `xla::megascale::runtime::MegaScaleRuntimeErrorOverlay` is now `xla::megascale::runtime::external::MegaScaleRuntimeError`. The wire format is unchanged. PiperOrigin-RevId: 988595845 --- third_party/xla/xla/megascale/BUILD | 3 + .../c_api_megascale_error_aggregator.cc | 2 +- .../c_api_megascale_error_aggregator.h | 2 +- .../c_api_client/megascale_c_api_client.cc | 15 +- .../c_api_client/megascale_c_api_client.h | 10 +- .../megascale_c_api_client_test.cc | 19 +- .../megascale_runtime_error_overlay.proto | 215 +++++++++++++++++- 7 files changed, 239 insertions(+), 27 deletions(-) diff --git a/third_party/xla/xla/megascale/BUILD b/third_party/xla/xla/megascale/BUILD index 39dde3a7bf946d..4540a271b94c2d 100644 --- a/third_party/xla/xla/megascale/BUILD +++ b/third_party/xla/xla/megascale/BUILD @@ -50,4 +50,7 @@ tf_proto_library( srcs = ["megascale_runtime_error_overlay.proto"], create_grpc_library = True, make_default_target_header_only = True, + protodeps = [ + "//xla/tsl/protobuf:status_proto", + ], ) diff --git a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc index 0168d09e5b957c..e75aa32c2a6bc8 100644 --- a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc +++ b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc @@ -93,7 +93,7 @@ CApiMegascaleErrorAggregator::~CApiMegascaleErrorAggregator() { void CApiMegascaleErrorAggregator::AddError( absl::string_view worker_id, - const runtime::MegaScaleRuntimeErrorOverlay& error) { + const runtime::external::MegaScaleRuntimeError& error) { std::string serialized_error = error.SerializeAsString(); PJRT_Megascale_ErrorAggregator_AddError_Args args{}; args.struct_size = PJRT_Megascale_ErrorAggregator_AddError_Args_STRUCT_SIZE; diff --git a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h index 1285cab2416e49..960291c636c6b3 100644 --- a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h +++ b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h @@ -58,7 +58,7 @@ class CApiMegascaleErrorAggregator { delete; void AddError(absl::string_view worker_id, - const runtime::MegaScaleRuntimeErrorOverlay& error); + const runtime::external::MegaScaleRuntimeError& error); ErrorDigest ProcessAndShutdown(); diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc index 8a26c9c477274a..5ac419154b0054 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc @@ -695,7 +695,7 @@ CreateMegascaleCollectives( } typedef absl::AnyInvocable + const runtime::external::MegaScaleRuntimeError& error)> MegaScaleErrorHandler; static absl::Mutex megascale_error_handlers_mutex(absl::kConstInit); @@ -718,7 +718,7 @@ absl::Status RegisterMegascaleErrorHandler(absl::string_view handler_name, args.handler_name_size = handler_name.size(); args.handler = +[](const char* serialized_error, size_t serialized_error_size, void* user_data) { - runtime::MegaScaleRuntimeErrorOverlay error; + runtime::external::MegaScaleRuntimeError error; error.ParseFromString( absl::string_view(serialized_error, serialized_error_size)); MegaScaleErrorHandler* handler = @@ -815,11 +815,12 @@ GetInterfaceAddressesHelper(absl::string_view megascale_port_name, return addresses; } -absl::StatusOr> +absl::StatusOr> GetOrCreateRuntimeError( - runtime::MegaScaleRuntimeErrorOverlay::ErrorType error_type, + runtime::external::MegaScaleRuntimeError::ErrorType error_type, absl::Time start_time, const absl::Status& status, int32_t launch_id, - std::optional + std::optional< + runtime::external::MegaScaleRuntimeError::UnrecoverableErrorType> unrecoverable_error_type) { ABSL_ASSIGN_OR_RETURN(const PJRT_Api* c_api, pjrt::PjrtApi(kTpuPjrtName)); ABSL_ASSIGN_OR_RETURN(PJRT_Megascale_Extension * extension, @@ -849,11 +850,11 @@ GetOrCreateRuntimeError( }; CHECK_NOTNULL(args.serialized_error); - runtime::MegaScaleRuntimeErrorOverlay error; + runtime::external::MegaScaleRuntimeError error; if (!error.ParseFromString(absl::string_view(args.serialized_error, args.serialized_error_size))) { return absl::InternalError( - "Failed to parse MegaScaleRuntimeErrorOverlay proto from serialized " + "Failed to parse MegaScaleRuntimeError proto from serialized " "error."); } diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h index 647d0765910188..8e763ac44c421d 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h @@ -84,7 +84,8 @@ CreateMegascaleCollectives( absl::Status RegisterMegascaleErrorHandler( absl::string_view handler_name, - absl::AnyInvocable + absl::AnyInvocable< + void(const runtime::external::MegaScaleRuntimeError& error)> handler); absl::Status UnregisterMegascaleErrorHandler(absl::string_view handler_name); @@ -96,11 +97,12 @@ GetInterfaceAddressesHelper(absl::string_view megascale_port_name, bool use_all_interfaces, bool limit_to_process_numa_local_interfaces); -absl::StatusOr> +absl::StatusOr> GetOrCreateRuntimeError( - runtime::MegaScaleRuntimeErrorOverlay::ErrorType error_type, + runtime::external::MegaScaleRuntimeError::ErrorType error_type, absl::Time start_time, const absl::Status& status, int32_t launch_id, - std::optional + std::optional< + runtime::external::MegaScaleRuntimeError::UnrecoverableErrorType> unrecoverable_error_type = std::nullopt); } // namespace c_api_client diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc index 968d98768562a9..98be9ba8171804 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc @@ -42,7 +42,7 @@ using ::xla::megascale::c_api_client::RegisterMegascaleErrorHandler; using ::xla::megascale::c_api_client::UnregisterMegascaleErrorHandler; using ::xla::megascale::runtime::DCNTopology; using ::xla::megascale::runtime::EndpointAddresses; -using ::xla::megascale::runtime::MegaScaleRuntimeErrorOverlay; +using ::xla::megascale::runtime::external::MegaScaleRuntimeError; TEST(MegaScaleCApiClientTest, CreateDefaultMegaScaleClientContext) { TF_ASSERT_OK_AND_ASSIGN(auto client_context, @@ -89,7 +89,7 @@ TEST(MegaScaleCApiClientTest, CreateMegascaleErrorAggregator) { TF_ASSERT_OK_AND_ASSIGN(auto aggregator, CreateMegascaleErrorAggregator("test_app")); EXPECT_NE(aggregator, nullptr); - MegaScaleRuntimeErrorOverlay error; + MegaScaleRuntimeError error; aggregator->AddError("worker1", error); EXPECT_EQ(aggregator->size(), 1); EXPECT_TRUE(aggregator->active()); @@ -99,7 +99,7 @@ TEST(MegaScaleCApiClientTest, CreateMegascaleErrorAggregator) { } TEST(MegaScaleCApiClientTest, RegisterAndUnregisterMegascaleErrorHandler) { - auto handler = [](const MegaScaleRuntimeErrorOverlay& error) {}; + auto handler = [](const MegaScaleRuntimeError& error) {}; EXPECT_OK(RegisterMegascaleErrorHandler("test_handler", handler)); EXPECT_OK(UnregisterMegascaleErrorHandler("test_handler")); } @@ -115,16 +115,15 @@ TEST(MegaScaleCApiClientTest, GetInterfaceAddressesHelper) { TEST(MegaScaleCApiClientTest, GetOrCreateRuntimeError) { TF_ASSERT_OK_AND_ASSIGN( - auto result, GetOrCreateRuntimeError( - MegaScaleRuntimeErrorOverlay::UNRECOVERABLE_ERROR, - absl::Now(), absl::OkStatus(), /*launch_id=*/0, - MegaScaleRuntimeErrorOverlay::HOST_TO_DEVICE_ERROR)); + auto result, + GetOrCreateRuntimeError(MegaScaleRuntimeError::UNRECOVERABLE_ERROR, + absl::Now(), absl::OkStatus(), /*launch_id=*/0, + MegaScaleRuntimeError::HOST_TO_DEVICE_ERROR)); const auto& [error, is_new] = result; - EXPECT_EQ(error.error_type(), - MegaScaleRuntimeErrorOverlay::UNRECOVERABLE_ERROR); + EXPECT_EQ(error.error_type(), MegaScaleRuntimeError::UNRECOVERABLE_ERROR); EXPECT_EQ(error.unrecoverable_error_type(), - MegaScaleRuntimeErrorOverlay::HOST_TO_DEVICE_ERROR); + MegaScaleRuntimeError::HOST_TO_DEVICE_ERROR); EXPECT_TRUE(is_new); } diff --git a/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto b/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto index be1bc8d66e11fa..7109e4294991ec 100644 --- a/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto +++ b/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto @@ -15,13 +15,201 @@ syntax = "proto3"; -package xla.megascale.runtime; +package xla.megascale.runtime.external; + +import "xla/tsl/protobuf/status.proto"; option java_multiple_files = true; option java_outer_classname = "Runtime"; -// Overlay for xla.megascale.runtime.MegaScaleRuntimeError -message MegaScaleRuntimeErrorOverlay { +enum TpuCoreTypeProto { + TPU_CORE_TYPE_INVALID = 0; + TPU_CORE_TYPE_TENSOR_CORE = 1; + TPU_CORE_TYPE_RESERVED = 2; + TPU_CORE_TYPE_SPARSE_CORE = 3; +} + +// LINT.IfChange(megascale_action) +message MegaScaleAction { + message DeviceToHost { + // Buffer produced as result of D2H transfer will be made available + // at done_key. + string done_key = 2; + } + + message HostToDevice {} + + message NetworkSend { + message Peer { + string key = 1; + } + + message Transfer { + Peer peer = 3; + } + + message TransferInfo { + Transfer transfer = 4; + } + + // Transfer buffer over network to destinations specified in transfer_info. + TransferInfo transfer_info = 1; + } + + message NetworkReceive { + // Buffer produced as a result of network received with this key will + // be available at done_key in callback registry. + string done_key = 1; + + // Expected size of the buffer. + int64 buffer_size = 2; + } + + message Compute {} + + // An input can be either a valid buffer to represent data dependency or a + // void buffer to indicate control dependency. + message Input { + // Index of action in MegaScaleActionGraph.actions that produces this + // buffer. + int32 action_index = 1; + } + + oneof action { + DeviceToHost device_to_host = 1; + HostToDevice host_to_device = 2; + NetworkSend network_send = 3; + NetworkReceive network_receive = 4; + Compute compute = 5; + } + + // List of inputs. Action is executed once all its inputs are available and + // any run_after actions are done. + repeated Input inputs = 6; + + // Index of this action in the graph. Mainly for debugging. + int32 index = 8; +} +// LINT.ThenChange() + +// LINT.IfChange(megascale_action_graph) +message MegaScaleActionGraph { + // Unique identifier for this graph. Typically this is set to unique + // identifiers generated by XLA compiler. + string id = 1; + + // All actions that belong to this Graph. Actions are executed once all their + // inputs are available (generated by predecessor actions) and any run_after + // actions are done. See MegaScaleAction for more details. The index of + // actions is referred to by dependent actions. + repeated MegaScaleAction actions = 2; + + // The HLO module name corresponding to the send op that triggered this graph. + string hlo_module = 7; +} +// LINT.ThenChange() + +// LINT.IfChange(executor_telemetry) +message ExecutorTelemetry { + message PendingGraphState { + MegaScaleActionGraph graph = 1; + int32 launch_id = 2; + repeated int32 completed_actions = 3; + repeated int32 pending_actions = 4; + repeated int32 active_actions = 6; + } + + repeated PendingGraphState pending_graphs = 1; +} +// LINT.ThenChange() + +// LINT.IfChange(program_run_info) +message ProgramRunInfo { + // run_id is only relevant within a TPU and shouldn't be compared across TPUs. + uint64 run_id = 1; + // launch_id should be consistent across TPUs running the same execution. + // (within and across slices) + int32 launch_id = 2; + // The name of the HLO module/program. + string name = 3; + // We store the fingerprint (as a hex string) just in case there are different + // modules that happen to have the same module_name, which generally shouldn't + // happen. + string fingerprint = 4; +} +// LINT.ThenChange() + +// LINT.IfChange(tpu_core_hlo) +message TpuCoreHLO { + // Chip ID local to the slice. + int32 chip_id = 1; + // Core index within its chip. + int32 core_idx = 2; + // Full HLO details. + string hlo = 3; + // Fingerprint of the running executable. + string executable_fingerprint = 4; + // Current program counter + optional int64 pc = 6; + // Tag information to be used with PC + optional int64 tag = 7; + string hostname = 8; + // Short name of the HLO. + string hlo_name = 9; + // Module name of the running executable. + string module_name = 10; + // Launch ID of the running executable. + int32 launch_id = 12; + // The type of the TPU core. + TpuCoreTypeProto tpu_core_type = 13; + // Short name of the computation where the HLO resides in. + string computation_name = 15; + // Basic debug info about recently completed programs runs. Sorted in reverse + // chronological order. (the first entry is the most recently completed run) + repeated ProgramRunInfo recent_completions = 16; + // True if the TPU is running a program actively. + optional bool active = 18; + // The error when obtaining the TPU core information. + repeated tensorflow.StatusProto errors = 19; +} +// LINT.ThenChange() + +// LINT.IfChange(rapid_eye_info) +message RapidEyeInfo { + ExecutorTelemetry telemetry = 1; + + repeated TpuCoreHLO tpu_core_hlos = 2; + + // Whether the RapidEyeInfo was truncated due to exceeding the size limit. + bool is_truncated = 3; + + // The number of slices, tasks per slice, local slice ID and task ID that are + // obtained from the workers' last used topology. This is useful for the error + // aggregator to determine if all error messages have been received. + int64 num_slices = 4; + int64 num_tasks_per_slice = 5; + int32 slice_id = 6; + int32 task_id = 7; + + // The incarnation ID for the workload shared across all workers. This will + // reset every time the workload restart. + int64 incarnation_id = 8; + + // Specifies the TPU platform type. + int32 tpu_version = 9; + string chip_config_name = 10; +} +// LINT.ThenChange() + +// LINT.IfChange(runtime_state_summary) +message RuntimeStateSummary { + // RapidEye info. + RapidEyeInfo rapid_eye_info = 4; +} +// LINT.ThenChange() + +// LINT.IfChange(megascale_runtime_error) +message MegaScaleRuntimeError { enum ErrorType { NO_ERROR = 0; HANG_DETECTED = 1; @@ -29,6 +217,18 @@ message MegaScaleRuntimeErrorOverlay { CANCELLED = 3; } + // How the hang error is detected. + enum HangErrorType { + // The default value. This indicates that the hang error type is not + // specified. + HANG_ERROR_UNSPECIFIED = 0; + // The hang error type is caused by a megascale timeout. + TIMEOUT_ERROR = 1; + // The hang error type is caused by no progress being made on all registered + // metrics like rpc, pjrt, etc. + NO_PROGRESS_ERROR = 2; + } + enum UnrecoverableErrorType { UNRECOVERABLE_ERROR_UNCLASSIFIED = 0; HOST_TO_DEVICE_ERROR = 1; @@ -36,9 +236,16 @@ message MegaScaleRuntimeErrorOverlay { } ErrorType error_type = 1; - UnrecoverableErrorType unrecoverable_error_type = 12; string error_message = 2; uint64 estimated_start_time_ms = 3; + RuntimeStateSummary runtime_state = 4; int32 triggering_launch_id = 5; string hostname = 6; + string job_name = 8; + int32 task_id = 9; + fixed64 event_id = 11; + UnrecoverableErrorType unrecoverable_error_type = 12; + fixed64 runtime_error_id = 13; + HangErrorType hang_error_type = 14; } +// LINT.ThenChange() From b4449a41869c3dbac7b4db333d019a941584e7a5 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 25 Sep 2026 18:04:01 -0700 Subject: [PATCH 06/17] Fix CompareToZero Android Lint error in GpuDelegateTest. PiperOrigin-RevId: 988599619 --- .../src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java index e1c8fd7fb5f40c..ee23a362acd639 100644 --- a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java +++ b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java @@ -214,7 +214,7 @@ private static ArrayList getTopKLabels(float[][] labels, int k) { @Override public int compare(Map.Entry o1, Map.Entry o2) { // Intentionally reversed to put high confidence at the head of the queue. - return o1.getValue().compareTo(o2.getValue()) * -1; + return o2.getValue().compareTo(o1.getValue()); } }); From f83e88165700f9dc01669daead9f21ea77f7ee9a Mon Sep 17 00:00:00 2001 From: Terry Heo Date: Fri, 25 Sep 2026 19:56:32 -0700 Subject: [PATCH 07/17] Build benchmark_model without protobuf in TFLite CMake The CMake benchmark_model build fetched the TensorFlow repository to compile tf.Example and TFLite profiling protos. Compile the proto-based outputs out instead, so benchmark_model builds without TensorFlow or protobuf. - Added the TFLITE_WITHOUT_PROTO define to benchmark_tflite_model.cc and profile_summary_formatter.{h,cc}, guarding the proto-based outputs (--output_proto_filepath, --result_file_path, --export_model_runtime_info and the 'proto' op profiling output mode). - Set TFLITE_WITHOUT_PROTO for benchmark_model and label_image in CMake and dropped their proto link dependencies. - Removed the CMake files for the benchmark result and profiling protos. - Downloaded stats_calculator and its forwarding headers in Findtensorflow_headers.cmake instead of taking them from the TensorFlow checkout. - Validated the field count of each Findtensorflow_headers.cmake entry. - Fetched TensorFlow only for the kernel tests. - Removed the unused TFLITE_GENERATED_HEADERS_DIR variable. PiperOrigin-RevId: 988641036 --- tensorflow/lite/CMakeLists.txt | 21 ++----- .../lite/examples/label_image/CMakeLists.txt | 8 ++- .../profiling/profile_summary_formatter.cc | 4 ++ .../profiling/profile_summary_formatter.h | 4 ++ .../lite/profiling/proto/CMakeLists.txt | 57 ------------------- .../lite/tools/benchmark/CMakeLists.txt | 20 ++----- .../tools/benchmark/benchmark_tflite_model.cc | 53 ++++++++++++++--- .../lite/tools/benchmark/proto/CMakeLists.txt | 36 ------------ .../modules/Findtensorflow_headers.cmake | 16 +++++- 9 files changed, 81 insertions(+), 138 deletions(-) delete mode 100644 tensorflow/lite/profiling/proto/CMakeLists.txt delete mode 100644 tensorflow/lite/tools/benchmark/proto/CMakeLists.txt diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index d6a7e460a32ef5..34c644f040523f 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -69,12 +69,12 @@ option(TFLITE_KERNEL_TEST "Enable tflite kernel unit test" OFF) # The TFLite runtime builds from this repository plus a few header-only TensorFlow # files (see tools/cmake/modules/Findtensorflow_headers.cmake). A full TensorFlow -# source tree is only needed by the benchmark_model tool and the kernel tests. +# source tree is only needed by the kernel tests. set(TENSORFLOW_SOURCE_DIR "" CACHE PATH "Directory that contains the TensorFlow project. Only needed when \ -TFLITE_ENABLE_BENCHMARK_MODEL or TFLITE_KERNEL_TEST is ON." +TFLITE_KERNEL_TEST is ON." ) -if(TFLITE_ENABLE_BENCHMARK_MODEL OR TFLITE_KERNEL_TEST) +if(TFLITE_KERNEL_TEST) if(NOT TENSORFLOW_SOURCE_DIR) message(STATUS "Downloading TensorFlow repository...") include(FetchContent) @@ -816,21 +816,8 @@ if(TFLITE_KERNEL_TEST) add_subdirectory(${TFLITE_SOURCE_DIR}/kernels) endif() -# Add the generated headers directory. Required for maintaining the -# tensorflow/lite directory structure for generated headers. -set(TFLITE_GENERATED_HEADERS_DIR ${CMAKE_BINARY_DIR}/tensorflow/lite) - +# The benchmark tool. if(TFLITE_ENABLE_BENCHMARK_MODEL) - # Add the profiling proto directory. - add_subdirectory(${TFLITE_SOURCE_DIR}/profiling/proto) - - # Add the benchmark result proto directory. - add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark/proto) - - # Add the tf example directory. - add_subdirectory(${TF_SOURCE_DIR}/core/example ${CMAKE_BINARY_DIR}/example_proto_generated) - - # The benchmark tool. add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark) endif() diff --git a/tensorflow/lite/examples/label_image/CMakeLists.txt b/tensorflow/lite/examples/label_image/CMakeLists.txt index 823d90dc044b94..6dda556abe8e67 100644 --- a/tensorflow/lite/examples/label_image/CMakeLists.txt +++ b/tensorflow/lite/examples/label_image/CMakeLists.txt @@ -21,7 +21,7 @@ populate_source_vars("${TFLITE_SOURCE_DIR}/examples/label_image" FILTER "_test\\.cc$" ) list(APPEND TFLITE_LABEL_IMAGE_SRCS - ${XLA_SOURCE_DIR}/xla/tsl/util/stats_calculator.cc + ${TENSORFLOW_HEADERS_INCLUDE_DIR}/xla/tsl/util/stats_calculator.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc ${TFLITE_SOURCE_DIR}/profiling/profile_buffer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summarizer.cc @@ -82,8 +82,10 @@ target_compile_options(label_image PRIVATE ${TFLITE_LABEL_IMAGE_CC_OPTIONS} ) +target_compile_definitions(label_image + PRIVATE + TFLITE_WITHOUT_PROTO +) target_link_libraries(label_image tensorflow-lite - profiling_info_proto - libprotobuf ) diff --git a/tensorflow/lite/profiling/profile_summary_formatter.cc b/tensorflow/lite/profiling/profile_summary_formatter.cc index afc9dc2db36fd2..20c3eb53a1bbf5 100644 --- a/tensorflow/lite/profiling/profile_summary_formatter.cc +++ b/tensorflow/lite/profiling/profile_summary_formatter.cc @@ -29,7 +29,9 @@ limitations under the License. #include "tensorflow/core/util/stat_summarizer_options.h" #include "tensorflow/core/util/stats_calculator.h" +#ifndef TFLITE_WITHOUT_PROTO #include "tensorflow/lite/profiling/proto/profiling_info.pb.h" +#endif // TFLITE_WITHOUT_PROTO #include "tensorflow/lite/tools/logging.h" namespace tflite { @@ -139,6 +141,7 @@ ProfileSummaryCSVFormatter::GetStatSummarizerOptions() const { return options; } +#ifndef TFLITE_WITHOUT_PROTO std::vector ProfileSummaryProtoFormatter::GetDetailsSortedByRunOrder( const tensorflow::StatsCalculator* stats_calculator) const { @@ -305,6 +308,7 @@ void ProfileSummaryProtoFormatter::HandleOutput( benchmark_profiling_data.SerializeToOstream(output_stream); } } +#endif // TFLITE_WITHOUT_PROTO } // namespace profiling } // namespace tflite diff --git a/tensorflow/lite/profiling/profile_summary_formatter.h b/tensorflow/lite/profiling/profile_summary_formatter.h index 62514eafecda93..5eab2f1539d572 100644 --- a/tensorflow/lite/profiling/profile_summary_formatter.h +++ b/tensorflow/lite/profiling/profile_summary_formatter.h @@ -28,7 +28,9 @@ limitations under the License. #include "tensorflow/core/util/stat_summarizer_options.h" #include "tensorflow/core/util/stats_calculator.h" +#ifndef TFLITE_WITHOUT_PROTO #include "tensorflow/lite/profiling/proto/profiling_info.pb.h" +#endif // TFLITE_WITHOUT_PROTO namespace tflite { namespace profiling { @@ -98,6 +100,7 @@ class ProfileSummaryCSVFormatter : public ProfileSummaryDefaultFormatter { tensorflow::StatSummarizerOptions GetStatSummarizerOptions() const override; }; +#ifndef TFLITE_WITHOUT_PROTO class ProfileSummaryProtoFormatter : public ProfileSummaryFormatter { public: std::string GetOutputString( @@ -139,6 +142,7 @@ class ProfileSummaryProtoFormatter : public ProfileSummaryFormatter { std::vector GetDetailsSortedByRunOrder( const tensorflow::StatsCalculator* stats_calculator) const; }; +#endif // TFLITE_WITHOUT_PROTO } // namespace profiling } // namespace tflite diff --git a/tensorflow/lite/profiling/proto/CMakeLists.txt b/tensorflow/lite/profiling/proto/CMakeLists.txt deleted file mode 100644 index 5738c992fc2839..00000000000000 --- a/tensorflow/lite/profiling/proto/CMakeLists.txt +++ /dev/null @@ -1,57 +0,0 @@ -# -# Copyright 2024 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 -# -# https://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. -find_package(Protobuf REQUIRED) - -add_library(profiling_info_proto profiling_info.proto) - -list(APPEND profiling_info_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/profiling_info.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/profiling_info.pb.h) - -# Generate profiling_info.pb.cc and profiling_info.pb.h from -# profiling_info.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${profiling_info_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/profiling_info.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/profiling_info.proto -) - -set_source_files_properties(${profiling_info_generated_files} PROPERTIES GENERATED TRUE) -target_sources(profiling_info_proto PRIVATE ${profiling_info_generated_files}) -target_link_libraries(profiling_info_proto protobuf::libprotobuf) -target_include_directories(profiling_info_proto PUBLIC ${CMAKE_BINARY_DIR}) - -add_library(model_runtime_info_proto model_runtime_info.proto) -list(APPEND model_runtime_info_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/model_runtime_info.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/model_runtime_info.pb.h -) - -# Generate model_runtime_info.pb.cc and model_runtime_info.pb.h from -# model_runtime_info.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${model_runtime_info_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/model_runtime_info.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/model_runtime_info.proto ${profiling_info_generated_files} -) - -set_source_files_properties(${model_runtime_info_generated_files} PROPERTIES GENERATED TRUE) -target_sources(model_runtime_info_proto PRIVATE ${model_runtime_info_generated_files}) -target_link_libraries(model_runtime_info_proto protobuf::libprotobuf) -target_include_directories(model_runtime_info_proto PUBLIC ${CMAKE_BINARY_DIR}) diff --git a/tensorflow/lite/tools/benchmark/CMakeLists.txt b/tensorflow/lite/tools/benchmark/CMakeLists.txt index aee420c6a52b9f..087d73ab41ce60 100644 --- a/tensorflow/lite/tools/benchmark/CMakeLists.txt +++ b/tensorflow/lite/tools/benchmark/CMakeLists.txt @@ -21,11 +21,10 @@ populate_source_vars("${TFLITE_SOURCE_DIR}/tools/benchmark" FILTER "(_test|_plus_flex_main|_performance_options.*)\\.cc$" ) list(APPEND TFLITE_BENCHMARK_SRCS - ${XLA_SOURCE_DIR}/xla/tsl/util/stats_calculator.cc + ${TENSORFLOW_HEADERS_INCLUDE_DIR}/xla/tsl/util/stats_calculator.cc ${TFLITE_SOURCE_DIR}/kernels/internal/utils/sparsity_format_converter.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc ${TFLITE_SOURCE_DIR}/profiling/memory_usage_monitor.cc - ${TFLITE_SOURCE_DIR}/profiling/model_runtime_info.cc ${TFLITE_SOURCE_DIR}/profiling/profile_buffer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summarizer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summary_formatter.cc @@ -44,15 +43,6 @@ list(APPEND TFLITE_BENCHMARK_SRCS list(APPEND TFLITE_BENCHMARK_LIBS tensorflow-lite -) - -list(APPEND TFLITE_BENCHMARK_LIBS - profiling_info_proto - feature_proto - example_proto - model_runtime_info_proto - benchmark_result_proto - protobuf::libprotobuf absl::log ) @@ -110,9 +100,11 @@ target_compile_options(benchmark_model PRIVATE ${TFLITE_BENCHMARK_CC_OPTIONS} ) -target_include_directories(benchmark_model - PUBLIC - ${CMAKE_BINARY_DIR} +# Proto-based outputs (e.g. --result_file_path, --export_model_runtime_info and +# the 'proto' op profiling output mode) are not supported in the CMake build. +target_compile_definitions(benchmark_model + PRIVATE + TFLITE_WITHOUT_PROTO ) target_link_libraries(benchmark_model ${TFLITE_BENCHMARK_LIBS} diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc index f6ccb18a963143..cfc7c7a65bc19c 100644 --- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc +++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc @@ -39,8 +39,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "ruy/profiler/profiler.h" // from @ruy -#include "tensorflow/core/example/example.pb.h" -#include "tensorflow/core/example/feature.pb.h" #include "tensorflow/lite/core/c/c_api_types.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/core/kernels/register.h" @@ -52,18 +50,25 @@ limitations under the License. #include "tensorflow/lite/kernels/cpu_backend_context.h" #include "tensorflow/lite/op_resolver.h" #include "tensorflow/lite/optional_debug_tools.h" -#include "tensorflow/lite/profiling/model_runtime_info.h" #include "tensorflow/lite/profiling/profile_summary_formatter.h" #include "tensorflow/lite/string_util.h" #include "tensorflow/lite/tools/benchmark/benchmark_params.h" #include "tensorflow/lite/tools/benchmark/benchmark_utils.h" #include "tensorflow/lite/tools/benchmark/profiling_listener.h" -#include "tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h" #include "tensorflow/lite/tools/delegates/delegate_provider.h" #include "tensorflow/lite/tools/logging.h" #include "tensorflow/lite/tools/model_loader.h" #include "tensorflow/lite/tools/utils.h" +// Proto-based outputs are compiled out when TFLITE_WITHOUT_PROTO is defined +// (e.g. the CMake build), which avoids a dependency on protobuf and TensorFlow. +#ifndef TFLITE_WITHOUT_PROTO +#include "tensorflow/core/example/example.pb.h" +#include "tensorflow/core/example/feature.pb.h" +#include "tensorflow/lite/profiling/model_runtime_info.h" +#include "tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h" +#endif // TFLITE_WITHOUT_PROTO + void RegisterSelectedOps(::tflite::MutableOpResolver* resolver); // Version with Weak linker attribute doing nothing: if someone links this @@ -75,7 +80,9 @@ RegisterSelectedOps(::tflite::MutableOpResolver* resolver) {} namespace tflite { namespace benchmark { namespace { +#ifndef TFLITE_WITHOUT_PROTO using ::tflite::tools::benchmark::BenchmarkResult; +#endif // TFLITE_WITHOUT_PROTO using utils::InputTensorData; using utils::VoidUniquePtr; @@ -89,12 +96,19 @@ constexpr bool kOpProfilingEnabledDefault = false; // Op profiling output modes. constexpr char kOpProfilingOutputModeStdout[] = "stdout"; constexpr char kOpProfilingOutputModeCsv[] = "csv"; +#ifndef TFLITE_WITHOUT_PROTO constexpr char kOpProfilingOutputModeProto[] = "proto"; +#endif // TFLITE_WITHOUT_PROTO + +const char* kOpProfilingOutputModes[] = { + kOpProfilingOutputModeStdout, + kOpProfilingOutputModeCsv, +#ifndef TFLITE_WITHOUT_PROTO + kOpProfilingOutputModeProto, +#endif // TFLITE_WITHOUT_PROTO +}; -const char* kOpProfilingOutputModes[] = {kOpProfilingOutputModeStdout, - kOpProfilingOutputModeCsv, - kOpProfilingOutputModeProto}; - +#ifndef TFLITE_WITHOUT_PROTO // Sets feature values in the tensorflow::Example proto from the tflite tensor. // Returns an error if the tensor type is not supported or the tensor dime is a // nullptr. @@ -137,6 +151,7 @@ TfLiteStatus MaybeSetFeatureValuesFromTensor(const TfLiteTensor& tensor, return kTfLiteError; } } +#endif // TFLITE_WITHOUT_PROTO // Dumps ruy profiling events if the ruy profiler is enabled. class RuyProfileListener : public BenchmarkListener { @@ -218,6 +233,7 @@ class OutputSaver : public BenchmarkListener { } } +#ifndef TFLITE_WITHOUT_PROTO // If the output_proto_filepath is specified, save the output tensors as // tensorflow::Example proto and serialize it to the file. const std::string output_proto_path = @@ -236,6 +252,12 @@ class OutputSaver : public BenchmarkListener { ofs.close(); } } +#else + if (!params_->Get("output_proto_filepath").empty()) { + TFLITE_LOG(WARN) + << "--output_proto_filepath is not supported in this build."; + } +#endif // TFLITE_WITHOUT_PROTO } private: @@ -243,6 +265,7 @@ class OutputSaver : public BenchmarkListener { const BenchmarkParams* params_ = nullptr; }; +#ifndef TFLITE_WITHOUT_PROTO // Dumps the Model Runtime Info if enabled when export_model_runtime_info is // set to true. class ModelRuntimeInfoListener : public BenchmarkListener { @@ -337,6 +360,7 @@ class ProtoBenchmarkReporter : public BenchmarkListener { private: std::string result_file_path_; }; +#endif // TFLITE_WITHOUT_PROTO std::vector Split(const std::string& str, const char delim) { if (str.empty()) { @@ -488,8 +512,10 @@ std::shared_ptr CreateProfileSummaryFormatter(const std::string& output_mode) { if (output_mode == kOpProfilingOutputModeCsv) { return std::make_shared(); +#ifndef TFLITE_WITHOUT_PROTO } else if (output_mode == kOpProfilingOutputModeProto) { return std::make_shared(); +#endif // TFLITE_WITHOUT_PROTO } else { return std::make_shared(); } @@ -1205,8 +1231,13 @@ TfLiteStatus BenchmarkTfLiteModel::Init() { new InterpreterStatePrinter(interpreter_.get()))); if (params_.Get("export_model_runtime_info")) { +#ifndef TFLITE_WITHOUT_PROTO AddOwnedListener(std::unique_ptr( new ModelRuntimeInfoListener(interpreter_.get()))); +#else + TFLITE_LOG(WARN) + << "--export_model_runtime_info is not supported in this build."; +#endif // TFLITE_WITHOUT_PROTO } interpreter_->SetAllowFp16PrecisionForFp32(params_.Get("allow_fp16")); @@ -1375,8 +1406,14 @@ TfLiteStatus BenchmarkTfLiteModel::Init() { AddOwnedListener( std::unique_ptr(new RuyProfileListener())); +#ifndef TFLITE_WITHOUT_PROTO AddOwnedListener( std::unique_ptr(new ProtoBenchmarkReporter())); +#else + if (!params_.Get("result_file_path").empty()) { + TFLITE_LOG(WARN) << "--result_file_path is not supported in this build."; + } +#endif // TFLITE_WITHOUT_PROTO AddOwnedListener(std::unique_ptr( new OutputSaver(interpreter_runner_.get()))); diff --git a/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt b/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt deleted file mode 100644 index 6a39f06e03c32b..00000000000000 --- a/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -# -# Copyright 2025 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 -# -# https://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. -find_package(Protobuf REQUIRED) - -add_library(benchmark_result_proto benchmark_result.proto) - -list(APPEND benchmark_result_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h) - -# Generate benchmark_result.pb.cc and benchmark_result.pb.h from -# benchmark_result.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${benchmark_result_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/benchmark_result.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/benchmark_result.proto -) - -set_source_files_properties(${benchmark_result_generated_files} PROPERTIES GENERATED TRUE) -target_sources(benchmark_result_proto PRIVATE ${benchmark_result_generated_files}) -target_link_libraries(benchmark_result_proto protobuf::libprotobuf) -target_include_directories(benchmark_result_proto PUBLIC ${CMAKE_BINARY_DIR}) diff --git a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake index 8a0514eaab93e5..7fef6f8c81c55e 100644 --- a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake +++ b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Downloads the few header-only files from the TensorFlow repository that the -# TFLite runtime includes (xla/tsl/... and third_party/fft2d/...), instead of -# cloning the whole TensorFlow repository. +# Downloads the few files from the TensorFlow repository that the TFLite runtime +# and the benchmark_model tool use (xla/tsl/..., tensorflow/core/util/... and +# third_party/fft2d/...), instead of cloning the whole TensorFlow repository. # # Sets TENSORFLOW_HEADERS_INCLUDE_DIR. @@ -36,10 +36,20 @@ set(_tensorflow_headers_files "third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|c0bec189723d52b4495dec62e197e2a4ffe7725bc313676984384e8b8fbb7a13" "third_party/xla/xla/tsl/lib/random/philox_random.h|xla/tsl/lib/random/philox_random.h|7a7659f95c59419373261af305736311bfd839b281259703b23f4df9b391a43e" "third_party/xla/xla/tsl/lib/random/random_distributions_utils.h|xla/tsl/lib/random/random_distributions_utils.h|1d65158a878510a1ec5835c5f26962b0ebe45bd56ecaec7e38314a8ac5e7a517" + # Used by the benchmark_model tool. + "tensorflow/core/util/stat_summarizer_options.h|tensorflow/core/util/stat_summarizer_options.h|e91f94b2fbb7e3a8d8b99664a1ff2443088ca1433b4372b1fcd2477b78492b3b" + "tensorflow/core/util/stats_calculator.h|tensorflow/core/util/stats_calculator.h|8d80ca00444162a55a288b3b7369f314e5cf65a3258387aef4af1df39db08e0a" + "third_party/xla/xla/tsl/util/stat_summarizer_options.h|xla/tsl/util/stat_summarizer_options.h|9f7d7cc5de38ae6e97a8982dc386a012ed70144c768ef2525fd25054de41c2b2" + "third_party/xla/xla/tsl/util/stats_calculator.cc|xla/tsl/util/stats_calculator.cc|704c6d22240a521a7cbee591cd932121f0f6dcbaf179037758a7f3bd4d851e82" + "third_party/xla/xla/tsl/util/stats_calculator.h|xla/tsl/util/stats_calculator.h|f5628ba1fbf39e7c4d2b280ff78c8daba40d9d462a97f8aa2abc0c5371b9f241" ) foreach(_entry IN LISTS _tensorflow_headers_files) string(REPLACE "|" ";" _entry "${_entry}") + list(LENGTH _entry _entry_len) + if(NOT _entry_len EQUAL 3) + message(FATAL_ERROR "Malformed entry in _tensorflow_headers_files: ${_entry}") + endif() list(GET _entry 0 _src) list(GET _entry 1 _path) list(GET _entry 2 _sha256) From 828ae62881c33c185402a597fd422322dffcd412 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 25 Sep 2026 21:10:16 -0700 Subject: [PATCH 08/17] Automated Code Change PiperOrigin-RevId: 988666007 --- tensorflow/lite/toco/tflite/BUILD | 1 + tensorflow/lite/toco/tflite/export.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/tensorflow/lite/toco/tflite/BUILD b/tensorflow/lite/toco/tflite/BUILD index ab265a0264ef4f..b665f95ada7714 100644 --- a/tensorflow/lite/toco/tflite/BUILD +++ b/tensorflow/lite/toco/tflite/BUILD @@ -142,6 +142,7 @@ cc_library( "//tensorflow/lite/toco:toco_port", "//tensorflow/lite/toco:tooling_util", "@com_google_absl//absl/log", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@flatbuffers", ], diff --git a/tensorflow/lite/toco/tflite/export.cc b/tensorflow/lite/toco/tflite/export.cc index d04d598dc62a6d..63889cb572c91b 100644 --- a/tensorflow/lite/toco/tflite/export.cc +++ b/tensorflow/lite/toco/tflite/export.cc @@ -23,6 +23,7 @@ limitations under the License. #include "flatbuffers/flexbuffers.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "flatbuffers/buffer.h" // from @flatbuffers From e1b7cbdd2cb63b2412686217b426e8df5e7e29b4 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Sat, 26 Sep 2026 03:09:34 -0700 Subject: [PATCH 09/17] Fix signed-integer-overflow in RaggedGather by validating output splits against Tsplits limits. PiperOrigin-RevId: 988777009 --- tensorflow/core/kernels/ragged_gather_op.cc | 25 +++++++++++++------ .../ops/ragged/ragged_gather_op_test.py | 19 ++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc index 72e5c840e86ad5..8e12616d26ad23 100644 --- a/tensorflow/core/kernels/ragged_gather_op.cc +++ b/tensorflow/core/kernels/ragged_gather_op.cc @@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include @@ -40,10 +41,10 @@ void WriteValueSlices( const auto& params_dense_values = params_dense_values_in.flat_outer_dims(); auto values = values_out->flat_outer_dims(); - int out_pos = 0; + SPLITS_TYPE out_pos = 0; for (const auto& slice : value_slices) { - for (int i = slice.first; i < slice.second; ++i) { - for (int j = 0; j < value_size; ++j) { + for (SPLITS_TYPE i = slice.first; i < slice.second; ++i) { + for (SPLITS_TYPE j = 0; j < value_size; ++j) { values(out_pos, j) = params_dense_values(i, j); } ++out_pos; @@ -171,17 +172,25 @@ class RaggedGatherOpBase : public OpKernel { // should add a new split point to out_splits that is 4 greater than the // previous split point in out_splits. for (int i = 0; i < indices.size(); ++i) { - int start = indices(i); - int limit = indices(i) + 1; + SPLITS_TYPE start = indices(i); + SPLITS_TYPE limit = indices(i) + 1; // Copy splits. for (int dim = 0; dim < params_nested_splits.size(); ++dim) { const auto& splits = params_nested_splits[dim]; int out_dim = dim + indices_in.dims() - 1; if (out_dim >= 0) { - SPLITS_TYPE delta = out_splits->at(out_dim).back() - splits(start); - for (int j = start; j < limit; ++j) { - out_splits->at(out_dim).push_back(splits(j + 1) + delta); + // Splits are validated to be non-negative and sorted, so `length` is + // in [0, max]. Check before adding so neither int32 nor int64 Tsplits + // can overflow. + const SPLITS_TYPE last = out_splits->at(out_dim).back(); + for (SPLITS_TYPE j = start; j < limit; ++j) { + const SPLITS_TYPE length = splits(j + 1) - splits(start); + if (length > std::numeric_limits::max() - last) { + return absl::InvalidArgumentError( + "Output splits value exceeds limits of Tsplits"); + } + out_splits->at(out_dim).push_back(last + length); } } start = splits(start); diff --git a/tensorflow/python/ops/ragged/ragged_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_op_test.py index d20cc3744060e1..06bdf06867ff15 100644 --- a/tensorflow/python/ops/ragged/ragged_gather_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_gather_op_test.py @@ -242,6 +242,25 @@ def testValueCountOverflowError(self): ) ) + def testSplitsOverflowErrorInt64(self): + # Zero-width values let int64 splits reach 2**62 without allocating memory; + # gathering the row twice would overflow the int64 output splits. + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'exceeds limits of Tsplits' + ): + self.evaluate( + gen_ragged_array_ops.ragged_gather( + params_nested_splits=[ + constant_op.constant([0, 2**62], dtype=dtypes.int64) + ], + params_dense_values=array_ops.zeros( + [2**62, 0], dtype=dtypes.uint8 + ), + indices=array_ops.zeros([2], dtype=dtypes.int32), + OUTPUT_RAGGED_RANK=1, + ) + ) + def testUnknownIndicesRankError(self): if context.executing_eagerly(): return From 422ab38b9e49d2d1b49c8a26f6c2fc645c7c26a5 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Sat, 26 Sep 2026 10:10:28 -0700 Subject: [PATCH 10/17] Fix Python 3 compatibility in huge_slice_op_test Replace long(2)**32 + 3 with 2**32 + 3 and add the missing test.main() entrypoint. PiperOrigin-RevId: 988876726 --- .../python/kernel_tests/array_ops/huge_slice_op_test.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py b/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py index 3df288faf21d46..bb702cca84e902 100644 --- a/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py +++ b/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py @@ -33,7 +33,11 @@ def testInt64Slicing(self): self.assertAllEqual([True, False, True], slice_val) slice_t = array_ops.slice( - a_large, constant_op.constant([long(2)**32 + 3], dtype=dtypes.int64), - [3]) + a_large, constant_op.constant([2**32 + 3], dtype=dtypes.int64), [3] + ) slice_val = self.evaluate(slice_t) self.assertAllEqual([True, False, True], slice_val) + + +if __name__ == "__main__": + test.main() From a2c628e9dc1f9f7230afe194947c86bb184ad846 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Sat, 26 Sep 2026 10:11:31 -0700 Subject: [PATCH 11/17] Cast bfloat16 logits to float32 on GPU in stateless_random_ops_test PiperOrigin-RevId: 988877018 --- .../python/kernel_tests/random/stateless_random_ops_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py index fad01517b67d2c..5d59970e3a0f9a 100644 --- a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py +++ b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py @@ -171,9 +171,7 @@ def multinomial_cases(): def wrap(op: ops.Operation, logits, logits_dtype, output_dtype, seed): device_type = get_device().device_type # Some dtypes are not supported on some devices - if (logits_dtype == dtypes.bfloat16 and device_type == 'GPU' and - not test_util.is_gpu_available( - cuda_only=True, min_cuda_compute_capability=(8, 0))): + if logits_dtype == dtypes.bfloat16 and device_type == 'GPU': logits_dtype = dtypes.float32 return op(seed=seed, logits=constant_op.constant(logits, dtype=logits_dtype), From 862ecb35c6f6704cda0cf7fa6a9da0106e99a6ec Mon Sep 17 00:00:00 2001 From: Terry Heo Date: Sat, 26 Sep 2026 12:21:53 -0700 Subject: [PATCH 12/17] Build TFLite kernel tests without TensorFlow in CMake The CMake kernel tests were the last user of the TensorFlow checkout. They now build from LiteRT sources only. - Took the schema and converter sources from tflite/converter and generated the mutable schema header in the build directory. - Dropped the unused compiler/mlir/tools/optimize/quantization_utils.cc and the tsl logging sources. - Replaced the tsl logging include in kernels/test_util with absl/log. - Added TFLITE_WITHOUT_MODEL_DUMP to compile out test model dumping, which needs tools/serialization, and set it in the CMake build. - Fixed stale sources and link deps (op_version, op_signature, absl check/log/random). - Removed the TensorFlow FetchContent block and TENSORFLOW_SOURCE_DIR. PiperOrigin-RevId: 988908594 --- tensorflow/lite/CMakeLists.txt | 46 -------------------------- tensorflow/lite/kernels/BUILD | 3 +- tensorflow/lite/kernels/CMakeLists.txt | 29 ++++++++-------- tensorflow/lite/kernels/test_util.cc | 9 ++++- tensorflow/lite/kernels/test_util.h | 3 +- 5 files changed, 27 insertions(+), 63 deletions(-) diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 34c644f040523f..39809306861f58 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -67,36 +67,6 @@ option(TFLITE_ENABLE_EXTERNAL_DELEGATE "Enable External Delegate backend" ON) option(TFLITE_KERNEL_TEST "Enable tflite kernel unit test" OFF) -# The TFLite runtime builds from this repository plus a few header-only TensorFlow -# files (see tools/cmake/modules/Findtensorflow_headers.cmake). A full TensorFlow -# source tree is only needed by the kernel tests. -set(TENSORFLOW_SOURCE_DIR "" CACHE PATH - "Directory that contains the TensorFlow project. Only needed when \ -TFLITE_KERNEL_TEST is ON." -) -if(TFLITE_KERNEL_TEST) - if(NOT TENSORFLOW_SOURCE_DIR) - message(STATUS "Downloading TensorFlow repository...") - include(FetchContent) - FetchContent_Declare( - tensorflow - GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git - GIT_TAG v2.21.0-rc0 - GIT_SHALLOW TRUE - GIT_PROGRESS TRUE - SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/tensorflow-src - ) - FetchContent_GetProperties(tensorflow) - if(NOT tensorflow_POPULATED) - FetchContent_Populate(tensorflow) - endif() - set(TENSORFLOW_SOURCE_DIR "${tensorflow_SOURCE_DIR}") - endif() - set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") - set(TSL_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/third_party/tsl") - set(XLA_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/") -endif() - if(${CMAKE_CROSSCOMPILING}) set(TFLITE_HOST_TOOLS_DIR "" CACHE PATH "Host tools directory") if(TFLITE_HOST_TOOLS_DIR STREQUAL "") @@ -168,14 +138,6 @@ macro(populate_tflite_source_vars RELATIVE_DIR SOURCES_VAR) "${TFLITE_SOURCE_DIR}/${RELATIVE_DIR}" ${SOURCES_VAR} ${ARGN} ) endmacro() -# Simplifies inclusion of non-test sources and headers from a directory -# relative to TF_SOURCE_DIR. See populate_source_vars() for the description of -# arguments including and following SOURCES_VAR. -macro(populate_tf_source_vars RELATIVE_DIR SOURCES_VAR) - populate_source_vars( - "${TF_SOURCE_DIR}/${RELATIVE_DIR}" ${SOURCES_VAR} ${ARGN} - ) -endmacro() # Make sure all repos have licenses. set(OVERRIDABLE_FETCH_CONTENT_LICENSE_CHECK ON) # Additional library dependencies based upon enabled features. @@ -197,14 +159,6 @@ find_package(ruy REQUIRED) # Header-only TensorFlow dependencies (xla/tsl/..., third_party/fft2d/...). find_package(tensorflow_headers REQUIRED) include_directories(${TENSORFLOW_HEADERS_INCLUDE_DIR}) -if(TENSORFLOW_SOURCE_DIR) - include_directories( - AFTER - ${TENSORFLOW_SOURCE_DIR} - ${TSL_SOURCE_DIR} - ${XLA_SOURCE_DIR} - ) -endif() # Download necessary dependencies. if(TFLITE_ENABLE_XNNPACK) # pthreadpool is used by XNNPACK. diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index d693d8f2aa11f9..d43aaa4fcbbdd6 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -174,8 +174,10 @@ cc_library( "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/synchronization", @@ -183,7 +185,6 @@ cc_library( "@com_google_googletest//:gtest", "@eigen_archive//:eigen3", "@flatbuffers", - "@tsl//tsl/platform:logging", ], ) diff --git a/tensorflow/lite/kernels/CMakeLists.txt b/tensorflow/lite/kernels/CMakeLists.txt index b9872323ac698d..ca9f48a5b3412f 100644 --- a/tensorflow/lite/kernels/CMakeLists.txt +++ b/tensorflow/lite/kernels/CMakeLists.txt @@ -18,10 +18,8 @@ find_package(google_benchmark REQUIRED) find_package(re2 REQUIRED) # Generate the mutable schema_generated.h header for tests. -set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") -set(SCHEMA_GENERATED_ROOT ${TF_SOURCE_DIR}/compiler/mlir/lite/schema) -set(SCHEMA_FILE ${SCHEMA_GENERATED_ROOT}/schema.fbs) -set(SCHEMA_GENERATED_PATH ${SCHEMA_GENERATED_ROOT}/mutable) +set(SCHEMA_FILE ${TFLITE_SOURCE_DIR}/converter/schema/schema.fbs) +set(SCHEMA_GENERATED_PATH ${PROJECT_BINARY_DIR}/tflite/converter/schema/mutable) set(SCHEMA_GENERATED_FILE ${SCHEMA_GENERATED_PATH}/schema_generated.h) # Use the util function in flatbuffer to generate the schema header. @@ -84,18 +82,15 @@ set(TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_list.cc ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_util.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc - ${TFLITE_SOURCE_DIR}/../compiler/mlir/lite/schema/schema_conversion_utils.cc + ${TFLITE_SOURCE_DIR}/converter/schema/schema_conversion_utils.cc ${TFLITE_SOURCE_DIR}/tools/command_line_flags.cc ${DELEGATE_PROVIDERS} ${TFLITE_SOURCE_DIR}/tools/optimize/model_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/tools/optimize/operator_property.cc - ${TF_SOURCE_DIR}/compiler/mlir/tools/optimize/quantization_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/kernels/internal/runtime_shape.cc + ${TFLITE_SOURCE_DIR}/converter/tools/optimize/operator_property.cc + ${TFLITE_SOURCE_DIR}/converter/kernels/internal/runtime_shape.cc ${TFLITE_SOURCE_DIR}/tools/optimize/quantization_utils.cc ${TFLITE_SOURCE_DIR}/tools/tool_params.cc - ${TFLITE_SOURCE_DIR}/tools/versioning/op_version.cc - ${XLA_SOURCE_DIR}/xla/tsl/platform/default/env_time.cc - ${XLA_SOURCE_DIR}/xla/tsl/platform/default/logging.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_version.cc internal/test_util.cc acceleration_test_util.cc acceleration_test_util_internal.cc @@ -111,10 +106,13 @@ endif() if(NOT TFLITE_ENABLE_GPU) list(APPEND TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/tools/versioning/op_signature.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_signature.cc ) endif() -set(TEST_FRAMEWORK_OPTIONS "") +# Model dumping (--dump_tflite_model_dir) needs tools/serialization, which is not +# built with CMake. +set(TEST_FRAMEWORK_OPTIONS "-DTFLITE_WITHOUT_MODEL_DUMP") if(TFLITE_ENABLE_XNNPACK) list(APPEND TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/tools/delegates/xnnpack_delegate_provider.cc @@ -130,9 +128,12 @@ target_link_libraries(tensorflow-lite-test-base re2 tensorflow-lite benchmark + absl::check + absl::log + absl::random_random ) add_dependencies(tensorflow-lite-test-base mutable_schema_file) -target_include_directories(tensorflow-lite-test-base PUBLIC ${SCHEMA_GENERATED_ROOT}) +target_include_directories(tensorflow-lite-test-base PUBLIC ${PROJECT_BINARY_DIR}) target_compile_options(tensorflow-lite-test-base PUBLIC ${TEST_FRAMEWORK_OPTIONS}) add_library(tensorflow-lite-test-gtest-main INTERFACE) @@ -163,7 +164,7 @@ macro(add_kernel_test TEST_SRC TEST_LIB) ${TEST_LIB} ) add_test(NAME ${TEST_NAME} COMMAND ${TEST_NAME} - WORKING_DIRECTORY ${TENSORFLOW_SOURCE_DIR} + WORKING_DIRECTORY ${TFLITE_SOURCE_DIR}/.. ) set_tests_properties(${TEST_NAME} PROPERTIES LABELS "plain") diff --git a/tensorflow/lite/kernels/test_util.cc b/tensorflow/lite/kernels/test_util.cc index 8596eae3639ac6..f66ea6ff0a49db 100644 --- a/tensorflow/lite/kernels/test_util.cc +++ b/tensorflow/lite/kernels/test_util.cc @@ -37,6 +37,8 @@ limitations under the License. #include "absl/base/const_init.h" #include "absl/base/no_destructor.h" #include "absl/base/thread_annotations.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_replace.h" @@ -65,11 +67,12 @@ limitations under the License. #include "tensorflow/lite/string_type.h" #include "tensorflow/lite/string_util.h" #include "tensorflow/lite/tools/logging.h" +#ifndef TFLITE_WITHOUT_MODEL_DUMP #include "tensorflow/lite/tools/serialization/writer_lib.h" +#endif // TFLITE_WITHOUT_MODEL_DUMP #include "tensorflow/lite/tools/versioning/op_version.h" #include "tensorflow/lite/types/fp16.h" // IWYU pragma: keep #include "tensorflow/lite/version.h" -#include "tsl/platform/logging.h" namespace tflite { @@ -164,6 +167,7 @@ MATCHER(Fp16Eq, "") { return AlmostEquals(actual, expected, fp16_ulps_in_fp32); } +#ifndef TFLITE_WITHOUT_MODEL_DUMP // Returns the name of the dumped model. The name is in the format of // DTS---.tflite. The model serial // number is used to distinguish different models dumped in the same test. @@ -281,6 +285,7 @@ std::unique_ptr ModifyDumpedModel( std::make_unique(std::move(data), fbb.GetSize()); return FlatBufferModel::VerifyAndBuildFromAllocation(std::move(allocation)); } +#endif // TFLITE_WITHOUT_MODEL_DUMP } // namespace @@ -703,6 +708,7 @@ int SingleOpModel::CountNumberOfDelegatedPartitions() const { } void SingleOpModel::MaybeDumpModel() { +#ifndef TFLITE_WITHOUT_MODEL_DUMP std::string dump_directory( tflite::KernelTestDelegateProviders::Get() ->ConstParams() @@ -750,6 +756,7 @@ void SingleOpModel::MaybeDumpModel() { output_file.write(reinterpret_cast(allocation->base()), allocation->bytes()); output_file.close(); +#endif // TFLITE_WITHOUT_MODEL_DUMP } SingleOpModel::~SingleOpModel() { diff --git a/tensorflow/lite/kernels/test_util.h b/tensorflow/lite/kernels/test_util.h index ce2ce190554a55..3b9121c24125bc 100644 --- a/tensorflow/lite/kernels/test_util.h +++ b/tensorflow/lite/kernels/test_util.h @@ -42,6 +42,8 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/types/span.h" #include "Eigen/Core" // from @eigen_archive #include "flatbuffers/buffer.h" // from @flatbuffers @@ -63,7 +65,6 @@ limitations under the License. #include "tensorflow/lite/types/fp16.h" #include "tensorflow/lite/types/half.h" #include "tensorflow/lite/util.h" -#include "tsl/platform/logging.h" namespace tflite { From 4dfd44407814ebcca8872ea3d599f0ec51728b85 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Sat, 26 Sep 2026 14:04:14 -0700 Subject: [PATCH 13/17] Validate contracting and batch dimensions in XLA Einsum and BatchDot to reject invalid operations with mismatched inner dimensions. Fixes #117771 PiperOrigin-RevId: 988931357 --- third_party/xla/xla/hlo/builder/lib/BUILD | 1 + third_party/xla/xla/hlo/builder/lib/matrix.cc | 15 ++++++-- .../xla/xla/hlo/builder/lib/matrix_test.cc | 35 +++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/hlo/builder/lib/BUILD b/third_party/xla/xla/hlo/builder/lib/BUILD index ad1f5002716126..5bce2511b0c6a0 100644 --- a/third_party/xla/xla/hlo/builder/lib/BUILD +++ b/third_party/xla/xla/hlo/builder/lib/BUILD @@ -324,6 +324,7 @@ xla_test( "//xla:array4d", "//xla:literal", "//xla:literal_util", + "//xla:shape_util", "//xla:types", "//xla/hlo/builder:xla_builder", "//xla/hlo/testlib:test", diff --git a/third_party/xla/xla/hlo/builder/lib/matrix.cc b/third_party/xla/xla/hlo/builder/lib/matrix.cc index 3587d6a447f396..1024575cf8ca6a 100644 --- a/third_party/xla/xla/hlo/builder/lib/matrix.cc +++ b/third_party/xla/xla/hlo/builder/lib/matrix.cc @@ -457,9 +457,14 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, } else if (x_shape.dimensions(i) == 1) { rhs_outer_dims.push_back(rhs_dim); lhs_delete_dims.push_back(i); - } else { + } else if (y_shape.dimensions(rhs_dim) == 1) { lhs_outer_dims.push_back(i); rhs_delete_dims.push_back(rhs_dim); + } else { + return InvalidArgument( + "Batch dimensions %d of LHS and %d of RHS must match or be 1. " + "Got %d and %d.", + i, rhs_dim, x_shape.dimensions(i), y_shape.dimensions(rhs_dim)); } } else if (is_contracting(dim_name)) { if (x_shape.dimensions(i) == y_shape.dimensions(rhs_dim)) { @@ -468,9 +473,15 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, } else if (x_shape.dimensions(i) == 1) { rhs_outer_dims.push_back(rhs_dim); lhs_delete_dims.push_back(i); - } else { + } else if (y_shape.dimensions(rhs_dim) == 1) { lhs_outer_dims.push_back(i); rhs_delete_dims.push_back(rhs_dim); + } else { + return InvalidArgument( + "Contracting dimensions %d of LHS and %d of RHS must match or be " + "1. " + "Got %d and %d.", + i, rhs_dim, x_shape.dimensions(i), y_shape.dimensions(rhs_dim)); } } else { lhs_outer_dims.push_back(i); diff --git a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc index b3bc3b87f953c5..598b5a3c0d459f 100644 --- a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc +++ b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/hlo/testlib/test.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/shape_util.h" #include "xla/tests/client_library_test_runner_mixin.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" #include "xla/tests/hlo_pjrt_test_base.h" @@ -329,5 +330,39 @@ TEST_F(MatrixTest, NormalizeEinsumString) { EXPECT_EQ(NormalizeEinsumString("...ba,ca..."), "...ba,ca...->...bc"); } +TEST_F(MatrixTest, EinsumIncompatibleContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0, 2.0, 3.0, 4.0}}); + auto b = + ConstantR2(&builder, {{1.0}, {2.0}, {3.0}, {4.0}, {5.0}, {6.0}}); + Einsum(a, b, "ab,bc->ac"); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, BatchDotIncompatibleContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0, 2.0, 3.0, 4.0}}); + auto b = + ConstantR2(&builder, {{1.0}, {2.0}, {3.0}, {4.0}, {5.0}, {6.0}}); + BatchDot(a, b); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, EinsumIncompatibleBatchDims) { + XlaBuilder builder(TestName()); + auto a = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {3, 2, 2}), "a"); + auto b = Parameter(&builder, 1, ShapeUtil::MakeShape(F32, {5, 2, 2}), "b"); + Einsum(a, b, "...ij,...jk->...ik"); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, EinsumBroadcastingContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0}}); + auto b = ConstantR2(&builder, {{1.0}, {2.0}, {3.0}}); + Einsum(a, b, "ab,bc->ac"); + EXPECT_TRUE(builder.GetCurrentStatus().ok()); +} + } // namespace } // namespace xla From 321852acecb0a3271b4faf667989c50cd451048c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 26 Sep 2026 17:41:27 -0700 Subject: [PATCH 14/17] Use absl::Span and absl::string_view in HloModule, HloInstruction, and HloCloneContext. Replace read-only `const std::vector&` and `const std::string&` parameters with `absl::Span` and `absl::string_view` in `HloModuleConfig::set_spmd_parameters_shardings`, `HloModule::set_profile_info`, `HloInstruction` cycle detection, and `HloCloneContext` to avoid unnecessary container copies and allow callers to pass slices or views directly. PiperOrigin-RevId: 988979844 --- third_party/xla/xla/hlo/ir/hlo_clone_context.h | 3 ++- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 2 +- third_party/xla/xla/hlo/ir/hlo_module.h | 9 ++++----- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_clone_context.h b/third_party/xla/xla/hlo/ir/hlo_clone_context.h index b1eaa31306c7cc..6cbe42966cf68e 100644 --- a/third_party/xla/xla/hlo/ir/hlo_clone_context.h +++ b/third_party/xla/xla/hlo/ir/hlo_clone_context.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" #include "xla/map_util.h" namespace xla { @@ -34,7 +35,7 @@ class HloCloneContext { // Creates a new HloCloneContext object to clone HloInstruction and // HloComputation objects to be added to the module specified as argument. // The suffix string will be appended to computation names. - explicit HloCloneContext(HloModule* module, const std::string& suffix = "") + explicit HloCloneContext(HloModule* module, absl::string_view suffix = "") : module_(module), suffix_(suffix) {} HloModule* module() const { return module_; } diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 8122426fc3d261..34719734a2fca1 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -4017,7 +4017,7 @@ std::string PrintCycle(const HloInstruction* child, DFSStack* dfs_stack, while (!dfs.empty() && result.empty()) { bool found_next_instr = false; auto process_users_or_successors = - [&](const std::vector& users_or_successors) { + [&](absl::Span users_or_successors) { for (const auto& user : users_or_successors) { if (user == child) { dfs.push_back(child); diff --git a/third_party/xla/xla/hlo/ir/hlo_module.h b/third_party/xla/xla/hlo/ir/hlo_module.h index 263db00361f89f..03d728da2d771f 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module.h +++ b/third_party/xla/xla/hlo/ir/hlo_module.h @@ -724,9 +724,8 @@ class HloModule { CHECK(spmd_parameters_shardings_.has_value()); return *spmd_parameters_shardings_; } - void set_spmd_parameters_shardings( - const std::vector& shardings) { - spmd_parameters_shardings_ = shardings; + void set_spmd_parameters_shardings(absl::Span shardings) { + spmd_parameters_shardings_.emplace(shardings.begin(), shardings.end()); } // Checks if this config has the entry computation output's HLO sharding for @@ -856,8 +855,8 @@ class HloModule { } void set_profile_info( - const std::vector& profile_info) { - profile_info_list_ = profile_info; + absl::Span profile_info) { + profile_info_list_.assign(profile_info.begin(), profile_info.end()); } const std::vector& profile_info() const { From 3eebeac433b341b8fb650cf742a5878ba833ec3d Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 26 Sep 2026 19:09:34 -0700 Subject: [PATCH 15/17] Use absl::string_view for HloModuleMetadata and HloPassPipeline parameters. Replace read-only `const std::string&` parameters with `absl::string_view` across `HloModuleMetadata`, `HloPassInterface`, and `HloPassPipeline` to avoid unnecessary `std::string` heap allocations when string literals or views are passed. PiperOrigin-RevId: 989002614 --- .../xla/xla/hlo/ir/hlo_module_metadata.cc | 3 ++- .../xla/xla/hlo/ir/hlo_module_metadata.h | 18 ++++++++---------- .../xla/xla/hlo/pass/hlo_pass_interface.h | 2 +- .../xla/xla/hlo/pass/hlo_pass_pipeline.cc | 8 ++++---- .../xla/xla/hlo/pass/hlo_pass_pipeline.h | 2 +- 5 files changed, 16 insertions(+), 17 deletions(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc index 2ede1b3ef6cf43..328b068c81407d 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc +++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/strings/string_view.h" #include "xla/service/hlo.pb.h" #include "xla/service/metrics.pb.h" #include "xla/tsl/platform/env.h" @@ -107,7 +108,7 @@ absl::Status HloModuleMetadata::set_custom_metadata( return absl::OkStatus(); } -absl::Status HloModuleMetadata::set_key_value_metric(const std::string& key, +absl::Status HloModuleMetadata::set_key_value_metric(absl::string_view key, int64_t value) { ABSL_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata, GetCurrentHloPassMetadata()); diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h index ec4d5be8926412..15944e6e5518f5 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h +++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h @@ -125,7 +125,7 @@ class HloModuleMetadata { const HloModuleMetadata& prepartitioning_metadata); // Setters for HloModuleMetadataProto. - void set_module_group_name(const std::string& name) { + void set_module_group_name(absl::string_view name) { module_metadata_.set_module_group_name(name); } void set_canonical_module_id(int64_t id) { @@ -137,7 +137,7 @@ class HloModuleMetadata { absl::Status set_custom_metadata(const ::tsl::protobuf::Message& message); // Adds a (key, value) pair metric if none was already set. Otherwise, it // updates the existing value. - absl::Status set_key_value_metric(const std::string& key, int64_t value); + absl::Status set_key_value_metric(absl::string_view key, int64_t value); absl::StatusOr current_pass_id() { ABSL_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata, @@ -146,23 +146,21 @@ class HloModuleMetadata { } // Setters for the current HloPassMetadata. - absl::Status set_current_pass_name(const std::string& pass_name) { + absl::Status set_current_pass_name(absl::string_view pass_name) { return MutateCurrentHloPassMetadata( - [&pass_name](HloPassMetadata* pass_metadata) { + [pass_name](HloPassMetadata* pass_metadata) { pass_metadata->set_pass_name(pass_name); }); } - absl::Status set_current_pass_pipeline_name( - const std::string& pipeline_name) { + absl::Status set_current_pass_pipeline_name(absl::string_view pipeline_name) { return MutateCurrentHloPassMetadata( - [&pipeline_name](HloPassMetadata* pass_metadata) { + [pipeline_name](HloPassMetadata* pass_metadata) { pass_metadata->set_pipeline_name(pipeline_name); }); } - absl::Status add_current_pass_dump_filename( - const std::string& dump_filename) { + absl::Status add_current_pass_dump_filename(absl::string_view dump_filename) { return MutateCurrentHloPassMetadata( - [&dump_filename](HloPassMetadata* pass_metadata) { + [dump_filename](HloPassMetadata* pass_metadata) { pass_metadata->add_dump_filenames(dump_filename); }); } diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h index a9f4d1e65a84b9..1fcdf4058c6c6c 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h +++ b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h @@ -118,7 +118,7 @@ class HloPassInterface { // If an HloPassMetadata has previously been created, it adds a (key, value) // pair metric if none was already set or updates the existing value. // If an HloPassMetadata doesn't exist, it simply returns. - static void SetKVMetric(HloModule* module, const std::string& key, + static void SetKVMetric(HloModule* module, absl::string_view key, int64_t value) { auto status = module->metadata()->set_key_value_metric(key, value); if (!status.ok()) { diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc index f493932b44acbd..5cf8c5712809cc 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc +++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc @@ -48,8 +48,8 @@ namespace xla { namespace { -void RecordPassStartMetadata(HloModule& module, const std::string& pass_name, - const std::string& pipeline_name) { +void RecordPassStartMetadata(HloModule& module, absl::string_view pass_name, + absl::string_view pipeline_name) { module.metadata()->RecordPassStart(); // An HloPassMetadata was just created so absl::Status should always be OK. CHECK_OK(module.metadata()->set_current_pass_name(pass_name)); @@ -57,7 +57,7 @@ void RecordPassStartMetadata(HloModule& module, const std::string& pass_name, } absl::Status AttemptRecordPassEndMetadata(HloModule& module, - const std::string& pass_name, + absl::string_view pass_name, bool module_changed) { // Module id is set here instead of RecordPassStartMetadata because it may // change in the middle of the pass, and we want the final id. @@ -69,7 +69,7 @@ absl::Status AttemptRecordPassEndMetadata(HloModule& module, return absl::OkStatus(); } -void RecordPassEndMetadata(HloModule& module, const std::string& pass_name, +void RecordPassEndMetadata(HloModule& module, absl::string_view pass_name, bool module_changed) { absl::Status status = AttemptRecordPassEndMetadata(module, pass_name, module_changed); diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h index 8d53cd9740a770..aa54edc8a4e69d 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h +++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h @@ -37,7 +37,7 @@ namespace xla { // Pipeline of HLO passes. class HloPassPipeline : public HloPassInterface { public: - explicit HloPassPipeline(const std::string& name, + explicit HloPassPipeline(absl::string_view name, CompilationStats* compilation_stats = nullptr) : name_(name), compilation_stats_(compilation_stats) { if (compilation_stats == nullptr) { From 5ce05c79adcdd8c2d0861a0aa57874dc4959e708 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 26 Sep 2026 21:40:57 -0700 Subject: [PATCH 16/17] Use absl::string_view for read-only string parameters in HloInstruction. Replace `const std::string&` parameters with `absl::string_view` in `HloInstruction` methods (`CreateInfeed`, `Clone`, `CloneWithNewShape`, `CloneWithNewOperands`, `add_frontend_attribute`, `set_infeed_config`, `set_outfeed_config`) and string-to-enum helpers (`StringToRandomAlgorithm`, `StringToRandomDistribution`, `StringToPrecision`, `StringToAlgorithm`) to avoid unnecessary `std::string` heap allocations. PiperOrigin-RevId: 989037408 --- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 28 ++++++++++--------- third_party/xla/xla/hlo/ir/hlo_instruction.h | 28 +++++++++---------- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 34719734a2fca1..cc60113b29b038 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -2106,9 +2106,9 @@ HloInstruction::CreateCollectivePermuteStart( /* static */ std::unique_ptr HloInstruction::CreateInfeed( const Shape& infeed_shape, HloInstruction* token_operand, - const std::string& config) { + absl::string_view config) { return std::make_unique(infeed_shape, token_operand, - config); + std::string(config)); } /* static */ std::unique_ptr HloInstruction::CreateOutfeed( @@ -2829,7 +2829,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( std::unique_ptr HloInstruction::CloneWithNewOperands( const Shape& shape, absl::Span new_operands, - const std::string& suffix, HloCloneContext* context) const { + absl::string_view suffix, HloCloneContext* context) const { VLOG(3) << "CloneWithNewOperands:\n " << ToString(); VLOG(3) << " new operands:"; for (const HloInstruction* new_operand : new_operands) { @@ -3100,7 +3100,7 @@ void HloInstruction::DetachFromOperandsAndUsersOutside( } std::unique_ptr HloInstruction::CloneWithNewShape( - const Shape& shape, const std::string& suffix, + const Shape& shape, absl::string_view suffix, HloCloneContext* context) const { std::unique_ptr clone = CloneWithNewOperands(shape, operands_, context); @@ -3113,7 +3113,7 @@ std::unique_ptr HloInstruction::CloneWithNewShape( } std::unique_ptr HloInstruction::Clone( - const std::string& suffix, HloCloneContext* context) const { + absl::string_view suffix, HloCloneContext* context) const { std::unique_ptr clone = CloneWithNewShape(shape(), suffix, context); return clone; @@ -5879,19 +5879,19 @@ absl::StatusOr StringToEnum(absl::string_view value_name, F enum_to_string, } // namespace absl::StatusOr StringToRandomAlgorithm( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, RandomAlgorithmToString, "algorithm"); } absl::StatusOr StringToRandomDistribution( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, RandomDistributionToString, "distribution"); } absl::StatusOr StringToPrecision( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, PrecisionToString, "precision"); } @@ -5913,7 +5913,7 @@ absl::StatusOr StringToResultAccuracy( } absl::StatusOr StringToAlgorithm( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, AlgorithmToString, "algorithm"); } @@ -6289,8 +6289,9 @@ std::string HloInstruction::infeed_config() const { return Cast(this)->infeed_config(); } -void HloInstruction::set_infeed_config(const std::string& config) { - return Cast(this)->set_infeed_config(config); +void HloInstruction::set_infeed_config(absl::string_view config) { + return Cast(this)->set_infeed_config( + std::string(config)); } const Shape& HloInstruction::outfeed_shape() const { @@ -6305,8 +6306,9 @@ const std::string& HloInstruction::outfeed_config() const { return Cast(this)->outfeed_config(); } -void HloInstruction::set_outfeed_config(const std::string& config) { - return Cast(this)->set_outfeed_config(config); +void HloInstruction::set_outfeed_config(absl::string_view config) { + return Cast(this)->set_outfeed_config( + std::string(config)); } const std::vector& HloInstruction::replica_groups() const { diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index b5fa706d78d90d..c233a089e30a74 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -875,7 +875,7 @@ class HloInstruction { // is a tuple containing the infeed_shape and the TOKEN. static std::unique_ptr CreateInfeed( const Shape& infeed_shape, HloInstruction* token_operand, - const std::string& config); + absl::string_view config); // Creates an outfeed instruction, which outputs data. outfeed_shape is the // shape of the data being outfed *not* the shape of the outfeed instruction @@ -1817,12 +1817,12 @@ class HloInstruction { // the instruction to form the name of the cloned instruction. // Ignores the control predecessors and successors of this HLO instruction. std::unique_ptr Clone( - const std::string& suffix = "clone", + absl::string_view suffix = "clone", HloCloneContext* context = nullptr) const; // Clones the HLO instruction as above but with new shape. std::unique_ptr CloneWithNewShape( - const Shape& shape, const std::string& suffix = "clone", + const Shape& shape, absl::string_view suffix = "clone", HloCloneContext* context = nullptr) const; // Clones the HLO instruction as above but with new shape and operands. @@ -1833,7 +1833,7 @@ class HloInstruction { // Clones the HLO instruction with new shape, operands and suffix. std::unique_ptr CloneWithNewOperands( const Shape& shape, absl::Span new_operands, - const std::string& suffix, HloCloneContext* context = nullptr) const; + absl::string_view suffix, HloCloneContext* context = nullptr) const; // Implementation for non-common logic of CloneWithNewOperands. // CloneWithNewOperands forwards to this method for some of the intstruction @@ -2011,10 +2011,9 @@ class HloInstruction { // Adds a single attribute only if it not already present in the // HloInstruction. Returns false if the attribute was already present. - bool add_frontend_attribute(const std::string& key, - const std::string& value) { - auto it = - mutable_rare()->frontend_attributes.mutable_map()->insert({key, value}); + bool add_frontend_attribute(absl::string_view key, absl::string_view value) { + auto it = mutable_rare()->frontend_attributes.mutable_map()->insert( + {std::string(key), std::string(value)}); return it.second; } @@ -2422,13 +2421,13 @@ class HloInstruction { std::string infeed_config() const; // Delegates to HloInfeedInstruction::set_infeed_config. - void set_infeed_config(const std::string& config); + void set_infeed_config(absl::string_view config); // Returns the config for the Outfeed instruction. const std::string& outfeed_config() const; // Delegates to HloOutfeedInstruction::set_outfeed_config. - void set_outfeed_config(const std::string& config); + void set_outfeed_config(absl::string_view config); // Returns the shape for the Outfeed instruction. const Shape& outfeed_shape() const; @@ -2984,16 +2983,15 @@ std::string SparsityConfigToString(const SparsityConfig& sparsity_config); std::string BlockScalingConfigToString( const BlockScalingConfig& block_scaling_config); -absl::StatusOr StringToRandomAlgorithm( - const std::string& name); +absl::StatusOr StringToRandomAlgorithm(absl::string_view name); absl::StatusOr StringToRandomDistribution( - const std::string& name); + absl::string_view name); absl::StatusOr StringToPrecision( - const std::string& name); + absl::string_view name); absl::StatusOr StringToShuffleMode( absl::string_view mode); absl::StatusOr StringToAlgorithm( - const std::string& name); + absl::string_view name); absl::StatusOr StringToResultAccuracy( absl::string_view name); absl::StatusOr StringToCustomCallSchedule( From 0c9cd330920b012f1ec542ca91c821f37a0ec253 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 26 Sep 2026 23:16:56 -0700 Subject: [PATCH 17/17] Use Abseil container algorithms in HLO IR classes. Replace iterator-based `std::find`, `std::replace`, `std::all_of`, and manual search loops with `absl::c_linear_search`, `absl::c_find`, `absl::c_any_of`, `absl::c_replace`, and `absl::c_all_of` across `HloInstruction`, `HloInstructionSequence`, and `OriginalValue`. PiperOrigin-RevId: 989057167 --- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 19 +++++++------------ .../xla/xla/hlo/ir/hlo_original_value.h | 6 +++--- third_party/xla/xla/hlo/ir/hlo_schedule.h | 13 ++++--------- 3 files changed, 14 insertions(+), 24 deletions(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index cc60113b29b038..3563727f001b7a 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -124,7 +124,7 @@ void HloInstruction::Users::Clear() { bool HloInstruction::Users::Contains(const HloInstruction* instruction) const { if (user_map_ == nullptr) { - return std::find(users_.begin(), users_.end(), instruction) != users_.end(); + return absl::c_linear_search(users_, instruction); } return user_map_->contains(instruction); } @@ -150,7 +150,7 @@ void HloInstruction::Users::AddUser(HloInstruction* user) { int64_t HloInstruction::Users::UserId(HloInstruction* user) { if (user_map_ == nullptr) { - auto it = std::find(users_.begin(), users_.end(), user); + auto it = absl::c_find(users_, user); CHECK(it != users_.end()); return it - users_.begin(); } @@ -3414,12 +3414,9 @@ void HloInstruction::RemoveOperandsAtAscendingIndices( } bool HloInstruction::HasConstantOperand() const { - for (const HloInstruction* operand : operands_) { - if (operand->IsConstant()) { - return true; - } - } - return false; + return absl::c_any_of(operands_, [](const HloInstruction* operand) { + return operand->IsConstant(); + }); } bool HloInstruction::IdenticalSlowPath( @@ -3608,8 +3605,7 @@ absl::Status HloInstruction::ReplaceUseWithDifferentShape( RemoveUser(user); TF_RET_CHECK(absl::c_count(user->operands_, this) >= 0); - std::replace(user->operands_.begin(), user->operands_.end(), this, - new_producer); + absl::c_replace(user->operands_, this, new_producer); new_producer->AddUser(user); // Custom fusions may not be able to handle deduplicated operands. if (user->opcode() == HloOpcode::kFusion) { @@ -3825,8 +3821,7 @@ absl::Status HloInstruction::ReplaceAllUsesWithDifferentShape( // graph. new_producer remains the only user of this instruction. new_producer_is_user = true; } else { - std::replace(user->operands_.begin(), user->operands_.end(), this, - new_producer); + absl::c_replace(user->operands_, this, new_producer); new_producer->AddUser(user); if (user->opcode() == HloOpcode::kFusion) { ABSL_RETURN_IF_ERROR( diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value.h b/third_party/xla/xla/hlo/ir/hlo_original_value.h index 20f2eadbb8b17e..7769db7599ba08 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value.h +++ b/third_party/xla/xla/hlo/ir/hlo_original_value.h @@ -117,9 +117,9 @@ class OriginalValue { } bool IsEmpty() const { - return std::all_of( - tree().leaves().begin(), tree().leaves().end(), - [](const auto& pair) { return !pair.second.has_value(); }); + return absl::c_all_of(tree().leaves(), [](const auto& pair) { + return !pair.second.has_value(); + }); } bool IsCompatibleWith(const Shape& shape) const; diff --git a/third_party/xla/xla/hlo/ir/hlo_schedule.h b/third_party/xla/xla/hlo/ir/hlo_schedule.h index aef51e1e44791e..49548f962a7387 100644 --- a/third_party/xla/xla/hlo/ir/hlo_schedule.h +++ b/third_party/xla/xla/hlo/ir/hlo_schedule.h @@ -73,12 +73,10 @@ class HloInstructionSequence { // Removes the instruction from the sequence. void remove_instruction(HloInstruction* instruction) { - auto instruction_it = std::find(instruction_sequence_.begin(), - instruction_sequence_.end(), instruction); + auto instruction_it = absl::c_find(instruction_sequence_, instruction); if (instruction_it != instruction_sequence_.end() && instruction->parent() != nullptr) { - auto id_it = std::find(id_sequence_.begin(), id_sequence_.end(), - instruction->unique_id()); + auto id_it = absl::c_find(id_sequence_, instruction->unique_id()); instruction_sequence_.erase(instruction_it); id_sequence_.erase(id_it); } @@ -87,11 +85,8 @@ class HloInstructionSequence { // Replaces the old instruction with the new instruction in the sequence. void replace_instruction(HloInstruction* old_instruction, HloInstruction* new_instruction) { - auto instruction_it = - std::find(instruction_sequence_.begin(), instruction_sequence_.end(), - old_instruction); - auto id_it = std::find(id_sequence_.begin(), id_sequence_.end(), - old_instruction->unique_id()); + auto instruction_it = absl::c_find(instruction_sequence_, old_instruction); + auto id_it = absl::c_find(id_sequence_, old_instruction->unique_id()); CHECK(instruction_it != instruction_sequence_.end()) << "Do not find instruction id " << old_instruction->unique_id(); CHECK(id_it != id_sequence_.end());