From 929901d903b97aae1367d6fc6033487180b6bed0 Mon Sep 17 00:00:00 2001 From: Tushar Darote Date: Mon, 2 Feb 2026 12:08:16 +0530 Subject: [PATCH 01/25] cmake: Exclude subdirectories from all builds Add EXCLUDE_FROM_ALL to add_subdirectory() calls in multiple CMake modules to prevent unnecessary compilation of third-party dependencies during default builds. This reduces build time and resource usage. - eigen.cmake - farmhash.cmake - fft2d.cmake - flatbuffers.cmake - gemmlowp.cmake - neon2sse.cmake - ruy.cmake Signed-off-by: Tushar Darote --- tensorflow/lite/tools/cmake/modules/cpuinfo.cmake | 1 + tensorflow/lite/tools/cmake/modules/eigen.cmake | 2 +- tensorflow/lite/tools/cmake/modules/farmhash.cmake | 1 + tensorflow/lite/tools/cmake/modules/fft2d.cmake | 1 + tensorflow/lite/tools/cmake/modules/flatbuffers.cmake | 1 + tensorflow/lite/tools/cmake/modules/gemmlowp.cmake | 1 + tensorflow/lite/tools/cmake/modules/neon2sse.cmake | 1 + tensorflow/lite/tools/cmake/modules/ruy.cmake | 1 + 8 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake b/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake index 52c0f67f61d4f6..ef177ad5caf27a 100644 --- a/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake +++ b/tensorflow/lite/tools/cmake/modules/cpuinfo.cmake @@ -41,4 +41,5 @@ set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "Disable cpuinfo micro-benchmarks") add_subdirectory( "${cpuinfo_SOURCE_DIR}" "${cpuinfo_BINARY_DIR}" + EXCLUDE_FROM_ALL ) diff --git a/tensorflow/lite/tools/cmake/modules/eigen.cmake b/tensorflow/lite/tools/cmake/modules/eigen.cmake index f03ae364d385e7..246ccb59e097e5 100644 --- a/tensorflow/lite/tools/cmake/modules/eigen.cmake +++ b/tensorflow/lite/tools/cmake/modules/eigen.cmake @@ -99,4 +99,4 @@ set(EIGEN_TEST_SYCL OFF CACHE BOOL "Disable Sycl test") set(EIGEN_SYCL_TRISYCL OFF CACHE BOOL "Disable triSYCL test") # Make sure only MPL2.0 or more permissively licensed code is included. add_compile_definitions(EIGEN_MPL2_ONLY) -add_subdirectory("${eigen_SOURCE_DIR}" "${eigen_BINARY_DIR}") +add_subdirectory("${eigen_SOURCE_DIR}" "${eigen_BINARY_DIR}" EXCLUDE_FROM_ALL) diff --git a/tensorflow/lite/tools/cmake/modules/farmhash.cmake b/tensorflow/lite/tools/cmake/modules/farmhash.cmake index 9c8e039e6feb24..7a5160fa7ebbac 100644 --- a/tensorflow/lite/tools/cmake/modules/farmhash.cmake +++ b/tensorflow/lite/tools/cmake/modules/farmhash.cmake @@ -44,4 +44,5 @@ set(FARMHASH_SOURCE_DIR "${farmhash_SOURCE_DIR}" CACHE PATH add_subdirectory( "${CMAKE_CURRENT_LIST_DIR}/farmhash" "${farmhash_BINARY_DIR}" + EXCLUDE_FROM_ALL ) diff --git a/tensorflow/lite/tools/cmake/modules/fft2d.cmake b/tensorflow/lite/tools/cmake/modules/fft2d.cmake index b4169104faee05..9d44a2c686530e 100644 --- a/tensorflow/lite/tools/cmake/modules/fft2d.cmake +++ b/tensorflow/lite/tools/cmake/modules/fft2d.cmake @@ -37,4 +37,5 @@ set(FFT2D_SOURCE_DIR "${fft2d_SOURCE_DIR}" CACHE PATH "fft2d source") add_subdirectory( "${CMAKE_CURRENT_LIST_DIR}/fft2d" "${fft2d_BINARY_DIR}" + EXCLUDE_FROM_ALL ) diff --git a/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake b/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake index ec0950f5e4dc5c..6133f2e596bf6b 100644 --- a/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake +++ b/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake @@ -44,6 +44,7 @@ add_definitions(-DNOMINMAX=1) add_subdirectory( "${flatbuffers_SOURCE_DIR}" "${flatbuffers_BINARY_DIR}" + EXCLUDE_FROM_ALL ) remove_definitions(-DNOMINMAX) diff --git a/tensorflow/lite/tools/cmake/modules/gemmlowp.cmake b/tensorflow/lite/tools/cmake/modules/gemmlowp.cmake index 76d9705475b05b..9c8e2c90c7cb1a 100644 --- a/tensorflow/lite/tools/cmake/modules/gemmlowp.cmake +++ b/tensorflow/lite/tools/cmake/modules/gemmlowp.cmake @@ -48,6 +48,7 @@ set(GEMMLOWP_SOURCE_DIR "${gemmlowp_SOURCE_DIR}" CACHE PATH "Source directory") add_subdirectory( "${gemmlowp_SOURCE_DIR}/contrib" "${gemmlowp_BINARY_DIR}" + EXCLUDE_FROM_ALL ) set(BUILD_TESTING ${BUILD_TESTING_TMP}) diff --git a/tensorflow/lite/tools/cmake/modules/neon2sse.cmake b/tensorflow/lite/tools/cmake/modules/neon2sse.cmake index c2612e34bd2e25..77bcf1acba3dbf 100644 --- a/tensorflow/lite/tools/cmake/modules/neon2sse.cmake +++ b/tensorflow/lite/tools/cmake/modules/neon2sse.cmake @@ -40,4 +40,5 @@ endif() add_subdirectory( "${neon2sse_SOURCE_DIR}" "${neon2sse_BINARY_DIR}" + EXCLUDE_FROM_ALL ) diff --git a/tensorflow/lite/tools/cmake/modules/ruy.cmake b/tensorflow/lite/tools/cmake/modules/ruy.cmake index 1c5965f94f501d..a6632d3247f4e5 100644 --- a/tensorflow/lite/tools/cmake/modules/ruy.cmake +++ b/tensorflow/lite/tools/cmake/modules/ruy.cmake @@ -37,4 +37,5 @@ set(RUY_SOURCE_DIR "${ruy_SOURCE_DIR}" CACHE PATH "RUY source directory") add_subdirectory( "${ruy_SOURCE_DIR}" "${ruy_BINARY_DIR}" + EXCLUDE_FROM_ALL ) From b31faa1567a8b4d62f83b1355db2aa3fa0311836 Mon Sep 17 00:00:00 2001 From: VibhorGautam Date: Mon, 9 Mar 2026 16:03:25 +0530 Subject: [PATCH 02/25] Document shuffle + zip interaction in tf.data.Dataset The shuffle() and zip() docstrings do not mention that zipping a shuffled dataset with an unshuffled one breaks element correspondence. Users hit this and think zip is re-triggering the shuffle. Added a section to shuffle() explaining the issue with a code example showing the two correct approaches (shuffle after zip, or same seed). Added a note to zip() pointing users to the shuffle docs. Fixes #70521 --- tensorflow/python/data/ops/dataset_ops.py | 26 +++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index 06ed4275dcbb2e..06db24abb7c1d3 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -1061,6 +1061,12 @@ def zip(*args, datasets=None, name=None) -> "DatasetV2": >>> [(i.item(), j.item()) for i, j in ds.as_numpy_iterator()] [(1, 13), (2, 14)] + Note: `zip` iterates its input datasets in parallel. If one dataset has + been shuffled and another has not, the element correspondence from the + original order will be lost. To keep elements aligned after shuffling, + apply `shuffle` after `zip` rather than before it. See + `tf.data.Dataset.shuffle` for more details. + Args: *args: Datasets or nested structures of datasets to zip together. This can't be set if `datasets` is set. @@ -1491,6 +1497,26 @@ def shuffle( # [18, 4, 9, 2, 17, 8, 5, 10, 0, 6, 16, 3, 19, 7, 14, 11, 15, 13, 12, 1] ``` + #### Using shuffle with zip + + When you `shuffle` a dataset and then `zip` it with an unshuffled dataset, + the element pairing will not match the original order because each input + dataset is iterated independently. If you need to shuffle multiple datasets + while preserving the correspondence between their elements, either shuffle + after zipping or use the same `seed` on all datasets: + + ```python + # Correct: shuffle after zipping to keep pairs aligned. + a = tf.data.Dataset.range(3) + b = tf.data.Dataset.range(3) + dataset = tf.data.Dataset.zip(a, b).shuffle(3) + + # Also correct: same seed + buffer_size on both datasets. + a = tf.data.Dataset.range(3).shuffle(3, seed=42) + b = tf.data.Dataset.range(3).shuffle(3, seed=42) + dataset = tf.data.Dataset.zip(a, b) + ``` + Args: buffer_size: An int or `tf.int64` scalar `tf.Tensor`, representing the number of elements from this dataset from which the new dataset will From 4ee6562a8a0b281c453cedd3e39e0c6092a57ef9 Mon Sep 17 00:00:00 2001 From: VibhorGautam Date: Sat, 4 Jul 2026 18:07:33 +0530 Subject: [PATCH 03/25] Remove fragile seed-matching example, recommend shuffle-after-zip only --- tensorflow/python/data/ops/dataset_ops.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index 06db24abb7c1d3..d886bbba9c267a 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -1502,19 +1502,14 @@ def shuffle( When you `shuffle` a dataset and then `zip` it with an unshuffled dataset, the element pairing will not match the original order because each input dataset is iterated independently. If you need to shuffle multiple datasets - while preserving the correspondence between their elements, either shuffle - after zipping or use the same `seed` on all datasets: + while preserving the correspondence between their elements, you should + shuffle after zipping: ```python # Correct: shuffle after zipping to keep pairs aligned. a = tf.data.Dataset.range(3) b = tf.data.Dataset.range(3) dataset = tf.data.Dataset.zip(a, b).shuffle(3) - - # Also correct: same seed + buffer_size on both datasets. - a = tf.data.Dataset.range(3).shuffle(3, seed=42) - b = tf.data.Dataset.range(3).shuffle(3, seed=42) - dataset = tf.data.Dataset.zip(a, b) ``` Args: From ed8d4c129b322ad7b1af4255e08b02413cad3858 Mon Sep 17 00:00:00 2001 From: deeven-seru Date: Fri, 24 Jul 2026 10:02:11 +0530 Subject: [PATCH 04/25] Fix signed-to-unsigned conversion bug in DatasetRandomAccessCache::Get --- .../core/kernels/data/cache_dataset_ops.cc | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc index 4ef46199b019fb..dc33b529b3add3 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc @@ -94,14 +94,14 @@ class DatasetRandomAccessCache { GetIteratorResourceFromDataset(ctx, input_)); TF_RETURN_IF_ERROR(iter_resource_->SetIteratorFromDataset(ctx, input_)); } - if (index >= cache_.size()) { - TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx)); - } if (index < 0) { return absl::InvalidArgumentError( absl::StrCat("Expected index >= 0; Received index: ", index)); } - *out_tensors = cache_.at(index); + if (static_cast(index) >= cache_.size()) { + TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx)); + } + *out_tensors = cache_.at(static_cast(index)); return absl::OkStatus(); } @@ -111,7 +111,7 @@ class DatasetRandomAccessCache { private: absl::Status ExtendTempCacheToIndex(int64_t index, OpKernelContext* ctx) { bool end_of_sequence; - while (cache_.size() <= index) { + while (cache_.size() <= static_cast(index)) { std::vector out_tensors; TF_RETURN_IF_ERROR( iter_resource_->GetNext(ctx, &out_tensors, &end_of_sequence)); @@ -151,18 +151,24 @@ class IteratorRandomAccessCache { explicit IteratorRandomAccessCache(const DatasetBase* input) : input_(input) {} - absl::Status Get(AnyContext ctx, size_t element_position, + absl::Status Get(AnyContext ctx, int64_t element_position, std::vector* out_tensors) { - if (element_position < cache_.size() && !cache_[element_position].empty()) { - *out_tensors = cache_[element_position]; + if (element_position < 0) { + return absl::InvalidArgumentError( + absl::StrCat("Element position must be non-negative; Received: ", + element_position)); + } + + if (static_cast(element_position) < cache_.size() && !cache_[static_cast(element_position)].empty()) { + *out_tensors = cache_[static_cast(element_position)]; return absl::OkStatus(); } TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors)); - if (element_position >= cache_.size()) { - cache_.resize(element_position + 1); + if (static_cast(element_position) >= cache_.size()) { + cache_.resize(static_cast(element_position) + 1); } - cache_[element_position] = *out_tensors; + cache_[static_cast(element_position)] = *out_tensors; return absl::OkStatus(); } From 3324146a29e555b679700efb83daabc353e7730b Mon Sep 17 00:00:00 2001 From: deeven-seru Date: Sat, 25 Jul 2026 15:59:02 +0530 Subject: [PATCH 05/25] Fix 32-bit index truncation issues per reviewer feedback --- .../core/kernels/data/cache_dataset_ops.cc | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc index dc33b529b3add3..624b2979f86684 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc @@ -98,10 +98,10 @@ class DatasetRandomAccessCache { return absl::InvalidArgumentError( absl::StrCat("Expected index >= 0; Received index: ", index)); } - if (static_cast(index) >= cache_.size()) { + if (index >= static_cast(cache_.size())) { TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx)); } - *out_tensors = cache_.at(static_cast(index)); + *out_tensors = cache_.at(index); return absl::OkStatus(); } @@ -111,7 +111,7 @@ class DatasetRandomAccessCache { private: absl::Status ExtendTempCacheToIndex(int64_t index, OpKernelContext* ctx) { bool end_of_sequence; - while (cache_.size() <= static_cast(index)) { + while (static_cast(cache_.size()) <= index) { std::vector out_tensors; TF_RETURN_IF_ERROR( iter_resource_->GetNext(ctx, &out_tensors, &end_of_sequence)); @@ -159,16 +159,16 @@ class IteratorRandomAccessCache { element_position)); } - if (static_cast(element_position) < cache_.size() && !cache_[static_cast(element_position)].empty()) { - *out_tensors = cache_[static_cast(element_position)]; + if (element_position < static_cast(cache_.size()) && !cache_[element_position].empty()) { + *out_tensors = cache_[element_position]; return absl::OkStatus(); } TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors)); - if (static_cast(element_position) >= cache_.size()) { - cache_.resize(static_cast(element_position) + 1); + if (element_position >= static_cast(cache_.size())) { + cache_.resize(element_position + 1); } - cache_[static_cast(element_position)] = *out_tensors; + cache_[element_position] = *out_tensors; return absl::OkStatus(); } From c46fe7ac35c478f070865bce1cafae9934be9895 Mon Sep 17 00:00:00 2001 From: Tirth Date: Wed, 29 Jul 2026 19:59:56 +0530 Subject: [PATCH 06/25] [tf.data] Validate `buffer_size` in shuffle to prevent process crash `_ShuffleDataset.__init__` converted `buffer_size` straight to a tensor with no upper-bound sanity check. The C++ kernel eagerly allocates a slot for every element up to `buffer_size` when the iterator is created (`std::vector>(buffer_size_)` in shuffle_dataset_op.cc), so a pathologically large value reaches that allocation and crashes the process instead of raising a catchable error. Add a sanity check in the Python wrapper that rejects `buffer_size` values above ~1 billion elements with a `ValueError` before they reach the op, and add a regression test covering the reported repro. Fixes #113167 --- .../python/data/kernel_tests/shuffle_test.py | 8 ++++++++ tensorflow/python/data/ops/shuffle_op.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/tensorflow/python/data/kernel_tests/shuffle_test.py b/tensorflow/python/data/kernel_tests/shuffle_test.py index 99266134f0f79d..3304993223ebae 100644 --- a/tensorflow/python/data/kernel_tests/shuffle_test.py +++ b/tensorflow/python/data/kernel_tests/shuffle_test.py @@ -178,6 +178,14 @@ def testDefaultArguments(self): for i in range(5): self.assertEqual(10, counts[i]) + @combinations.generate(test_base.default_test_combinations()) + def testExcessiveBufferSize(self): + dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5]) + with self.assertRaisesRegex(ValueError, "buffer_size"): + dataset.shuffle(buffer_size=sys.maxsize) + with self.assertRaisesRegex(ValueError, "buffer_size"): + dataset.shuffle(buffer_size=2**31 - 1) + @combinations.generate( combinations.times( test_base.default_test_combinations(), diff --git a/tensorflow/python/data/ops/shuffle_op.py b/tensorflow/python/data/ops/shuffle_op.py index 474e146eeffb3f..be0ecd96103538 100644 --- a/tensorflow/python/data/ops/shuffle_op.py +++ b/tensorflow/python/data/ops/shuffle_op.py @@ -21,6 +21,13 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import gen_dataset_ops +# Sanity limit on the number of elements in the shuffle buffer. The C++ +# kernel eagerly allocates a slot for every element up to `buffer_size` when +# the iterator is created, so a pathologically large value (e.g. +# `sys.maxsize`) reaches that allocation and crashes the process instead of +# raising a catchable error. +_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS = 1 << 30 # ~1 billion elements + def _shuffle( # pylint: disable=unused-private-name input_dataset, @@ -45,6 +52,14 @@ def __init__( name=None, ): """See `Dataset.shuffle()` for details.""" + if (isinstance(buffer_size, int) and + buffer_size > _MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS): + raise ValueError( + f"`buffer_size` must not exceed " + f"{_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS} elements, but got " + f"{buffer_size}. Requesting a shuffle buffer this large would " + "cause the dataset to abort the process instead of raising a " + "catchable error.") self._input_dataset = input_dataset self._buffer_size = ops.convert_to_tensor( buffer_size, dtype=dtypes.int64, name="buffer_size") From 66bf27947ee2edb6524b44227fddbb9b82c4c99e Mon Sep 17 00:00:00 2001 From: Tirth Date: Fri, 31 Jul 2026 23:07:09 +0530 Subject: [PATCH 07/25] Address review: validate buffer_size via constant_value, not isinstance isinstance(buffer_size, int) misses NumPy integers and constant Tensor inputs, both of which bypass the check and can still reach the crashing C++ allocation. Use tensor_util.constant_value on the converted tensor instead, and add tests for np.int64 and constant Tensor buffer_size values. --- .../python/data/kernel_tests/shuffle_test.py | 5 +++++ tensorflow/python/data/ops/shuffle_op.py | 18 ++++++++++-------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/tensorflow/python/data/kernel_tests/shuffle_test.py b/tensorflow/python/data/kernel_tests/shuffle_test.py index 3304993223ebae..97e4342e6c7d21 100644 --- a/tensorflow/python/data/kernel_tests/shuffle_test.py +++ b/tensorflow/python/data/kernel_tests/shuffle_test.py @@ -185,6 +185,11 @@ def testExcessiveBufferSize(self): dataset.shuffle(buffer_size=sys.maxsize) with self.assertRaisesRegex(ValueError, "buffer_size"): dataset.shuffle(buffer_size=2**31 - 1) + with self.assertRaisesRegex(ValueError, "buffer_size"): + dataset.shuffle(buffer_size=np.int64(2**31 - 1)) + with self.assertRaisesRegex(ValueError, "buffer_size"): + dataset.shuffle( + buffer_size=constant_op.constant(2**31 - 1, dtype=dtypes.int64)) @combinations.generate( combinations.times( diff --git a/tensorflow/python/data/ops/shuffle_op.py b/tensorflow/python/data/ops/shuffle_op.py index be0ecd96103538..5858ee98ad048a 100644 --- a/tensorflow/python/data/ops/shuffle_op.py +++ b/tensorflow/python/data/ops/shuffle_op.py @@ -19,6 +19,7 @@ from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.ops import gen_dataset_ops # Sanity limit on the number of elements in the shuffle buffer. The C++ @@ -52,17 +53,18 @@ def __init__( name=None, ): """See `Dataset.shuffle()` for details.""" - if (isinstance(buffer_size, int) and - buffer_size > _MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS): - raise ValueError( - f"`buffer_size` must not exceed " - f"{_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS} elements, but got " - f"{buffer_size}. Requesting a shuffle buffer this large would " - "cause the dataset to abort the process instead of raising a " - "catchable error.") self._input_dataset = input_dataset self._buffer_size = ops.convert_to_tensor( buffer_size, dtype=dtypes.int64, name="buffer_size") + constant_buffer_size = tensor_util.constant_value(self._buffer_size) + if (constant_buffer_size is not None and + constant_buffer_size > _MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS): + raise ValueError( + f"`buffer_size` must not exceed " + f"{_MAX_SHUFFLE_BUFFER_SIZE_ELEMENTS} elements, but got " + f"{constant_buffer_size}. Requesting a shuffle buffer this large " + "would cause the dataset to abort the process instead of raising " + "a catchable error.") self._seed, self._seed2 = random_seed.get_seed(seed) self._reshuffle_each_iteration = reshuffle_each_iteration self._name = name From 7d433840b201380fc3a3e8ad28036cf4df1242b0 Mon Sep 17 00:00:00 2001 From: abhijeet117 Date: Sun, 23 Aug 2026 23:17:27 +0530 Subject: [PATCH 08/25] Allow a Tensor buffer_size in Dataset.prefetch prefetch compared buffer_size with AUTOTUNE to set the legacy_autotune op attribute, but for a Tensor buffer_size that comparison produces a Tensor and the op construction fails with TypeError even though the docs allow an int64 scalar tf.Tensor. Decide from the statically known value of the converted tensor instead; symbolic tensors use the non-legacy path. --- .../python/data/kernel_tests/prefetch_test.py | 13 +++++++++++++ tensorflow/python/data/ops/prefetch_op.py | 10 +++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/data/kernel_tests/prefetch_test.py b/tensorflow/python/data/kernel_tests/prefetch_test.py index 7f075c079b1490..d9e18b6afa1bab 100644 --- a/tensorflow/python/data/kernel_tests/prefetch_test.py +++ b/tensorflow/python/data/kernel_tests/prefetch_test.py @@ -25,6 +25,7 @@ from tensorflow.python.data.ops import options as options_lib from tensorflow.python.data.ops import prefetch_op from tensorflow.python.framework import combinations +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.ops import script_ops @@ -40,6 +41,18 @@ def testBufferSize(self, buffer_size): dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size) self.assertDatasetProduces(dataset, expected_output=range(10)) + @combinations.generate(test_base.default_test_combinations()) + def testTensorBufferSize(self): + # `buffer_size` is documented as an int64 scalar `tf.Tensor`. + dataset = dataset_ops.Dataset.range(10).prefetch( + buffer_size=constant_op.constant(2, dtypes.int64)) + self.assertDatasetProduces(dataset, expected_output=range(10)) + + # A tensor holding the AUTOTUNE value keeps autotuning enabled. + dataset = dataset_ops.Dataset.range(10).prefetch( + buffer_size=constant_op.constant(-1, dtypes.int64)) + self.assertDatasetProduces(dataset, expected_output=range(10)) + @combinations.generate( combinations.times(test_base.eager_only_combinations(), combinations.combine(buffer_size=[0, 1, 2, 42]))) diff --git a/tensorflow/python/data/ops/prefetch_op.py b/tensorflow/python/data/ops/prefetch_op.py index 49ae904d83d8aa..cecbf21ea90038 100644 --- a/tensorflow/python/data/ops/prefetch_op.py +++ b/tensorflow/python/data/ops/prefetch_op.py @@ -18,6 +18,7 @@ from tensorflow.python.data.ops import debug_mode from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.ops import gen_dataset_ops @@ -39,6 +40,13 @@ def __init__(self, input_dataset, buffer_size, slack_period=None, name=None): self._buffer_size = ops.convert_to_tensor( buffer_size, dtype=dtypes.int64, name="buffer_size") self._name = name + # `legacy_autotune` must be a Python bool, so decide from the statically + # known value of the converted `buffer_size`; it is None for symbolic + # tensors, which then use the non-legacy path. + buffer_size_constant = tensor_util.constant_value(self._buffer_size) + legacy_autotune = ( + buffer_size_constant is not None + and int(buffer_size_constant) == dataset_ops.AUTOTUNE) # pylint: disable=protected-access # We colocate the prefetch dataset with its input as this collocation only # happens automatically in graph mode. @@ -47,6 +55,6 @@ def __init__(self, input_dataset, buffer_size, slack_period=None, name=None): input_dataset._variant_tensor, buffer_size=self._buffer_size, slack_period=slack_period, - legacy_autotune=(buffer_size == dataset_ops.AUTOTUNE), + legacy_autotune=legacy_autotune, **self._common_args) super().__init__(input_dataset, variant_tensor) From 2bbfb12a2b3441b728c83102d3673e87232446f7 Mon Sep 17 00:00:00 2001 From: lakshit verma Date: Thu, 27 Aug 2026 23:36:27 +0530 Subject: [PATCH 09/25] Reject non-empty tensors missing tensor_content in ParseFast Enforce shape.num_elements() == 0 when tensor_content is omitted in ParseTensorSubmessage to prevent allocating uninitialized heap memory. --- .../core/distributed_runtime/tensor_coding.cc | 3 +++ .../distributed_runtime/tensor_coding_test.cc | 17 +++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/tensorflow/core/distributed_runtime/tensor_coding.cc b/tensorflow/core/distributed_runtime/tensor_coding.cc index e26e3653a22a7a..3f0caa745f08ad 100644 --- a/tensorflow/core/distributed_runtime/tensor_coding.cc +++ b/tensorflow/core/distributed_runtime/tensor_coding.cc @@ -167,6 +167,9 @@ bool TensorResponse::ParseTensorSubmessage( .ok()) { return false; } + if (shape.num_elements() != 0) { + return false; + } Tensor t(allocator_, tensor_meta->dtype(), shape); tensor_ = std::move(t); } diff --git a/tensorflow/core/distributed_runtime/tensor_coding_test.cc b/tensorflow/core/distributed_runtime/tensor_coding_test.cc index 1cb8817bbf8318..0aa6f82119520b 100644 --- a/tensorflow/core/distributed_runtime/tensor_coding_test.cc +++ b/tensorflow/core/distributed_runtime/tensor_coding_test.cc @@ -180,6 +180,23 @@ TEST_F(TensorResponseTest, InitPartialOverflow) { EXPECT_TRUE(absl::IsInvalidArgument(s)); } +TEST_F(TensorResponseTest, NonEmptyTensorMissingContentRejected) { + RecvTensorResponse proto; + proto.set_is_dead(false); + proto.set_send_start_micros(123456); + TensorProto* tensor_proto = proto.mutable_tensor(); + tensor_proto->set_dtype(DT_FLOAT); + tensor_proto->mutable_tensor_shape()->add_dim()->set_size(10); + + std::string encoded; + proto.AppendToString(&encoded); + StringSource source(&encoded, 1024); + TensorResponse response; + DummyDevice cpu_device(Env::Default()); + response.InitAlloc(&cpu_device, AllocatorAttributes()); + EXPECT_FALSE(response.ParseFrom(&source).ok()); +} + std::string MakeFloatTensorTestCase(int num_elems) { std::vector v(num_elems); for (int i = 0; i < num_elems; i++) { From 8aebb14cc55fb3298c691ea1af2cf128d27c49e2 Mon Sep 17 00:00:00 2001 From: lakshit verma Date: Fri, 28 Aug 2026 13:36:23 +0530 Subject: [PATCH 10/25] Update tensor_coding test assertions for ParseSlow fallback and amplification --- .../distributed_runtime/tensor_coding_test.cc | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/distributed_runtime/tensor_coding_test.cc b/tensorflow/core/distributed_runtime/tensor_coding_test.cc index 0aa6f82119520b..423a27c8caf10e 100644 --- a/tensorflow/core/distributed_runtime/tensor_coding_test.cc +++ b/tensorflow/core/distributed_runtime/tensor_coding_test.cc @@ -180,7 +180,25 @@ TEST_F(TensorResponseTest, InitPartialOverflow) { EXPECT_TRUE(absl::IsInvalidArgument(s)); } -TEST_F(TensorResponseTest, NonEmptyTensorMissingContentRejected) { +TEST_F(TensorResponseTest, ZeroLengthTensorMissingContentAccepted) { + RecvTensorResponse proto; + proto.set_is_dead(false); + proto.set_send_start_micros(123456); + TensorProto* tensor_proto = proto.mutable_tensor(); + tensor_proto->set_dtype(DT_FLOAT); + tensor_proto->mutable_tensor_shape()->add_dim()->set_size(0); + + std::string encoded; + proto.AppendToString(&encoded); + StringSource source(&encoded, 1024); + TensorResponse response; + DummyDevice cpu_device(Env::Default()); + response.InitAlloc(&cpu_device, AllocatorAttributes()); + EXPECT_TRUE(response.ParseFrom(&source).ok()); + EXPECT_EQ(response.tensor().NumElements(), 0); +} + +TEST_F(TensorResponseTest, NonEmptyTensorMissingContentZeroInitializedViaSlowPath) { RecvTensorResponse proto; proto.set_is_dead(false); proto.set_send_start_micros(123456); @@ -188,6 +206,28 @@ TEST_F(TensorResponseTest, NonEmptyTensorMissingContentRejected) { tensor_proto->set_dtype(DT_FLOAT); tensor_proto->mutable_tensor_shape()->add_dim()->set_size(10); + std::string encoded; + proto.AppendToString(&encoded); + StringSource source(&encoded, 1024); + TensorResponse response; + DummyDevice cpu_device(Env::Default()); + response.InitAlloc(&cpu_device, AllocatorAttributes()); + EXPECT_TRUE(response.ParseFrom(&source).ok()); + EXPECT_EQ(response.tensor().NumElements(), 10); + for (int i = 0; i < 10; ++i) { + EXPECT_EQ(response.tensor().flat()(i), 0.0f); + } +} + +TEST_F(TensorResponseTest, AmplificationProtoWithoutContentRejected) { + RecvTensorResponse proto; + proto.set_is_dead(false); + proto.set_send_start_micros(123456); + TensorProto* tensor_proto = proto.mutable_tensor(); + tensor_proto->set_dtype(DT_FLOAT); + // 1 billion floats = 4GB (> 2GB safe limit) with no tensor_content + tensor_proto->mutable_tensor_shape()->add_dim()->set_size(1000000000); + std::string encoded; proto.AppendToString(&encoded); StringSource source(&encoded, 1024); From 5121c32d2b595011d519a857e6acc872c1b8750f Mon Sep 17 00:00:00 2001 From: Shyamli Agrawal Date: Mon, 31 Aug 2026 11:45:12 -0700 Subject: [PATCH 11/25] Use the new XLA GPU autotune cache format by default. PiperOrigin-RevId: 973997845 --- third_party/xla/xla/debug_options_flags.cc | 2 +- .../xla/xla/tools/xla_gpu_compile_lib_test.cc | 46 ------------------- 2 files changed, 1 insertion(+), 47 deletions(-) diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 66a1e7eb3767c5..66c478dd3879da 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -479,7 +479,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_per_fusion_autotune_cache_dir(""); - opts.set_xla_gpu_use_new_autotune_cache_format(false); + opts.set_xla_gpu_use_new_autotune_cache_format(true); opts.set_xla_compile_all_supported_configs(false); diff --git a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc index c850c02df1b3f9..a574b33fefa67b 100644 --- a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc +++ b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc @@ -60,29 +60,6 @@ class XlaCompileLibTest : public HloTestBase { } std::unique_ptr module_; - - void ComputeAutotuneResults(AutotuneResults& results) { - static constexpr absl::string_view kHloText = R"( -HloModule t -ENTRY e { - p0 = f16[1,16,17,3] parameter(0) - p1 = f16[16,17,3] parameter(1) - ROOT _ = f16[1,16,16] dot(p0, p1), - lhs_contracting_dims={2,3}, rhs_contracting_dims={1,2} -})"; - - HloModuleConfig config = GetModuleConfigForTest(); - DebugOptions opts = config.debug_options(); - opts.set_xla_gpu_autotune_level(3); - config.set_debug_options(opts); - - gpu::AutotunerCache::ClearAutotuneResults(); - ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(kHloText, config)); - (void)CreateExecutable(std::move(module), /*run_hlo_passes=*/true); - - ASSERT_OK(gpu::AutotunerCache::SerializeAutotuneResults(&results)); - } }; TEST_F(XlaCompileLibTest, CompilesForGpuWithDevice) { @@ -193,34 +170,11 @@ TEST_F(XlaCompileLibTest, MainForGpu) { EXPECT_EQ(result.status().code(), tensorflow::error::OK); } -TEST_F(XlaCompileLibTest, LoadAutotuneDataGpuDataPresentAndAutotuningEnabled) { - gpu::AutotunerCache::ClearAutotuneResults(); - - HloModuleAndMetadata mod; - mod.hlo_module = std::move(module_); - auto data = std::make_unique(); - ComputeAutotuneResults(data->autotune_results.emplace()); - gpu::AutotunerCache::ClearAutotuneResults(); - mod.backend_specific_data = std::move(data); - - DebugOptions opts = mod.hlo_module->config().debug_options(); - opts.set_xla_gpu_autotune_level(3); - mod.hlo_module->mutable_config().set_debug_options(opts); - - EXPECT_THAT(internal::LoadAutotuneDataFromModule(&mod, BackendType::kGpu), - absl_testing::IsOkAndHolds(true)); - EXPECT_FALSE(gpu::AutotunerCache::ResultCacheIsEmpty()); -} - TEST_F(XlaCompileLibTest, LoadAutotuneDataGpuDataPresentAndAutotuningDisabled) { gpu::AutotunerCache::ClearAutotuneResults(); HloModuleAndMetadata mod; mod.hlo_module = std::move(module_); - auto data = std::make_unique(); - ComputeAutotuneResults(data->autotune_results.emplace()); - gpu::AutotunerCache::ClearAutotuneResults(); - mod.backend_specific_data = std::move(data); DebugOptions opts = mod.hlo_module->config().debug_options(); opts.set_xla_gpu_autotune_level(0); From ad3cdcf9a6adf2a3a41a89d647a68e497e073ba8 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Mon, 31 Aug 2026 11:50:58 -0700 Subject: [PATCH 12/25] Remove StreamExecutorGpuClient and update all references to use CommonPjRtClient instead. PiperOrigin-RevId: 974000504 --- third_party/xla/xla/pjrt/c/BUILD | 2 + .../xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc | 16 ++-- third_party/xla/xla/pjrt/gpu/BUILD | 5 + .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc | 6 +- .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h | 6 -- .../gpu/se_gpu_pjrt_client_multi_gpu_test.cc | 18 ++-- .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc | 18 ++-- .../xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc | 14 ++- .../pjrt/gpu/se_gpu_pjrt_compiler_aot_test.cc | 92 +++++++++---------- .../xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc | 11 ++- .../xla_gpu/xla_gpu_pjrt_client_test.cc | 6 +- 11 files changed, 103 insertions(+), 91 deletions(-) diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD index c0d66b6772925d..34b28d4728ce8a 100644 --- a/third_party/xla/xla/pjrt/c/BUILD +++ b/third_party/xla/xla/pjrt/c/BUILD @@ -832,6 +832,7 @@ xla_test( "//xla/pjrt:pjrt_compiler", "//xla/pjrt/distributed:in_memory_key_value_store", "//xla/pjrt/gpu:se_gpu_pjrt_client", + "//xla/pjrt/se:pjrt_stream_executor_client", "//xla/service:custom_call_target_registry", "//xla/service:device_assignment", "//xla/stream_executor/gpu:gpu_init", @@ -839,6 +840,7 @@ xla_test( "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base", "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc index e58556a18ede63..1ddc92bf54458d 100644 --- a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc +++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc @@ -31,6 +31,7 @@ limitations under the License. #include #include #include "absl/algorithm/container.h" +#include "absl/base/casts.h" #include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/log/check.h" @@ -65,6 +66,7 @@ limitations under the License. #include "xla/pjrt/gpu/se_gpu_pjrt_client.h" #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/se/pjrt_stream_executor_client.h" #include "xla/service/custom_call_target_registry.h" #include "xla/service/device_assignment.h" #include "xla/shape.h" @@ -968,10 +970,11 @@ TEST(PjrtCApiGpuExtensionTest, PJRT_Error* error = api->PJRT_Client_Create(&create_arg); EXPECT_EQ(error, nullptr) << GetErrorMessage(error, api); - xla::PjRtClient* cpp_client = create_arg.client->client.get(); - auto* gpu_client = absl::down_cast(cpp_client); + auto* gpu_client = absl::down_cast( + absl::down_cast(create_arg.client->client.get()) + ->raw_client()); std::vector data(4, 0.0f); - EXPECT_TRUE(gpu_client->raw_client()->ShouldStageHostToDeviceTransfers( + EXPECT_TRUE(gpu_client->ShouldStageHostToDeviceTransfers( data.data(), sizeof(float) * data.size())); PJRT_Client_Destroy_Args destroy_args; @@ -1008,10 +1011,11 @@ TEST(PjrtCApiGpuExtensionTest, PJRT_Error* error = api->PJRT_Client_Create(&create_arg); EXPECT_EQ(error, nullptr) << GetErrorMessage(error, api); - xla::PjRtClient* cpp_client = create_arg.client->client.get(); - auto* gpu_client = absl::down_cast(cpp_client); + auto* gpu_client = absl::down_cast( + absl::down_cast(create_arg.client->client.get()) + ->raw_client()); std::vector data(4, 0.0f); - EXPECT_FALSE(gpu_client->raw_client()->ShouldStageHostToDeviceTransfers( + EXPECT_FALSE(gpu_client->ShouldStageHostToDeviceTransfers( data.data(), sizeof(float) * data.size())); PJRT_Client_Destroy_Args destroy_args; diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD index 4b331c67bc86ca..91833397fdf4a8 100644 --- a/third_party/xla/xla/pjrt/gpu/BUILD +++ b/third_party/xla/xla/pjrt/gpu/BUILD @@ -743,6 +743,7 @@ cc_library( "//xla/hlo/builder:xla_computation", "//xla/hlo/ir:hlo", "//xla/mlir_hlo:mhlo_passes", + "//xla/pjrt:common_pjrt_client", "//xla/pjrt:layout_mode", "//xla/pjrt:maybe_owning_mlir_module", "//xla/pjrt:mlir_to_hlo", @@ -752,6 +753,7 @@ cc_library( "//xla/pjrt:pjrt_compiler", "//xla/pjrt:pjrt_executable", "//xla/pjrt:utils", + "//xla/pjrt/se:pjrt_stream_executor_client", "//xla/pjrt/se:stream_executor_executable", "//xla/service:compiled_module", "//xla/service:compiler", @@ -906,6 +908,7 @@ xla_test( "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:test", "//xla/mlir_hlo", + "//xla/pjrt:common_pjrt_client", "//xla/pjrt:maybe_owning_mlir_module", "//xla/pjrt:mock_pjrt_client", "//xla/pjrt:pjrt_abi_version", @@ -924,6 +927,7 @@ xla_test( "//xla/stream_executor/cuda:cuda_platform_id", "//xla/tests:literal_test_util", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/base", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:status_matchers", @@ -949,6 +953,7 @@ xla_test( "//xla:literal_util", "//xla:shape_util", "//xla:xla_data_proto_cc", + "//xla/client:local_client", "//xla/hlo/builder:xla_computation", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc index 27333f76cdc530..ca8451d79ad792 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc @@ -672,7 +672,7 @@ void StreamExecutorGpuRawClient::ScheduleTransfersOnLocalDevice( tsl::profiler::TraceMe trace([&] { return tsl::profiler::TraceMeEncode( absl::StrFormat( - "[%v] StreamExecutorGpuClient::ScheduleTransfersOnLocalDevice", + "[%v] StreamExecutorGpuRawClient::ScheduleTransfersOnLocalDevice", local_device_state->local_device_id()), {{"num_buffers", transfer_specs.size()}}); }); @@ -1887,7 +1887,7 @@ const int StreamExecutorGpuHbmMemorySpace::kKindId = []() { return static_cast(kind_id); }(); -std::unique_ptr MakeStreamExecutorGpuClient( +std::unique_ptr MakeStreamExecutorGpuClient( std::string platform_name, std::vector> devices, int process_index, std::unique_ptr raw_client, @@ -1900,7 +1900,7 @@ std::unique_ptr MakeStreamExecutorGpuClient( attrs.pjrt_c_api_minor_version = 0; attrs.attributes["serialize_with_sdy"] = true; attrs.attributes["supports_cross_host_transfers"] = PjRtValueType(true); - auto result = std::make_unique( + auto result = std::make_unique( tsl::Fingerprint64(platform_name), platform_name, platform_version, process_index, std::move(topology), std::move(raw_client), std::move(kv_store), std::move(attrs)); diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h index 8339e25e5c8271..e79b15e8b6ed91 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h @@ -205,12 +205,6 @@ class StreamExecutorGpuRawClient : public PjRtStreamExecutorRawClient { std::shared_ptr memory_registration_; }; -// A custom PjRtClient that overrides the device assignment method. -class StreamExecutorGpuClient : public xla::PjRtStreamExecutorClient { - public: - using PjRtStreamExecutorClient::PjRtStreamExecutorClient; -}; - absl::StatusOr> GetStreamExecutorGpuClient( const GpuClientOptions& options); diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc index 2801598a403a7a..b3c4969501fa5f 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc @@ -312,10 +312,10 @@ TEST(StreamExecutorGpuClientTest, return; } std::unique_ptr& client = *client_status; - auto* gpu_client = - absl::down_cast(client.get()); const gpu::GpuExecutableRunOptions* run_options = - gpu_client->gpu_run_options(); + absl::down_cast( + absl::down_cast(client.get())->raw_client()) + ->gpu_run_options(); if (run_options == nullptr || !run_options->execution_timeout_handler()) { statuses[i] = absl::InternalError( @@ -345,9 +345,10 @@ TEST(StreamExecutorGpuClientTest, options.abort_collectives_on_failure = true; ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options)); - auto* gpu_client = absl::down_cast(client.get()); const gpu::GpuExecutableRunOptions* run_options = - gpu_client->gpu_run_options(); + absl::down_cast( + absl::down_cast(client.get())->raw_client()) + ->gpu_run_options(); ASSERT_NE(run_options, nullptr); ASSERT_TRUE(run_options->execution_timeout_handler()); @@ -414,10 +415,11 @@ TEST(StreamExecutorGpuClientTest, ASSERT_OK(status); } - auto* gpu_client0 = - absl::down_cast(pjrt_clients[0].get()); const gpu::GpuExecutableRunOptions* run_options = - gpu_client0->gpu_run_options(); + absl::down_cast( + absl::down_cast(pjrt_clients[0].get()) + ->raw_client()) + ->gpu_run_options(); ASSERT_NE(run_options, nullptr); ASSERT_TRUE(run_options->execution_timeout_handler()); diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index c6fb70e01c7bd3..593f3cb71a3733 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -1487,11 +1487,10 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToTrue) { std::vector data(1024, 1.0f); Shape shape = ShapeUtil::MakeShape(F32, {1024}); - // TODO(b/b/482307468) Switch to absl::down_cast after upgrade. - [[deprecated("remove after absl upgrade")]] auto* staging_client = - absl::down_cast(client_staging.get()); + auto* staging_client = absl::down_cast( + absl::down_cast(client_staging.get())->raw_client()); - EXPECT_TRUE(staging_client->raw_client()->ShouldStageHostToDeviceTransfers( + EXPECT_TRUE(staging_client->ShouldStageHostToDeviceTransfers( data.data(), sizeof(float) * data.size())); TF_ASSERT_OK_AND_ASSIGN( @@ -1518,13 +1517,12 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToFalse) { std::vector data(1024, 1.0f); Shape shape = ShapeUtil::MakeShape(F32, {1024}); - // TODO(b/b/482307468) Switch to absl::down_cast after upgrade. - [[deprecated("remove after absl upgrade")]] auto* no_staging_client = - absl::down_cast(client_no_staging.get()); + auto* no_staging_client = absl::down_cast( + absl::down_cast(client_no_staging.get()) + ->raw_client()); - EXPECT_FALSE( - no_staging_client->raw_client()->ShouldStageHostToDeviceTransfers( - data.data(), sizeof(float) * data.size())); + EXPECT_FALSE(no_staging_client->ShouldStageHostToDeviceTransfers( + data.data(), sizeof(float) * data.size())); TF_ASSERT_OK_AND_ASSIGN( auto buffer, diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc index ef50ea82d194ad..8fe78aaaafb752 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc @@ -38,6 +38,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/layout_util.h" #include "xla/mlir_hlo/mhlo/transforms/passes.h" +#include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/gpu/se_gpu_pjrt_client.h" #include "xla/pjrt/gpu/se_gpu_pjrt_runtime_abi_version.h" #include "xla/pjrt/gpu/se_gpu_topology_description.h" @@ -49,6 +50,7 @@ limitations under the License. #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/se/pjrt_stream_executor_client.h" #include "xla/pjrt/se/stream_executor_executable.h" #include "xla/pjrt/utils.h" #include "xla/primitive_util.h" @@ -110,14 +112,18 @@ absl::StatusOr> GetCompilerForPlatform( absl::StatusOr GetStreamExecutor( PjRtClient* client) { - const StreamExecutorGpuClient* gpu_client = - dynamic_cast(client); + const auto* gpu_client = dynamic_cast(client); if (gpu_client != nullptr) { - return gpu_client->client()->backend().default_stream_executor(); + if (const auto* raw_gpu_client = + dynamic_cast( + gpu_client->raw_client())) { + return raw_gpu_client->client()->backend().default_stream_executor(); + } } return absl::InvalidArgumentError( - "Given PjRtClient is not a StreamExecutorGpuClient."); + "Given PjRtClient does not contain a xla::LocalClient needed for " + "autotuning."); } } // namespace diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_aot_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_aot_test.cc index c90fff7cc02914..da728f79ecd4ca 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_aot_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_aot_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "mlir/Parser/Parser.h" #include "riegeli/bytes/string_reader.h" +#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/parser/hlo_parser.h" @@ -95,21 +96,25 @@ void ValidateResult( LiteralTestUtil::Equal(LiteralUtil::CreateR0(2), *result_literal)); } +xla::LocalClient* GetLocalClient(PjRtClient* client) { + return absl::down_cast( + absl::down_cast(client)->raw_client()) + ->client(); +} + TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileMlirAndLoad) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); Compiler::GpuTargetConfig gpu_target_config = xla::Compiler::GpuTargetConfig( - se_client->client()->backend().default_stream_executor()); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + GetLocalClient(client.get())->backend().default_stream_executor()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); auto context = std::make_unique(); context->loadDialect(); auto mlir_module = mlir::parseSourceString(mlir_str, context.get()); - TF_ASSERT_OK_AND_ASSIGN(auto topology, se_client->GetTopologyDescription()); + TF_ASSERT_OK_AND_ASSIGN(auto topology, client->GetTopologyDescription()); xla::CompileOptions opts; opts.gpu_target_config = gpu_target_config; @@ -120,9 +125,8 @@ TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileMlirAndLoad) { MaybeOwningMlirModule(std::move(context), std::move(mlir_module)), *topology, /*client=*/nullptr)); EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(SizeIs(1))); - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_executable, - se_client->Load(std::move(executable), LoadOptions())); + TF_ASSERT_OK_AND_ASSIGN(auto loaded_executable, + client->Load(std::move(executable), LoadOptions())); TF_ASSERT_OK_AND_ASSIGN( std::vector>> result, @@ -133,18 +137,16 @@ TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileMlirAndLoad) { TEST(StreamExecutorGpuCompilerTest, AotCompileDeserializeRoundTrip) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); Compiler::GpuTargetConfig gpu_target_config = xla::Compiler::GpuTargetConfig( - se_client->client()->backend().default_stream_executor()); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + GetLocalClient(client.get())->backend().default_stream_executor()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); auto context = std::make_unique(); context->loadDialect(); auto mlir_module = mlir::parseSourceString(mlir_str, context.get()); - TF_ASSERT_OK_AND_ASSIGN(auto topology, se_client->GetTopologyDescription()); + TF_ASSERT_OK_AND_ASSIGN(auto topology, client->GetTopologyDescription()); xla::CompileOptions opts; opts.gpu_target_config = gpu_target_config; @@ -171,17 +173,15 @@ TEST(StreamExecutorGpuCompilerTest, AotCompileDeserializeRoundTrip) { TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileXlaAndLoad) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); Compiler::GpuTargetConfig gpu_target_config{ - se_client->client()->backend().default_stream_executor()}; - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + GetLocalClient(client.get())->backend().default_stream_executor()}; + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram)); TF_ASSERT_OK_AND_ASSIGN(const PjRtTopologyDescription* topology, - se_client->GetTopologyDescription()); + client->GetTopologyDescription()); xla::CompileOptions opts; opts.gpu_target_config = gpu_target_config; @@ -191,7 +191,7 @@ TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileXlaAndLoad) { EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(SizeIs(1))); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr loaded_executable, - se_client->Load(std::move(executable), LoadOptions())); + client->Load(std::move(executable), LoadOptions())); TF_ASSERT_OK_AND_ASSIGN( std::vector>> result, loaded_executable->Execute(/*argument_handles=*/{{}}, {})); @@ -201,18 +201,16 @@ TEST(StreamExecutorGpuCompilerTest, SuccessAotCompileXlaAndLoad) { TEST(StreamExecutorGpuCompilerTest, SuccessLoadFromSerializedExecutable) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); xla::CompileOptions opts; opts.gpu_target_config = Compiler::GpuTargetConfig( - se_client->client()->backend().default_stream_executor()); + GetLocalClient(client.get())->backend().default_stream_executor()); TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram)); TF_ASSERT_OK_AND_ASSIGN(const PjRtTopologyDescription* topology, - se_client->GetTopologyDescription()); + client->GetTopologyDescription()); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, compiler.Compile(opts, computation, *topology, /*client=*/nullptr)); @@ -223,8 +221,8 @@ TEST(StreamExecutorGpuCompilerTest, SuccessLoadFromSerializedExecutable) { executable->SerializeExecutable()); TF_ASSERT_OK_AND_ASSIGN( auto loaded_executable, - se_client->LoadSerializedExecutable(serialized_executable, std::nullopt, - LoadOptions())); + client->LoadSerializedExecutable(serialized_executable, std::nullopt, + LoadOptions())); TF_ASSERT_OK_AND_ASSIGN( auto result, loaded_executable->Execute(/*argument_handles=*/{{}}, {})); @@ -240,33 +238,31 @@ ENTRY main { TEST(StreamExecutorGpuCompilerTest, SuccessSerializeDeserialize) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); xla::CompileOptions opts; opts.gpu_target_config = Compiler::GpuTargetConfig( - se_client->client()->backend().default_stream_executor()); + GetLocalClient(client.get())->backend().default_stream_executor()); TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgramIdentity)); TF_ASSERT_OK_AND_ASSIGN(const PjRtTopologyDescription* topology, - se_client->GetTopologyDescription()); + client->GetTopologyDescription()); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, compiler.Compile(opts, computation, *topology, /*client=*/nullptr)); EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(SizeIs(1))); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr loaded_executable, - se_client->Load(std::move(executable), LoadOptions())); + client->Load(std::move(executable), LoadOptions())); // Serialize the executable and deserialize it without failure. TF_ASSERT_OK_AND_ASSIGN(std::string serialized_executable, loaded_executable->SerializeExecutable()); TF_ASSERT_OK_AND_ASSIGN( auto deserialized_executable, - se_client->LoadSerializedExecutable(serialized_executable, std::nullopt, - LoadOptions())); + client->LoadSerializedExecutable(serialized_executable, std::nullopt, + LoadOptions())); EXPECT_EQ(deserialized_executable->GetExecutable()->name(), "Identity"); } @@ -285,13 +281,11 @@ constexpr char const* kD2HProgramTupleOutput = R"( TEST(StreamExecutorGpuCompilerTest, UnloadedExecutableMemoryStats) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); xla::CompileOptions options; options.gpu_target_config = Compiler::GpuTargetConfig( - se_client->client()->backend().default_stream_executor()); + GetLocalClient(client.get())->backend().default_stream_executor()); // Build the output shape with the correct memory space set. Shape shape = ShapeUtil::MakeShapeWithDenseLayout(S32, {4}, {0}); @@ -303,7 +297,7 @@ TEST(StreamExecutorGpuCompilerTest, UnloadedExecutableMemoryStats) { TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kD2HProgramTupleOutput)); TF_ASSERT_OK_AND_ASSIGN(const PjRtTopologyDescription* topology, - se_client->GetTopologyDescription()); + client->GetTopologyDescription()); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, compiler.Compile(options, computation, *topology, /*client=*/nullptr)); @@ -337,12 +331,10 @@ TEST(StreamExecutorGpuCompilerTest, AutoLayoutIsSupported) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_client = absl::WrapUnique( - absl::down_cast(client.release())); - StreamExecutorGpuCompiler compiler(se_client->platform_id(), - se_client->client()->platform()->id()); + StreamExecutorGpuCompiler compiler( + client->platform_id(), GetLocalClient(client.get())->platform()->id()); TF_ASSERT_OK_AND_ASSIGN(const PjRtTopologyDescription* topology, - se_client->GetTopologyDescription()); + client->GetTopologyDescription()); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr m, ParseAndReturnUnverifiedModule( diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc index cdb2f122b4b003..06d2fdedb24d13 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/status_matchers.h" @@ -42,6 +43,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" +#include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/gpu/se_gpu_pjrt_client.h" #include "xla/pjrt/gpu/se_gpu_topology_description.h" #include "xla/pjrt/maybe_owning_mlir_module.h" @@ -382,11 +384,14 @@ TEST(StreamExecutorGpuCompilerTest, CrossCompilation) { ASSERT_OK_AND_ASSIGN(std::unique_ptr client, GetStreamExecutorGpuClient(GpuClientOptions())); - auto se_gpu_client = dynamic_cast(client.get()); - ASSERT_NE(se_gpu_client, nullptr); + auto common_client = dynamic_cast(client.get()); + ASSERT_NE(common_client, nullptr); se::StreamExecutor* stream_executor = - se_gpu_client->client()->backend().default_stream_executor(); + absl::down_cast(common_client->raw_client()) + ->client() + ->backend() + .default_stream_executor(); auto hlo_module = std::make_shared("name", HloModuleConfig()); diff --git a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client_test.cc index 1b393842b52d80..5e581c61cf256d 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client_test.cc @@ -24,7 +24,11 @@ namespace xla { TEST(XlaCpuPjrtClientTest, GetXlaPjrtGpuClient) { ASSERT_OK_AND_ASSIGN(auto client, GetXlaPjrtGpuClient({})); EXPECT_EQ(client->platform_name(), "cuda"); - EXPECT_NE(dynamic_cast(client.get()), nullptr); + auto* common_client = dynamic_cast(client.get()); + ASSERT_NE(common_client, nullptr); + ASSERT_NE( + dynamic_cast(common_client->raw_client()), + nullptr); } } // namespace xla From fa0bc06d695e8eab2ec4a3b930440a17e06e40ea Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Mon, 31 Aug 2026 11:55:47 -0700 Subject: [PATCH 13/25] Prevent fusing AllGather instructions into GEMM fusions. Currently, this is not supported by the autotuner and we get the following error: RET_CHECK failure (xla/backends/gpu/runtime/collective_thunk.cc:378) params.collective_params && params.collective_params->device_assn Collective parameters and device assignment are required for collective thunk execution PiperOrigin-RevId: 974002561 --- third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc index e16ab8303037ea..cd8591f85691f6 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc @@ -709,7 +709,9 @@ bool AllowedInGemmFusion(const HloInstruction& instr) { // the moment. Includes kFusion, kReduce, kAllReduce, etc. return false; } - return HloPredicateIsNotOp(&instr); + return HloPredicateIsNotOp(&instr); } // Returns true if we should consider fusing the instruction into the GEMM From 75cafafc65635855f436e90ce8ddd2ff5d811185 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 11:56:33 -0700 Subject: [PATCH 14/25] Automated Code Change PiperOrigin-RevId: 974002876 --- tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc index 211c385ef1b53f..78d0ba7d52f182 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc @@ -18,7 +18,6 @@ limitations under the License. #include #include #include -#include #include #include #include From e0595310294074bffe6ea1e26d26a7f2d7e27010 Mon Sep 17 00:00:00 2001 From: Ezekiel Calubaquib Date: Mon, 31 Aug 2026 12:02:08 -0700 Subject: [PATCH 15/25] Allow visibility of profile_interface and profile_factory prior torch_tpu migration layout PiperOrigin-RevId: 974005318 --- third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD index 83d63330c8a04b..e0ffcd0018792d 100644 --- a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD +++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD @@ -87,7 +87,7 @@ cc_library( "@xla//xla/python:__pkg__", "//learning/brain/tfrc/executor/api:__pkg__", "//net/grpc/internal/src/core/ext/xprof_profiler:__pkg__", - "//third_party/py/torch_tpu/pjrt:__pkg__", + "//third_party/py/torch_tpu:__subpackages__", ]), deps = [ ":profiler_interface", @@ -144,7 +144,7 @@ cc_library( "@xla//xla/tsl/profiler:internal", "@xla//xla/tsl/profiler:xla_profiler_backends", "//net/grpc/internal/src/core/ext/xprof_profiler:__pkg__", - "//third_party/py/torch_tpu/pjrt:__pkg__", + "//third_party/py/torch_tpu:__subpackages__", ]), deps = [ "//tsl/profiler/protobuf:xplane_proto_cc", From dbff6667f34b09897a915392ffebdad643cc79c4 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 12:08:33 -0700 Subject: [PATCH 16/25] Fix heap buffer overflow in DynamicUpdateSliceInt4 by clamping memcpy length. Problem is that DynamicUpdateSliceInt4 copied input->bytes into the output buffer without bounding it against `output->bytes`. PiperOrigin-RevId: 974008610 --- tensorflow/lite/kernels/dynamic_update_slice.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/kernels/dynamic_update_slice.cc b/tensorflow/lite/kernels/dynamic_update_slice.cc index a2c04baa07f7d3..5278053c3e39f1 100644 --- a/tensorflow/lite/kernels/dynamic_update_slice.cc +++ b/tensorflow/lite/kernels/dynamic_update_slice.cc @@ -278,15 +278,23 @@ void DynamicUpdateSliceInt4(const TfLiteTensor* input, // If the update is the entirety of the output, then simply copy it and // return. if (input_shape.FlatSize() == update_shape.FlatSize()) { - memcpy(output_data, update_data, input->bytes); + // Clamp the copy length to the destination size. `input->bytes` is derived + // from the raw FlatBuffer buffer size for constant tensors and may be + // inflated relative to the logical shape, so it must be bounded by + // `output->bytes` to avoid an out-of-bounds write. + memcpy(output_data, update_data, std::min(input->bytes, output->bytes)); return; } RuntimeShape clamped_start_indices = ClampStartIndices(input_dims, indices_data, input_shape, update_shape); // If the operation is not done in-place, copy the input data to the output. + // Clamp the copy length to the destination size, since `input->bytes` may be + // larger than the output allocation for constant tensors loaded from a + // FlatBuffer model. + size_t bytes = std::min(input->bytes, output->bytes); if (input->data.data != output->data.data) { - memcpy(output->data.data, input->data.data, input->bytes); + memcpy(output->data.data, input->data.data, bytes); } // Update tensor has no elements. Skip. From 94f18a7b7e9324248069de5a3e0ae0884e7a539a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 12:57:50 -0700 Subject: [PATCH 17/25] Automated Code Change PiperOrigin-RevId: 974030983 --- tensorflow/compiler/mlir/tensorflow/BUILD | 1 + .../compiler/mlir/tensorflow/analysis/side_effect_analysis.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD index ef25dc5626ad42..36f022a83171a6 100644 --- a/tensorflow/compiler/mlir/tensorflow/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/BUILD @@ -655,6 +655,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:node_hash_map", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:vlog_is_on", "@llvm-project//llvm:Support", "@llvm-project//mlir:Analysis", "@llvm-project//mlir:FuncDialect", diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc index 6d148209c27abf..4a9d75f01d8a90 100644 --- a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc +++ b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/container/node_hash_map.h" #include "absl/log/log.h" +#include "absl/log/vlog_is_on.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" From 3f3bfd6ae08a85728bdb506dbc0ebc2d9b3af422 Mon Sep 17 00:00:00 2001 From: Bryan Massoth Date: Mon, 31 Aug 2026 13:39:40 -0700 Subject: [PATCH 18/25] Add FourSigFigs for rounding step time outputs. PiperOrigin-RevId: 974053683 --- third_party/xla/xla/tsl/profiler/utils/format_utils.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/third_party/xla/xla/tsl/profiler/utils/format_utils.h b/third_party/xla/xla/tsl/profiler/utils/format_utils.h index 583c68842e5bd8..9e78f75c8cee1b 100644 --- a/third_party/xla/xla/tsl/profiler/utils/format_utils.h +++ b/third_party/xla/xla/tsl/profiler/utils/format_utils.h @@ -57,6 +57,10 @@ inline std::string MaxPrecision(double d) { return internal::FormatDouble("%.17g", d); } +inline std::string FourSigFigs(double d) { + return internal::FormatDouble("%.4g", d); +} + } // namespace profiler } // namespace tsl From 8a1b178025dc0d49e1ed4539811df6e00f60c828 Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Mon, 31 Aug 2026 14:03:33 -0700 Subject: [PATCH 19/25] PR #47666: [XLA:GPU][oneAPI] Add XPU target check in Triton ThreadDim extraction Imported from GitHub PR https://github.com/openxla/xla/pull/47666 Like ROCm, Intel-XPU Triton targets do not support warp specialization and do not annotate functions with `ttg.total-num-warps` attribute. Hence, this PR (like ROCm) computes `ThreadDim` from the other launch information attributes. Copybara import of the project: -- e94243603c9674f6955a3331fcb3be202fcdaf9e by Akhil Goel : Compute ThreadDim for xpu triton Merging this change closes #47666 PiperOrigin-RevId: 974067049 --- .../xla/xla/backends/gpu/codegen/triton/lowering_util.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc b/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc index 2e680d9eb43dc8..6fae64d06bfc52 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc @@ -73,7 +73,7 @@ absl::StatusOr ExtractThreadDims( if (!num_warps_attr) { return absl::InternalError("ttg.num-warps attribute not found."); } - // AMD/ROCm Triton backend does not support warp specialization. + // AMD/ROCm and Intel XPU Triton backends do not support warp specialization. // Consequently, `ttg.total-num-warps` and `nvvm.reqntid` are not added // to triton module/function. // ThreadDim is therefore calculated from the Module attributes and not @@ -82,7 +82,8 @@ absl::StatusOr ExtractThreadDims( if (!target) { return absl::InternalError("ttg.target attribute not found."); } - if (target.getValue().find("gfx") != std::string::npos) { + if (target.getValue().find("gfx") != std::string::npos || + target.getValue().find("xpu") != std::string::npos) { stream_executor::ThreadDim thread_dims( num_warps_attr.getInt() * threads_per_warp_attr.getInt(), 1, 1); return thread_dims; From afc9e6af3f38fdc312c0b1592f218c9a563425fa Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Mon, 31 Aug 2026 14:51:05 -0700 Subject: [PATCH 20/25] Propagate actual failure status to definition event promises in AsyncHostToDeviceTransferManager. Previously, if `CommonAsyncHostToDeviceTransferManager::Create` failed during buffer creation (e.g., due to OOM or invalid arguments), definition event promises for any previously allocated buffers were fulfilled with a generic `absl::UnknownError`. This obscured the actual failure reason. Update `Create` to capture the failure `absl::Status` from the buffer creation loop and set each definition event promise with that actual status. PiperOrigin-RevId: 974091435 --- .../pjrt/host_to_device_transfer_manager.cc | 153 +++++++++--------- 1 file changed, 77 insertions(+), 76 deletions(-) diff --git a/third_party/xla/xla/pjrt/host_to_device_transfer_manager.cc b/third_party/xla/xla/pjrt/host_to_device_transfer_manager.cc index a062c9c856fa68..4fc880aff30a11 100644 --- a/third_party/xla/xla/pjrt/host_to_device_transfer_manager.cc +++ b/third_party/xla/xla/pjrt/host_to_device_transfer_manager.cc @@ -110,91 +110,92 @@ class CommonAsyncHostToDeviceTransferManager definition_events.reserve(shape_specs.size()); device_shapes.reserve(shape_specs.size()); - absl::Cleanup fail_definitions = [&] { - for (PjRtDeviceEventPromiseRef& definition_promise : definition_events) { - definition_promise.SetError(absl::UnknownError( - "Failed to create host to device transfer manager.")); - } - }; - - for (int i = 0; i < shape_specs.size(); ++i) { - const PjRtClient::ShapeSpec& shape_spec = shape_specs[i]; - if (shape_spec.element_type == TUPLE) { - return Unimplemented( - "Async buffer transfer of tuples not implemented."); - } + auto create_buffers = [&]() -> absl::Status { + for (int i = 0; i < shape_specs.size(); ++i) { + const PjRtClient::ShapeSpec& shape_spec = shape_specs[i]; + if (shape_spec.element_type == TUPLE) { + return Unimplemented( + "Async buffer transfer of tuples not implemented."); + } - auto allocation_event = - client->CreateAllocationEventForTransfers(memory_space, debug_info); - if (allocation_event) { - allocation_events.push_back( - std::make_unique(allocation_event)); - } else { - allocation_events.push_back({}); - } + auto allocation_event = + client->CreateAllocationEventForTransfers(memory_space, debug_info); + if (allocation_event) { + allocation_events.push_back( + std::make_unique(allocation_event)); + } else { + allocation_events.push_back({}); + } - ABSL_ASSIGN_OR_RETURN( - Shape device_shape, - client->MakeDefaultShapeForMemorySpace( - memory_space, - shape_spec.element_type == xla::TOKEN - ? xla::ShapeUtil::MakeTokenShape() - : xla::ShapeUtil::MakeShape(shape_spec.element_type, - shape_spec.dims), - device_layouts.has_value() && (*device_layouts)[i].has_value() - ? &(*(*device_layouts)[i]) - : nullptr)); - auto shared_device_shape = - std::make_shared(std::move(device_shape)); - ABSL_ASSIGN_OR_RETURN( - int64_t on_device_bytes_count, - client->GetOnDeviceBytesCount(memory_space, *shared_device_shape)); - PjRtRawBufferRef raw_buffer; - if (donated_buffer_refs.has_value()) { - raw_buffer = (*donated_buffer_refs)[i]; - if (!raw_buffer) { - return InvalidArgument("Donated buffer ref at index %d is null", i); + ABSL_ASSIGN_OR_RETURN( + Shape device_shape, + client->MakeDefaultShapeForMemorySpace( + memory_space, + shape_spec.element_type == xla::TOKEN + ? xla::ShapeUtil::MakeTokenShape() + : xla::ShapeUtil::MakeShape(shape_spec.element_type, + shape_spec.dims), + device_layouts.has_value() && (*device_layouts)[i].has_value() + ? &(*(*device_layouts)[i]) + : nullptr)); + auto shared_device_shape = + std::make_shared(std::move(device_shape)); + ABSL_ASSIGN_OR_RETURN( + int64_t on_device_bytes_count, + client->GetOnDeviceBytesCount(memory_space, *shared_device_shape)); + PjRtRawBufferRef raw_buffer; + if (donated_buffer_refs.has_value()) { + raw_buffer = (*donated_buffer_refs)[i]; + if (!raw_buffer) { + return InvalidArgument("Donated buffer ref at index %d is null", i); + } + if (raw_buffer->GetOnDeviceSizeInBytes() != on_device_bytes_count) { + return InvalidArgument( + "Donated buffer size %d does not match target buffer size %d", + raw_buffer->GetOnDeviceSizeInBytes(), on_device_bytes_count); + } + } else { + ABSL_ASSIGN_OR_RETURN(raw_buffer, + client->AllocateRawBuffer( + memory_space, on_device_bytes_count, + /*retry_on_oom=*/true, allocation_event)); } - if (raw_buffer->GetOnDeviceSizeInBytes() != on_device_bytes_count) { - return InvalidArgument( - "Donated buffer size %d does not match target buffer size %d", - raw_buffer->GetOnDeviceSizeInBytes(), on_device_bytes_count); + + // We make an event that will become available when the final transfer + // is complete. + PjRtDeviceEventPromiseRef definition_event_promise; + PjRtDeviceEventRef definition_event; + if (client->event_tracking_enabled()) { + ABSL_ASSIGN_OR_RETURN( + std::tie(definition_event_promise, definition_event), + client->CreateLinkedEventPromise( + memory_space, + absl::StrCat("AsyncHostToDeviceTransferManager Op:", + debug_info.value_or("")))); + } else { + ABSL_ASSIGN_OR_RETURN(std::tie(definition_event_promise, definition_event), + client->CreateLinkedEventPromise(memory_space, "")); } - } else { + definition_events.push_back(std::move(definition_event_promise)); + ABSL_ASSIGN_OR_RETURN( - raw_buffer, - client->AllocateRawBuffer(memory_space, on_device_bytes_count, - /*retry_on_oom=*/true, allocation_event)); + auto buffer, + client->DefineBuffer(shared_device_shape, memory_space, raw_buffer, + {std::move(definition_event)})); + device_shapes.push_back(std::move(shared_device_shape)); + buffers.push_back(std::move(buffer)); + undispatched_buffer_refs.push_back(raw_buffer); + buffer_sizes.push_back(on_device_bytes_count); } + return absl::OkStatus(); + }; - // We make an event that will become available when the final transfer - // is complete. - PjRtDeviceEventPromiseRef definition_event_promise; - PjRtDeviceEventRef definition_event; - if (client->event_tracking_enabled()) { - ABSL_ASSIGN_OR_RETURN( - std::tie(definition_event_promise, definition_event), - client->CreateLinkedEventPromise( - memory_space, - absl::StrCat("AsyncHostToDeviceTransferManager Op:", - debug_info.value_or("")))); - } else { - ABSL_ASSIGN_OR_RETURN(std::tie(definition_event_promise, definition_event), - client->CreateLinkedEventPromise(memory_space, "")); + if (absl::Status status = create_buffers(); !status.ok()) { + for (PjRtDeviceEventPromiseRef& definition_promise : definition_events) { + definition_promise.SetError(status); } - definition_events.push_back(std::move(definition_event_promise)); - - ABSL_ASSIGN_OR_RETURN( - auto buffer, - client->DefineBuffer(shared_device_shape, memory_space, raw_buffer, - {std::move(definition_event)})); - device_shapes.push_back(std::move(shared_device_shape)); - buffers.push_back(std::move(buffer)); - undispatched_buffer_refs.push_back(raw_buffer); - buffer_sizes.push_back(on_device_bytes_count); + return status; } - - std::move(fail_definitions).Cancel(); return std::unique_ptr( new CommonAsyncHostToDeviceTransferManager( std::move(buffers), std::move(undispatched_buffer_refs), From a24bc3540f607a6304321bc19b56adbbbcd08b59 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Mon, 31 Aug 2026 15:04:11 -0700 Subject: [PATCH 21/25] Don't DCE tokens in opt barriers. This allows expressing an after relation like so: token = create_token() token, x = opt_barrier((token, x)) ys = [opt_barrier((token, y))[1] for y in ys] without needing to add dce-sinks after all the second opt-barriers. PiperOrigin-RevId: 974098204 --- .../simplifiers/algebraic_simplifier.cc | 5 +++ .../simplifiers/algebraic_simplifier_test.cc | 35 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 137e6599691540..00480dc82bf636 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -5382,6 +5382,11 @@ absl::Status AlgebraicSimplifierVisitor::HandleOptimizationBarrier( } used_elements[use->tuple_index()] = true; } + for (size_t i = 0; i < barrier->shape().tuple_shapes().size(); ++i) { + if (barrier->shape().tuple_shapes()[i].IsToken()) { + used_elements[i] = true; + } + } HloInstruction* operand = barrier->mutable_operand(0); if (operand->opcode() == HloOpcode::kTuple) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index 1d7021f9c48b83..294e6653bb57a9 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -12658,6 +12658,41 @@ TEST_F(AlgebraicSimplifierTest, DoNotSimplifyOptimizationBarrierSideEffects) { ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).value()); } +TEST_F(AlgebraicSimplifierTest, SimplifyOptimizationBarrierTokens) { + constexpr absl::string_view kModuleStr = R"( + HloModule m + + ENTRY entry { + param.0 = f32[] parameter(0) + param.1 = f32[] parameter(1) + add.0 = f32[] add(param.0, param.1) + sub.0 = f32[] subtract(param.0, param.1) + mul.0 = f32[] multiply(param.0, param.1) + tok = token[] after-all() + tuple.0 = (f32[], f32[], f32[], token[]) tuple(mul.0, sub.0, add.0, tok) + b = (f32[], f32[], f32[], token[]) opt-barrier(tuple.0) + gte.0 = f32[] get-tuple-element(b), index=1 + ROOT t = (f32[], f32[]) tuple(mul.0,gte.0) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + EXPECT_EQ(m->entry_computation() + ->root_instruction() + ->operand(1) + ->operand(0) + ->operand(0) + ->operand_count(), + 4); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value()); + EXPECT_EQ(m->entry_computation() + ->root_instruction() + ->operand(1) + ->operand(0) + ->operand(0) + ->operand_count(), + 3); +} + TEST_F(AlgebraicSimplifierTest, GTETupleShardingLoss) { // Verify the gte(tuple) folding does not happen if it loses sharding info. constexpr absl::string_view kModuleStr = R"( From 66420a1b734566ea23f9c7454528875140602fbc Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 31 Aug 2026 15:41:23 -0700 Subject: [PATCH 22/25] Add BlockScalingConfig to xla_data.proto and hlo.proto to describe block scaling and zero point configurations for convolution operands. Update HloInstruction::CreateConvolve and HloConvolutionInstruction to support block scaling and custom scale operands. Update HLO parser and HloCreationUtils (MakeConvolveHlo) for block-scaled convolutions. Migrate HLO passes (conv_operand_swapper, convolution_group_converter, dot_as_convolution_util, space_to_batch_converter, spmd convolution_handler) to the new CreateConvolve interface. PiperOrigin-RevId: 974117253 --- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 70 +++++++++-- third_party/xla/xla/hlo/ir/hlo_instruction.h | 8 ++ .../xla/xla/hlo/ir/hlo_instruction_test.cc | 54 +++++++++ .../xla/xla/hlo/ir/hlo_instructions.cc | 19 ++- third_party/xla/xla/hlo/ir/hlo_instructions.h | 9 ++ third_party/xla/xla/hlo/parser/hlo_parser.cc | 110 +++++++++++++++++- .../xla/xla/hlo/parser/hlo_parser_test.cc | 65 +++++++++++ .../simplifiers/conv_operand_swapper.cc | 9 +- .../convolution_group_converter.cc | 19 ++- .../xla/service/dot_as_convolution_util.cc | 3 +- .../xla/xla/service/gpu/conv_utils_test.cc | 11 +- third_party/xla/xla/service/hlo.proto | 5 +- .../xla/xla/service/hlo_creation_utils.cc | 6 +- .../xla/xla/service/hlo_creation_utils.h | 4 +- third_party/xla/xla/service/hlo_verifier.cc | 74 ++++++++---- .../xla/xla/service/hlo_verifier_test.cc | 64 ++++++++++ .../xla/service/space_to_batch_converter.cc | 8 +- .../xla/service/spmd/convolution_handler.cc | 11 +- third_party/xla/xla/xla_data.proto | 18 +++ 19 files changed, 515 insertions(+), 52 deletions(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index be8d15cbc75aee..73012d47597a92 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -1022,12 +1022,13 @@ absl::StatusOr> HloInstruction::CreateFromProto( PrecisionConfig precision_config = proto.precision_config(); precision_config.mutable_operand_precision()->Resize( proto.operand_ids_size(), PrecisionConfig::DEFAULT); - instruction = CreateConvolve( - shape, all_operands(), - std::max(proto.feature_group_count(), 1), - std::max(proto.batch_group_count(), 1), proto.window(), - proto.convolution_dimension_numbers(), precision_config, - proto.sparsity_config(), proto.conv_kind()); + instruction = + CreateConvolve(shape, all_operands(), + std::max(proto.feature_group_count(), 1), + std::max(proto.batch_group_count(), 1), + proto.window(), proto.convolution_dimension_numbers(), + precision_config, proto.sparsity_config(), + proto.block_scaling_config(), proto.conv_kind()); break; } case HloOpcode::kReduceWindow: @@ -1685,10 +1686,13 @@ HloInstruction::CreateRngBitGenerator(const Shape& shape, HloInstruction* state, int64_t feature_group_count, int64_t batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, - const SparsityConfig& sparsity_config, ConvolutionKind convolution_kind) { + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config, + ConvolutionKind convolution_kind) { return std::make_unique( shape, operands, feature_group_count, batch_group_count, window, - dimension_numbers, precision_config, sparsity_config, convolution_kind); + dimension_numbers, precision_config, sparsity_config, + block_scaling_config, convolution_kind); } /* static */ std::unique_ptr HloInstruction::CreateFft( @@ -5726,6 +5730,46 @@ std::string SparsityConfigToString(const SparsityConfig& sparsity_config) { return StrJoin(result, " "); } +std::string BlockScalingConfigToString( + const BlockScalingConfig& block_scaling_config) { + std::vector result; + if (block_scaling_config.has_lhs()) { + std::string lhs_str = + StrCat("lhs={scale_idx=", block_scaling_config.lhs().scale_idx()); + if (block_scaling_config.lhs().has_zero_idx()) { + StrAppend(&lhs_str, " zero_idx=", block_scaling_config.lhs().zero_idx()); + } + if (!block_scaling_config.lhs().strides().empty()) { + StrAppend(&lhs_str, " strides=", + StrJoin(block_scaling_config.lhs().strides(), "x")); + } + if (!block_scaling_config.lhs().steps().empty()) { + StrAppend(&lhs_str, + " steps=", StrJoin(block_scaling_config.lhs().steps(), "x")); + } + StrAppend(&lhs_str, "}"); + result.push_back(lhs_str); + } + if (block_scaling_config.has_rhs()) { + std::string rhs_str = + StrCat("rhs={scale_idx=", block_scaling_config.rhs().scale_idx()); + if (block_scaling_config.rhs().has_zero_idx()) { + StrAppend(&rhs_str, " zero_idx=", block_scaling_config.rhs().zero_idx()); + } + if (!block_scaling_config.rhs().strides().empty()) { + StrAppend(&rhs_str, " strides=", + StrJoin(block_scaling_config.rhs().strides(), "x")); + } + if (!block_scaling_config.rhs().steps().empty()) { + StrAppend(&rhs_str, + " steps=", StrJoin(block_scaling_config.rhs().steps(), "x")); + } + StrAppend(&rhs_str, "}"); + result.push_back(rhs_str); + } + return StrJoin(result, " "); +} + std::string ConvolutionDimensionNumbersToString( const ConvolutionDimensionNumbers& dnums) { auto len_required = [](int64_t a, int64_t b, absl::Span cs) { @@ -6379,6 +6423,16 @@ void HloInstruction::set_sparsity_config( Cast(this)->set_sparsity_config(sparsity_config); } +const BlockScalingConfig& HloInstruction::block_scaling_config() const { + return Cast(this)->block_scaling_config(); +} + +void HloInstruction::set_block_scaling_config( + const BlockScalingConfig& block_scaling_config) { + Cast(this)->set_block_scaling_config( + block_scaling_config); +} + const DomainMetadata& HloInstruction::operand_side_metadata() const { return Cast(this)->operand_side_metadata(); } diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index 308507b026a324..b0a90e9fe7fd3f 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -440,6 +440,7 @@ class HloInstruction { const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, const SparsityConfig& sparsity_config = SparsityConfig(), + const BlockScalingConfig& block_scaling_config = BlockScalingConfig(), ConvolutionKind convolution_kind = CONVOLUTION_KIND_UNSET); // Creates an FFT op, of the type indicated by fft_type. @@ -2534,6 +2535,10 @@ class HloInstruction { const SparsityConfig& sparsity_config() const; void set_sparsity_config(const SparsityConfig& config); + // Delegates to HloConvolutionInstruction::block_scaling_config. + const BlockScalingConfig& block_scaling_config() const; + void set_block_scaling_config(const BlockScalingConfig& config); + // Returns true if the instruction is an async-start, async-update, or // async-done. bool IsAsynchronous() const { return HloOpcodeIsAsync(opcode_); } @@ -2958,6 +2963,9 @@ std::string ConvolutionDimensionNumbersToString( const ConvolutionDimensionNumbers& dnums); std::string SparsityConfigToString(const SparsityConfig& sparsity_config); +std::string BlockScalingConfigToString( + const BlockScalingConfig& block_scaling_config); + absl::StatusOr StringToRandomAlgorithm( const std::string& name); absl::StatusOr StringToRandomDistribution( diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc index 0db5924abf1ed3..fb5a4904ebadd7 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc @@ -150,6 +150,60 @@ TEST_F(HloInstructionTest, SparsityConfigToString_LHSAndRHS) { R"(%convolution = bf16[256,256]{1,0} convolution(%lhs, %rhs, %lhs_indices, %rhs_indices), dim_labels=bf_io->bf, sparsity_config={lhs={sparsity=1x4 dimension=0 stride=1 idx=2} rhs={sparsity=1x4 dimension=0 stride=1 idx=3}})"); } +TEST_F(HloInstructionTest, BlockScalingConfigToString) { + { + BlockScalingConfig config; + EXPECT_EQ(BlockScalingConfigToString(config), ""); + } + { + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + // Unpopulated zero_idx, strides, steps + EXPECT_EQ(BlockScalingConfigToString(config), "lhs={scale_idx=2}"); + } + { + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + config.mutable_lhs()->set_zero_idx(0); + // zero_idx set to 0 should be printed when has_zero_idx() is true. + EXPECT_EQ(BlockScalingConfigToString(config), + "lhs={scale_idx=2 zero_idx=0}"); + } + { + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + config.mutable_lhs()->set_zero_idx(3); + config.mutable_lhs()->add_strides(1); + config.mutable_lhs()->add_strides(4); + config.mutable_lhs()->add_steps(1); + config.mutable_lhs()->add_steps(2); + EXPECT_EQ(BlockScalingConfigToString(config), + "lhs={scale_idx=2 zero_idx=3 strides=1x4 steps=1x2}"); + } + { + BlockScalingConfig config; + config.mutable_rhs()->set_scale_idx(3); + config.mutable_rhs()->add_strides(2); + config.mutable_rhs()->add_strides(4); + EXPECT_EQ(BlockScalingConfigToString(config), + "rhs={scale_idx=3 strides=2x4}"); + } + { + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + config.mutable_lhs()->set_zero_idx(0); + config.mutable_lhs()->add_strides(1); + config.mutable_rhs()->set_scale_idx(3); + config.mutable_rhs()->set_zero_idx(1); + config.mutable_rhs()->add_strides(2); + config.mutable_rhs()->add_steps(4); + EXPECT_EQ( + BlockScalingConfigToString(config), + "lhs={scale_idx=2 zero_idx=0 strides=1} rhs={scale_idx=3 zero_idx=1 " + "strides=2 steps=4}"); + } +} + TEST_F(HloInstructionTest, GetStackTraceStringFromStackFrameId) { auto module = CreateNewVerifiedModule(); HloComputation::Builder builder("main"); diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index f86be70c5b0a41..0ed3f092b02857 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -3541,7 +3541,9 @@ HloConvolutionInstruction::HloConvolutionInstruction( int64_t feature_group_count, int64_t batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, - const SparsityConfig& sparsity_config, ConvolutionKind convolution_kind) + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config, + ConvolutionKind convolution_kind) : HloInstruction(HloOpcode::kConvolution, shape), feature_group_count_(feature_group_count), batch_group_count_(batch_group_count), @@ -3549,6 +3551,7 @@ HloConvolutionInstruction::HloConvolutionInstruction( convolution_dimension_numbers_(dimension_numbers), precision_config_(precision_config), sparsity_config_(sparsity_config), + block_scaling_config_(block_scaling_config), convolution_kind_(convolution_kind) { if (window_util::HasBaseDilation(window)) { SetAndSanitizeName(StrCat(name(), "-base-dilated")); @@ -3584,6 +3587,7 @@ void HloConvolutionInstruction::ToProto(HloInstructionProto* proto) const { proto->set_conv_kind(convolution_kind_); } *proto->mutable_sparsity_config() = sparsity_config_; + *proto->mutable_block_scaling_config() = block_scaling_config_; } void HloConvolutionInstruction::PrintExtraAttributesImpl( @@ -3624,6 +3628,13 @@ void HloConvolutionInstruction::PrintExtraAttributesImpl( printer->Append("}"); }); } + if (block_scaling_config_.has_lhs() || block_scaling_config_.has_rhs()) { + printer.Next([this](Printer* printer) { + printer->Append("block_scaling_config={"); + printer->Append(BlockScalingConfigToString(block_scaling_config_)); + printer->Append("}"); + }); + } } bool HloConvolutionInstruction::IdenticalSlowPath( @@ -3649,7 +3660,9 @@ bool HloConvolutionInstruction::IdenticalSlowPath( protobuf_util::HaveSameSerialization( precision_config(), casted_other.precision_config()) && protobuf_util::HaveSameSerialization(sparsity_config(), - casted_other.sparsity_config()); + casted_other.sparsity_config()) && + protobuf_util::HaveSameSerialization( + block_scaling_config(), casted_other.block_scaling_config()); } std::unique_ptr @@ -3659,7 +3672,7 @@ HloConvolutionInstruction::CloneWithNewOperandsImpl( return std::make_unique( shape, new_operands, feature_group_count_, batch_group_count_, window(), convolution_dimension_numbers_, precision_config_, sparsity_config_, - convolution_kind_); + block_scaling_config_, convolution_kind_); } HloReduceWindowInstruction::HloReduceWindowInstruction( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 1f14120b5e6f6d..7d1495baa9244f 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -2068,6 +2068,7 @@ class HloConvolutionInstruction : public HloInstruction { const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config, ConvolutionKind convolution_kind = CONVOLUTION_KIND_UNSET); const Window& window() const override { return window_; } void set_window(const Window& window) override { window_ = window; } @@ -2110,6 +2111,13 @@ class HloConvolutionInstruction : public HloInstruction { sparsity_config_ = sparsity_config; } + const BlockScalingConfig& block_scaling_config() const { + return block_scaling_config_; + } + void set_block_scaling_config(const BlockScalingConfig& config) { + block_scaling_config_ = config; + } + std::string ToCategory() const override; void ToProto(HloInstructionProto* proto) const override; @@ -2143,6 +2151,7 @@ class HloConvolutionInstruction : public HloInstruction { // The sparsity configuration used for the convolution. SparsityConfig sparsity_config_; // Convolution block scaling config. + BlockScalingConfig block_scaling_config_; // Conv type (fprop, dgrad, wgrad) ConvolutionKind convolution_kind_ = CONVOLUTION_KIND_UNSET; }; diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 46f3a12490c101..a5ec279715662d 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -303,6 +303,9 @@ class HloParserImpl : public HloParser { ParseCollectiveDeviceListBaseOnly(); bool ParseSparsityConfig(SparsityConfig* result); bool ParseTensorSparsityConfig(SparsityConfig::TensorSparsityConfig* result); + bool ParseBlockScalingConfig(BlockScalingConfig* result); + bool ParseTensorBlockScalingConfig( + BlockScalingConfig::TensorBlockScalingConfig* result); private: // Types of attributes. @@ -359,6 +362,7 @@ class HloParserImpl : public HloParser { kMode, kConvKind, kSparsityConfig, + kBlockScalingConfig, kDebugAttributesTable, }; @@ -2758,6 +2762,10 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT optional parsed_sparsity_config; attrs["sparsity_config"] = {/*required=*/false, AttrTy::kSparsityConfig, &parsed_sparsity_config}; + optional parsed_block_scaling_config; + attrs["block_scaling_config"] = {/*required=*/false, + AttrTy::kBlockScalingConfig, + &parsed_block_scaling_config}; if ((!preset_operands && !ParseOperands(&operands, builder)) || !ParseAttributes(attrs, allow_attributes, shape)) { return nullptr; @@ -2788,6 +2796,8 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT } SparsityConfig sparsity_config = parsed_sparsity_config.value_or(SparsityConfig()); + BlockScalingConfig block_scaling_config = + parsed_block_scaling_config.value_or(BlockScalingConfig()); if (!maybe_infer_shape([&] { return ShapeInference::InferConvolveShape( operands[0]->shape(), operands[1]->shape(), @@ -2800,7 +2810,7 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT return builder->AddInstruction(HloInstruction::CreateConvolve( *shape, operands, feature_group_count.value(), batch_group_count.value(), *window, *dnums, precision_config, - sparsity_config, kind)); + sparsity_config, block_scaling_config, kind)); } case HloOpcode::kFft: { optional fft_type; @@ -6014,6 +6024,7 @@ bool HloParserImpl::ParseAttributeHelper( AttrTy attr_type = attr_it->second.attr_type; void* attr_out_ptr = attr_it->second.result; bool success = [&] { + // Dispatch parsing logic based on the attribute type. LocTy attr_loc = lexer_.GetLoc(); switch (attr_type) { case AttrTy::kBool: { @@ -6473,13 +6484,22 @@ bool HloParserImpl::ParseAttributeHelper( static_cast*>(attr_out_ptr)->emplace(result); return true; } + case AttrTy::kBlockScalingConfig: { + BlockScalingConfig result; + if (!ParseBlockScalingConfig(&result)) { + return false; + } + static_cast*>(attr_out_ptr) + ->emplace(result); + return true; + } } }(); if (!success) { return Error(loc, StrFormat("error parsing attribute %s", name)); } return true; -} +} // NOLINT(readability/fn_size) bool HloParserImpl::CopyAttributeToProtoMessage( absl::flat_hash_set non_proto_attrs, @@ -7688,6 +7708,92 @@ bool HloParserImpl::ParseTensorSparsityConfig( "expects '}' at the end of TensorSparsityConfig"); } +bool HloParserImpl::ParseBlockScalingConfig(BlockScalingConfig* result) { + VLOG(kDebugLevel) << "ParseBlockScalingConfig"; + if (!ParseToken(TokKind::kLbrace, + "expected '{' to start BlockScalingConfig")) { + return false; + } + if (lexer_.GetKind() == TokKind::kRbrace) { + // empty + } else { + do { + std::string attribute; + if (!ParseAttributeName(&attribute)) { + return false; + } + if (attribute == "lhs" || attribute == "rhs") { + BlockScalingConfig::TensorBlockScalingConfig* tensor_config; + if (attribute == "lhs") { + tensor_config = result->mutable_lhs(); + } else { + tensor_config = result->mutable_rhs(); + } + if (!ParseTensorBlockScalingConfig(tensor_config)) { + return false; + } + } else { + return Error(lexer_.GetLoc(), "unknown attribute"); + } + } while (lexer_.GetKind() != TokKind::kRbrace); + } + return ParseToken(TokKind::kRbrace, + "expects '}' at the end of BlockScalingConfig"); +} + +bool HloParserImpl::ParseTensorBlockScalingConfig( + BlockScalingConfig::TensorBlockScalingConfig* result) { + VLOG(kDebugLevel) << "ParseTensorBlockScalingConfig"; + CHECK(result != nullptr); + if (!ParseToken(TokKind::kLbrace, + "expected '{' to start TensorBlockScalingConfig")) { + return false; + } + if (lexer_.GetKind() == TokKind::kRbrace) { + // empty + } else { + do { + std::string attribute; + if (!ParseAttributeName(&attribute)) { + return false; + } + if (attribute == "scale_idx") { + int64_t scale_idx; + if (!ParseInt64(&scale_idx)) { + return Error(lexer_.GetLoc(), "expects int64_t"); + } + result->set_scale_idx(scale_idx); + } else if (attribute == "zero_idx") { + int64_t zero_idx; + if (!ParseInt64(&zero_idx)) { + return Error(lexer_.GetLoc(), "expects int64_t"); + } + result->set_zero_idx(zero_idx); + } else if (attribute == "strides") { + std::vector strides; + if (!ParseDxD("strides", &strides)) { + return Error(lexer_.GetLoc(), "expects strides in form NxMxK"); + } + for (int64_t s : strides) { + result->add_strides(s); + } + } else if (attribute == "steps") { + std::vector steps; + if (!ParseDxD("steps", &steps)) { + return Error(lexer_.GetLoc(), "expects steps in form NxMxK"); + } + for (int64_t s : steps) { + result->add_steps(s); + } + } else { + return Error(lexer_.GetLoc(), "unknown attribute"); + } + } while (lexer_.GetKind() != TokKind::kRbrace); + } + return ParseToken(TokKind::kRbrace, + "expects '}' at the end of TensorBlockScalingConfig"); +} + bool HloParserImpl::ParseDxD(const std::string& name, std::vector* result) { LocTy loc = lexer_.GetLoc(); diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index 2d7d10392df8c6..ae3857c563526f 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -7315,6 +7315,71 @@ TEST_F(HloParserTest, SparsityConfig_Both) { EXPECT_EQ(config.rhs().stride(), 1); } +TEST_F(HloParserTest, BlockScalingConfig_RHSOnly) { + const char* const hlo_string = R"( + HloModule BlockScalingConfigModule + ENTRY BlockScalingConfig { + %input = f32[1,2] parameter(0) + %filter = f32[2,2] parameter(1) + %scale = f32[2,1] parameter(2) + ROOT %convolution = f32[1,2] convolution(%input, %filter, %scale), dim_labels=bf_io->bf, + block_scaling_config={rhs={scale_idx=2 strides=1x4 steps=1x1}} + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + auto* conv = module->entry_computation()->root_instruction(); + auto config = conv->block_scaling_config(); + EXPECT_EQ(config.rhs().scale_idx(), 2); + EXPECT_THAT(config.rhs().strides(), ::testing::ElementsAre(1, 4)); + EXPECT_THAT(config.rhs().steps(), ::testing::ElementsAre(1, 1)); +} + +TEST_F(HloParserTest, BlockScalingConfig_Both) { + const char* const hlo_string = R"( + HloModule BlockScalingConfigModule + ENTRY BlockScalingConfig { + %input = f32[1,2] parameter(0) + %filter = f32[2,2] parameter(1) + %lhs_scale = f32[1,1] parameter(2) + %rhs_scale = f32[2,1] parameter(3) + ROOT %convolution = f32[1,2] convolution(%input, %filter, %lhs_scale, %rhs_scale), dim_labels=bf_io->bf, + block_scaling_config={lhs={scale_idx=2 strides=1x4 steps=1x1} rhs={scale_idx=3 strides=1x4 steps=1x1}} + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + auto* conv = module->entry_computation()->root_instruction(); + auto config = conv->block_scaling_config(); + EXPECT_EQ(config.lhs().scale_idx(), 2); + EXPECT_THAT(config.lhs().strides(), ::testing::ElementsAre(1, 4)); + EXPECT_THAT(config.lhs().steps(), ::testing::ElementsAre(1, 1)); + EXPECT_EQ(config.rhs().scale_idx(), 3); + EXPECT_THAT(config.rhs().strides(), ::testing::ElementsAre(1, 4)); + EXPECT_THAT(config.rhs().steps(), ::testing::ElementsAre(1, 1)); +} + +TEST_F(HloParserTest, BlockScalingConfig_RoundTrip) { + const char* const hlo_string = R"( +HloModule BlockScalingConfigModule +ENTRY BlockScalingConfig { + %input = f32[1,2] parameter(0) + %filter = f32[2,2] parameter(1) + %lhs_scale = f32[1,1] parameter(2) + %rhs_scale = f32[2,1] parameter(3) + ROOT %convolution = f32[1,2] convolution(%input, %filter, %lhs_scale, %rhs_scale), window={size=1x1}, dim_labels=bf_io->bf, block_scaling_config={lhs={scale_idx=2 zero_idx=3 strides=1x4 steps=1x1} rhs={scale_idx=3 strides=1x4 steps=1x1}} +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + auto config_before = + module->entry_computation()->root_instruction()->block_scaling_config(); + std::string printed = module->ToString(); + ASSERT_OK_AND_ASSIGN(auto parsed_module, + ParseAndReturnUnverifiedModule(printed)); + auto config_after = parsed_module->entry_computation() + ->root_instruction() + ->block_scaling_config(); + EXPECT_EQ(config_after.DebugString(), config_before.DebugString()); +} + TEST_F(HloParserTest, DesugarParsingTest_DotStart) { const char* const hlo = R"( HloModule async_dot_example diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/conv_operand_swapper.cc b/third_party/xla/xla/hlo/transforms/simplifiers/conv_operand_swapper.cc index da3675a37dd19b..08b80bc6551810 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/conv_operand_swapper.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/conv_operand_swapper.cc @@ -46,6 +46,12 @@ absl::StatusOr SwapConvolutionOperandsIfBeneficial( return false; } + // Do not swap operands for convolutions with block scaling. + if (convolution->block_scaling_config().has_lhs() || + convolution->block_scaling_config().has_rhs()) { + return false; + } + const auto& dnums = convolution->convolution_dimension_numbers(); const auto& window_dims = convolution->window().dimensions(); Window swapped_window; @@ -150,7 +156,8 @@ absl::StatusOr SwapConvolutionOperandsIfBeneficial( /*batch_group_count=*/1, swapped_window, swapped_dnums, precision_config, /*preferred_element_type=*/convolution->shape().element_type(), - /*sparsity_config=*/convolution->sparsity_config())); + /*sparsity_config=*/convolution->sparsity_config(), + /*block_scaling_config=*/convolution->block_scaling_config())); if (conv_is_lowerable_callback && !conv_is_lowerable_callback(new_convolution)) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/convolution_group_converter.cc b/third_party/xla/xla/hlo/transforms/simplifiers/convolution_group_converter.cc index b340e122d86be7..778be054e2987c 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/convolution_group_converter.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/convolution_group_converter.cc @@ -219,6 +219,14 @@ absl::Status ConvolutionVisitor::HandleBatchGroupCount( return absl::OkStatus(); } + // Do not expand if we have block scaling or structured sparsity. + if (convolution->block_scaling_config().has_lhs() || + convolution->block_scaling_config().has_rhs() || + convolution->sparsity_config().has_lhs() || + convolution->sparsity_config().has_rhs()) { + return absl::OkStatus(); + } + VLOG(2) << "Dealing with batch_group_count " << batch_group_count << " for convolution " << convolution->ToString() << "\n"; @@ -316,7 +324,7 @@ absl::Status ConvolutionVisitor::HandleBatchGroupCount( /*batch_group_count=*/1, window, dim_numbers, convolution->precision_config(), /*preferred_element_type=*/convolution->shape().element_type(), - convolution->sparsity_config()) + convolution->sparsity_config(), convolution->block_scaling_config()) .value(); convolution->SetupDerivedInstruction(new_convolution); CHECK_OK(computation_->ReplaceInstruction( @@ -444,6 +452,13 @@ absl::Status ConvolutionVisitor::HandleConvolution( return absl::OkStatus(); } + if (convolution->block_scaling_config().has_lhs() || + convolution->block_scaling_config().has_rhs() || + convolution->sparsity_config().has_lhs() || + convolution->sparsity_config().has_rhs()) { + return absl::OkStatus(); + } + ConvolutionDimensionNumbers dim_numbers = convolution->convolution_dimension_numbers(); auto filter = convolution->mutable_operand(1); @@ -677,7 +692,7 @@ absl::Status ConvolutionVisitor::HandleConvolution( /*batch_group_count=*/1, window, dim_numbers, convolution->precision_config(), /*preferred_element_type=*/convolution->shape().element_type(), - convolution->sparsity_config()) + convolution->sparsity_config(), convolution->block_scaling_config()) .value(); convolution->SetupDerivedInstruction(new_convolution); changed_ = true; diff --git a/third_party/xla/xla/service/dot_as_convolution_util.cc b/third_party/xla/xla/service/dot_as_convolution_util.cc index 9c4fb8ec1a2526..5cc9ae4649112c 100644 --- a/third_party/xla/xla/service/dot_as_convolution_util.cc +++ b/third_party/xla/xla/service/dot_as_convolution_util.cc @@ -188,7 +188,8 @@ CreateShardedConvForDotGeneralConvolution( sharded_conv_shape, operands, /*feature_group_count=*/conv.feature_group_count(), /*batch_group_count=*/conv.batch_group_count(), window, conv_dnums, - conv.precision_config(), conv.sparsity_config()); + conv.precision_config(), conv.sparsity_config(), + conv.block_scaling_config()); } DotConvolutionDimsInfo ParseDotGeneralFromDot(const HloInstruction* dot) { diff --git a/third_party/xla/xla/service/gpu/conv_utils_test.cc b/third_party/xla/xla/service/gpu/conv_utils_test.cc index 2ab50709edead3..da6aef6a584104 100644 --- a/third_party/xla/xla/service/gpu/conv_utils_test.cc +++ b/third_party/xla/xla/service/gpu/conv_utils_test.cc @@ -102,7 +102,8 @@ TEST_F(ConvUtilsTest, BackwardFilterConvolveWithPaddedActivations) { ShapeUtil::MakeShape(F32, {3, 3, 32, 32}), {activations, gradients}, /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, dnums_for_backward_filter_, DefaultPrecisionConfig(2), - /*sparsity_config=*/{}, CONVOLUTION_KIND_WGRAD)); + /*sparsity_config=*/{}, /*block_scaling_config=*/{}, + CONVOLUTION_KIND_WGRAD)); auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = @@ -168,7 +169,7 @@ TEST_F(ConvUtilsTest, BackwardInputConvolveEvenPadding) { /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, conv_dnums, DefaultPrecisionConfig(2), /*sparsity_config=*/{}, - CONVOLUTION_KIND_WGRAD)); + /*block_scaling_config=*/{}, CONVOLUTION_KIND_WGRAD)); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( conv->shape(), ShapeInference::InferConvolveShape( @@ -224,7 +225,8 @@ TEST_F(ConvUtilsTest, BackwardInputConvolveUnevenPaddingOnGradients) { ShapeUtil::MakeShape(F32, {20, 10, 10, 192}), {output, reverse_kernel}, /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, dnums_for_backward_input_, DefaultPrecisionConfig(2), - /*sparsity_config=*/{}, CONVOLUTION_KIND_DGRAD)); + /*sparsity_config=*/{}, /*block_scaling_config=*/{}, + CONVOLUTION_KIND_DGRAD)); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( conv->shape(), @@ -278,7 +280,8 @@ TEST_F(ConvUtilsTest, BackwardInputConvolveUnevenPaddingOnActivations) { ShapeUtil::MakeShape(F32, {1, 1, 14, 1}), {output, reverse_kernel}, /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, dnums_for_backward_input_, DefaultPrecisionConfig(2), - /*sparsity_config=*/{}, CONVOLUTION_KIND_DGRAD)); + /*sparsity_config=*/{}, /*block_scaling_config=*/{}, + CONVOLUTION_KIND_DGRAD)); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( conv->shape(), diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index ef2e1bfcbf9812..b2df4f33b6ec37 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -128,7 +128,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 101 +// Next ID: 102 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -436,6 +436,9 @@ message HloInstructionProto { // For collective-broadcast: indicates the last operand is a scalar // root rank index rather than a data tensor to broadcast. bool has_dynamic_root = 100; + + // Convolution block scaling config. + BlockScalingConfig block_scaling_config = 101; } // Serialization of HloComputation. diff --git a/third_party/xla/xla/service/hlo_creation_utils.cc b/third_party/xla/xla/service/hlo_creation_utils.cc index aefec4318ecb50..0b5dcca89d7c33 100644 --- a/third_party/xla/xla/service/hlo_creation_utils.cc +++ b/third_party/xla/xla/service/hlo_creation_utils.cc @@ -132,7 +132,8 @@ absl::StatusOr MakeConvolveHlo( const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, std::optional preferred_element_type, - const SparsityConfig& sparsity_config, const OpMetadata* metadata, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config, const OpMetadata* metadata, const FrontendAttributes* frontend_attributes) { HloComputation* computation = lhs->parent(); CHECK_EQ(computation, rhs->parent()); @@ -144,7 +145,8 @@ absl::StatusOr MakeConvolveHlo( return computation->AddInstruction( HloInstruction::CreateConvolve( convolve_shape, {lhs, rhs}, feature_group_count, batch_group_count, - window, dimension_numbers, precision_config, sparsity_config), + window, dimension_numbers, precision_config, sparsity_config, + block_scaling_config), metadata, frontend_attributes); } diff --git a/third_party/xla/xla/service/hlo_creation_utils.h b/third_party/xla/xla/service/hlo_creation_utils.h index bf5a2a6d40bada..ea4805ec25deb7 100644 --- a/third_party/xla/xla/service/hlo_creation_utils.h +++ b/third_party/xla/xla/service/hlo_creation_utils.h @@ -86,7 +86,9 @@ absl::StatusOr MakeConvolveHlo( const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config, std::optional preferred_element_type, - const SparsityConfig& sparsity_config, const OpMetadata* metadata = nullptr, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config = BlockScalingConfig(), + const OpMetadata* metadata = nullptr, const FrontendAttributes* frontend_attributes = nullptr); // Creates a transpose HLO instruction and adds it to the computation containing diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index 62a54a223899fe..d79888ecabdbfc 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -303,36 +303,66 @@ absl::Status ShapeVerifier::HandleScaledDot(HloInstruction* scaled_dot) { } absl::Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { - if (convolution->sparsity_config().has_lhs()) { - int32_t idx = convolution->sparsity_config().lhs().idx(); - if (idx < 2 || idx >= convolution->operand_count()) { - return InvalidArgument("Sparsity idx %d out of bounds for lhs", idx); - } - if (!convolution->operand(idx)->shape().IsArray()) { - return InvalidArgument( - "Expected array argument for lhs sparsity at index %d, but got %s", - idx, ShapeUtil::HumanString(convolution->operand(idx)->shape())); + auto check_idx_in_range = [&](int32_t idx, int32_t low, int32_t high, + absl::string_view desc) -> absl::Status { + if (idx < low || idx >= high) { + return InvalidArgument("%s %d out of bounds", desc, idx); } + return absl::OkStatus(); + }; + + if (convolution->sparsity_config().has_lhs()) { + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->sparsity_config().lhs().idx(), 2, + convolution->operand_count(), "Sparsity idx for lhs")); } if (convolution->sparsity_config().has_rhs()) { - int32_t idx = convolution->sparsity_config().rhs().idx(); - if (idx < 2 || idx >= convolution->operand_count()) { - return InvalidArgument("Sparsity idx %d out of bounds for rhs", idx); + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->sparsity_config().rhs().idx(), 2, + convolution->operand_count(), "Sparsity idx for rhs")); + } + if (convolution->block_scaling_config().has_lhs()) { + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->block_scaling_config().lhs().scale_idx(), 2, + convolution->operand_count(), "Block scaling scale_idx for lhs")); + if (convolution->block_scaling_config().lhs().has_zero_idx()) { + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->block_scaling_config().lhs().zero_idx(), 2, + convolution->operand_count(), "Block scaling zero_idx for lhs")); + if (convolution->block_scaling_config().lhs().scale_idx() == + convolution->block_scaling_config().lhs().zero_idx()) { + return InvalidArgument( + "LHS block scaling scale_idx and zero_idx cannot be the same (%d)", + convolution->block_scaling_config().lhs().scale_idx()); + } } - if (!convolution->operand(idx)->shape().IsArray()) { - return InvalidArgument( - "Expected array argument for rhs sparsity at index %d, but got %s", - idx, ShapeUtil::HumanString(convolution->operand(idx)->shape())); + } + if (convolution->block_scaling_config().has_rhs()) { + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->block_scaling_config().rhs().scale_idx(), 2, + convolution->operand_count(), "Block scaling scale_idx for rhs")); + if (convolution->block_scaling_config().rhs().has_zero_idx()) { + ABSL_RETURN_IF_ERROR(check_idx_in_range( + convolution->block_scaling_config().rhs().zero_idx(), 2, + convolution->operand_count(), "Block scaling zero_idx for rhs")); + if (convolution->block_scaling_config().rhs().scale_idx() == + convolution->block_scaling_config().rhs().zero_idx()) { + return InvalidArgument( + "RHS block scaling scale_idx and zero_idx cannot be the same (%d)", + convolution->block_scaling_config().rhs().scale_idx()); + } } } - if (convolution->sparsity_config().has_lhs() && - convolution->sparsity_config().has_rhs()) { - if (convolution->sparsity_config().lhs().idx() == - convolution->sparsity_config().rhs().idx()) { - return InvalidArgument("LHS and RHS sparsity idx cannot be the same (%d)", - convolution->sparsity_config().lhs().idx()); + if (convolution->block_scaling_config().has_lhs() && + convolution->block_scaling_config().has_rhs()) { + if (convolution->block_scaling_config().lhs().scale_idx() == + convolution->block_scaling_config().rhs().scale_idx()) { + return InvalidArgument( + "LHS and RHS block scaling scale_idx cannot be the same (%d)", + convolution->block_scaling_config().lhs().scale_idx()); } } + ABSL_ASSIGN_OR_RETURN( Shape expected, ShapeInference::InferConvolveShape( diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index 117e53583e19bb..4981119fcfea9e 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -946,6 +946,70 @@ TEST_F(HloVerifierTest, ConvNegativeBaseDilationNotAllowed) { HasSubstr("non-positive base area dilation factor")); } +TEST_F(HloVerifierTest, ConvBlockScalingConfigScaleIdxOutOfBoundsNotAllowed) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnUnverifiedModule(kConvHloString)); + auto* conv = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_rhs()->set_scale_idx(5); + conv->set_block_scaling_config(config); + + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Block scaling scale_idx for rhs 5 out of bounds")); +} + +TEST_F(HloVerifierTest, ConvBlockScalingConfigScaleIdxZeroNotAllowed) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnUnverifiedModule(kConvHloString)); + auto* conv = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(0); + conv->set_block_scaling_config(config); + + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Block scaling scale_idx for lhs 0 out of bounds")); +} + +static const char* const kConvWith4OperandsHloString = R"( +HloModule module +ENTRY entry_computation { + param0 = f16[128,128,56,56] parameter(0) + param1 = f16[3,3,128,128] parameter(1) + param2 = f8e8m0fnu[128,128,56,56] parameter(2) + param3 = f8e8m0fnu[3,3,128,128] parameter(3) + ROOT conv = f16[128,128,28,28] convolution(param0, param1, param2, param3), + window={size=3x3 stride=2x2}, dim_labels=bf01_01io->bf01 +})"; + +TEST_F(HloVerifierTest, ConvBlockScalingConfigSameScaleAndZeroIdxNotAllowed) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnUnverifiedModule(kConvWith4OperandsHloString)); + auto* conv = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + config.mutable_lhs()->set_zero_idx(2); + conv->set_block_scaling_config(config); + + EXPECT_THAT( + verifier().Run(module.get()).status().message(), + HasSubstr( + "LHS block scaling scale_idx and zero_idx cannot be the same (2)")); +} + +TEST_F(HloVerifierTest, ConvBlockScalingConfigLhsAndRhsSameScaleIdxNotAllowed) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnUnverifiedModule(kConvWith4OperandsHloString)); + auto* conv = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + config.mutable_rhs()->set_scale_idx(2); + conv->set_block_scaling_config(config); + + EXPECT_THAT( + verifier().Run(module.get()).status().message(), + HasSubstr("LHS and RHS block scaling scale_idx cannot be the same (2)")); +} + static const char* const kAddWithLayoutChangeHlo = R"( HloModule AddWithLayoutChange ENTRY AddWithLayoutChange { diff --git a/third_party/xla/xla/service/space_to_batch_converter.cc b/third_party/xla/xla/service/space_to_batch_converter.cc index b0e4c37afa4690..1637100b052add 100644 --- a/third_party/xla/xla/service/space_to_batch_converter.cc +++ b/third_party/xla/xla/service/space_to_batch_converter.cc @@ -3006,7 +3006,7 @@ absl::Status ConvolutionVisitor::PropagateOnConv(HloInstruction* convolution) { convolution->feature_group_count(), convolution->batch_group_count(), new_window, new_dim_numbers, convolution->precision_config(), /*preferred_element_type=*/convolution->shape().element_type(), - convolution->sparsity_config())); + convolution->sparsity_config(), convolution->block_scaling_config())); convolution->SetupDerivedInstruction(new_conv); old_to_new_instrs_[convolution] = new_conv; @@ -3758,7 +3758,7 @@ absl::Status ConvolutionVisitor::PropagateOnBackpropFilterConv( convolution->batch_group_count(), new_window, new_dim_numbers, convolution->precision_config(), /*preferred_element_type=*/convolution->shape().element_type(), - convolution->sparsity_config())); + convolution->sparsity_config(), convolution->block_scaling_config())); convolution->SetupDerivedInstruction(new_conv); VLOG(2) << "New backprop filter convolution " << new_conv->ToString(); @@ -4125,8 +4125,8 @@ absl::Status ConvolutionVisitor::PerformSpaceToBatchOnConvolution( convolution->feature_group_count(), convolution->batch_group_count(), new_window, new_dim_numbers, convolution->precision_config(), /*preferred_element_type=*/convolution->shape().element_type(), - convolution->sparsity_config(), &convolution->metadata(), - &convolution->frontend_attributes())); + convolution->sparsity_config(), convolution->block_scaling_config(), + &convolution->metadata(), &convolution->frontend_attributes())); convolution->SetupDerivedInstruction(new_conv); // If the activations were to be batch-to-spaced again, simply use the diff --git a/third_party/xla/xla/service/spmd/convolution_handler.cc b/third_party/xla/xla/service/spmd/convolution_handler.cc index a990f77cc9b8b3..7bd6f0dff685d7 100644 --- a/third_party/xla/xla/service/spmd/convolution_handler.cc +++ b/third_party/xla/xla/service/spmd/convolution_handler.cc @@ -993,10 +993,13 @@ absl::StatusOr> CreateShardedConvolution( /*preferred_element_type=*/conv.shape().element_type())); *sharded_conv_shape.mutable_layout() = conv.shape().layout(); CHECK(!conv.sparsity_config().has_lhs() && !conv.sparsity_config().has_rhs()); + CHECK(!conv.block_scaling_config().has_lhs() && + !conv.block_scaling_config().has_rhs()); return HloInstruction::CreateConvolve( sharded_conv_shape, {sharded_lhs_hlo, sharded_rhs_hlo}, feature_group_count, batch_group_count, window, conv_dnums, - conv.precision_config(), conv.sparsity_config()); + conv.precision_config(), conv.sparsity_config(), + conv.block_scaling_config()); } // Partition convolution. @@ -1026,6 +1029,12 @@ absl::Status SpmdPartitioningVisitor::HandleConvolution(HloInstruction* hlo) { if (hlo->sharding().IsSingleDevice()) { return DefaultAction(hlo); } + // TODO(b/535773961): Support sharding for scaled / sparse convolutions. + if (hlo->block_scaling_config().has_lhs() || + hlo->block_scaling_config().has_rhs() || + hlo->sparsity_config().has_lhs() || hlo->sparsity_config().has_rhs()) { + return DefaultAction(hlo); + } const dot_as_convolution_util::DotConvolutionDimsInfo dims_info = dot_as_convolution_util::ParseConvolutionDimsInfo(hlo); diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto index 82efb42bec2453..730fc669e01b27 100644 --- a/third_party/xla/xla/xla_data.proto +++ b/third_party/xla/xla/xla_data.proto @@ -892,6 +892,24 @@ message SparsityConfig { TensorSparsityConfig rhs = 2; } +// Describes the block scaling configuration for a convolution's operands. +message BlockScalingConfig { + message TensorBlockScalingConfig { + // The operand index for the scale argument. + int32 scale_idx = 1; + // The operand index for the zero point argument. + optional int32 zero_idx = 2; + // The strides for the block scaling. + repeated int64 strides = 3; + // The steps for the block scaling. + repeated int64 steps = 4; + } + // Block scaling config for the LHS operand. + TensorBlockScalingConfig lhs = 1; + // Block scaling config for the RHS operand. + TensorBlockScalingConfig rhs = 2; +} + enum PaddingType { PADDING_INVALID = 0; PADDING_VALID = 1; // Only valid portion of the base are covered. From eafe265156f003a68580a2d90dfdb024903dc393 Mon Sep 17 00:00:00 2001 From: Yue Sheng Date: Mon, 31 Aug 2026 15:42:23 -0700 Subject: [PATCH 23/25] [Mosaic TPU] Support packed types strided load with static indices by emulation. PiperOrigin-RevId: 974117693 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 32 +++++++++++++------ .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 1 + 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 92dcb8c92f7d13..57afd88cfad38c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -591,7 +591,9 @@ struct MemRefSqueezeFoldCast : public OpRewritePattern { } for (auto [source_dim, result_dim] : llvm::zip(cast_source_type.getShape(), cast_result_type.getShape())) { - if (source_dim == result_dim) continue; + if (source_dim == result_dim) { + continue; + } if (ShapedType::isDynamic(source_dim) && !ShapedType::isDynamic(result_dim)) { // The result type must be more dynamic than the source type. @@ -847,6 +849,16 @@ LogicalResult StridedLoadOp::verify() { /*min_stride=*/0); } +OpFoldResult StridedLoadOp::fold(FoldAdaptor adaptor) { + if (llvm::all_of(getStrides(), [](int32_t s) { return s == 1; })) { + OpBuilder builder(*this); + return tpu::VectorLoadOp::create(builder, getLoc(), getType(), getBase(), + getIndices()) + .getResult(); + } + return nullptr; +} + LogicalResult StridedStoreOp::verify() { return verifyStridedOp(*this, getBase().getType(), getValueToStore().getType(), @@ -869,10 +881,11 @@ LogicalResult verifyStoreOp(Op op) { return op.emitError( "Not implemented: masked store with non-32-bit element type"); } - if (value_ty.getShape() != op.getMask().getType().getShape()) + if (value_ty.getShape() != op.getMask().getType().getShape()) { return op.emitOpError("Expected mask shape to match result shape: (") << value_ty.getShape() << "). Got: (" << op.getMask().getType().getShape() << ")."; + } } return success(); } @@ -1529,15 +1542,16 @@ LogicalResult ScanOp::verify() { if (input_ty.getElementType().isInteger(1) && getKind() != ReductionKind::kSum) { return emitOpError("Only sum reduction is supported for i1 vector inputs."); - } else if (getKind() != ReductionKind::kSum && - getKind() != ReductionKind::kMax && - getKind() != ReductionKind::kMin) { + } + if (getKind() != ReductionKind::kSum && getKind() != ReductionKind::kMax && + getKind() != ReductionKind::kMin) { return emitOpError("Only sum, max and min reductions are supported."); } if (getMask() == nullptr) { return success(); - } else if (input_ty.getElementType().isInteger(1)) { + } + if (input_ty.getElementType().isInteger(1)) { return emitOpError("Mask is not supported for i1 vector inputs."); } @@ -1798,11 +1812,11 @@ LogicalResult EnqueueIndirectDMAOp::verify() { if (is_gather) { return verifyGather(getOperation(), /*operand_shape=*/source_ty.getShape(), /*offsets_shape=*/offsets_shape, - /*results_memory_space=*/target_ty.getShape()); + /*result_shape=*/target_ty.getShape()); } - return verifyScatter(getOperation(), /*updates_ty=*/source_ty.getShape(), + return verifyScatter(getOperation(), /*updates_shape=*/source_ty.getShape(), /*offsets_shape=*/offsets_shape, - /*operand_ty=*/target_ty.getShape()); + /*operand_shape=*/target_ty.getShape()); } void WaitDMA2Op::build(OpBuilder& builder, OperationState& state, diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index 0c2391d5fb776b..ffbc06b052b25e 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -335,6 +335,7 @@ def TPU_StridedLoadOp : TPU_Op<"strided_load", [DefaultMemRead]> { $base `[` $indices `]` attr-dict `:` type($base) `,` type($result) }]; let hasVerifier = 1; + let hasFolder = 1; } def TPU_StridedStoreOp : TPU_Op<"strided_store", [DefaultMemWrite]> { From 505c29c0e4cb0e0ad55b4e05d0259c9444ab1e0e Mon Sep 17 00:00:00 2001 From: Zac Mustin Date: Mon, 31 Aug 2026 16:23:28 -0700 Subject: [PATCH 24/25] Switch PJRT C API GPU client to use StreamExecutor client directly. This change should be a no-op. PiperOrigin-RevId: 974137702 --- third_party/xla/xla/pjrt/c/BUILD | 1 - third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc | 7 +++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD index 34b28d4728ce8a..100ab0c5e42804 100644 --- a/third_party/xla/xla/pjrt/c/BUILD +++ b/third_party/xla/xla/pjrt/c/BUILD @@ -699,7 +699,6 @@ cc_library( "//xla/pjrt/gpu:se_gpu_topology_description", "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config", "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options", - "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client", "//xla/python:custom_call_batch_partitioner", "//xla/python:custom_partition_callback", "//xla/python:debug_callback_partitioner", # To register "DebugCallbackCustomCallPartitioner" custom partitioning handler. diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc index b52b03b51fb6c3..2845e9ff6a4f75 100644 --- a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc +++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc @@ -58,6 +58,7 @@ limitations under the License. #include "xla/pjrt/extensions/abi_version/gpu_abi_version_extension.h" #include "xla/pjrt/extensions/cross_host_transfers/pjrt_c_api_cross_host_transfers_extension.h" #include "xla/pjrt/gpu/gpu_helpers.h" +#include "xla/pjrt/gpu/se_gpu_pjrt_client.h" #include "xla/pjrt/gpu/se_gpu_topology_description.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_common.h" @@ -66,7 +67,6 @@ limitations under the License. #include "xla/pjrt/pjrt_executable.h" #include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h" #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h" -#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h" #include "xla/python/custom_call_batch_partitioner.h" #include "xla/python/custom_partition_callback.h" #include "xla/service/compiler.h" @@ -239,8 +239,11 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) { } options.max_inflight_computations = static_cast(v); } + if (options.use_tfrt_gpu_client) { + options.use_async_dispatch = true; + } PJRT_ASSIGN_OR_RETURN(std::unique_ptr client, - xla::GetXlaPjrtGpuClient(options)); + xla::GetStreamExecutorGpuClient(options)); args->client = pjrt::CreateWrapperClient(GetGpuPjrtApi(), std::move(client)); return nullptr; } From 3e1e9b9d6a3ccf96d9f56add2b12565627c42daf Mon Sep 17 00:00:00 2001 From: Ionel Gog Date: Mon, 31 Aug 2026 17:06:14 -0700 Subject: [PATCH 25/25] [IFRT IR] Add Fingerprint method to IfrtIRProgram PiperOrigin-RevId: 974156581 --- third_party/xla/xla/python/ifrt/hlo/BUILD | 6 +- .../xla/xla/python/ifrt/hlo/hlo_program.cc | 102 +--------- .../xla/python/ifrt/hlo/hlo_program_test.cc | 84 +-------- third_party/xla/xla/python/ifrt/ir/BUILD | 6 +- .../xla/xla/python/ifrt/ir/ifrt_ir_program.cc | 13 ++ .../xla/xla/python/ifrt/ir/ifrt_ir_program.h | 5 +- .../python/ifrt/ir/ifrt_ir_program_test.cc | 177 +++++++++++++++++- .../xla/xla/python/ifrt/ir/transforms/BUILD | 3 +- .../ifrt/ir/transforms/ifrt_to_dot_pass.cc | 14 +- .../xla/python/ifrt/ir/transforms/utils.cc | 12 -- .../xla/xla/python/ifrt/ir/transforms/utils.h | 5 +- third_party/xla/xla/python/ifrt/ir/utils.cc | 1 - third_party/xla/xla/python/ifrt/mlir/BUILD | 47 +++++ .../xla/python/ifrt/mlir/fingerprint_utils.cc | 134 +++++++++++++ .../xla/python/ifrt/mlir/fingerprint_utils.h | 34 ++++ .../ifrt/mlir/fingerprint_utils_test.cc | 131 +++++++++++++ 16 files changed, 561 insertions(+), 213 deletions(-) create mode 100644 third_party/xla/xla/python/ifrt/mlir/BUILD create mode 100644 third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.cc create mode 100644 third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.h create mode 100644 third_party/xla/xla/python/ifrt/mlir/fingerprint_utils_test.cc diff --git a/third_party/xla/xla/python/ifrt/hlo/BUILD b/third_party/xla/xla/python/ifrt/hlo/BUILD index 8c6033598c97b6..80e23d2eda034a 100644 --- a/third_party/xla/xla/python/ifrt/hlo/BUILD +++ b/third_party/xla/xla/python/ifrt/hlo/BUILD @@ -23,6 +23,7 @@ cc_library( "//xla/pjrt:mlir_to_hlo", "//xla/python/ifrt", "//xla/python/ifrt:rtti", + "//xla/python/ifrt/mlir:fingerprint_utils", "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler", "//xla/tsl/platform:errors", "@com_google_absl//absl/log:check", @@ -30,11 +31,7 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", - "@highwayhash", - "@highwayhash//:arch_specific", - "@highwayhash//:hh_types", "@llvm-project//llvm:Support", - "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:BytecodeWriter", "@llvm-project//mlir:IR", "@llvm-project//mlir:Parser", @@ -49,7 +46,6 @@ xla_cc_test( ":hlo_program", "//xla/pjrt:maybe_owning_mlir_module", "//xla/pjrt:mlir_to_hlo", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:string_view", diff --git a/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc b/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc index 7c655d124bf3aa..6f357ed926a1db 100644 --- a/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc +++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc @@ -15,9 +15,7 @@ limitations under the License. #include "xla/python/ifrt/hlo/hlo_program.h" -#include #include -#include #include #include #include @@ -26,18 +24,11 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "highwayhash/arch_specific.h" -#include "highwayhash/hh_types.h" -#include "highwayhash/highwayhash.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" -#include "mlir/Bytecode/BytecodeImplementation.h" #include "mlir/Bytecode/BytecodeWriter.h" -#include "mlir/Bytecode/Encoding.h" -#include "mlir/IR/Attributes.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/DialectRegistry.h" -#include "mlir/IR/Location.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OperationSupport.h" #include "mlir/IR/OwningOpRef.h" @@ -45,6 +36,7 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "xla/pjrt/maybe_owning_mlir_module.h" #include "xla/pjrt/mlir_to_hlo.h" +#include "xla/python/ifrt/mlir/fingerprint_utils.h" #include "xla/python/ifrt/rtti.h" #include "xla/status_macros.h" #include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h" @@ -95,99 +87,15 @@ absl::StatusOr> HloProgram::FromBytes( std::move(module)); } -namespace { - -// Calculates a HighwayHash fingerprint in a streaming manner. -class HighwayHashStream final : public llvm::raw_ostream { - public: - HighwayHashStream() : hash_(kHighwayHashKey), pos_(0) { SetUnbuffered(); } - - ~HighwayHashStream() override { CHECK_EQ(buffer_pos_, 0); } - - // Destructively calculates the fingerprint of the data consumed so far. - uint64_t fingerprint() && { - flush_buffer(); - highwayhash::HHResult64 result; - hash_.Finalize(&result); - return result; - } - - private: - // `HighwayHashCatT::Append` updates the hash directly from the source buffer - // if the provided buffer size is larger than `sizeof(HHPacket)`. - static constexpr size_t kBufferSize = sizeof(highwayhash::HHPacket); - - // Arbitrarily chosen, forever-unchanging hash key required by HighwayHash. - static constexpr highwayhash::HHKey kHighwayHashKey = { - 0x4451e30f87db9609ULL, - 0xca7358a1fd2737f8ULL, - 0x4b2c991fcee4fdeaULL, - 0x0b2658e18326f6baULL, - }; - - void write_impl(const char* Ptr, size_t Size) final { - // For tiny writes, it is more efficient to accumulate the data to a buffer - // first and flush it since `HighwayHashCatT::Append` is optimized for - // large writes. - static constexpr size_t kSmallWriteSize = 4; - static_assert(kSmallWriteSize <= kBufferSize); - - if (Size <= kSmallWriteSize) { - if (buffer_pos_ + Size > kBufferSize) { - flush_buffer(); - } - std::memcpy(buffer_ + buffer_pos_, Ptr, Size); - buffer_pos_ += Size; - } else { - flush_buffer(); - hash_.Append(Ptr, Size); - } - pos_ += Size; - } - - uint64_t current_pos() const final { return pos_; } - - void flush_buffer() { - if (buffer_pos_ > 0) { - hash_.Append(buffer_, buffer_pos_); - buffer_pos_ = 0; - } - } - - highwayhash::HighwayHashCatT hash_; - uint64_t pos_; - - char buffer_[kBufferSize]; - uint64_t buffer_pos_ = 0; -}; - -} // namespace - absl::StatusOr HloProgram::Fingerprint() const { - tsl::StatusScopedDiagnosticHandler diag_handler(mlir_module_->getContext()); - - mlir::BytecodeWriterConfig config; - config.setElideLocations(true); - - // Use a version before `kUseListOrdering` due to an MLIR bug where use list - // ordering is not stable. - // - // TODO(b/503120525): Remove this workaround once - // https://github.com/llvm/llvm-project/pull/191942 lands. - config.setDesiredBytecodeVersion( - mlir::bytecode::BytecodeVersion::kLazyLoading); - - HighwayHashStream os; - mlir::LogicalResult result = - mlir::writeBytecodeToFile(mlir_module_, os, config); - absl::Status status = diag_handler.consumeStatus(); - if (!status.ok()) { + absl::StatusOr fingerprint = FingerprintModuleOp(mlir_module_); + if (!fingerprint.ok()) { + absl::Status status = fingerprint.status(); tsl::errors::AppendToMessage( &status, "Failed while calculating HloProgram fingerprint"); return status; } - TF_RET_CHECK(mlir::succeeded(result)); - return std::move(os).fingerprint(); + return *fingerprint; } xla::MaybeOwningMlirModule HloProgram::ToMaybeOwningMlirModule() && { diff --git a/third_party/xla/xla/python/ifrt/hlo/hlo_program_test.cc b/third_party/xla/xla/python/ifrt/hlo/hlo_program_test.cc index 11ddf4de6bdd98..56c5600c2b6306 100644 --- a/third_party/xla/xla/python/ifrt/hlo/hlo_program_test.cc +++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program_test.cc @@ -28,7 +28,6 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/pjrt/maybe_owning_mlir_module.h" #include "xla/pjrt/mlir_to_hlo.h" -#include "xla/tsl/platform/statusor.h" namespace xla::ifrt { namespace { @@ -54,77 +53,6 @@ absl::StatusOr> ParseHloProgramString( std::move(module)); } -TEST(HloProgramTest, Fingerprint) { - static constexpr absl::string_view kModule1 = R"( -module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { - func.func @main(%arg0: tensor) -> tensor { - %0 = mhlo.constant dense<1.000000e+00> : tensor - %1 = mhlo.add %arg0, %0 : tensor - return %1 : tensor - } -} -)"; - TF_ASSERT_OK_AND_ASSIGN(auto program1, ParseHloProgramString(kModule1)); - - static constexpr absl::string_view kModule2 = R"( -module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { - func.func @main(%arg0: tensor) -> tensor { - %0 = mhlo.constant dense<2.000000e+00> : tensor - %1 = mhlo.add %arg0, %0 : tensor - return %1 : tensor - } -} -)"; - TF_ASSERT_OK_AND_ASSIGN(auto program2, ParseHloProgramString(kModule2)); - - EXPECT_EQ(program1->Fingerprint(), program1->Fingerprint()); - EXPECT_NE(program1->Fingerprint(), program2->Fingerprint()); -} - -TEST(HloProgramTest, FingerprintIgnoresDebugInfo) { - TF_ASSERT_OK_AND_ASSIGN( - const std::unique_ptr hlo_program1, - ParseHloProgramString(R"( -module @foo { - func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> { - return %arg0 : tensor<2x3xi32> loc("foo") - } -})")); - TF_ASSERT_OK_AND_ASSIGN( - const std::unique_ptr hlo_program2, - ParseHloProgramString(R"( -module @foo { - func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> { - return %arg0 : tensor<2x3xi32> loc("bar") - } -})")); - - EXPECT_EQ(hlo_program1->Fingerprint(), hlo_program2->Fingerprint()); -} - -TEST(HloProgramTest, FingerprintIgnoresDebugInfoStructure) { - TF_ASSERT_OK_AND_ASSIGN( - const std::unique_ptr hlo_program1, - ParseHloProgramString(R"( -module @foo { - func.func @main(%arg0: tensor<2x3xi32> loc("foo")) -> tensor<2x3xi32> { - return %arg0 : tensor<2x3xi32> loc("foo") - } loc("foo") -} loc("foo") -)")); - TF_ASSERT_OK_AND_ASSIGN( - const std::unique_ptr hlo_program2, - ParseHloProgramString(R"( -module @foo { - func.func @main(%arg0: tensor<2x3xi32> loc("bar")) -> tensor<2x3xi32> { - return %arg0 : tensor<2x3xi32> loc("baz") - } loc("qux") -} loc("quux") -)")); - - EXPECT_EQ(hlo_program1->Fingerprint(), hlo_program2->Fingerprint()); -} - TEST(HloProgramTest, BytesRoundTrip) { static constexpr absl::string_view kModule = R"( module @hlo_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { @@ -135,9 +63,9 @@ module @hlo_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas } } )"; - TF_ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); - TF_ASSERT_OK_AND_ASSIGN(auto serialized, program->ToBytes()); - TF_ASSERT_OK_AND_ASSIGN(auto deserialized, HloProgram::FromBytes(serialized)); + ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); + ASSERT_OK_AND_ASSIGN(auto serialized, program->ToBytes()); + ASSERT_OK_AND_ASSIGN(auto deserialized, HloProgram::FromBytes(serialized)); EXPECT_EQ(program->Fingerprint(), deserialized->Fingerprint()); } @@ -151,7 +79,7 @@ module @hlo_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas } } )"; - TF_ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); + ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); mlir::ModuleOp mlir_module = program->mlir_module(); xla::MaybeOwningMlirModule module = @@ -169,7 +97,7 @@ module @hlo_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas } } )"; - TF_ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); + ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); EXPECT_EQ(program->name(), "hlo_module"); } @@ -183,7 +111,7 @@ module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { } } )"; - TF_ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); + ASSERT_OK_AND_ASSIGN(auto program, ParseHloProgramString(kModule)); EXPECT_THAT(program->name(), ContainsRegex(R"(unnamed_[0-9a-f]+)")); } diff --git a/third_party/xla/xla/python/ifrt/ir/BUILD b/third_party/xla/xla/python/ifrt/ir/BUILD index 7b7d1a3d734e7d..76e3e968b407c0 100644 --- a/third_party/xla/xla/python/ifrt/ir/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/BUILD @@ -209,6 +209,7 @@ cc_library( "//xla/python/ifrt:serdes", "//xla/python/ifrt:serdes_default_version_accessor", "//xla/python/ifrt:serdes_version", + "//xla/python/ifrt/mlir:fingerprint_utils", "//xla/python/pjrt_ifrt:xla_ifrt", "//xla/tsl/platform:errors", "@com_google_absl//absl/container:flat_hash_map", @@ -247,11 +248,13 @@ xla_cc_test( ":ifrt_ir_program", "//xla/client:executable_build_options", "//xla/pjrt:pjrt_executable", + "//xla/python/ifrt/ir/support:module_parsing", + "//xla/python/ifrt/mlir:fingerprint_utils", "//xla/service:device_assignment", - "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", + "@llvm-project//mlir:IR", ], ) @@ -946,7 +949,6 @@ cc_library( "//xla/python/ifrt/hlo:hlo_program", "//xla/service:hlo_module_config", "//xla/service:hlo_proto_cc", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.cc index 8632b76c3e9f66..5a6d63d5a22cdd 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.cc +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.cc @@ -43,9 +43,11 @@ limitations under the License. #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/executable.h" #include "xla/python/ifrt/ir/ifrt_ir_compile_options.pb.h" +#include "xla/python/ifrt/mlir/fingerprint_utils.h" #include "xla/python/ifrt/rtti.h" #include "xla/python/ifrt/serdes_version.h" #include "xla/python/pjrt_ifrt/xla_compiler.h" +#include "xla/tsl/platform/errors.h" #include "tsl/platform/human_readable_json.h" namespace xla { @@ -56,6 +58,17 @@ char SerializeIfrtIRProgramOptions::ID = 0; char DeserializeIfrtIRProgramOptions::ID = 0; char IfrtIRCompileOptions::ID = 0; +absl::StatusOr IfrtIRProgram::Fingerprint() const { + absl::StatusOr fingerprint = FingerprintModuleOp(mlir_module); + if (!fingerprint.ok()) { + absl::Status status = fingerprint.status(); + tsl::errors::AppendToMessage( + &status, "Failed while calculating IfrtIRProgram fingerprint"); + return status; + } + return *fingerprint; +} + absl::StatusOr> GetIfrtIRCompileOptions( std::unique_ptr options) { if (!isa(options.get())) { diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.h b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.h index fdf7f99f86e60b..d7fed522a1ca7f 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.h +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program.h @@ -48,7 +48,6 @@ limitations under the License. #include "xla/python/ifrt/serdes.h" #include "xla/python/ifrt/serdes_default_version_accessor.h" #include "xla/python/ifrt/serdes_version.h" -#include "xla/tsl/platform/errors.h" namespace xla { namespace ifrt { @@ -70,6 +69,10 @@ struct IfrtIRProgram : RTTIExtends { // Returns true if the program exclusively owns the MLIR context. bool OwnsMlirContext() const { return mlir_context != nullptr; } + // Returns a fingerprint of the IFRT IR program. Two IFRT IR programs are + // equivalent if their fingerprints are the same. May ignore debug info. + absl::StatusOr Fingerprint() const; + // Key for the `fill_all_statuses` attribute in the custom_options attribute // map. If set to true, all executables will have their status filled if // `options.fill_status` is set. Otherwise, only leaf executables will have diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program_test.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program_test.cc index 0684d381d93c5a..64545e0b59b639 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program_test.cc +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_program_test.cc @@ -15,17 +15,23 @@ limitations under the License. #include "xla/python/ifrt/ir/ifrt_ir_program.h" +#include #include #include #include #include #include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" #include "xla/client/executable_build_options.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/python/ifrt/ir/ifrt_ir_compile_options.pb.h" +#include "xla/python/ifrt/ir/support/module_parsing.h" +#include "xla/python/ifrt/mlir/fingerprint_utils.h" #include "xla/service/device_assignment.h" -#include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" namespace xla { @@ -43,19 +49,19 @@ TEST(IfrtIRCompileOptionsTest, ToFromProto) { xla::ExecutableBuildOptions build_option; build_option.set_device_assignment(xla::DeviceAssignment(2, 4)); src.executable_build_options = build_option; - TF_ASSERT_OK_AND_ASSIGN(CompileOptionsProto compile_options_proto, - src.ToProto()); + ASSERT_OK_AND_ASSIGN(CompileOptionsProto compile_options_proto, + src.ToProto()); proto.mutable_compile_option_overrides()->insert( {absl::StrCat("key", i), compile_options_proto}); } - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr options, - IfrtIRCompileOptions::FromProto(proto)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr options, + IfrtIRCompileOptions::FromProto(proto)); EXPECT_EQ(options->compile_options_overrides->size(), 4); EXPECT_EQ(options->device_assignments.size(), num_devices); - TF_ASSERT_OK_AND_ASSIGN(IfrtIrCompileOptionsProto from_to_proto, - options->ToProto()); + ASSERT_OK_AND_ASSIGN(IfrtIrCompileOptionsProto from_to_proto, + options->ToProto()); for (int i = 0; i < 4; ++i) { std::string key = absl::StrCat("key", i); @@ -70,6 +76,163 @@ TEST(IfrtIRCompileOptionsTest, ToFromProto) { std::vector(proto.device_ids().begin(), proto.device_ids().end())); } +TEST(IfrtIRProgramTest, IfrtIrModuleSameFingerprint) { + static constexpr absl::string_view kIfrtModule = R"( +!array = !ifrt.array, #ifrt.sharding_param<1 to [0] on 1>, [0]> +module { + func.func @main(%arg0: !array) -> !array attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @add_one::@main(%arg0) on devices [0] + : (!array) -> !array + return %0 : !array + } + + module @add_one { + func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> { + %0 = stablehlo.constant dense<1> : tensor<2xi32> + %1 = stablehlo.add %arg0, %0 : tensor<2xi32> + return %1 : tensor<2xi32> + } + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module, + support::ParseMlirModuleString(kIfrtModule, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module)); + EXPECT_EQ(fp1, fp2); +} + +TEST(IfrtIRProgramTest, IfrtIrModuleDifferentDevicesDifferentFingerprints) { + static constexpr absl::string_view kIfrtModuleDevice0 = R"( +!array0 = !ifrt.array, #ifrt.sharding_param<1 to [0] on 1>, [0]> +module { + func.func @main(%arg0: !array0) -> !array0 attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @add_one::@main(%arg0) on devices [0] + : (!array0) -> !array0 + return %0 : !array0 + } + + module @add_one { + func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> { + return %arg0 : tensor<2xi32> + } + } +} +)"; + static constexpr absl::string_view kIfrtModuleDevice1 = R"( +!array1 = !ifrt.array, #ifrt.sharding_param<1 to [0] on 1>, [1]> +module { + func.func @main(%arg0: !array1) -> !array1 attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @add_one::@main(%arg0) on devices [1] + : (!array1) -> !array1 + return %0 : !array1 + } + + module @add_one { + func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> { + return %arg0 : tensor<2xi32> + } + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN( + mlir::OwningOpRef module0, + support::ParseMlirModuleString(kIfrtModuleDevice0, context)); + ASSERT_OK_AND_ASSIGN( + mlir::OwningOpRef module1, + support::ParseMlirModuleString(kIfrtModuleDevice1, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp0, FingerprintModuleOp(*module0)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + EXPECT_NE(fp0, fp1); +} + +TEST(IfrtIRProgramTest, IfrtIrModuleDifferentShardingDifferentFingerprints) { + static constexpr absl::string_view kIfrtModuleSharding1 = R"( +!array = !ifrt.array, #ifrt.sharding_param<1 to [0] on 2>, [0, 1]> +module { + func.func @main(%arg0: !array) -> !array attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @identity::@main(%arg0) on devices [0, 1] + : (!array) -> !array + return %0 : !array + } + + module @identity { + func.func @main(%arg0: tensor<4xi32>) -> tensor<4xi32> { + return %arg0 : tensor<4xi32> + } + } +} +)"; + static constexpr absl::string_view kIfrtModuleSharding2 = R"( +!array = !ifrt.array, #ifrt.sharding_param<2 to [0] on 2>, [0, 1]> +module { + func.func @main(%arg0: !array) -> !array attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @identity::@main(%arg0) on devices [0, 1] + : (!array) -> !array + return %0 : !array + } + + module @identity { + func.func @main(%arg0: tensor<4xi32>) -> tensor<4xi32> { + return %arg0 : tensor<4xi32> + } + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN( + mlir::OwningOpRef module1, + support::ParseMlirModuleString(kIfrtModuleSharding1, context)); + ASSERT_OK_AND_ASSIGN( + mlir::OwningOpRef module2, + support::ParseMlirModuleString(kIfrtModuleSharding2, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module2)); + EXPECT_NE(fp1, fp2); +} + +TEST(IfrtIRProgramTest, IfrtIrModuleIgnoresDebugInfo) { + static constexpr absl::string_view kIfrtModule1 = R"( +!array = !ifrt.array, #ifrt.sharding_param<1 to [0] on 1>, [0]> +module @ifrt_mod { + func.func @main(%arg0: !array loc("arg_loc1")) -> !array + attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @identity(%arg0) on devices [0] + : (!array) -> !array loc("call_loc1") + return %0 : !array loc("return_loc1") + } loc("func_loc1") + + func.func @identity(%arg0: tensor<2xi32>) -> tensor<2xi32> { + return %arg0 : tensor<2xi32> + } +} loc("module_loc1") +)"; + static constexpr absl::string_view kIfrtModule2 = R"( +!array = !ifrt.array, #ifrt.sharding_param<1 to [0] on 1>, [0]> +module @ifrt_mod { + func.func @main(%arg0: !array loc("arg_loc2")) -> !array attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.Call @identity(%arg0) on devices [0] + : (!array) -> !array loc("call_loc2") + return %0 : !array loc("return_loc2") + } loc("func_loc2") + + func.func @identity(%arg0: tensor<2xi32>) -> tensor<2xi32> { + return %arg0 : tensor<2xi32> + } +} loc("module_loc2") +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module1, + support::ParseMlirModuleString(kIfrtModule1, context)); + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module2, + support::ParseMlirModuleString(kIfrtModule2, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module2)); + EXPECT_EQ(fp1, fp2); +} + } // namespace } // namespace ifrt } // namespace xla diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/BUILD b/third_party/xla/xla/python/ifrt/ir/transforms/BUILD index 475195dc5d4148..da90f12f134efa 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/transforms/BUILD @@ -83,6 +83,7 @@ cc_library( "//xla/python/ifrt/ir:version", "//xla/python/ifrt/ir:vifrt", "//xla/python/ifrt/ir/support:sharding_conversions", + "//xla/python/ifrt/mlir:fingerprint_utils", "//xla/python/pjrt_ifrt", "//xla/python/pjrt_ifrt:xla_ifrt", "//xla/service:hlo_proto_cc", @@ -151,7 +152,6 @@ cc_library( "//xla/python/pjrt_ifrt:xla_ifrt", "//xla/service:device_assignment", "//xla/service/spmd/shardy:utils", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -165,7 +165,6 @@ cc_library( "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", - "@tsl//tsl/platform:fingerprint", ], ) diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_to_dot_pass.cc b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_to_dot_pass.cc index ea4a4f43cfc9e6..6d267f3cb1a01e 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_to_dot_pass.cc +++ b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_to_dot_pass.cc @@ -47,6 +47,7 @@ limitations under the License. #include "xla/python/ifrt/ir/transforms/debug.h" #include "xla/python/ifrt/ir/transforms/passes.h" #include "xla/python/ifrt/ir/transforms/utils.h" +#include "xla/python/ifrt/mlir/fingerprint_utils.h" #include "xla/python/ifrt/shape.h" #include "xla/service/hlo.pb.h" #include "xla/tsl/platform/env.h" @@ -277,10 +278,15 @@ void IfrtToDotPass::runOnOperation() { std::string module_name = module_op.getName().value_or("unknown").str(); // Include the module fingerprint in the file name to avoid exporting a // module multiple times. - std::string file_path = - tsl::io::JoinPath(dot_graph_dump_to, - absl::StrCat("ifrt_", module_name, "_", - MlirModuleFingerprint(module_op), ".dot")); + absl::StatusOr fingerprint = FingerprintModuleOp(module_op); + if (!fingerprint.ok()) { + LOG(WARNING) << "Failed to get fingerprint for module " << module_name + << ": " << fingerprint.status(); + return; + } + std::string file_path = tsl::io::JoinPath( + dot_graph_dump_to, + absl::StrCat("ifrt_", module_name, "_", *fingerprint, ".dot")); std::unique_ptr f; if (const absl::Status status = tsl::Env::Default()->NewWritableFile(file_path, &f); diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/utils.cc b/third_party/xla/xla/python/ifrt/ir/transforms/utils.cc index 653c6cf42761a2..8a65583c278447 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/utils.cc +++ b/third_party/xla/xla/python/ifrt/ir/transforms/utils.cc @@ -15,7 +15,6 @@ limitations under the License. #include "xla/python/ifrt/ir/transforms/utils.h" -#include #include #include #include @@ -74,9 +73,7 @@ limitations under the License. #include "xla/service/device_assignment.h" #include "xla/service/spmd/shardy/utils.h" #include "xla/status_macros.h" -#include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" -#include "tsl/platform/fingerprint.h" namespace xla { namespace ifrt { @@ -335,15 +332,6 @@ absl::StatusOr> ExpandPlatformNames( return expanded_platform_names; } -uint64_t MlirModuleFingerprint(mlir::ModuleOp module) { - std::string s; - llvm::raw_string_ostream os(s); - mlir::OpPrintingFlags flags; - flags.enableDebugInfo(false); - module.print(os, flags); - return tsl::Fingerprint64(os.str()); -} - absl::StatusOr GetModuleXlaCompileOverrides( mlir::StringAttr compile_options_key, std::shared_ptr< diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/utils.h b/third_party/xla/xla/python/ifrt/ir/transforms/utils.h index 9fd06205679c44..5cf8d7c9277ed1 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/utils.h +++ b/third_party/xla/xla/python/ifrt/ir/transforms/utils.h @@ -16,7 +16,6 @@ limitations under the License. #ifndef XLA_PYTHON_IFRT_IR_TRANSFORMS_UTILS_H_ #define XLA_PYTHON_IFRT_IR_TRANSFORMS_UTILS_H_ -#include #include #include #include @@ -36,6 +35,7 @@ limitations under the License. #include "mlir/Pass/Pass.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/python/ifrt/array_spec.h" +#include "xla/python/ifrt/compiler.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/dtype.h" @@ -86,9 +86,6 @@ absl::StatusOr> ExpandPlatformNames( // Returns a pretty string representation of the location. std::string GetPrettyLocation(mlir::Location loc); -// Returns a fingerprint of the provided module. -uint64_t MlirModuleFingerprint(mlir::ModuleOp module); - // Extracts the XlaCompileOptions overrides for the given atom program module. // Returns nullptr if no overrides are found. absl::StatusOr GetModuleXlaCompileOverrides( diff --git a/third_party/xla/xla/python/ifrt/ir/utils.cc b/third_party/xla/xla/python/ifrt/ir/utils.cc index 5c78fabf5aa361..f4a453781b5d1f 100644 --- a/third_party/xla/xla/python/ifrt/ir/utils.cc +++ b/third_party/xla/xla/python/ifrt/ir/utils.cc @@ -41,7 +41,6 @@ limitations under the License. #include "xla/service/hlo.pb.h" #include "xla/service/hlo_module_config.h" #include "xla/shape.h" -#include "xla/tsl/platform/statusor.h" namespace xla { namespace ifrt { diff --git a/third_party/xla/xla/python/ifrt/mlir/BUILD b/third_party/xla/xla/python/ifrt/mlir/BUILD new file mode 100644 index 00000000000000..55e8ff0613a6d1 --- /dev/null +++ b/third_party/xla/xla/python/ifrt/mlir/BUILD @@ -0,0 +1,47 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") +load("//xla:xla.default.bzl", "xla_cc_test") +load("//xla/tsl:tsl.bzl", "internal_visibility") +load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable") + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + default_visibility = internal_visibility([ + "//xla/python/ifrt:internal", + ]), +) + +cc_library( + name = "fingerprint_utils", + srcs = ["fingerprint_utils.cc"], + hdrs = ["fingerprint_utils.h"], + compatible_with = get_compatible_with_portable(), + deps = [ + "//xla:status_macros", + "//xla/service:hlo_proto_cc", + "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@highwayhash", + "@highwayhash//:arch_specific", + "@highwayhash//:hh_types", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:BytecodeOpInterface", + "@llvm-project//mlir:BytecodeWriter", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + ], +) + +xla_cc_test( + name = "fingerprint_utils_test", + srcs = ["fingerprint_utils_test.cc"], + deps = [ + ":fingerprint_utils", + "//xla/pjrt:mlir_to_hlo", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + "@llvm-project//mlir:IR", + ], +) diff --git a/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.cc b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.cc new file mode 100644 index 00000000000000..a97392f7b8db18 --- /dev/null +++ b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.cc @@ -0,0 +1,134 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/python/ifrt/mlir/fingerprint_utils.h" + +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "highwayhash/arch_specific.h" +#include "highwayhash/hh_types.h" +#include "highwayhash/highwayhash.h" +#include "llvm/Support/raw_ostream.h" +#include "mlir/Bytecode/BytecodeWriter.h" +#include "mlir/Bytecode/Encoding.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Support/LLVM.h" +#include "xla/service/hlo.pb.h" +#include "xla/status_macros.h" +#include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h" + +namespace xla { +namespace ifrt { + +namespace { + +// Calculates a HighwayHash fingerprint in a streaming manner. +class HighwayHashStream final : public llvm::raw_ostream { + public: + HighwayHashStream() : hash_(kHighwayHashKey), pos_(0) { SetUnbuffered(); } + + ~HighwayHashStream() override { CHECK_EQ(buffer_pos_, 0); } + + // Destructively calculates the fingerprint of the data consumed so far. + uint64_t fingerprint() && { + flush_buffer(); + highwayhash::HHResult64 result; + hash_.Finalize(&result); + return result; + } + + private: + // `HighwayHashCatT::Append` updates the hash directly from the source buffer + // if the provided buffer size is larger than `sizeof(HHPacket)`. + static constexpr size_t kBufferSize = sizeof(highwayhash::HHPacket); + + // Arbitrarily chosen, forever-unchanging hash key required by HighwayHash. + static constexpr highwayhash::HHKey kHighwayHashKey = { + 0x4451e30f87db9609ULL, + 0xca7358a1fd2737f8ULL, + 0x4b2c991fcee4fdeaULL, + 0x0b2658e18326f6baULL, + }; + + void write_impl(const char* Ptr, size_t Size) final { + // For tiny writes, it is more efficient to accumulate the data to a buffer + // first and flush it since `HighwayHashCatT::Append` is optimized for + // large writes. + static constexpr size_t kSmallWriteSize = 4; + static_assert(kSmallWriteSize <= kBufferSize); + + if (Size <= kSmallWriteSize) { + if (buffer_pos_ + Size > kBufferSize) { + flush_buffer(); + } + std::memcpy(buffer_ + buffer_pos_, Ptr, Size); + buffer_pos_ += Size; + } else { + flush_buffer(); + hash_.Append(Ptr, Size); + } + pos_ += Size; + } + + uint64_t current_pos() const final { return pos_; } + + void flush_buffer() { + if (buffer_pos_ > 0) { + hash_.Append(buffer_, buffer_pos_); + buffer_pos_ = 0; + } + } + + highwayhash::HighwayHashCatT hash_; + uint64_t pos_; + + char buffer_[kBufferSize]; + uint64_t buffer_pos_ = 0; +}; + +} // namespace + +absl::StatusOr FingerprintModuleOp(mlir::ModuleOp module_op) { + tsl::StatusScopedDiagnosticHandler diag_handler(module_op->getContext()); + + mlir::BytecodeWriterConfig config; + config.setElideLocations(true); + + // Use a version before `kUseListOrdering` due to an MLIR bug where use list + // ordering is not stable. + // + // TODO(b/503120525): Remove this workaround once + // https://github.com/llvm/llvm-project/pull/191942 lands. + config.setDesiredBytecodeVersion( + mlir::bytecode::BytecodeVersion::kLazyLoading); + + HighwayHashStream os; + mlir::LogicalResult result = mlir::writeBytecodeToFile(module_op, os, config); + ABSL_RETURN_IF_ERROR(diag_handler.consumeStatus()); + TF_RET_CHECK(mlir::succeeded(result)); + return std::move(os).fingerprint(); +} + +} // namespace ifrt +} // namespace xla diff --git a/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.h b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.h new file mode 100644 index 00000000000000..be573c4e429704 --- /dev/null +++ b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils.h @@ -0,0 +1,34 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_PYTHON_IFRT_MLIR_FINGERPRINT_UTILS_H_ +#define XLA_PYTHON_IFRT_MLIR_FINGERPRINT_UTILS_H_ + +#include + +#include "absl/status/statusor.h" +#include "mlir/IR/BuiltinOps.h" + +namespace xla { +namespace ifrt { + +// Returns a fingerprint of the given MLIR module. Two MLIR modules are +// equivalent if their fingerprints are the same. May ignore debug info. +absl::StatusOr FingerprintModuleOp(mlir::ModuleOp module_op); + +} // namespace ifrt +} // namespace xla + +#endif // XLA_PYTHON_IFRT_MLIR_FINGERPRINT_UTILS_H_ diff --git a/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils_test.cc b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils_test.cc new file mode 100644 index 00000000000000..a3192a1cfd20db --- /dev/null +++ b/third_party/xla/xla/python/ifrt/mlir/fingerprint_utils_test.cc @@ -0,0 +1,131 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/python/ifrt/mlir/fingerprint_utils.h" + +#include + +#include +#include +#include "absl/strings/string_view.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "xla/pjrt/mlir_to_hlo.h" + +namespace xla { +namespace ifrt { +namespace { + +TEST(FingerprintUtilsTest, IdenticalModulesHaveSameFingerprint) { + static constexpr absl::string_view kModule = R"( +module { + func.func @main(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<1.000000e+00> : tensor + %1 = stablehlo.add %arg0, %0 : tensor + return %1 : tensor + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module, + xla::ParseMlirModuleString(kModule, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module)); + EXPECT_EQ(fp1, fp2); +} + +TEST(FingerprintUtilsTest, DistinctStablehloModulesHaveDifferentFingerprints) { + static constexpr absl::string_view kModule1 = R"( +module { + func.func @main(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<1.000000e+00> : tensor + %1 = stablehlo.add %arg0, %0 : tensor + return %1 : tensor + } +} +)"; + static constexpr absl::string_view kModule2 = R"( +module { + func.func @main(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor + %1 = stablehlo.add %arg0, %0 : tensor + return %1 : tensor + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module1, + xla::ParseMlirModuleString(kModule1, context)); + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module2, + xla::ParseMlirModuleString(kModule2, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module2)); + EXPECT_NE(fp1, fp2); +} + +TEST(FingerprintUtilsTest, IgnoresDebugLocations) { + static constexpr absl::string_view kModule1 = R"( +module @foo { + func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> { + return %arg0 : tensor<2x3xi32> loc("foo") + } +} +)"; + static constexpr absl::string_view kModule2 = R"( +module @foo { + func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> { + return %arg0 : tensor<2x3xi32> loc("bar") + } +} +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module1, + xla::ParseMlirModuleString(kModule1, context)); + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module2, + xla::ParseMlirModuleString(kModule2, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module2)); + EXPECT_EQ(fp1, fp2); +} + +TEST(FingerprintUtilsTest, IgnoresDebugLocationStructure) { + static constexpr absl::string_view kModule1 = R"( +module @foo { + func.func @main(%arg0: tensor<2x3xi32> loc("foo")) -> tensor<2x3xi32> { + return %arg0 : tensor<2x3xi32> loc("foo") + } loc("foo") +} loc("foo") +)"; + static constexpr absl::string_view kModule2 = R"( +module @foo { + func.func @main(%arg0: tensor<2x3xi32> loc("bar")) -> tensor<2x3xi32> { + return %arg0 : tensor<2x3xi32> loc("baz") + } loc("qux") +} loc("quux") +)"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module1, + xla::ParseMlirModuleString(kModule1, context)); + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module2, + xla::ParseMlirModuleString(kModule2, context)); + ASSERT_OK_AND_ASSIGN(uint64_t fp1, FingerprintModuleOp(*module1)); + ASSERT_OK_AND_ASSIGN(uint64_t fp2, FingerprintModuleOp(*module2)); + EXPECT_EQ(fp1, fp2); +} + +} // namespace +} // namespace ifrt +} // namespace xla