From 3ba5685d8b1142fa7bab4d965b06064ddf4cb6cd Mon Sep 17 00:00:00 2001 From: kokol16 Date: Sat, 5 Sep 2026 07:37:46 -0400 Subject: [PATCH 01/32] Fix TopK GPU loop index overflow TopK used a 32-bit loop index. Adding the thread stride near INT32_MAX could overflow the index and make the kernel read outside the tensor. I changed the loop index to int64. It is converted back to int when TopK stores the result because the output indices are int32. --- tensorflow/core/kernels/topk_op_gpu.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/topk_op_gpu.h b/tensorflow/core/kernels/topk_op_gpu.h index abd07645df3b6d..ab853e196ac45c 100644 --- a/tensorflow/core/kernels/topk_op_gpu.h +++ b/tensorflow/core/kernels/topk_op_gpu.h @@ -238,12 +238,12 @@ __device__ void heapTopK(const T* __restrict__ input, int length, int k, // Now iterate over the remaining items. // If an item is smaller than the min element, it is not amongst the top k. // Otherwise, replace the min element with it and push upwards. - for (int index = heap_end_index; index < length; index += step_size) { + for (int64_t index = heap_end_index; index < length; index += step_size) { // We prefer elements with lower indices. This is given here. // Later elements automatically have higher indices, so can be discarded. if (input[index] > heap.root().value) { // This element should replace the min. - heap.replace_root({index, input[index]}, k); + heap.replace_root({static_cast(index), input[index]}, k); } } From 93fbfd4464d61c10badf9f905458be998c1f5334 Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Thu, 17 Sep 2026 06:24:59 +0530 Subject: [PATCH 02/32] Fix argument parsing for toco_compatible flag --- tensorflow/python/tools/optimize_for_inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/tools/optimize_for_inference.py b/tensorflow/python/tools/optimize_for_inference.py index 148511cd1651c1..8d7b409391bffd 100644 --- a/tensorflow/python/tools/optimize_for_inference.py +++ b/tensorflow/python/tools/optimize_for_inference.py @@ -153,7 +153,9 @@ def parse_args(): """) parser.add_argument( "--toco_compatible", - type=bool, + nargs="?", + const=True, + type="bool", default=False, help="""\ If true, only use ops compatible with Tensorflow From f1fbb409492c6dfde2738f4c6ab47a55c96adf7f Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Thu, 17 Sep 2026 06:26:47 +0530 Subject: [PATCH 03/32] Update tensorflow/python/tools/optimize_for_inference.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tensorflow/python/tools/optimize_for_inference.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorflow/python/tools/optimize_for_inference.py b/tensorflow/python/tools/optimize_for_inference.py index 8d7b409391bffd..331f9b071da1fe 100644 --- a/tensorflow/python/tools/optimize_for_inference.py +++ b/tensorflow/python/tools/optimize_for_inference.py @@ -153,8 +153,6 @@ def parse_args(): """) parser.add_argument( "--toco_compatible", - nargs="?", - const=True, type="bool", default=False, help="""\ From 521fad051180e3dc4f84fd223db16a01d7d42b61 Mon Sep 17 00:00:00 2001 From: Jiayang Liu Date: Thu, 17 Sep 2026 12:47:20 +0800 Subject: [PATCH 04/32] Fix tf.signal.kaiser_window losing its second derivative at beta == 0 kaiser_window called bessel_i0e with abs(beta). Since the second derivative of abs vanishes at the origin, the resulting cusp made the second derivative of the window with respect to beta evaluate to zero at beta == 0, even though the window is I0(beta * s) / I0(beta), an even analytic function of beta. bessel_i0e is an even function, so beta is now passed to it with its original sign. The exponential scaling factor keeps using abs(beta), so negative beta values still produce the same result as positive ones and the computation stays numerically stable for large beta. Adds a unit test for the second derivative at beta == 0. Verified that the forward values are bit-for-bit unchanged for beta >= 0 in float16/float32/ float64, and that negative beta now matches positive beta exactly. Fixes #127245 --- .../kernel_tests/signal/window_ops_test.py | 40 +++++++++++++++++++ tensorflow/python/ops/signal/window_ops.py | 18 ++++++--- 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/tensorflow/python/kernel_tests/signal/window_ops_test.py b/tensorflow/python/kernel_tests/signal/window_ops_test.py index d09cfea991e62d..0eea6763f6ea5e 100644 --- a/tensorflow/python/kernel_tests/signal/window_ops_test.py +++ b/tensorflow/python/kernel_tests/signal/window_ops_test.py @@ -20,11 +20,14 @@ from absl.testing import parameterized import numpy as np +from tensorflow.python.eager import backprop from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.kernel_tests.signal import test_util +from tensorflow.python.ops import math_ops from tensorflow.python.ops.signal import window_ops from tensorflow.python.platform import test @@ -119,6 +122,43 @@ def test_kaiser_window_negative_beta(self, window_length, beta, tf_dtype_tol): pos_beta_win, neg_beta_win, tf_dtype_tol[1], tf_dtype_tol[1] ) + def test_kaiser_window_second_derivative_at_zero_beta(self): + """Check that kaiser_window is twice differentiable at beta == 0. + + The window equals I0(beta * s) / I0(beta), where s does not depend on + beta and I0 is an even entire function. The window is therefore an even + analytic function of beta and must have a well defined second derivative + at beta == 0. Applying abs() to beta inside the implementation puts a + cusp at the origin, where the second derivative of abs vanishes, and + makes this second derivative evaluate to zero. + """ + window_length = 10 + weights = np.arange(1, window_length + 1, dtype=np.float64)**2 + beta = constant_op.constant(0.0, dtypes.float64) + + def target(beta_value): + window = window_ops.kaiser_window( + window_length, beta_value, dtypes.float64) + return math_ops.reduce_sum(window * weights) + + with backprop.GradientTape() as tape_outer: + tape_outer.watch(beta) + with backprop.GradientTape() as tape_inner: + tape_inner.watch(beta) + value = target(beta) + first = tape_inner.gradient(value, beta) + second = tape_outer.gradient(first, beta) + + # The window is even in beta, so the first derivative vanishes at zero. + self.assertAllClose(0.0, first, atol=1e-12) + # For window(beta) = I0(beta * s) / I0(beta), the second derivative at + # beta == 0 is (s^2 - 1) / 2. + halflen = (window_length - 1) / 2.0 + arg = np.arange(-halflen, halflen + 0.1, dtype=np.float64) + shape = np.sqrt(np.maximum(0.0, 1.0 - (arg / halflen)**2)) + expected = np.sum(weights * (shape**2 - 1.0) / 2.0) + self.assertAllClose(expected, second, rtol=1e-6, atol=1e-9) + @parameterized.parameters( itertools.product(_WINDOW_LENGTHS, (False, True), _TF_DTYPE_TOLERANCE) ) diff --git a/tensorflow/python/ops/signal/window_ops.py b/tensorflow/python/ops/signal/window_ops.py index 37d09ce0ea7ed5..d52e4e83f15741 100644 --- a/tensorflow/python/ops/signal/window_ops.py +++ b/tensorflow/python/ops/signal/window_ops.py @@ -79,15 +79,21 @@ def kaiser_window(window_length, beta=12., dtype=dtypes.float32, name=None): dtype=dtypes.float32) # Convert everything into given dtype which can be float16. arg = math_ops.cast(arg, dtype=dtype) - # I0 (modified Bessel function of the first kind) is an even function, - # i.e., I0(-x) = I0(x). Take abs(beta) so that negative beta values - # produce the same (correct) result as positive ones. - beta = math_ops.abs(math_ops.cast(beta, dtype=dtype)) one = math_ops.cast(1.0, dtype=dtype) halflen_float = math_ops.cast(halflen_float, dtype=dtype) - num = beta * math_ops.sqrt(nn_ops.relu( + shape = math_ops.sqrt(nn_ops.relu( one - math_ops.square(arg / halflen_float))) - window = math_ops.exp(num - beta) * ( + # `bessel_i0e` is an even function, i.e. i0e(-x) = i0e(x), so `beta` is + # passed to it with its original sign. Applying `abs` to `beta` before + # calling `bessel_i0e` would introduce a cusp at beta == 0 (the second + # derivative of `abs` vanishes at the origin), which makes the second + # derivative of the window with respect to `beta` evaluate to zero. The + # exponential factor still uses `abs(beta)`, so that negative `beta` + # values produce the same result as positive ones. + beta = math_ops.cast(beta, dtype=dtype) + beta_abs = math_ops.abs(beta) + num = beta * shape + window = math_ops.exp(beta_abs * shape - beta_abs) * ( special_math_ops.bessel_i0e(num) / special_math_ops.bessel_i0e(beta)) return window From c0f15896aefd3a6dee3ccc74f102cc959154c077 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Wed, 16 Sep 2026 22:13:26 -0700 Subject: [PATCH 05/32] Reject growing TensorLists in XLA instead of dropping the write A tf.TensorArray created with dynamic_size=True lowers to TensorListSetItem with resize_if_index_out_of_bounds set, and the non-XLA kernel grows the list when the index is past the end. The tf2xla lowering never read that attribute, so under jit_compile=True the write was silently discarded and the program returned a truncated result with no error: ta = tf.TensorArray(tf.bool, size=0, dynamic_size=True) # writes at 0 and 1 inside a while_loop, then stack() eager -> [True, True] tf.function -> [True, True] tf.function(jit_compile=True) -> [] The loss is not specific to size=0. XLA freezes the buffer at the size TensorListReserve allocated and drops every write beyond it, so size=1 with two writes returns [True] just as quietly. ExecuteTensorListSetItem lowers to a DynamicUpdateSlice, which clamps its start indices rather than failing, so the out-of-bounds element lands back inside the buffer. XLA cannot support the growth itself, since that needs a reallocation and the buffer shape is static, and it cannot bounds check at runtime either, since a compiled program cannot raise a data-dependent error. Reject the op at compile time instead, so the failure is a clear message rather than a wrong answer. This also makes the two lowerings of the same user-level API agree. The older TensorArray path already rejects dynamic_size unconditionally in TensorArrayOp in tensor_array_ops.cc; only the TensorList path that TF2 actually emits was missing the check. Fixes #127528 --- tensorflow/compiler/tests/BUILD | 2 + .../compiler/tests/tensor_list_ops_test.py | 82 +++++++++++++++++++ .../tf2xla/kernels/tensor_list_ops.cc | 25 +++++- 3 files changed, 108 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 56b3fc99746d7d..c1ee3e5164b5c2 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -1820,6 +1820,8 @@ tf_xla_py_strict_test( "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:list_ops", "//tensorflow/python/ops:map_fn", + "//tensorflow/python/ops:tensor_array_ops", + "//tensorflow/python/ops:while_loop", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 05afb8f781c451..059c7f2120affb 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -28,6 +28,8 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import list_ops from tensorflow.python.ops import map_fn +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.ops import while_loop from tensorflow.python.platform import test @@ -302,6 +304,86 @@ def testZerosLikeForTensorList(self): self.assertAllEqual(z.shape.as_list(), [None]) self.assertAllEqual(z, [0.0, 0.0]) + def testSetItemResizeIsRejected(self): + # Regression test for GitHub issue 127528. A TensorListSetItem that is + # allowed to grow the list cannot be compiled, because the buffer that + # TensorListReserve allocated has a static shape. It must be rejected at + # compile time rather than silently dropping the write: the lowering ends + # in a DynamicUpdateSlice, which clamps its start indices instead of + # failing, so an out-of-bounds element lands back inside the buffer. + @def_function.function(jit_compile=True) + def grows(): + l = list_ops.tensor_list_reserve( + element_shape=[], element_dtype=dtypes.float32, num_elements=1 + ) + l = list_ops.tensor_list_set_item( + input_handle=l, + index=0, + item=constant_op.constant(1.0), + resize_if_index_out_of_bounds=True, + ) + # Stack inside the compiled function so that no TensorList crosses the + # XLA/TF boundary, which would raise a different error and stop this + # test from covering the rejection above. + return list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + + with self.session(): + with self.assertRaisesRegex( + errors.UnimplementedError, + "TensorLists that grow on an out-of-bounds write", + ): + self.evaluate(grows()) + + def testSetItemWithoutResizeStillCompiles(self): + # The rejection above keys off the attribute rather than the index, so an + # ordinary fixed-size write must still compile and run. + @def_function.function(jit_compile=True) + def fixed(): + l = list_ops.tensor_list_reserve( + element_shape=[], element_dtype=dtypes.float32, num_elements=1 + ) + l = list_ops.tensor_list_set_item( + input_handle=l, + index=0, + item=constant_op.constant(1.0), + resize_if_index_out_of_bounds=False, + ) + return list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + + with self.session(): + self.assertAllEqual(self.evaluate(fixed()), [1.0]) + + def testDynamicSizeTensorArrayIsRejected(self): + # End to end case for GitHub issue 127528. A dynamic_size TensorArray + # written past its initial size used to compile and return a truncated + # result with no error, so this asserts the compile-time failure instead. + # The whole function is jit-compiled so the list stays inside XLA rather + # than crossing the XLA/TF boundary at the stack op. + @def_function.function(jit_compile=True) + def f(): + ta = tensor_array_ops.TensorArray( + dtypes.float32, size=0, dynamic_size=True + ) + + def cond(i, ta): + del ta # Unused. + return i < 2 + + def body(i, ta): + return i + 1, ta.write(i, constant_op.constant(1.0)) + + _, out = while_loop.while_loop( + cond, body, [constant_op.constant(0), ta] + ) + return out.stack() + + with self.session(): + with self.assertRaisesRegex( + errors.UnimplementedError, + "TensorLists that grow on an out-of-bounds write", + ): + self.evaluate(f()) + def testInvalidSplitLength(self): with self.session(), self.test_scope(): tensor_list_split = list_ops.tensor_list_split( diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index 1471e53d63d6c7..cb38810442ad03 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -649,7 +649,30 @@ REGISTER_XLA_OP( class TensorListSetItemOp : public XlaOpKernel { public: - explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + // A tf.TensorArray built with dynamic_size=True lowers to TensorListSetItem + // with this attribute set, and the non-XLA kernel grows the list when the + // index is past the end. XLA needs static shapes, so the buffer that + // TensorListReserve allocated cannot grow. + // + // Without this check the write is silently dropped: + // ExecuteTensorListSetItem lowers to a DynamicUpdateSlice, which clamps + // its start indices rather than failing, so the element lands back inside + // the existing buffer and the program returns a truncated result with no + // error at all. + // + // The older TensorArray lowering already rejects the same thing, see + // TensorArrayOp in tensor_array_ops.cc. + bool resize_if_index_out_of_bounds; + OP_REQUIRES_OK(ctx, ctx->GetAttr("resize_if_index_out_of_bounds", + &resize_if_index_out_of_bounds)); + OP_REQUIRES(ctx, !resize_if_index_out_of_bounds, + errors::Unimplemented( + "TensorLists that grow on an out-of-bounds write are not " + "supported by XLA. This is typically a tf.TensorArray " + "created with dynamic_size=True; give it a static size " + "large enough to hold every write instead.")); + } void Compile(XlaOpKernelContext* ctx) override { xla::XlaOp list = ctx->Input(0); From a0e7c96a8c394a7a6f08d104736a9188b90088b0 Mon Sep 17 00:00:00 2001 From: Jiayang Liu Date: Thu, 17 Sep 2026 14:05:46 +0800 Subject: [PATCH 06/32] Avoid catastrophic cancellation in kaiser_window's exponential factor The scaling factor was computed as `beta_abs * shape - beta_abs`. Near the centre of the window `shape` is close to one, so the two products nearly cancel and the difference loses precision. For large `beta` the result can round to exactly zero, making the factor evaluate to 1.0 instead of the correct value. Factor `beta_abs` out so that the subtraction is performed first: `beta_abs * (shape - 1)`. The two forms are algebraically identical and differ only in floating-point rounding. Measured against a higher precision reference, the new form is the more accurate one, by up to ~400x in float64 for large `beta`. Forward values therefore change slightly, by at most ~1.6e-2 relative in float16 at `beta = 60`, ~2.0e-6 in float32 and ~3.5e-15 in float64, where the new values are closer to the exact result. The second derivative at `beta == 0` and the treatment of negative `beta` are unaffected. --- tensorflow/python/ops/signal/window_ops.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/signal/window_ops.py b/tensorflow/python/ops/signal/window_ops.py index d52e4e83f15741..1177455cc9ff2d 100644 --- a/tensorflow/python/ops/signal/window_ops.py +++ b/tensorflow/python/ops/signal/window_ops.py @@ -93,7 +93,11 @@ def kaiser_window(window_length, beta=12., dtype=dtypes.float32, name=None): beta = math_ops.cast(beta, dtype=dtype) beta_abs = math_ops.abs(beta) num = beta * shape - window = math_ops.exp(beta_abs * shape - beta_abs) * ( + # `beta_abs` is factored out so that the subtraction is performed first. + # Evaluating `beta_abs * shape - beta_abs` instead loses precision to + # catastrophic cancellation when `beta_abs` is large and `shape` is close + # to one, which happens near the center of the window. + window = math_ops.exp(beta_abs * (shape - one)) * ( special_math_ops.bessel_i0e(num) / special_math_ops.bessel_i0e(beta)) return window From 4a383a3eb61c717dfa33fbd52a14a3305118c653 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 17 Sep 2026 07:03:27 -0700 Subject: [PATCH 07/32] Correct the comment describing how the write is dropped The comment credited the loss entirely to DynamicUpdateSlice clamping, but that only covers the case where the element fits and the index is past the end. When the element does not fit at all, which is what a list reserved with size 0 hits, ExecuteTensorListSetItem takes an earlier return that logs a warning and returns the list unchanged with an OK status. Describe both, since the reported repro goes through the first one. No functional change. --- .../compiler/tf2xla/kernels/tensor_list_ops.cc | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index cb38810442ad03..b2a2a61e39c165 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -655,11 +655,14 @@ class TensorListSetItemOp : public XlaOpKernel { // index is past the end. XLA needs static shapes, so the buffer that // TensorListReserve allocated cannot grow. // - // Without this check the write is silently dropped: - // ExecuteTensorListSetItem lowers to a DynamicUpdateSlice, which clamps - // its start indices rather than failing, so the element lands back inside - // the existing buffer and the program returns a truncated result with no - // error at all. + // Without this check the write is silently dropped, in one of two ways + // depending on the shapes. When the element does not fit at all, as for a + // list reserved with size 0, ExecuteTensorListSetItem takes its early + // return, logs a warning and hands back the list unchanged with an OK + // status. When the element does fit but the index is past the end, the + // DynamicUpdateSlice it lowers to clamps its start indices rather than + // failing, so the element overwrites an earlier one. Either way the + // program returns a truncated result with no error at all. // // The older TensorArray lowering already rejects the same thing, see // TensorArrayOp in tensor_array_ops.cc. From c9b1de7c64fb5ce016323a4b6e6d3ff255fb5850 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 17 Sep 2026 07:10:58 -0700 Subject: [PATCH 08/32] Rename the list variable in the new tests and fix a stale comment Uses tensor_list rather than l in the two tests added here, per review feedback. The rest of this file uses l under its g-bad-name pragma, so this leaves the new tests inconsistent with their neighbours; flagging that for the reviewer to settle either way. Also corrects the same inaccuracy fixed in the kernel comment: the write is dropped either by the early return when the element does not fit, or by DynamicUpdateSlice clamping an out-of-range index, not by clamping alone. No functional change; both tests were re-checked against an unpatched 2.21.0 wheel and still fail there as intended. --- .../compiler/tests/tensor_list_ops_test.py | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 059c7f2120affb..2f0e2535a37fcc 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -308,16 +308,17 @@ def testSetItemResizeIsRejected(self): # Regression test for GitHub issue 127528. A TensorListSetItem that is # allowed to grow the list cannot be compiled, because the buffer that # TensorListReserve allocated has a static shape. It must be rejected at - # compile time rather than silently dropping the write: the lowering ends - # in a DynamicUpdateSlice, which clamps its start indices instead of - # failing, so an out-of-bounds element lands back inside the buffer. + # compile time rather than silently dropping the write, which the lowering + # does either by returning the list unchanged when the element does not + # fit, or by letting DynamicUpdateSlice clamp an out-of-range index so the + # element overwrites an earlier one. @def_function.function(jit_compile=True) def grows(): - l = list_ops.tensor_list_reserve( + tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 ) - l = list_ops.tensor_list_set_item( - input_handle=l, + tensor_list = list_ops.tensor_list_set_item( + input_handle=tensor_list, index=0, item=constant_op.constant(1.0), resize_if_index_out_of_bounds=True, @@ -325,7 +326,9 @@ def grows(): # Stack inside the compiled function so that no TensorList crosses the # XLA/TF boundary, which would raise a different error and stop this # test from covering the rejection above. - return list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + return list_ops.tensor_list_stack( + tensor_list, element_dtype=dtypes.float32 + ) with self.session(): with self.assertRaisesRegex( @@ -339,16 +342,18 @@ def testSetItemWithoutResizeStillCompiles(self): # ordinary fixed-size write must still compile and run. @def_function.function(jit_compile=True) def fixed(): - l = list_ops.tensor_list_reserve( + tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 ) - l = list_ops.tensor_list_set_item( - input_handle=l, + tensor_list = list_ops.tensor_list_set_item( + input_handle=tensor_list, index=0, item=constant_op.constant(1.0), resize_if_index_out_of_bounds=False, ) - return list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + return list_ops.tensor_list_stack( + tensor_list, element_dtype=dtypes.float32 + ) with self.session(): self.assertAllEqual(self.evaluate(fixed()), [1.0]) From 2d57c5771ccd3b950d1732eefe481355252c6a1e Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 17 Sep 2026 07:14:35 -0700 Subject: [PATCH 09/32] Reject growing TensorLists in the MLIR bridge as well The classic tf2xla lowering is not the only path that compiles these ops. When the MLIR bridge runs, CreateTensorListOpsDecompositionPass is added in compile_mlir_util.cc under lower_to_xla_hlo, and HandleTensorListSetItemOp decomposed the op without reading resize_if_index_out_of_bounds, so the same write was dropped just as quietly there. Reject it in the pass too, with the same message the kernel uses, and add a lit case covering it. The buffer the pass builds has a static shape for the same reason the kernel's does, and cutil::SetElement lowers to a DynamicUpdateSlice that clamps an out-of-range index rather than failing. The Python tests now accept InvalidArgumentError as well as UnimplementedError, since emitOpError surfaces as the former, and pin the assertion on the message instead of the type so either path satisfies it. --- .../tests/tensor_list_ops_decomposition.mlir | 18 ++++++++++++++++++ .../tensor_list_ops_decomposition.cc | 17 +++++++++++++++++ .../compiler/tests/tensor_list_ops_test.py | 12 ++++++++++-- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir index 51d631c8515b84..3d4badf472af5d 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir @@ -628,3 +628,21 @@ func.func @main(%arg0: tensor<*xi32>) -> () { %push = "tf.TensorListPushBack"(%tl, %elem) : (tensor>>, tensor) -> tensor>> func.return } + +// ----- + +// Tests that the pass reports an error on a set item that is allowed to grow +// the list, which a tf.TensorArray with dynamic_size=True emits. The buffer has +// a static shape and cannot grow, so decomposing it would silently drop the +// write. + +func.func @main() -> () { + %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + %num = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + %index = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor) -> tensor>> + %elem = "tf._SomeOp"() : () -> tensor + // expected-error @+1 {{TensorLists that grow on an out-of-bounds write are not supported by XLA}} + %set = "tf.TensorListSetItem"(%tl, %index, %elem) {resize_if_index_out_of_bounds = true} : (tensor>>, tensor, tensor) -> tensor>> + func.return +} diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc index ea6b06b5dfe2fb..b94576fdc05a64 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc @@ -714,6 +714,23 @@ LogicalResult HandleTensorListGetItemOp( LogicalResult HandleTensorListSetItemOp( TF::TensorListSetItemOp set_item, llvm::SmallDenseMap* buffer_to_size) { + // A tf.TensorArray built with dynamic_size=True sets this attribute, and the + // non-XLA kernel grows the list when the index is past the end. The buffer + // built here has a static shape and cannot grow, and cutil::SetElement lowers + // to a DynamicUpdateSlice, which clamps an out-of-range index rather than + // failing. Decomposing the op regardless would silently drop the write and + // leave the program returning a truncated result. + // + // The classic tf2xla lowering rejects this too, see TensorListSetItemOp in + // tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc. + if (set_item.getResizeIfIndexOutOfBounds()) { + return set_item.emitOpError( + "TensorLists that grow on an out-of-bounds write are not supported by " + "XLA. This is typically a tf.TensorArray created with " + "dynamic_size=True; give it a static size large enough to hold every " + "write instead."); + } + auto buffer = set_item.getInputHandle(); auto it = buffer_to_size->find(buffer); if (it == buffer_to_size->end()) { diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 2f0e2535a37fcc..9b06659cdb18b8 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -331,8 +331,12 @@ def grows(): ) with self.session(): + # Either bridge may compile this. The classic tf2xla kernel raises + # UnimplementedError; the MLIR pass reports the same message through + # emitOpError, which surfaces as InvalidArgumentError. The message is + # what pins this to the rejection under test. with self.assertRaisesRegex( - errors.UnimplementedError, + (errors.UnimplementedError, errors.InvalidArgumentError), "TensorLists that grow on an out-of-bounds write", ): self.evaluate(grows()) @@ -383,8 +387,12 @@ def body(i, ta): return out.stack() with self.session(): + # Either bridge may compile this. The classic tf2xla kernel raises + # UnimplementedError; the MLIR pass reports the same message through + # emitOpError, which surfaces as InvalidArgumentError. The message is + # what pins this to the rejection under test. with self.assertRaisesRegex( - errors.UnimplementedError, + (errors.UnimplementedError, errors.InvalidArgumentError), "TensorLists that grow on an out-of-bounds write", ): self.evaluate(f()) From e474e9dd559f82c7a88818d73992e825c4355857 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 17 Sep 2026 09:20:27 -0700 Subject: [PATCH 10/32] Use test_scope so the new tests reach the XLA kernel The new tests wrapped their bodies in def_function.function with jit_compile=True and ran them under self.session(). In that harness the ops are instantiated as ordinary CPU kernels and XlaCompiler never runs, so neither rejection fired and both tests failed with the error not raised. Switch them to the idiom the rest of this file uses, building the ops under self.session() with self.test_scope() and evaluating the result. That is how testInvalidSplitLength asserts UnimplementedError out of the same lowering, and how testGetItemFromEmptyList reaches these kernels. No test in tensorflow/compiler/tests asserts a compile error through jit_compile, which is the hint that the previous shape was wrong. Drops the end to end TensorArray test. There is no established way here to make a dynamic_size TensorArray in a while_loop reach the XLA kernel under this harness, and a test that cannot reach the code it covers is worse than no test. The op level test covers the kernel rejection and the new lit case covers the MLIR one. Removes the two imports and the two BUILD deps that test had needed. --- tensorflow/compiler/tests/BUILD | 2 - .../compiler/tests/tensor_list_ops_test.py | 73 +++---------------- 2 files changed, 12 insertions(+), 63 deletions(-) diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index c1ee3e5164b5c2..56b3fc99746d7d 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -1820,8 +1820,6 @@ tf_xla_py_strict_test( "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:list_ops", "//tensorflow/python/ops:map_fn", - "//tensorflow/python/ops:tensor_array_ops", - "//tensorflow/python/ops:while_loop", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 9b06659cdb18b8..58cb724cedd95e 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -28,8 +28,6 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import list_ops from tensorflow.python.ops import map_fn -from tensorflow.python.ops import tensor_array_ops -from tensorflow.python.ops import while_loop from tensorflow.python.platform import test @@ -312,90 +310,43 @@ def testSetItemResizeIsRejected(self): # does either by returning the list unchanged when the element does not # fit, or by letting DynamicUpdateSlice clamp an out-of-range index so the # element overwrites an earlier one. - @def_function.function(jit_compile=True) - def grows(): + with self.session(), self.test_scope(): tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 ) - tensor_list = list_ops.tensor_list_set_item( + set_item = list_ops.tensor_list_set_item( input_handle=tensor_list, index=0, item=constant_op.constant(1.0), resize_if_index_out_of_bounds=True, ) - # Stack inside the compiled function so that no TensorList crosses the - # XLA/TF boundary, which would raise a different error and stop this - # test from covering the rejection above. - return list_ops.tensor_list_stack( - tensor_list, element_dtype=dtypes.float32 - ) - - with self.session(): - # Either bridge may compile this. The classic tf2xla kernel raises - # UnimplementedError; the MLIR pass reports the same message through - # emitOpError, which surfaces as InvalidArgumentError. The message is - # what pins this to the rejection under test. + # The classic lowering raises UnimplementedError from the kernel, and + # the MLIR bridge reports the same message through emitOpError, which + # surfaces as InvalidArgumentError. The message is what pins this to + # the rejection under test. with self.assertRaisesRegex( (errors.UnimplementedError, errors.InvalidArgumentError), "TensorLists that grow on an out-of-bounds write", ): - self.evaluate(grows()) + self.evaluate(set_item) def testSetItemWithoutResizeStillCompiles(self): # The rejection above keys off the attribute rather than the index, so an # ordinary fixed-size write must still compile and run. - @def_function.function(jit_compile=True) - def fixed(): + with self.session(), self.test_scope(): tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 ) - tensor_list = list_ops.tensor_list_set_item( + set_item = list_ops.tensor_list_set_item( input_handle=tensor_list, index=0, item=constant_op.constant(1.0), resize_if_index_out_of_bounds=False, ) - return list_ops.tensor_list_stack( - tensor_list, element_dtype=dtypes.float32 - ) - - with self.session(): - self.assertAllEqual(self.evaluate(fixed()), [1.0]) - - def testDynamicSizeTensorArrayIsRejected(self): - # End to end case for GitHub issue 127528. A dynamic_size TensorArray - # written past its initial size used to compile and return a truncated - # result with no error, so this asserts the compile-time failure instead. - # The whole function is jit-compiled so the list stays inside XLA rather - # than crossing the XLA/TF boundary at the stack op. - @def_function.function(jit_compile=True) - def f(): - ta = tensor_array_ops.TensorArray( - dtypes.float32, size=0, dynamic_size=True + stacked = list_ops.tensor_list_stack( + set_item, element_dtype=dtypes.float32 ) - - def cond(i, ta): - del ta # Unused. - return i < 2 - - def body(i, ta): - return i + 1, ta.write(i, constant_op.constant(1.0)) - - _, out = while_loop.while_loop( - cond, body, [constant_op.constant(0), ta] - ) - return out.stack() - - with self.session(): - # Either bridge may compile this. The classic tf2xla kernel raises - # UnimplementedError; the MLIR pass reports the same message through - # emitOpError, which surfaces as InvalidArgumentError. The message is - # what pins this to the rejection under test. - with self.assertRaisesRegex( - (errors.UnimplementedError, errors.InvalidArgumentError), - "TensorLists that grow on an out-of-bounds write", - ): - self.evaluate(f()) + self.assertAllEqual(self.evaluate(stacked), [1.0]) def testInvalidSplitLength(self): with self.session(), self.test_scope(): From cdece3ba5e413dd777aee698339322cb5ae86be0 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 17 Sep 2026 12:43:06 -0700 Subject: [PATCH 11/32] Downgrade the rejection to a warning on both paths Per review: making this a hard compile error rejects every graph that sets resize_if_index_out_of_bounds, which is every dynamic_size TensorArray, including the many that never actually write past the end and compile correctly today. That is too broad to land. Warn instead and let compilation proceed: - the tf2xla kernel logs the message rather than failing construction - the MLIR pass uses emitWarning rather than emitOpError, so decomposition continues - the lit case expects a warning - the Python test evaluates the op and checks the result instead of asserting a raise, with a counterpart pinning that the attribute alone changes nothing but the logging Note this leaves the truncation itself in place, so the wrong result reported in the issue still occurs; the change only adds a log trail. Adds the logging.h include the kernel now needs. --- .../tests/tensor_list_ops_decomposition.mlir | 11 +++--- .../tensor_list_ops_decomposition.cc | 11 +++--- .../compiler/tests/tensor_list_ops_test.py | 35 +++++++++--------- .../tf2xla/kernels/tensor_list_ops.cc | 36 +++++++++---------- 4 files changed, 48 insertions(+), 45 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir index 3d4badf472af5d..df57fc3c420169 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir @@ -631,10 +631,11 @@ func.func @main(%arg0: tensor<*xi32>) -> () { // ----- -// Tests that the pass reports an error on a set item that is allowed to grow -// the list, which a tf.TensorArray with dynamic_size=True emits. The buffer has -// a static shape and cannot grow, so decomposing it would silently drop the -// write. +// Tests that the pass warns on a set item that is allowed to grow the list, +// which a tf.TensorArray with dynamic_size=True emits. The buffer has a static +// shape and cannot grow, so the write is dropped. Decomposition still proceeds, +// since the attribute is set by every dynamic_size TensorArray including those +// that never write past the end. func.func @main() -> () { %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> @@ -642,7 +643,7 @@ func.func @main() -> () { %index = "tf.Const"() {value = dense<0> : tensor} : () -> tensor %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor) -> tensor>> %elem = "tf._SomeOp"() : () -> tensor - // expected-error @+1 {{TensorLists that grow on an out-of-bounds write are not supported by XLA}} + // expected-warning @+1 {{TensorLists that grow on an out-of-bounds write are not supported by XLA}} %set = "tf.TensorListSetItem"(%tl, %index, %elem) {resize_if_index_out_of_bounds = true} : (tensor>>, tensor, tensor) -> tensor>> func.return } diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc index b94576fdc05a64..f358940a37cbf7 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc @@ -718,13 +718,16 @@ LogicalResult HandleTensorListSetItemOp( // non-XLA kernel grows the list when the index is past the end. The buffer // built here has a static shape and cannot grow, and cutil::SetElement lowers // to a DynamicUpdateSlice, which clamps an out-of-range index rather than - // failing. Decomposing the op regardless would silently drop the write and - // leave the program returning a truncated result. + // failing, so the write is dropped and the program returns a truncated + // result. // - // The classic tf2xla lowering rejects this too, see TensorListSetItemOp in + // This is only a warning rather than an error, because the attribute is set + // by every dynamic_size TensorArray, including the many that never actually + // write past the end and compile correctly today. The classic tf2xla + // lowering warns in the same case, see TensorListSetItemOp in // tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc. if (set_item.getResizeIfIndexOutOfBounds()) { - return set_item.emitOpError( + set_item.emitWarning( "TensorLists that grow on an out-of-bounds write are not supported by " "XLA. This is typically a tf.TensorArray created with " "dynamic_size=True; give it a static size large enough to hold every " diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 58cb724cedd95e..e282e9278bad09 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -302,14 +302,17 @@ def testZerosLikeForTensorList(self): self.assertAllEqual(z.shape.as_list(), [None]) self.assertAllEqual(z, [0.0, 0.0]) - def testSetItemResizeIsRejected(self): + def testSetItemResizeWarnsAndTruncates(self): # Regression test for GitHub issue 127528. A TensorListSetItem that is - # allowed to grow the list cannot be compiled, because the buffer that - # TensorListReserve allocated has a static shape. It must be rejected at - # compile time rather than silently dropping the write, which the lowering - # does either by returning the list unchanged when the element does not - # fit, or by letting DynamicUpdateSlice clamp an out-of-range index so the - # element overwrites an earlier one. + # allowed to grow the list cannot grow it under XLA, because the buffer + # that TensorListReserve allocated has a static shape. The write is + # dropped, either by returning the list unchanged when the element does + # not fit, or by letting DynamicUpdateSlice clamp an out-of-range index so + # the element overwrites an earlier one. + # + # The lowering logs a warning rather than failing, because the attribute is + # set by every dynamic_size TensorArray including the many that never write + # past the end, so compilation still has to succeed here. with self.session(), self.test_scope(): tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 @@ -320,19 +323,15 @@ def testSetItemResizeIsRejected(self): item=constant_op.constant(1.0), resize_if_index_out_of_bounds=True, ) - # The classic lowering raises UnimplementedError from the kernel, and - # the MLIR bridge reports the same message through emitOpError, which - # surfaces as InvalidArgumentError. The message is what pins this to - # the rejection under test. - with self.assertRaisesRegex( - (errors.UnimplementedError, errors.InvalidArgumentError), - "TensorLists that grow on an out-of-bounds write", - ): - self.evaluate(set_item) + stacked = list_ops.tensor_list_stack( + set_item, element_dtype=dtypes.float32 + ) + self.assertAllEqual(self.evaluate(stacked), [1.0]) def testSetItemWithoutResizeStillCompiles(self): - # The rejection above keys off the attribute rather than the index, so an - # ordinary fixed-size write must still compile and run. + # Counterpart to the test above: the warning is keyed off the attribute, + # so this pins that an ordinary fixed-size write compiles to the same + # result and that the attribute alone changes nothing but the logging. with self.session(), self.test_scope(): tensor_list = list_ops.tensor_list_reserve( element_shape=[], element_dtype=dtypes.float32, num_elements=1 diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index b2a2a61e39c165..30bec2790f7b3c 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -41,6 +41,7 @@ limitations under the License. #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -653,28 +654,27 @@ class TensorListSetItemOp : public XlaOpKernel { // A tf.TensorArray built with dynamic_size=True lowers to TensorListSetItem // with this attribute set, and the non-XLA kernel grows the list when the // index is past the end. XLA needs static shapes, so the buffer that - // TensorListReserve allocated cannot grow. + // TensorListReserve allocated cannot grow, and a write past the end is + // dropped in one of two ways depending on the shapes. When the element does + // not fit at all, as for a list reserved with size 0, + // ExecuteTensorListSetItem takes its early return and hands back the list + // unchanged with an OK status. When the element does fit but the index is + // past the end, the DynamicUpdateSlice it lowers to clamps its start + // indices rather than failing, so the element overwrites an earlier one. + // Either way the program returns a truncated result. // - // Without this check the write is silently dropped, in one of two ways - // depending on the shapes. When the element does not fit at all, as for a - // list reserved with size 0, ExecuteTensorListSetItem takes its early - // return, logs a warning and hands back the list unchanged with an OK - // status. When the element does fit but the index is past the end, the - // DynamicUpdateSlice it lowers to clamps its start indices rather than - // failing, so the element overwrites an earlier one. Either way the - // program returns a truncated result with no error at all. - // - // The older TensorArray lowering already rejects the same thing, see - // TensorArrayOp in tensor_array_ops.cc. + // This is only a warning rather than a compile error, because the attribute + // is set by every dynamic_size TensorArray, including the many that never + // actually write past the end and compile correctly today. bool resize_if_index_out_of_bounds; OP_REQUIRES_OK(ctx, ctx->GetAttr("resize_if_index_out_of_bounds", &resize_if_index_out_of_bounds)); - OP_REQUIRES(ctx, !resize_if_index_out_of_bounds, - errors::Unimplemented( - "TensorLists that grow on an out-of-bounds write are not " - "supported by XLA. This is typically a tf.TensorArray " - "created with dynamic_size=True; give it a static size " - "large enough to hold every write instead.")); + if (resize_if_index_out_of_bounds) { + LOG(WARNING) << "TensorLists that grow on an out-of-bounds write are not " + "supported by XLA. This is typically a tf.TensorArray " + "created with dynamic_size=True; give it a static size " + "large enough to hold every write instead."; + } } void Compile(XlaOpKernelContext* ctx) override { From 874cfd5699925a844fbbaff54497ab28c63e2ae7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:15:56 -0700 Subject: [PATCH 12/32] PR #48926: Bump github/codeql-action/upload-sarif from 4.37.9 to 4.38.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48926 Bumps [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) from 4.37.9 to 4.38.0.
Release notes

Sourced from github/codeql-action/upload-sarif's releases.

v4.38.0

  • On GitHub-hosted runners, the CodeQL Action now deletes unused CodeQL bundles from the toolcache before downloading a different bundle, which frees up disk space for the analysis. We expect to roll this change out to everyone in September. #4124
  • The CodeQL Action now supports CodeQL releases that are compatible with Linux Arm64 and downloads the native linux-arm64 CodeQL bundle when available. #4072
  • Update default CodeQL bundle version to 2.27.0. #4129
Changelog

Sourced from github/codeql-action/upload-sarif's changelog.

CodeQL Action Changelog

See the releases page for the relevant changes to the CodeQL CLI and language packs.

[UNRELEASED]

No user facing changes.

4.38.0 - 09 Sept 2026

  • On GitHub-hosted runners, the CodeQL Action now deletes unused CodeQL bundles from the toolcache before downloading a different bundle, which frees up disk space for the analysis. We expect to roll this change out to everyone in September. #4124
  • The CodeQL Action now supports CodeQL releases that are compatible with Linux Arm64 and downloads the native linux-arm64 CodeQL bundle when available. #4072
  • Update default CodeQL bundle version to 2.27.0. #4129

4.37.9 - 26 Aug 2026

  • Update default CodeQL bundle version to 2.26.4. #4106

4.37.8 - 21 Aug 2026

No user facing changes.

4.37.7 - 13 Aug 2026

  • Update default CodeQL bundle version to 2.26.3. #4085

4.37.6 - 04 Aug 2026

  • Changed the default filepath for the new remote file address format that was introduced in CodeQL Action 4.37.0 / 3.37.0 to .github/codeql-config.yml to align it with the suggested path that is used elsewhere. #4070

4.37.5 - 03 Aug 2026

  • Fixed a bug where a network error while streaming the download of the CodeQL bundle could terminate the init Action instead of falling back to downloading the bundle before extracting it. #4061

4.37.4 - 29 Jul 2026

  • This version of the CodeQL Action adds support for the tools input for the codeql-action/init step to be specified using a github-codeql-tools repository property. This feature will gradually be rolled out following the release of this version. Once rolled out, this allows for the CodeQL CLI version that is used in GitHub-managed workflows, such as Default Setup, to be set to a custom value. For example, customers who run into issues with rate limits when a new CodeQL CLI version is released can set the value to toolcache to always use the CodeQL CLI version that is available in the runner toolcache. For Advanced Setup workflows, the value provided for tools in the workflow definition always takes precedence unless the value of the repository property starts with !. #4037
  • Update default CodeQL bundle version to 2.26.2. #4051

4.37.3 - 22 Jul 2026

No user facing changes.

4.37.2 - 21 Jul 2026

  • The new address format for the config-file input that was introduced in CodeQL Action 4.37.0 is now enabled by default. In addition to the format described there, the remote= prefix can now be used to explicitly indicate that the input refers to a remote file. All previous input formats continue to be accepted as well. #4023
  • The CodeQL Action can now make use of configured private registries in Default Setup to retrieve CodeQL configuration files from remote repositories that require authentication. This will allow customers to store their CodeQL configuration in a single repository that can then be referenced by Default Setup workflows in other repositories. We expect to roll this and other, related changes out to everyone in July. #4007

4.37.1 - 16 Jul 2026

... (truncated)

Commits
  • b96794f Merge pull request #4131 from github/update-v4.38.0-7e08580a9
  • 02d5093 Update changelog for v4.38.0
  • 7e08580 Merge pull request #4130 from github/henrymercer/workflow-runner-sizing
  • bfcc52b Run slow macOS checks on larger runners
  • 8c251e7 Merge pull request #4129 from github/update-bundle/codeql-bundle-v2.27.0
  • 0b7ca40 Add changelog note
  • 40484b3 Update default bundle to codeql-bundle-v2.27.0
  • 977e6ce Merge pull request #4124 from github/henrymercer/toolcache-bundle-cleanup
  • 40a6b38 Address toolcache cleanup review feedback
  • deece8f Apply suggestion from @​henrymercer
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=github/codeql-action/upload-sarif&package-manager=github_actions&previous-version=4.37.9&new-version=4.38.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Copybara import of the project: -- 64555e142a81a3f779a0078e93b9e597d04712dd by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>: Bump github/codeql-action/upload-sarif from 4.37.9 to 4.38.0 Bumps [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) from 4.37.9 to 4.38.0. - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/github/codeql-action/compare/cdf488f595d80d6e07e03d4674febd5ab45fa938...b96794f015dfd88f77b49b1c93e0fa7110f94c63) --- updated-dependencies: - dependency-name: github/codeql-action/upload-sarif dependency-version: 4.38.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Merging this change closes #48926 PiperOrigin-RevId: 983804916 --- third_party/xla/.github/workflows/scorecards-analysis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/.github/workflows/scorecards-analysis.yml b/third_party/xla/.github/workflows/scorecards-analysis.yml index ec2c190fd2ef51..f1f84938ab300c 100644 --- a/third_party/xla/.github/workflows/scorecards-analysis.yml +++ b/third_party/xla/.github/workflows/scorecards-analysis.yml @@ -67,6 +67,6 @@ jobs: # Upload the results to GitHub's code scanning dashboard (optional). # Commenting out will disable upload of results to your repo's Code Scanning dashboard - name: "Upload to code-scanning" - uses: github/codeql-action/upload-sarif@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + uses: github/codeql-action/upload-sarif@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: sarif_file: results.sarif From cce353619eb4d2a0a6a5b79e0e8c58e8c327552f Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Fri, 18 Sep 2026 05:17:51 -0700 Subject: [PATCH 13/32] Fix recursive DWYU rule With my previous DWYU change I accidentally registered a loop in the macro expansion which makes DWYU segfault. This is the fix for that. Also improves testing. PiperOrigin-RevId: 983805520 --- third_party/xla/.bant-macros | 1 - third_party/xla/BUILD.bazel | 2 ++ third_party/xla/build_tools/lint/BUILD | 2 ++ .../xla/build_tools/lint/run_dwyu_test.py | 30 ++++++++++++++----- 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/third_party/xla/.bant-macros b/third_party/xla/.bant-macros index 47774a7f8fe9e8..d3fb19952adabf 100644 --- a/third_party/xla/.bant-macros +++ b/third_party/xla/.bant-macros @@ -33,7 +33,6 @@ exhaustive_xla_test = bant_forward_args(cc_test()) # Binary rules xla_cc_binary = bant_forward_args(cc_binary()) -cc_binary = bant_forward_args(cc_binary()) # Library rules onednn_cc_library = bant_forward_args(cc_library()) diff --git a/third_party/xla/BUILD.bazel b/third_party/xla/BUILD.bazel index 5a6d40abe0a80a..3a2613f00bbf71 100644 --- a/third_party/xla/BUILD.bazel +++ b/third_party/xla/BUILD.bazel @@ -28,6 +28,8 @@ license( license_kinds = ["@rules_license//licenses/spdx:Apache-2.0"], ) +exports_files([".bant-macros"]) + filegroup( name = "clang_tidy_config", srcs = [".clang-tidy"], diff --git a/third_party/xla/build_tools/lint/BUILD b/third_party/xla/build_tools/lint/BUILD index bf17ce87f87517..f9ea2203214e5a 100644 --- a/third_party/xla/build_tools/lint/BUILD +++ b/third_party/xla/build_tools/lint/BUILD @@ -117,7 +117,9 @@ pytype_strict_library( py_test( name = "run_dwyu_test", srcs = ["run_dwyu_test.py"], + data = ["//:.bant-macros"], # copybara:comment_replace data = ["//third_party/xla:.bant-macros"], deps = [ + ":check_dwyu", ":run_dwyu", ], ) diff --git a/third_party/xla/build_tools/lint/run_dwyu_test.py b/third_party/xla/build_tools/lint/run_dwyu_test.py index 32f767327bde87..16d0ba4d0c6db7 100644 --- a/third_party/xla/build_tools/lint/run_dwyu_test.py +++ b/third_party/xla/build_tools/lint/run_dwyu_test.py @@ -22,9 +22,11 @@ import tempfile import unittest +from build_tools.lint import check_dwyu from build_tools.lint import run_dwyu _BANT = shutil.which(os.environ.get("BANT_BIN", "bant")) +_BANT_MACROS = pathlib.Path(run_dwyu.__file__).parents[2] / ".bant-macros" class BantWorkspaceTest(unittest.TestCase): @@ -42,7 +44,7 @@ def setUp(self): self.external.mkdir(parents=True) (self.external / "_main").symlink_to(self.root, target_is_directory=True) self.write_file("MODULE.bazel", 'module(name = "xla")\n') - self.write_file(".bant-macros", "") + self.write_file(".bant-macros", _BANT_MACROS.read_text()) self.write_file("third_party/tsl/WORKSPACE", "") def write_file(self, path, contents): @@ -138,7 +140,7 @@ def setUp(self): def check_targets(self, *targets): return subprocess.run( [ - sys.executable, + sys.executable or shutil.which("python3"), run_dwyu.__file__, "--bant", _BANT, @@ -229,10 +231,6 @@ def test_keeps_implementation_edge_but_reports_other_unused_dependency(self): ) def test_conditional_deps_macro_resolves_wrapped_dependency(self): - self.write_file( - ".bant-macros", - "if_cuda_is_configured = _arg_0\nif_static = _arg_0\n", - ) self.write_file( "xla/consumer/BUILD", """\ @@ -248,7 +246,25 @@ def test_conditional_deps_macro_resolves_wrapped_dependency(self): self.assertEqual(result.stdout, "") self.assertIn("Checked DWYU on 1 targets.", result.stderr) + def test_all_allowed_rules_supported_by_bant_macros(self): + rules = [] + labels = [] + for index, rule in enumerate(check_dwyu.DEFAULT_ALLOWED_RULES): + target = f"target_{index}_{rule}" + rules.append( + f'{rule}(name = "{target}", srcs = ["consumer.cc"],' + ' deps = ["//xla/platform:errors"])' + ) + labels.append(f"//xla/consumer:{target}") + self.write_file("xla/consumer/BUILD", "\n".join(rules) + "\n") + result = self.check_targets(*labels) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertEqual(result.stdout, "") + self.assertIn( + f"Checked DWYU on {len(labels)} targets.", + result.stderr, + ) + if __name__ == "__main__": unittest.main() - From 14057ffcfdb7ceca7504297e2785805457298fee Mon Sep 17 00:00:00 2001 From: zoranjovanovic-ns <126815388+zoranjovanovic-ns@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:18:47 -0700 Subject: [PATCH 14/32] PR #48981: [ROCm] Skip NumericTestsForBlas.Infinity TF32x3 on MI350 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48981 📝 Summary of Changes Skip NumericTestsForBlas.Infinity/dot_tf32_tf32_f32_x3 on MI350 (gfx950). 🎯 Justification hipBLASLt emulates TF32 on gfx950 and returns NaN for inf * 1. TF32x3 relies on that high×high product after zeroing the low order terms, so the test fails on MI350. It still passes on gfx942 and with Triton, so the skip is scoped to gfx950 only. 🚀 Kind of Contribution Please remove what does not apply: 🧪 Tests 🧪 Unit Tests: GTEST_SKIP() on ROCm gfx950: xla/backends/gpu/codegen/triton/dot_algorithms_test.cc — NumericTestsForBlas/NumericTestsForBlas.Infinity/dot_tf32_tf32_f32_x3 Copybara import of the project: -- 6e41451edc164f240c9bed076324a5adfed7b6cb by Zoran Jovanovic : [ROCm] Skip NumericTestsForBlas.Infinity TF32x3 on MI350 Merging this change closes #48981 PiperOrigin-RevId: 983805760 --- .../xla/backends/gpu/codegen/triton/dot_algorithms_test.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc index 26997ce8df3b22..81a06dc5487b9f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc @@ -1142,6 +1142,11 @@ class NumericTestsForTriton : public TritonAlgorithmTest, }; TEST_P(NumericTestsForBlas, Infinity) { + // hipBLASLt emulates TF32 on gfx950 and returns NaN for inf*1. + if (GetParam() == PC::ALG_DOT_TF32_TF32_F32_X3 && GpuComputeComp().IsRocm() && + GpuComputeComp().rocm_compute_capability()->gfx9_mi350()) { + GTEST_SKIP() << "hipBLASLt FAST_TF32 inf*1 returns NaN on MI350 "; + } std::string hlo_text = HloText(); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, GetOptimizedModule(hlo_text)); From 8f1c038046eb54cfaa090d3729cc6e4f49299a9e Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 18 Sep 2026 05:26:44 -0700 Subject: [PATCH 15/32] PR #48110: [ROCm] Pick the rocBLAS complex GEMM fallback at runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48110 ## 📝 Summary of Changes Replace the compile-time `TF_ROCM_VERSION` gate on the C64/C128 rocBLAS redirect from #44047 with a runtime capability check. Complex matmuls take the hipBLASLt plan when the descriptor builds and the heuristic offers an algorithm, and fall back to rocBLAS (`rocblas_cgemm`/`zgemm`) otherwise. ## 🎯 Justification hipBLASLt can support complex C64/C128 GEMM kernels in ROCm 10, but a build-time macro describes the ROCm that XLA was compiled against, not the hipBLASLt loaded at runtime. so making the check to runtime so legacy rocms will still use rocblas and newer one switch to hipblasLT. ## 🚀 Kind of Contribution ♻️ Cleanup Copybara import of the project: -- 8e63c8691ee87cd538cf16ec811faeb5bbbb56f7 by magaonka : [ROCm] Pick the rocBLAS complex GEMM fallback at runtime hipBLASLt gained complex GEMM kernels in ROCm 10, but whether they are usable depends on the hipBLASLt loaded at runtime and on the device, not on the ROCm version XLA was built against. A compile-time TF_ROCM_VERSION gate strands mismatched build/runtime combinations on the wrong path in both directions. Try the hipBLASLt plan for C64/C128 and keep it only if the descriptor builds and the heuristic offers an algorithm; otherwise redirect to the openxla/xla#44047 rocBLAS path. Both conditions are needed: older hipBLASLt rejects the complex matmul descriptor outright, while a newer one could still be missing Tensile kernels for a given architecture. Merging this change closes #48110 PiperOrigin-RevId: 983807962 --- .../xla/stream_executor/rocm/hip_blas_lt.cc | 50 +++++++++++++------ .../xla/stream_executor/rocm/hip_blas_lt.h | 8 ++- 2 files changed, 42 insertions(+), 16 deletions(-) diff --git a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc index cf063977e54d87..9c2543e98d5fcb 100644 --- a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc +++ b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc @@ -91,6 +91,11 @@ void GroupGemmUpdateArgs( int8_t bias_type, bool has_matrix_bias); namespace { +// Upper bound for the heuristic query only, never allocated. Generous on +// purpose: a low ceiling would hide complex kernels that need scratch and +// divert the matmul to rocBLAS. +constexpr size_t kComplexProbeMaxWorkspaceBytes = 64 * 1024 * 1024; + template absl::Status SetAttr(hipblasLtMatrixLayout_t handle, hipblasLtMatrixLayoutAttribute_t attr, T value) { @@ -311,21 +316,8 @@ auto BlasLt::RegularMatmulPlan::GetAlgorithms(size_t max_algorithm_count, return algorithms; } -absl::StatusOr BlasLt::GetMatmulPlan( +absl::StatusOr BlasLt::GetHipBlasLtMatmulPlan( const gpu::GemmConfig& cfg, Epilogue epilogue) const { - // hipBLASLt has no complex GEMM kernels; redirect C64/C128 to rocBLAS while - // still consuming the cublasLt matmul custom call emitted by GemmRewriter. - // TODO(magaonka-amd): Once we get a hipBLASLt that supports complex GEMMs, - // clean up this routing code. - if (cfg.output_layout.dtype == xla::C64 || - cfg.output_layout.dtype == xla::C128) { - TF_RET_CHECK(epilogue == Epilogue::kDefault) - << "rocBLAS complex fallback does not support epilogues (bias, GELU, " - "etc.); got epilogue=" - << static_cast(epilogue); - return std::make_unique(*this, cfg); - } - auto lhs_layout = cfg.lhs_layout, rhs_layout = cfg.rhs_layout, output_layout = cfg.output_layout, c_layout = cfg.c_layout; @@ -566,6 +558,36 @@ absl::StatusOr BlasLt::GetMatmulPlan( return plan; } +absl::StatusOr BlasLt::GetMatmulPlan( + const gpu::GemmConfig& cfg, Epilogue epilogue) const { + if (cfg.output_layout.dtype != xla::C64 && + cfg.output_layout.dtype != xla::C128) { + return GetHipBlasLtMatmulPlan(cfg, epilogue); + } + + // Older hipBLASLt rejects the complex matmul descriptor outright, and a newer + // one can still have no kernels for this device, so test both rather than the + // ROCm version XLA was built against. + auto plan = GetHipBlasLtMatmulPlan(cfg, epilogue); + if (plan.ok()) { + auto algorithms = (*plan)->GetAlgorithms(/*max_algorithm_count=*/1, + kComplexProbeMaxWorkspaceBytes); + if (algorithms.ok() && !algorithms->empty()) { + return plan; + } + } + + TF_RET_CHECK(epilogue == Epilogue::kDefault) + << "rocBLAS complex fallback does not support epilogues (bias, GELU, " + "etc.); got epilogue=" + << static_cast(epilogue); + VLOG(1) << "hipBLASLt offers no complex GEMM algorithm here; routing " + << xla::primitive_util::LowercasePrimitiveTypeName( + cfg.output_layout.dtype) + << " matmul to rocBLAS"; + return std::make_unique(*this, cfg); +} + absl::Status BlasLt::RocBlasGemmPlan::ExecuteOnStream( Stream* stream, const gpu::BlasLt::MemoryArgs& args, blas::ProfileResult* profile_result) const { diff --git a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h index b71f97e6341bed..8a08f42b855408 100644 --- a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h +++ b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h @@ -193,8 +193,8 @@ class BlasLt : public gpu::BlasLt { int8_t bias_type_ = 0; }; // class GroupedMatmulPlan - // Executes complex (C64/C128) matmuls via rocBLAS (rocblas_cgemm/zgemm), - // since hipBLASLt has no complex GEMM kernels in current ROCm releases. + // Executes complex (C64/C128) matmuls via rocBLAS (rocblas_cgemm/zgemm) when + // the hipBLASLt loaded at runtime has no algorithm for them. class RocBlasGemmPlan : public gpu::BlasLt::MatmulPlan { public: friend class BlasLt; @@ -239,6 +239,10 @@ class BlasLt : public gpu::BlasLt { ~BlasLt() override = default; private: + // Fails if hipBLASLt does not handle the requested types. + absl::StatusOr GetHipBlasLtMatmulPlan( + const gpu::GemmConfig& cfg, Epilogue epilogue) const; + StreamExecutor* executor_; mutable absl::Mutex mu_; Owned handle_ ABSL_GUARDED_BY(mu_); From befa3dbaefc1feee42937c3dcbd416656a2719f7 Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Fri, 18 Sep 2026 05:27:19 -0700 Subject: [PATCH 16/32] PR #48711: [ROCm]Enable local_client_execute_test and gpu_bandwidth_benchmark_test on ROCm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48711 📝 Summary of Changes - Enable `local_client_execute_test` on ROCm by dropping the stale `cuda-only` tag (keep `optonly`). - Enable `gpu_bandwidth_benchmark_test` on ROCm by linking `rocm_platform` and switching to `xla_test`. - Keep NVIDIA on H100 (`requires-gpu-sm90-only`) and add `amdgpu_any` for ROCm. 🎯 Justification Add more unit test coverage on ROCm platform. 🚀 Kind of Contribution 🧪 Tests Copybara import of the project: -- 3686a353cc8bc068e29da882093b7e5bb435ae04 by linchen1 : Enable local_client_execute_test and gpu_bandwidth_benchmark_test on ROCm. Signed-off-by: linchen1 Merging this change closes #48711 PiperOrigin-RevId: 983808132 --- third_party/xla/xla/tests/BUILD | 5 +---- third_party/xla/xla/tools/cost_model/BUILD | 16 ++++++++-------- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 05ada6ce5eb56f..717d98c173580c 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -3013,10 +3013,7 @@ xla_test( "interpreter", ], shard_count = 30, - tags = [ - "cuda-only", - "optonly", - ], + tags = ["optonly"], use_legacy_runtime = True, deps = [ ":literal_test_util", diff --git a/third_party/xla/xla/tools/cost_model/BUILD b/third_party/xla/xla/tools/cost_model/BUILD index 3c8026b2be1cea..6517af28e5af9c 100644 --- a/third_party/xla/xla/tools/cost_model/BUILD +++ b/third_party/xla/xla/tools/cost_model/BUILD @@ -1,8 +1,6 @@ +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@rules_cc//cc:cc_library.bzl", "cc_library") -load( - "//xla:xla.default.bzl", - "xla_cc_test", -) +load("//xla/tests:build_defs.bzl", "xla_test") load( "//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", @@ -30,15 +28,17 @@ cc_library( "@com_google_absl//absl/types:span", ] + if_cuda_is_configured([ "//xla/stream_executor:cuda_platform", + ]) + if_rocm_is_configured([ + "//xla/stream_executor:rocm_platform", ]), ) -xla_cc_test( +xla_test( name = "gpu_bandwidth_benchmark_test", srcs = ["gpu_bandwidth_benchmark_test.cc"], - tags = [ - "cuda-only", - "requires-gpu-sm90", + backends = [ + "h100", + "amdgpu_any", ], deps = [ ":gpu_bandwidth_benchmark", From 50ebcb5c84921b29749333c0df08f2a9bd32f3e7 Mon Sep 17 00:00:00 2001 From: Shaogang Wang Date: Fri, 18 Sep 2026 05:33:08 -0700 Subject: [PATCH 17/32] PR #48789: [XLA:GPU] Optimize VMM allocator overlap lookups with ordered address indexes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48789 ## 📝 Summary of Changes Reduce the host overhead of device_address_vmm_allocator with below changes: 1. Replace the VMM allocator's address hash maps with ordered `absl::btree_map` indexes for targeted overlap checks. 2. Store deferred deallocations in an intrusive queue so reused allocations can cancel their pending frees directly, while preserving mapping reuse and deferred teardown behavior. ## 🎯 Justification Avoid scanning all tracked ranges during overlap checks and searching the pending queue when reusing a known allocation. This reduces CPU bookkeeping for workloads with many live or deferred VMM allocations. ## 🚀 Kind of Contribution ⚡️ Performance Improvement, 🧪 Tests ## 🧪 Unit Tests: Added coverage in `//xla/stream_executor:device_address_vmm_allocator_test`: - `OverlapLookupPreservesRangeBoundariesForActiveAndStaleAddresses` - `ManyReservationAliasesReuseAndRemoveOrderedRecords` - `StaleAliasReuseValidatesSourceAndFullReservationRange` - `MixedPendingOperationsCanBeCancelledAndRequeuedInAnyOrder` Copybara import of the project: -- 44cc7cb82c2c704a78e91cf641530fa7b3391ab2 by Shawn Wang : Optimize VMM allocator overlap lookups with ordered address indexes -- cf69d2304c5bb34736c6c51501d110f931fbc557 by Shawn Wang : Reuse VMM aliases before overlap checks and unlink pending operations Reactivate an exact stale reservation alias directly from its validated source record before searching either destination address index. Preserve the existing validation and waiting path for new or changed mappings. Embed intrusive pending-operation nodes in allocation and alias records. Reuse now cancels the known operation in constant time without scanning or shifting queued entries, or allocating separate queue nodes. Preserve FIFO stream ordering, batch accounting, and deferred teardown. Add tests for full-range/source validation, repeated cancellation and requeueing of mixed operations, and draining uncancelled aliases. Validation: 11 CPU allocator tests and 34 CUDA allocator tests passed, including the two-device cases with CUDA_VISIBLE_DEVICES=0,1. -- 740911322f61c8ccebade5106ff542c299620846 by Shawn Wang : Simplify VMM allocator reuse and deferred completion Build deferred-operation keys from queue nodes and centralize unlinking and metadata snapshots in completion. Remove the redundant address-index lookup when reactivating an allocation and flatten the stale-alias branch made redundant by early reuse. Preserve the early alias reuse path, constant-time queue cancellation, and existing validation and synchronization behavior. Validation: 11 CPU and 34 CUDA allocator tests passed with no skips in the maxtext-2026-09-12 container, using at most two B200 GPUs. Merging this change closes #48789 PiperOrigin-RevId: 983809891 --- third_party/xla/xla/stream_executor/BUILD | 2 + .../device_address_vmm_allocator.cc | 336 ++++++++++-------- .../device_address_vmm_allocator.h | 121 ++++--- .../device_address_vmm_allocator_test.cc | 265 ++++++++++++++ 4 files changed, 522 insertions(+), 202 deletions(-) diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index cb5d49d469c1ea..15ba54aaf48e28 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -370,6 +370,7 @@ cc_library( ":stream_executor_h", "//xla/service:device_assignment", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", @@ -389,6 +390,7 @@ xla_cc_test( srcs = ["device_address_vmm_allocator_test.cc"], deps = [ ":device_address", + ":device_address_allocator", ":device_address_vmm_allocator", ":memory_allocation", ":memory_reservation", diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc index 417518b1477768..9013f87dc5553c 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc @@ -132,6 +132,7 @@ void DeviceAddressVmmAllocator::AllocationRecord::MarkAllocatorStale( void DeviceAddressVmmAllocator::AllocationRecord::ReactivateAllocator( uint64_t new_size) { CHECK(allocator_stale()); + CHECK_EQ(allocator_deallocation_.pending.seqno, 0); CHECK(!allocator_address_mapping_.is_null()); allocator_address_ = DeviceAddressBase(allocator_address_.opaque(), new_size); allocator_stale_seqno_ = 0; @@ -154,14 +155,22 @@ void DeviceAddressVmmAllocator::AllocationRecord::MarkReservationStale( void DeviceAddressVmmAllocator::AllocationRecord::ReactivateReservation() { CHECK(reservation_stale()); + CHECK_EQ(reservation_alias_->deallocation.pending.seqno, 0); reservation_alias_->stale_seqno = 0; } void DeviceAddressVmmAllocator::AllocationRecord::CompleteStaleReservation() { CHECK(reservation_stale()); + CHECK_EQ(reservation_alias_->deallocation.pending.seqno, 0); reservation_alias_.reset(); } +DeviceAddressVmmAllocator::PendingDeallocationNode& +DeviceAddressVmmAllocator::AllocationRecord::reservation_deallocation() { + CHECK(has_reservation_alias()); + return reservation_alias_->deallocation; +} + // Interval between CPU polls of the GPU-written deallocation timeline while // waiting for deferred frees to become safe. The 50us value is a conservative // initial tradeoff: long enough to avoid busy-spinning a CPU core and short @@ -265,10 +274,8 @@ DeviceAddressVmmAllocator::~DeviceAddressVmmAllocator() { // no longer callable after the derived destructor has run, and the stream // is allowed to have been destroyed already. state->open_deallocation_batch_seqno = 0; - while (!state->pending_deallocations.empty()) { - PendingDeallocation pending = state->pending_deallocations.front(); - state->pending_deallocations.pop_front(); - CompletePendingDeallocation(*state, pending); + while (state->pending_head != nullptr) { + CompletePendingDeallocation(*state, *state->pending_head); } } @@ -319,16 +326,14 @@ absl::StatusOr DeviceAddressVmmAllocator::GetStream( absl::Status DeviceAddressVmmAllocator::DrainPendingDeallocations( PerDeviceState& state) { ABSL_RETURN_IF_ERROR(FlushOpenDeallocationBatch(state)); - if (state.pending_deallocations.empty()) { + if (state.pending_head == nullptr) { return absl::OkStatus(); } - uint64_t target_seqno = state.pending_deallocations.back().seqno; + uint64_t target_seqno = state.pending_tail->pending.seqno; ABSL_RETURN_IF_ERROR(WaitUntilSeqno(state, target_seqno)); - while (!state.pending_deallocations.empty() && - state.pending_deallocations.front().seqno <= target_seqno) { - PendingDeallocation pending = state.pending_deallocations.front(); - state.pending_deallocations.pop_front(); - CompletePendingDeallocation(state, pending); + while (state.pending_head != nullptr && + state.pending_head->pending.seqno <= target_seqno) { + CompletePendingDeallocation(state, *state.pending_head); } return absl::OkStatus(); } @@ -358,7 +363,8 @@ MemoryAllocation* DeviceAddressVmmAllocator::GetRawAllocation( // Allocator addresses are keyed directly by their VA. Stale records remain in // this map until deferred teardown completes, so expose their backing // allocation for diagnostics/reuse checks until the pending operation drains. - auto allocation_it = state->records_by_allocator_address.find(addr.opaque()); + auto allocation_it = + state->records_by_allocator_address.find(AddressStart(addr)); if (allocation_it != state->records_by_allocator_address.end() && allocation_it->second->allocator_matches(addr)) { return allocation_it->second->raw_allocation(); @@ -367,7 +373,7 @@ MemoryAllocation* DeviceAddressVmmAllocator::GetRawAllocation( // Reservation aliases created by Map() share one index. Only active aliases // are exposed; stale or already-unmapped aliases intentionally return // nullptr. - auto reservation_it = state->reservation_records.find(addr.opaque()); + auto reservation_it = state->reservation_records.find(AddressStart(addr)); if (reservation_it != state->reservation_records.end() && reservation_it->second->reservation_active() && reservation_it->second->reservation_matches(addr)) { @@ -385,7 +391,8 @@ MemoryReservation* DeviceAddressVmmAllocator::GetReservation( PerDeviceState* state = *state_or; absl::MutexLock lock(state->mu); - auto allocation_it = state->records_by_allocator_address.find(addr.opaque()); + auto allocation_it = + state->records_by_allocator_address.find(AddressStart(addr)); if (allocation_it != state->records_by_allocator_address.end() && allocation_it->second->allocator_active() && allocation_it->second->allocator_matches(addr)) { @@ -447,7 +454,7 @@ DeviceAddressVmmAllocator::TrackAllocatorAddressMappedAllocation( std::unique_ptr reservation, MemoryReservation::ScopedMapping mapping, int64_t memory_space, bool multi_device) { - void* va_ptr = allocator_address.opaque(); + uintptr_t va_ptr = AddressStart(allocator_address); CHECK(raw_allocation != nullptr); uint64_t physical_size = raw_allocation->address().size(); auto record = std::make_unique( @@ -465,7 +472,7 @@ std::optional DeviceAddressVmmAllocator::TryReuseMappedAllocation( PerDeviceState& state, const MappedAllocateRequest& request) { auto record_it = state.records_by_allocator_address.find( - request.reservation_address.opaque()); + AddressStart(request.reservation_address)); if (record_it == state.records_by_allocator_address.end()) { return std::nullopt; } @@ -479,20 +486,9 @@ DeviceAddressVmmAllocator::TryReuseMappedAllocation( return std::nullopt; } - auto pending_it = state.pending_deallocations.end(); - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == PendingDeallocationKind::kAllocation && - it->addr.IsSameAs(record->allocator_address())) { - pending_it = it; - break; - } - } - CHECK(pending_it != state.pending_deallocations.end()); - DeviceAddressBase reused_mem(record->allocator_key(), request.size); - MoveAllocatorRecordToActive(state, *record, request.size); - state.pending_deallocations.erase(pending_it); + ErasePendingDeallocation(state, record->allocator_deallocation()); + record->ReactivateAllocator(request.size); return reused_mem; } @@ -520,16 +516,9 @@ DeviceAddressVmmAllocator::EnsureReservationAvailableForFreshMapping( if (stale_overlap.has_value()) { CHECK(!stale_overlap->is_active); AllocationRecord& record = *stale_overlap->record; - if (stale_overlap->is_allocator) { - pending_completion_key = PendingDeallocationKey{ - PendingDeallocationKind::kAllocation, - record.allocator_stale_seqno(), record.allocator_address()}; - } else { - CHECK(record.reservation_stale()); - pending_completion_key = PendingDeallocationKey{ - PendingDeallocationKind::kMap, record.reservation_stale_seqno(), - record.reservation_address()}; - } + pending_completion_key = stale_overlap->is_allocator + ? record.allocator_deallocation().key() + : record.reservation_deallocation().key(); } } if (!pending_completion_key.has_value()) { @@ -623,12 +612,14 @@ DeviceAddressVmmAllocator::TryWithPendingReclaim(PerDeviceState& state, // Target 1.1x the requested size to provide some headroom. uint64_t target_size = rounded_size + rounded_size / 10; - for (const PendingDeallocation& pending : state.pending_deallocations) { + for (const PendingDeallocationNode* node = state.pending_head; + node != nullptr; node = node->next) { + const PendingDeallocation& pending = node->pending; if (pending.kind == PendingDeallocationKind::kMap) { continue; } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state.records_by_allocator_address.end()); CHECK(record_it->second->allocator_stale()); if (record_it->second->memory_space() == reclaim_exempt_memory_space_) { @@ -639,8 +630,7 @@ DeviceAddressVmmAllocator::TryWithPendingReclaim(PerDeviceState& state, CHECK_GT(reclaimable_bytes, 0); accumulated_size += reclaimable_bytes; target_seqno = std::max(target_seqno, pending.seqno); - selected.push_back( - PendingDeallocationKey{pending.kind, pending.seqno, pending.addr}); + selected.push_back(node->key()); if (accumulated_size >= target_size) { break; } @@ -679,17 +669,18 @@ DeviceAddressVmmAllocator::Allocate(int device_ordinal, uint64_t size, auto try_reuse = [&]() -> std::optional { state->mu.AssertHeld(); uint64_t rounded_size = RoundUpToGranularity(*state, size); - for (auto it = state->pending_deallocations.begin(); - it != state->pending_deallocations.end(); ++it) { - if (it->kind != PendingDeallocationKind::kAllocation) { + for (PendingDeallocationNode* node = state->pending_head; node != nullptr; + node = node->next) { + const PendingDeallocation& pending = node->pending; + if (pending.kind != PendingDeallocationKind::kAllocation) { continue; } auto record_it = - state->records_by_allocator_address.find(it->addr.opaque()); + state->records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state->records_by_allocator_address.end()); AllocationRecord& record = *record_it->second; CHECK(record.allocator_stale()); - CHECK(record.allocator_matches(it->addr)); + CHECK(record.allocator_matches(pending.addr)); if (record.kind() != AllocationRecord::Kind::kAllocate) { continue; } @@ -706,8 +697,8 @@ DeviceAddressVmmAllocator::Allocate(int device_ordinal, uint64_t size, } DeviceAddressBase reused_mem(record.allocator_key(), size); - MoveAllocatorRecordToActive(*state, record, size); - state->pending_deallocations.erase(it); + ErasePendingDeallocation(*state, *node); + record.ReactivateAllocator(size); return reused_mem; } @@ -815,7 +806,7 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, absl::MutexLock lock(state->mu); - auto record_it = state->records_by_allocator_address.find(mem.opaque()); + auto record_it = state->records_by_allocator_address.find(AddressStart(mem)); if (record_it == state->records_by_allocator_address.end() || !record_it->second->allocator_active() || !record_it->second->allocator_matches(mem)) { @@ -825,7 +816,7 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, mem.opaque())); } AllocationRecord& record = *record_it->second; - auto reservation_it = state->reservation_records.find(mem.opaque()); + auto reservation_it = state->reservation_records.find(AddressStart(mem)); CHECK(reservation_it == state->reservation_records.end() || !reservation_it->second->reservation_active()); if (record.reservation_active()) { @@ -853,7 +844,8 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, // Move the returned allocator address out of active ownership and keep its // mapping alive as stale state until the stream reaches `seqno`. record.MarkAllocatorStale(seqno); - state->pending_deallocations.push_back( + EnqueuePendingDeallocation( + *state, record.allocator_deallocation(), PendingDeallocation{PendingDeallocationKind::kAllocation, seqno, record.allocator_address(), reclaimable_bytes}); return absl::OkStatus(); @@ -902,46 +894,73 @@ DeviceAddressVmmAllocator::FindOverlappingRecord( } }; - auto check_record = [&](AllocationRecord* record, - DeviceAddressBase tracked_address, bool is_allocator, - bool is_active) -> std::optional { - if (matches(tracked_address)) { - return OverlappingRecord{record, tracked_address, is_allocator, - is_active}; + auto find_in_index = + [&](const auto& records, auto get_record, + bool is_allocator) -> std::optional { + auto check_record = + [&](const auto& entry) -> std::optional { + AllocationRecord* record = get_record(entry.second); + CHECK(is_allocator || record->has_reservation_alias()); + bool is_active = is_allocator ? record->allocator_active() + : record->reservation_active(); + if (!(is_active ? include_active : include_stale)) { + return std::nullopt; + } + DeviceAddressBase tracked_address = is_allocator + ? record->allocator_address() + : record->reservation_address(); + if (matches(tracked_address)) { + return OverlappingRecord{record, tracked_address, is_allocator, + is_active}; + } + return std::nullopt; + }; + + if (overlap_kind == OverlapKind::kExact) { + auto it = records.find(AddressStart(address)); + return it == records.end() ? std::nullopt : check_record(*it); + } + if (address.is_null() || address.size() == 0) { + return std::nullopt; } - return std::nullopt; - }; - if (include_allocator) { - for (const auto& [_, record_owner] : state.records_by_allocator_address) { - AllocationRecord* record = record_owner.get(); - bool include_record = (include_active && record->allocator_active()) || - (include_stale && record->allocator_stale()); - if (!include_record) { - continue; - } - if (auto overlap = - check_record(record, record->allocator_address(), - /*is_allocator=*/true, - /*is_active=*/record->allocator_active())) { + auto it = records.lower_bound(AddressStart(address)); + // Only the immediate predecessor can extend into the query: tracked ranges + // are disjoint, even while stale. Read the current range from the record + // because reusing an allocation can change its logical size. + if (it != records.begin()) { + auto previous = it; + --previous; + if (auto overlap = check_record(*previous)) { return overlap; } } - } - if (include_reservation) { - for (const auto& [_, record] : state.reservation_records) { - CHECK(record->has_reservation_alias()); - bool include_record = (include_active && record->reservation_active()) || - (include_stale && record->reservation_stale()); - if (!include_record) { - continue; - } - if (auto overlap = check_record( - record, record->reservation_address(), /*is_allocator=*/false, - /*is_active=*/record->reservation_active())) { + // With kBoth, the first successor either overlaps, exactly matches (so no + // other range overlaps), or starts beyond the query. Filtering by state can + // require checking additional successors, but never unrelated ranges. + for (; it != records.end() && it->first < AddressEnd(address); ++it) { + if (auto overlap = check_record(*it)) { return overlap; } } + return std::nullopt; + }; + + if (include_allocator) { + if (auto overlap = find_in_index( + state.records_by_allocator_address, + [](const std::unique_ptr& record) { + return record.get(); + }, + /*is_allocator=*/true)) { + return overlap; + } + } + if (include_reservation) { + return find_in_index( + state.reservation_records, + [](AllocationRecord* record) { return record; }, + /*is_allocator=*/false); } return std::nullopt; @@ -952,7 +971,7 @@ DeviceAddressVmmAllocator::ResolveMapSourceRecord( PerDeviceState& state, DeviceAddressBase source_address, uint64_t size) const { auto allocation_it = - state.records_by_allocator_address.find(source_address.opaque()); + state.records_by_allocator_address.find(AddressStart(source_address)); if (allocation_it == state.records_by_allocator_address.end() || !allocation_it->second->allocator_active() || !allocation_it->second->allocator_matches(source_address)) { @@ -1004,6 +1023,19 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( ABSL_ASSIGN_OR_RETURN( source_record, ResolveMapSourceRecord(state, request.source_address, request.size)); + + // The source already identifies its one reservation alias. An exact + // stale alias still owns this entire VA range, so no other tracked range + // can overlap it. Reactivate it before any destination-index searches. + // Source validation above and all state changes remain under state.mu. + if (source_record->reservation_stale() && + source_record->reservation_matches(request.reservation_address)) { + ErasePendingDeallocation(state, + source_record->reservation_deallocation()); + source_record->ReactivateReservation(); + return absl::OkStatus(); + } + ABSL_RETURN_IF_ERROR( CheckNoPartialReservationOverlap(state, request.reservation_address)); @@ -1027,13 +1059,10 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( } if (source_record->reservation_stale()) { - CHECK(source_record->has_reservation_alias()); - if (!source_record->reservation_matches(request.reservation_address)) { - pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kMap, - source_record->reservation_stale_seqno(), - source_record->reservation_address()}; - } + // An exact alias was reused above; any remaining source alias must be + // retired before mapping the source at a different destination. + pending_completion_key = + source_record->reservation_deallocation().key(); } if (!pending_completion_key.has_value()) { @@ -1043,22 +1072,12 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( if (stale_reservation_overlap.has_value()) { AllocationRecord& stale_record = *stale_reservation_overlap->record; - // A deferred UnMap() leaves the old mapping valid. Reuse it when it - // aliases the requested physical allocation; otherwise wait before - // overwriting it. + // Matching source aliases were reactivated above. This stale + // destination belongs to another source; wait before overwriting it. CHECK(stale_record.has_reservation_alias()); CHECK(stale_record.reservation_matches(request.reservation_address)); - if (stale_record.raw_allocation() == - source_record->raw_allocation()) { - stale_record.ReactivateReservation(); - ErasePendingDeallocation(state, PendingDeallocationKind::kMap, - request.reservation_address); - return absl::OkStatus(); - } pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kMap, - stale_record.reservation_stale_seqno(), - stale_record.reservation_address()}; + stale_record.reservation_deallocation().key(); } } @@ -1067,11 +1086,8 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( state, request.reservation_address, AddressRole::kAllocator, RecordState::kStale, OverlapKind::kExact); if (stale_allocator_overlap.has_value()) { - AllocationRecord& stale_record = *stale_allocator_overlap->record; pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kAllocation, - stale_record.allocator_stale_seqno(), - stale_record.allocator_address()}; + stale_allocator_overlap->record->allocator_deallocation().key(); } } @@ -1090,8 +1106,8 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( << ", actual=" << mapped.opaque(); source_record->AddActiveReservationAlias(std::move(mapping)); - auto mapping_insert_result = - state.reservation_records.emplace(mapped.opaque(), source_record); + auto mapping_insert_result = state.reservation_records.emplace( + AddressStart(mapped), source_record); CHECK(mapping_insert_result.second); return absl::OkStatus(); } @@ -1144,17 +1160,18 @@ DeviceAddressVmmAllocator::OpenBatchSize(const PerDeviceState& state) const { // Entries are appended in non-decreasing seqno order, so the open batch is // the trailing run carrying the open seqno. Walking it costs at most // kMaxOpenDeallocationBatchEntries steps. - for (auto it = state.pending_deallocations.rbegin(); - it != state.pending_deallocations.rend(); ++it) { - if (it->seqno != state.open_deallocation_batch_seqno) { + for (const PendingDeallocationNode* node = state.pending_tail; + node != nullptr; node = node->previous) { + const PendingDeallocation& pending = node->pending; + if (pending.seqno != state.open_deallocation_batch_seqno) { break; } ++size.entries; if (std::numeric_limits::max() - size.bytes < - it->reclaimable_bytes) { + pending.reclaimable_bytes) { size.bytes = std::numeric_limits::max(); } else { - size.bytes += it->reclaimable_bytes; + size.bytes += pending.reclaimable_bytes; } } return size; @@ -1216,25 +1233,38 @@ absl::Status DeviceAddressVmmAllocator::FlushOpenDeallocationBatch( return absl::OkStatus(); } -void DeviceAddressVmmAllocator::ErasePendingDeallocation( - PerDeviceState& state, PendingDeallocationKind kind, - DeviceAddressBase addr) { - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == kind && it->addr.IsSameAs(addr)) { - state.pending_deallocations.erase(it); - return; - } +void DeviceAddressVmmAllocator::EnqueuePendingDeallocation( + PerDeviceState& state, PendingDeallocationNode& node, + PendingDeallocation pending) { + CHECK_EQ(node.pending.seqno, 0); + CHECK_GT(pending.seqno, 0); + node.pending = pending; + node.previous = state.pending_tail; + node.next = nullptr; + if (state.pending_tail != nullptr) { + state.pending_tail->next = &node; + } else { + state.pending_head = &node; } + state.pending_tail = &node; } -void DeviceAddressVmmAllocator::MoveAllocatorRecordToActive( - PerDeviceState& state, AllocationRecord& record, uint64_t new_size) { - void* allocator_va = record.allocator_key(); - auto record_it = state.records_by_allocator_address.find(allocator_va); - CHECK(record_it != state.records_by_allocator_address.end()); - CHECK_EQ(record_it->second.get(), &record); - record.ReactivateAllocator(new_size); +void DeviceAddressVmmAllocator::ErasePendingDeallocation( + PerDeviceState& state, PendingDeallocationNode& node) { + CHECK_GT(node.pending.seqno, 0); + if (node.previous != nullptr) { + node.previous->next = node.next; + } else { + CHECK_EQ(state.pending_head, &node); + state.pending_head = node.next; + } + if (node.next != nullptr) { + node.next->previous = node.previous; + } else { + CHECK_EQ(state.pending_tail, &node); + state.pending_tail = node.previous; + } + node = {}; } absl::Status DeviceAddressVmmAllocator::WaitUntilSeqno(PerDeviceState& state, @@ -1259,19 +1289,20 @@ absl::Status DeviceAddressVmmAllocator::WaitUntilSeqno(PerDeviceState& state, void DeviceAddressVmmAllocator::CompleteReadyAllocatorDeallocationsForReclaim( PerDeviceState& state, uint64_t completed_seqno) { std::vector selected; - for (const PendingDeallocation& pending : state.pending_deallocations) { + for (const PendingDeallocationNode* node = state.pending_head; + node != nullptr; node = node->next) { + const PendingDeallocation& pending = node->pending; if (pending.seqno > completed_seqno || pending.kind == PendingDeallocationKind::kMap) { continue; } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); if (record_it != state.records_by_allocator_address.end() && record_it->second->memory_space() == reclaim_exempt_memory_space_) { continue; } - selected.push_back( - PendingDeallocationKey{pending.kind, pending.seqno, pending.addr}); + selected.push_back(node->key()); } for (const PendingDeallocationKey& key : selected) { CompletePendingDeallocationByKey(state, key); @@ -1280,22 +1311,22 @@ void DeviceAddressVmmAllocator::CompleteReadyAllocatorDeallocationsForReclaim( void DeviceAddressVmmAllocator::CompletePendingDeallocationByKey( PerDeviceState& state, const PendingDeallocationKey& key) { - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == key.kind && it->seqno == key.seqno && - it->addr.IsSameAs(key.addr)) { - PendingDeallocation pending = *it; - state.pending_deallocations.erase(it); - CompletePendingDeallocation(state, pending); + for (PendingDeallocationNode* node = state.pending_head; node != nullptr; + node = node->next) { + if (node->pending.kind == key.kind && node->pending.seqno == key.seqno && + node->pending.addr.IsSameAs(key.addr)) { + CompletePendingDeallocation(state, *node); return; } } } void DeviceAddressVmmAllocator::CompletePendingDeallocation( - PerDeviceState& state, const PendingDeallocation& pending) { + PerDeviceState& state, PendingDeallocationNode& node) { + const PendingDeallocation pending = node.pending; + ErasePendingDeallocation(state, node); if (pending.kind == PendingDeallocationKind::kMap) { - auto record_it = state.reservation_records.find(pending.addr.opaque()); + auto record_it = state.reservation_records.find(AddressStart(pending.addr)); CHECK(record_it != state.reservation_records.end()); AllocationRecord& record = *record_it->second; CHECK(record.reservation_stale()); @@ -1308,7 +1339,7 @@ void DeviceAddressVmmAllocator::CompletePendingDeallocation( } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state.records_by_allocator_address.end()); CHECK_EQ(pending.kind, PendingDeallocationKind::kAllocation); CHECK(record_it->second->allocator_stale()); @@ -1329,10 +1360,8 @@ void DeviceAddressVmmAllocator::CompletePendingDeallocation( // sequence numbers are assigned in that same order. Without this invariant // the alias could sit in a still-open batch with no stream marker at all. CHECK_LE(record.reservation_stale_seqno(), pending.seqno); - CompletePendingDeallocationByKey( - state, PendingDeallocationKey{PendingDeallocationKind::kMap, - record.reservation_stale_seqno(), - record.reservation_address()}); + CompletePendingDeallocationByKey(state, + record.reservation_deallocation().key()); CHECK(!record.has_reservation_alias()); } uint64_t physical_size = record.raw_allocation()->address().size(); @@ -1362,7 +1391,7 @@ absl::Status DeviceAddressVmmAllocator::UnMap(int device_ordinal, // UnMap() only accepts the exact active reservation range previously created // by Map(). Allocator addresses and subranges are not valid UnMap() inputs. auto reservation_it = - state->reservation_records.find(reservation_address.opaque()); + state->reservation_records.find(AddressStart(reservation_address)); if (reservation_it == state->reservation_records.end()) { return absl::NotFoundError(absl::StrFormat( "UnMap() requires an exact active reservation range created by Map(): " @@ -1397,7 +1426,8 @@ absl::Status DeviceAddressVmmAllocator::UnMap(int device_ordinal, // stream marker will be enqueued for the whole batch when it is flushed. uint64_t seqno = GetOrCreateOpenDeallocationBatchSeqno(*state); record->MarkReservationStale(seqno); - state->pending_deallocations.push_back( + EnqueuePendingDeallocation( + *state, record->reservation_deallocation(), PendingDeallocation{PendingDeallocationKind::kMap, seqno, reservation_address, /*reclaimable_bytes=*/0}); return absl::OkStatus(); diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h index 2b1891aa304c08..2e4df779996dd7 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h @@ -17,13 +17,13 @@ limitations under the License. #define XLA_STREAM_EXECUTOR_DEVICE_ADDRESS_VMM_ALLOCATOR_H_ #include -#include #include #include #include #include #include "absl/base/thread_annotations.h" +#include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -292,6 +292,43 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { kMap, }; + // Stable identity for one pending operation. A batch shares one sequence + // number, so the operation kind and address are also needed to select an + // entry after WaitUntilSeqno() temporarily releases state.mu. + struct PendingDeallocationKey { + PendingDeallocationKind kind; + uint64_t seqno; + DeviceAddressBase addr; + }; + + // Snapshot of a stream-ordered deferred operation. Copy before unlinking its + // node: completing the operation can destroy the record that owns the node. + struct PendingDeallocation { + PendingDeallocationKind kind; + uint64_t seqno; + // Allocator address for allocation deallocations; reservation address for + // kMap. + DeviceAddressBase addr; + // Physical bytes released when this entry completes; 0 for kMap. Store the + // bytes charged when enqueuing so open-batch accounting remains consistent. + uint64_t reclaimable_bytes; + }; + + // Intrusive queue node owned by an allocation or its reservation alias. + // Records stay at stable addresses while queued, allowing cancellation + // without searching the queue or allocating separate queue storage. + struct PendingDeallocationNode { + // A zero seqno means this node is not queued. + PendingDeallocation pending{PendingDeallocationKind::kAllocation, 0, + DeviceAddressBase(), 0}; + PendingDeallocationNode* previous = nullptr; + PendingDeallocationNode* next = nullptr; + + PendingDeallocationKey key() const { + return {pending.kind, pending.seqno, pending.addr}; + } + }; + // Lifetime record for one raw physical allocation. // // The record is owned by records_by_allocator_address while either the @@ -347,6 +384,11 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { } uint64_t reservation_stale_seqno() const; + PendingDeallocationNode& allocator_deallocation() { + return allocator_deallocation_; + } + PendingDeallocationNode& reservation_deallocation(); + void MarkAllocatorStale(uint64_t seqno); void ReactivateAllocator(uint64_t new_size); @@ -361,6 +403,7 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { MemoryReservation::ScopedMapping mapping; // Zero while active; the deferred UnMap() sequence number while stale. uint64_t stale_seqno = 0; + PendingDeallocationNode deallocation = {}; }; Kind kind_; @@ -379,33 +422,7 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { std::optional reservation_alias_; uint64_t allocator_stale_seqno_ = 0; - }; - - // Queue entry for a stream-ordered deferred operation. The heavy resources - // live in AllocationRecord; this entry only says which stale address becomes - // safe to complete when the GPU timeline reaches `seqno`. - struct PendingDeallocation { - PendingDeallocationKind kind; - // GPU stream sequence number recorded at deallocation time. When the - // pinned_timeline value reaches this seqno, the memory is safe to free. - uint64_t seqno; - // Allocator address for allocation deallocations; reservation address for - // kMap. - DeviceAddressBase addr; - // Physical bytes released when this entry completes; 0 for kMap, which only - // drops a reservation alias. Stored here rather than re-derived from the - // AllocationRecord so that open-batch accounting reads exactly the value - // that was accounted when the entry was queued. - uint64_t reclaimable_bytes; - }; - - // Stable identity for one pending operation. A batch shares one sequence - // number, so the operation kind and address are also needed to select an - // entry after WaitUntilSeqno() temporarily releases state.mu. - struct PendingDeallocationKey { - PendingDeallocationKind kind; - uint64_t seqno; - DeviceAddressBase addr; + PendingDeallocationNode allocator_deallocation_; }; struct PerDeviceState { @@ -455,25 +472,30 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { // The next deferred operation opens a new batch with seqno 2. // // Only the sequence number is stored. The batch's entry and byte totals are - // the trailing run of `pending_deallocations` carrying this seqno (entries + // the trailing run of the pending queue carrying this seqno (entries // are appended in non-decreasing seqno order, so that run is contiguous), // and are computed on demand by OpenDeallocationBatchSize(). Keeping them // derived means cancelling an entry -- which happens whenever a stale // record is reused -- needs no batch bookkeeping at all. uint64_t open_deallocation_batch_seqno ABSL_GUARDED_BY(mu) = 0; - std::deque pending_deallocations ABSL_GUARDED_BY(mu); - // Owns AllocationRecord objects. Key is the allocator address pointer - // (`AllocationRecord::allocator_address().opaque()`), including the - // reservation-derived allocator address returned by the mapped Allocate() - // overload. Allocator-address active/stale state is stored in + PendingDeallocationNode* pending_head ABSL_GUARDED_BY(mu) = nullptr; + PendingDeallocationNode* pending_tail ABSL_GUARDED_BY(mu) = nullptr; + // Owns AllocationRecord objects. Key is the allocator address as uintptr_t, + // including the reservation-derived address returned by the mapped + // Allocate() overload. Allocator-address active/stale state is stored in // AllocationRecord::allocator_active()/allocator_stale(). - absl::flat_hash_map> + // Both address indexes are ordered by range start. Ranges within each + // index are disjoint, including stale ranges retained until deferred + // teardown. This permits overlap checks starting at lower_bound() instead + // of scanning every allocation. The pointed-to records stay valid when + // insertion or erasure invalidates btree iterators. + absl::btree_map> records_by_allocator_address ABSL_GUARDED_BY(mu); - // Reservation-address index. Keys are reservation alias pointers - // (`AllocationRecord::reservation_address().opaque()`) created by Map(). + // Reservation-address index. Keys are reservation alias addresses as + // uintptr_t, created by Map(). // Active/stale state is stored in the record. - absl::flat_hash_map reservation_records + absl::btree_map reservation_records ABSL_GUARDED_BY(mu); }; @@ -677,14 +699,15 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { absl::Status DrainPendingDeallocations(PerDeviceState& state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - // Removes the matching pending entry when a stale record is reused. - void ErasePendingDeallocation(PerDeviceState& state, - PendingDeallocationKind kind, - DeviceAddressBase addr) + // Appends a record-owned node in stream sequence order. + void EnqueuePendingDeallocation(PerDeviceState& state, + PendingDeallocationNode& node, + PendingDeallocation pending) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - void MoveAllocatorRecordToActive(PerDeviceState& state, - AllocationRecord& record, uint64_t new_size) + // Unlinks a queued node in constant time, before reusing or completing it. + void ErasePendingDeallocation(PerDeviceState& state, + PendingDeallocationNode& node) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); // Waits for the device timeline to reach `target_seqno`. Temporarily releases @@ -699,12 +722,12 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { uint64_t completed_seqno) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - // Completes a pending operation whose stream sequence has passed by dropping - // its ScopedMappings, allocator-owned reservation, and raw allocation - // reference. This is where VA unmap, reservation release, and PA budget - // accounting happen. + // Unlinks and completes a pending operation whose stream sequence has passed. + // Copies its metadata before releasing the mappings and owning record, which + // can destroy the node. This is where VA unmap, reservation release, and PA + // budget accounting happen. void CompletePendingDeallocation(PerDeviceState& state, - const PendingDeallocation& pending) + PendingDeallocationNode& node) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); // Finds, erases, and completes the selected pending entry if it is still diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc index a4417c2886db08..9dac5e156ebb8b 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc @@ -18,7 +18,10 @@ limitations under the License. #include #include #include +#include #include +#include +#include #include #include @@ -27,6 +30,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/memory_allocation.h" #include "xla/stream_executor/memory_reservation.h" #include "xla/stream_executor/mock_platform.h" @@ -72,6 +76,7 @@ class TestMemoryReservation final : public MemoryReservation { } int active_mapping_count() const { return active_mapping_count_; } + int mapping_count() const { return mapping_count_; } private: absl::Status Map(size_t reservation_offset, size_t allocation_offset, @@ -82,6 +87,7 @@ class TestMemoryReservation final : public MemoryReservation { return absl::InvalidArgumentError("mapping range is out of bounds"); } ++active_mapping_count_; + ++mapping_count_; return absl::OkStatus(); } @@ -101,6 +107,7 @@ class TestMemoryReservation final : public MemoryReservation { std::unique_ptr storage_; uint64_t size_; int active_mapping_count_ = 0; + int mapping_count_ = 0; }; class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { @@ -120,6 +127,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { } int allocation_count() const { return allocation_count_; } + int timeline_write_count() const { return timeline_write_count_; } protected: absl::Status InitializeDeviceState(PerDeviceState& state) override { @@ -152,6 +160,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { absl::Status EnqueueDeferredDeallocation(PerDeviceState& state, uint64_t seqno) override { + ++timeline_write_count_; __atomic_store_n(state.pinned_timeline, seqno, __ATOMIC_RELEASE); return absl::OkStatus(); } @@ -167,6 +176,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { uint64_t physical_size_padding_; std::function on_device_destroy_; int allocation_count_ = 0; + int timeline_write_count_ = 0; }; class DeviceAddressVmmAllocatorTest : public ::testing::Test { @@ -404,5 +414,260 @@ TEST_F(DeviceAddressVmmAllocatorTest, EXPECT_EQ(allocator->allocation_count(), 2); } +TEST_F(DeviceAddressVmmAllocatorTest, + OverlapLookupPreservesRangeBoundariesForActiveAndStaleAddresses) { + for (bool alias : {false, true}) { + for (bool stale : {false, true}) { + SCOPED_TRACE(testing::Message() + << "alias=" << alias << " stale=" << stale); + TestMemoryReservation reservation(16 * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, + TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + ASSERT_OK_AND_ASSIGN( + auto tracked, + alias ? allocator->Allocate(/*device_ordinal=*/0, 4 * kGranularity) + : allocator->Allocate( + /*device_ordinal=*/0, 4 * kGranularity, + /*retry_on_failure=*/false, /*memory_space=*/0, + &reservation, 4 * kGranularity, 4 * kGranularity)); + if (alias) { + ASSERT_THAT(allocator->Map(0, tracked.cref(), &reservation, + 4 * kGranularity, 4 * kGranularity), + absl_testing::IsOk()); + } + if (stale) { + if (alias) { + ASSERT_THAT(allocator->UnMap(0, &reservation, 4 * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); + } else { + ASSERT_THAT(allocator->Deallocate(0, tracked.Release()), + absl_testing::IsOk()); + } + } + ASSERT_OK_AND_ASSIGN(auto source, + allocator->Allocate(0, 16 * kGranularity)); + struct Range { + uint64_t offset; + uint64_t size; + }; + // Intersect the successor, lie inside the predecessor, intersect its + // end, share its start with a different size, or contain the whole range. + const Range partial_overlaps[] = {{2, 3}, {5, 1}, {7, 2}, + {4, 2}, {4, 5}, {2, 8}}; + for (Range range : partial_overlaps) { + SCOPED_TRACE(testing::Message() + << "offset=" << range.offset << " size=" << range.size); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, + range.offset * kGranularity, + range.size * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_THAT(allocator->Allocate( + 0, range.size * kGranularity, false, 0, &reservation, + range.offset * kGranularity, range.size * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + // Rejecting an overlap must not drain the stale mapping. + EXPECT_EQ(reservation.active_mapping_count(), 1); + } + + // Exactly adjacent ranges are disjoint on either side of the record. + ASSERT_THAT( + allocator->Map(0, source.cref(), &reservation, 0, 4 * kGranularity), + absl_testing::IsOk()); + ASSERT_OK_AND_ASSIGN( + auto right, + allocator->Allocate(0, 4 * kGranularity, false, 0, &reservation, + 8 * kGranularity, 4 * kGranularity)); + auto exact = + allocator->Allocate(0, 4 * kGranularity, false, 0, &reservation, + 4 * kGranularity, 4 * kGranularity); + if (stale) { + ASSERT_THAT(exact, absl_testing::IsOk()); + } else { + EXPECT_THAT(exact, StatusIs(absl::StatusCode::kAlreadyExists)); + } + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, 4 * kGranularity), + absl_testing::IsOk()); + if (alias && !stale) { + ASSERT_THAT(allocator->UnMap(0, &reservation, 4 * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); + } + } + } +} + +TEST_F(DeviceAddressVmmAllocatorTest, + ManyReservationAliasesReuseAndRemoveOrderedRecords) { + constexpr int kCount = 512; + TestMemoryReservation reservation((2 * kCount + 1) * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + std::vector> sources; + for (int i = 0; i < kCount; ++i) { + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, kGranularity)); + sources.push_back(std::move(source)); + } + // Insert in an order that interleaves low and high addresses, leaving gaps + // between aliases. Records must remain valid across index rebalancing. + for (int n = 0; n < kCount; ++n) { + int i = (n % 2 == 0) ? n / 2 : kCount - 1 - n / 2; + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (2 * i + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (2 * i + 1) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + for (int i = kCount - 1; i >= 0; --i) { + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (2 * i + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + EXPECT_EQ(reservation.active_mapping_count(), kCount); + EXPECT_EQ(reservation.mapping_count(), kCount); + EXPECT_EQ(allocator->allocation_count(), kCount); + + // A query starting in a gap must still find the next occupied range. + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, 4 * kGranularity)); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, + kCount * kGranularity, 4 * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (2 * i + 1) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); + // The removed alias ranges can now be covered by a larger mapping. + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, + kCount * kGranularity, 4 * kGranularity), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, kCount * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); +} + +TEST_F(DeviceAddressVmmAllocatorTest, + StaleAliasReuseValidatesSourceAndFullReservationRange) { + TestMemoryReservation reservation(2 * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, 2 * kGranularity)); + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, kGranularity), + absl_testing::IsOk()); + + // Matching the start is insufficient: a larger range must still be rejected, + // without retiring the stale alias or its pending operation. + EXPECT_THAT( + allocator->Map(0, source.cref(), &reservation, 0, 2 * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_EQ(reservation.active_mapping_count(), 1); + EXPECT_EQ(allocator->timeline_write_count(), 0); + + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + absl_testing::IsOk()); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + StatusIs(absl::StatusCode::kAlreadyExists)); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.mapping_count(), 1); + EXPECT_EQ(reservation.active_mapping_count(), 1); + EXPECT_EQ(allocator->timeline_write_count(), 0); + + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, kGranularity), + absl_testing::IsOk()); + DeviceAddressBase stale_source = source.Release(); + ASSERT_THAT(allocator->Deallocate(0, stale_source), absl_testing::IsOk()); + EXPECT_THAT(allocator->Map(0, stale_source, &reservation, 0, kGranularity), + StatusIs(absl::StatusCode::kNotFound)); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); +} + +TEST_F(DeviceAddressVmmAllocatorTest, + MixedPendingOperationsCanBeCancelledAndRequeuedInAnyOrder) { + constexpr int kCount = 4; + TestMemoryReservation reservation(2 * kCount * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + std::vector> temps; + std::vector> sources; + for (int i = 0; i < kCount; ++i) { + ASSERT_OK_AND_ASSIGN( + auto temp, allocator->Allocate(0, kGranularity, false, 0, &reservation, + i * kGranularity, kGranularity)); + temps.push_back(std::move(temp)); + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, kGranularity)); + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + sources.push_back(std::move(source)); + } + + for (int step = 0; step < 3; ++step) { + SCOPED_TRACE(step); + // Interleave allocator deallocations and alias unmaps in the same batch. + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->Deallocate(0, temps[i].Release()), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + // Cancel from the middle, head, tail, and finally the remaining entries. + // The following iteration requeues the same record-owned nodes. + for (int i : {1, 0, 3, 2}) { + ASSERT_OK_AND_ASSIGN( + temps[i], allocator->Allocate(0, kGranularity, false, 0, &reservation, + i * kGranularity, kGranularity)); + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), + absl_testing::IsOk()); + EXPECT_EQ(allocator->timeline_write_count(), 0); + EXPECT_EQ(allocator->allocation_count(), 2 * kCount); + EXPECT_EQ(reservation.mapping_count(), 2 * kCount); + EXPECT_EQ(reservation.active_mapping_count(), 2 * kCount); + } + + // Leave one alias pending between cancelled entries. Draining must release + // only that alias, then permit a fresh mapping at the same address. + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + for (int i : {3, 0, 2}) { + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 2 * kCount - 1); + EXPECT_EQ(allocator->timeline_write_count(), 1); + ASSERT_THAT(allocator->Map(0, sources[1].cref(), &reservation, + (kCount + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + temps.clear(); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); + EXPECT_EQ(allocator->timeline_write_count(), 2); +} + } // namespace } // namespace stream_executor From a34fdf2e24c036471eef239125ca50d972ada4ae Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Fri, 18 Sep 2026 05:57:58 -0700 Subject: [PATCH 18/32] Let `KernelArgsPackingSpec` be validated before launch time. A packing spec's `AddAddressArgument(i)` indices refer to positions in the device buffer list that the runtime builds from `emitters::KernelArguments`. Nothing checks that those indices are in range until `BuildArguments` runs, at which point the kernel is already being launched, so an off-by-one between a kernel's signature and the emitter that set up its arguments only shows up as a runtime error on the machine executing the program. `Validate(num_available_args)` performs that check up front, so emitters can fail at compile time with a message that points at the emitter rather than at the launch. `size()`, `MaxArgumentIndex()` and `KernelArgPackingSpec::relocation_argument_index()` expose the information it needs and are useful for diagnostics on their own. `Identity(num_args)` covers the common "pass every buffer pointer straight through, in order" case. `ptx_custom_kernel.cc` had a private copy of it; it now uses the shared one. PiperOrigin-RevId: 983818570 --- .../gpu/codegen/kernels/ptx_custom_kernel.cc | 20 ++----- .../kernel_args_packing_spec.cc | 33 +++++++++++ .../kernel_args_packing_spec.h | 28 ++++++++++ .../kernel_args_packing_spec_test.cc | 56 +++++++++++++++++++ 4 files changed, 123 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc index 9dd19b68021369..44883901699c6d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc +++ b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc @@ -33,16 +33,6 @@ limitations under the License. namespace xla::gpu::kernel { namespace se = ::stream_executor; -namespace { -se::KernelArgsPackingSpec IdentityPackingSpec(int num_args) { - se::KernelArgsPackingSpec packing_spec; - for (int i = 0; i < num_args; i++) { - packing_spec.AddAddressArgument(i); - } - return packing_spec; -} -} // namespace - // Note: Make sure that the kernel_name matches the kernel name in the ptx, // otherwise you will get a "CUDA_ERROR_NOT_FOUND: named symbol not found.". // E.g. `.visible .entry AddI32(...)` would have a kernel name of "AddI32". @@ -54,7 +44,8 @@ absl::StatusOr GetPtxCustomKernel(std::string kernel_name, size_t shared_memory_bytes) { se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - ptx, kernel_name, /*arity=*/num_args, IdentityPackingSpec(num_args)); + ptx, kernel_name, /*arity=*/num_args, + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, shared_memory_bytes); } @@ -65,7 +56,8 @@ absl::StatusOr GetPtxCustomKernel( se::ClusterDim cluster_dim, size_t shared_memory_bytes) { se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - ptx, kernel_name, /*arity=*/num_args, IdentityPackingSpec(num_args)); + ptx, kernel_name, /*arity=*/num_args, + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, cluster_dim, shared_memory_bytes); } @@ -77,7 +69,7 @@ absl::StatusOr GetOwnedPtxCustomKernel( se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateOwningCudaPtxInMemorySpec( std::move(ptx), kernel_name, /*arity=*/num_args, - IdentityPackingSpec(num_args)); + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, shared_memory_bytes); } @@ -89,7 +81,7 @@ absl::StatusOr CreateOwnedCubinCustomKernel( se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateOwningCudaCubinInMemorySpec( std::move(cubin), kernel_name, /*arity=*/num_args, - IdentityPackingSpec(num_args)); + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), std::move(kernel_spec), block_dim, thread_dim, shared_memory_bytes); } diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc index d174e3ac179384..d52f92f94957d0 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc @@ -90,6 +90,39 @@ KernelArgPackingSpec KernelArgPackingSpec::BuildArgRelocation( KernelArgPackingRelocation::Kind::kBits64Absolute, argument_index)}); } +KernelArgsPackingSpec KernelArgsPackingSpec::Identity(int num_args) { + CHECK_GE(num_args, 0); + std::vector kernel_arguments; + kernel_arguments.reserve(num_args); + for (int i = 0; i < num_args; ++i) { + kernel_arguments.push_back(KernelArgPackingSpec::BuildArgRelocation(i)); + } + return KernelArgsPackingSpec(std::move(kernel_arguments)); +} + +std::optional KernelArgsPackingSpec::MaxArgumentIndex() const { + std::optional max_index; + for (const KernelArgPackingSpec& kernel_argument : kernel_arguments_) { + std::optional index = kernel_argument.relocation_argument_index(); + if (index.has_value() && (!max_index.has_value() || *index > *max_index)) { + max_index = index; + } + } + return max_index; +} + +absl::Status KernelArgsPackingSpec::Validate(size_t num_available_args) const { + std::optional max_index = MaxArgumentIndex(); + if (max_index.has_value() && + static_cast(*max_index) >= num_available_args) { + return absl::InvalidArgumentError(absl::StrFormat( + "Kernel arguments packing spec refers to device buffer %d, but only " + "%zu buffers are available", + *max_index, num_available_args)); + } + return absl::OkStatus(); +} + absl::StatusOr> KernelArgsPackingSpec::BuildArguments( absl::Span> args, diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h index 586fcfc963eeca..ceec408db33ec3 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/stream_executor/kernel_args.h" @@ -66,6 +67,14 @@ class KernelArgPackingSpec { // Builds KernelArgPackingSpec that refers to the given arg number. static KernelArgPackingSpec BuildArgRelocation(int argument_index); + // Returns the index of the argument this spec relocates, or `std::nullopt` + // if this spec holds a constant. + std::optional relocation_argument_index() const { + return relocation_.has_value() + ? std::make_optional(relocation_->argument_index()) + : std::nullopt; + } + // Builds KernelArgPackingSpec that refers to a constant. The value must be // trivially copyable. template @@ -146,6 +155,11 @@ class KernelArgsPackingSpec { std::vector kernel_arguments) : kernel_arguments_(std::move(kernel_arguments)) {} + // Builds the packing spec that passes the first `num_args` device buffer + // pointers to the kernel unchanged, in order. This is the right spec for + // kernels whose signature is exactly "all inputs, then all outputs". + static KernelArgsPackingSpec Identity(int num_args); + // Adds a single argument packing spec to the kernel arguments packing spec. void AddArgument(KernelArgPackingSpec spec) { kernel_arguments_.push_back(std::move(spec)); @@ -163,6 +177,20 @@ class KernelArgsPackingSpec { kernel_arguments_.push_back(KernelArgPackingSpec::BuildFor(value)); } + // Number of arguments the kernel will be launched with. + size_t size() const { return kernel_arguments_.size(); } + + // Returns the largest device buffer index referenced by any relocation, or + // `std::nullopt` if the spec contains no relocations. + std::optional MaxArgumentIndex() const; + + // Returns an error if any relocation refers to a device buffer index that is + // out of range for `num_available_args` buffers. `BuildArguments` performs + // the same check, but only once the kernel is actually launched; calling + // `Validate` while emitting lets a mismatch between the packing spec and the + // kernel arguments surface at compile time instead. + absl::Status Validate(size_t num_available_args) const; + // Materializes the argument buffers for this packing spec. The `args` span // must contain at least the number of arguments referenced in the packing // spec, otherwise an error will be returned. diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc b/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc index 281d455fae6359..36dac56a368bb0 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include @@ -232,5 +233,60 @@ TEST(KernelArgsPackingSpecTest, FromProto) { ElementsAre(0x34, 0x12, 0x00, 0x00)); } +TEST(KernelArgsPackingSpecTest, Identity) { + KernelArgsPackingSpec spec = KernelArgsPackingSpec::Identity(3); + + EXPECT_EQ(spec.size(), 3); + EXPECT_THAT(spec.ToProto(), IsOkAndHolds(EqualsProto(R"pb( + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 0 } + } + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 1 } + } + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 2 } + } + )pb"))); +} + +TEST(KernelArgsPackingSpecTest, IdentityWithZeroArguments) { + KernelArgsPackingSpec spec = KernelArgsPackingSpec::Identity(0); + + EXPECT_EQ(spec.size(), 0); + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); +} + +TEST(KernelArgsPackingSpecTest, MaxArgumentIndex) { + KernelArgsPackingSpec spec; + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); + + spec.AddConstantArgument(0x1234); + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); + + spec.AddAddressArgument(/*argument_index=*/5); + spec.AddAddressArgument(/*argument_index=*/2); + EXPECT_EQ(spec.MaxArgumentIndex(), 5); +} + +TEST(KernelArgsPackingSpecTest, Validate) { + KernelArgsPackingSpec spec; + spec.AddAddressArgument(/*argument_index=*/0); + spec.AddConstantArgument(0x1234); + spec.AddAddressArgument(/*argument_index=*/2); + + EXPECT_OK(spec.Validate(/*num_available_args=*/3)); + EXPECT_OK(spec.Validate(/*num_available_args=*/4)); + EXPECT_THAT(spec.Validate(/*num_available_args=*/2), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST(KernelArgsPackingSpecTest, ValidateConstantsOnlySpecNeedsNoArguments) { + KernelArgsPackingSpec spec; + spec.AddConstantArgument(0x1234); + + EXPECT_OK(spec.Validate(/*num_available_args=*/0)); +} + } // namespace } // namespace stream_executor From ce721307cd4ecd353d26e82704d3a82e19de98bf Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Fri, 18 Sep 2026 05:58:12 -0700 Subject: [PATCH 19/32] Upgrade rules_cc to 0.2.25 and bazel_features to 1.50.0 in XLA and TensorFlow - Bump rules_cc from 0.2.20 to 0.2.25 in MODULE.bazel and workspace3.bzl for both XLA and TensorFlow. - Bump bazel_features from 1.25.0 to 1.50.0 in workspace3.bzl for both XLA and TensorFlow to provide bazel_features.cc.cc_common_is_in_rules_cc used by rules_cc 0.2.25 in WORKSPACE mode. PiperOrigin-RevId: 983818634 --- MODULE.bazel | 2 +- tensorflow/workspace3.bzl | 12 ++++++------ third_party/xla/MODULE.bazel | 2 +- third_party/xla/workspace3.bzl | 12 ++++++------ 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/MODULE.bazel b/MODULE.bazel index 539d8107d2fa2f..61ea33e7a9a983 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -35,7 +35,7 @@ single_version_override( bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_github_googlecloudplatform_google_cloud_cpp") bazel_dep(name = "crc32c", version = "1.1.0", repo_name = "com_github_google_crc32c") bazel_dep(name = "brotli", version = "1.1.0", repo_name = "org_brotli") -bazel_dep(name = "rules_cc", version = "0.2.20") +bazel_dep(name = "rules_cc", version = "0.2.25") bazel_dep(name = "curl", version = "8.11.0.bcr.5") bazel_dep(name = "rules_webtesting", version = "0.4.1", repo_name = "io_bazel_rules_webtesting") bazel_dep(name = "rules_closure", version = "0.15.0", repo_name = "io_bazel_rules_closure") diff --git a/tensorflow/workspace3.bzl b/tensorflow/workspace3.bzl index 06809b6207741a..bf557b54aa9769 100644 --- a/tensorflow/workspace3.bzl +++ b/tensorflow/workspace3.bzl @@ -60,19 +60,19 @@ def workspace(): tf_http_archive( name = "bazel_features", - sha256 = "4fd9922d464686820ffd8fcefa28ccffa147f7cdc6b6ac0d8b07fde565c65d66", - strip_prefix = "bazel_features-1.25.0", + sha256 = "094367e732ece23f334eaf84089a720b861d053beeba6a6a68356d3aee1dc32b", + strip_prefix = "bazel_features-1.50.0", urls = tf_mirror_urls( - "https://github.com/bazel-contrib/bazel_features/releases/download/v1.25.0/bazel_features-v1.25.0.tar.gz", + "https://github.com/bazel-contrib/bazel_features/releases/download/v1.50.0/bazel_features-v1.50.0.tar.gz", ), ) tf_http_archive( name = "rules_cc", - sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", - strip_prefix = "rules_cc-0.2.20", + sha256 = "bd7124a844d0403b4b353bcea34d6c8b2ba88dc26881c26c9ee668da89b71846", + strip_prefix = "rules_cc-0.2.25", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.25/rules_cc-0.2.25.tar.gz", ), ) diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 432b20da54537f..9b994cf3db0c72 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -23,7 +23,7 @@ bazel_dep(name = "pybind11_abseil", version = "202402.0") bazel_dep(name = "pybind11_bazel", version = "3.0.0") bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7") bazel_dep(name = "re2", version = "2025-11-05.bcr.1", repo_name = "com_googlesource_code_re2") -bazel_dep(name = "rules_cc", version = "0.2.20") +bazel_dep(name = "rules_cc", version = "0.2.25") bazel_dep(name = "rules_java", version = "8.16.1") bazel_dep(name = "rules_license", version = "1.0.0") bazel_dep(name = "rules_python", version = "2.2.0") diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index 4ba3a8a3de1c9b..3f9a1f215739ac 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -56,17 +56,17 @@ def workspace(): tf_http_archive( name = "bazel_features", - sha256 = "4fd9922d464686820ffd8fcefa28ccffa147f7cdc6b6ac0d8b07fde565c65d66", - strip_prefix = "bazel_features-1.25.0", - urls = tf_mirror_urls("https://github.com/bazel-contrib/bazel_features/releases/download/v1.25.0/bazel_features-v1.25.0.tar.gz"), + sha256 = "094367e732ece23f334eaf84089a720b861d053beeba6a6a68356d3aee1dc32b", + strip_prefix = "bazel_features-1.50.0", + urls = tf_mirror_urls("https://github.com/bazel-contrib/bazel_features/releases/download/v1.50.0/bazel_features-v1.50.0.tar.gz"), ) tf_http_archive( name = "rules_cc", - sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", - strip_prefix = "rules_cc-0.2.20", + sha256 = "bd7124a844d0403b4b353bcea34d6c8b2ba88dc26881c26c9ee668da89b71846", + strip_prefix = "rules_cc-0.2.25", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.25/rules_cc-0.2.25.tar.gz", ), ) From e6b054e42c57fc7d256d4a7e04f5942854b12699 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Fri, 18 Sep 2026 06:44:35 -0700 Subject: [PATCH 20/32] Add missing `//third_party/absl/strings:string_view` dependency to `dot_algorithms_test`. `dot_algorithms_test.cc` directly includes `absl/strings/string_view.h`, which DWYU (`bant`) requires as an explicit dependency now that `xla_test` targets are checked. PiperOrigin-RevId: 983835500 --- third_party/xla/xla/backends/gpu/codegen/triton/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD index 35be416a6d9edc..0449cd64363a0b 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD @@ -457,6 +457,7 @@ xla_test( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", From 82c8d8b323045ee35391132df3eb7a6683d06819 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Fri, 18 Sep 2026 06:46:11 -0700 Subject: [PATCH 21/32] Let `emitters::KernelArguments` be created from a custom slice provider and report argument positions. Two changes that let out-of-tree emitters (in particular XLA:GPU native custom-call handlers) build `KernelArguments` correctly instead of assembling `KernelArgument`s by hand: 1. A `Create` overload that takes a `SliceProvider` callback instead of a `BufferAssignment`. `ThunkEmitter` resolves buffers through `GetAllocationSlice`, which first consults per-instruction allocation overrides and only then falls back to the buffer assignment. Emitters with such overrides cannot use the `BufferAssignment` overloads without silently dropping the override and launching the kernel on the wrong buffers. 2. `OperandIndex` and `ResultIndex`, which report where an operand or a result subshape ended up in `args()`. Those positions are the index space of `stream_executor::KernelArgsPackingSpec` relocations, so callers previously had to hardcode them and keep them in sync with the argument layout by hand. This is particularly error-prone for the interleaved overload. Both return `FailedPrecondition` when the arguments were assembled by hand rather than by a `Create` factory. All overloads now share a single private `CreateInternal` implementation, including the one that interleaves results with operands. That overload used to duplicate the construction logic, and had already drifted from it: it ignored the `xla.no_invariant_operands` frontend attribute, so an operand the producer marked as non-invariant was passed to the kernel as invariant anyway. It now also rejects interleaving indices that are out of order or that do not cover every result, rather than reporting a confusing count mismatch afterwards. PiperOrigin-RevId: 983836076 --- third_party/xla/xla/codegen/emitters/BUILD | 3 +- .../xla/codegen/emitters/kernel_arguments.cc | 224 +++++++++------ .../xla/codegen/emitters/kernel_arguments.h | 70 +++++ .../codegen/emitters/kernel_arguments_test.cc | 256 ++++++++++++++++++ 4 files changed, 476 insertions(+), 77 deletions(-) diff --git a/third_party/xla/xla/codegen/emitters/BUILD b/third_party/xla/xla/codegen/emitters/BUILD index 50e9a0d16946bb..b3e3a248ebc56e 100644 --- a/third_party/xla/xla/codegen/emitters/BUILD +++ b/third_party/xla/xla/codegen/emitters/BUILD @@ -234,12 +234,13 @@ cc_library( "//xla/service:shaped_slice", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", - "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc index 03ec8c97f81517..41d946c8b7b25e 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc @@ -16,17 +16,16 @@ limitations under the License. #include #include -#include #include #include #include -#include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "llvm/ADT/STLExtras.h" @@ -125,13 +124,15 @@ void FillKernelArgumentAttributes( struct OutputArguments { std::vector output_arguments; + // Shape index of each entry of `output_arguments`, in the same order. + std::vector output_shape_indices; absl::flat_hash_set buffers_written; }; // Extract output arguments from an instruction's shape and return both // the arguments and the set of written buffer slices absl::StatusOr ExtractOutputArguments( - const BufferAssignment& buffer_assignment, + KernelArguments::SliceProvider slice_provider, const HloInstruction* hlo_instruction) { OutputArguments result; ABSL_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( @@ -139,56 +140,144 @@ absl::StatusOr ExtractOutputArguments( [&](const Shape& subshape, const ShapeIndex& index) { if (!subshape.IsArray()) return absl::OkStatus(); - ABSL_ASSIGN_OR_RETURN( - BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(hlo_instruction, index)); + ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, + slice_provider(*hlo_instruction, index)); result.output_arguments.emplace_back(KernelArgument(subshape, slice)); + result.output_shape_indices.push_back(index); result.buffers_written.insert(slice); return absl::OkStatus(); })); return result; } -absl::StatusOr CreateKernelArguments( - const BufferAssignment& buffer_assignment, - const KernelArguments::BufferAlignment& buffer_alignment, + +// Validates `interleaved_output_indices` and returns the position that each +// result must take in the argument list, in result order. +// +// An empty `interleaved_output_indices` means "no interleaving": the results +// follow the operands. +absl::StatusOr> ResultPositions( + absl::Span interleaved_output_indices, size_t num_operands, + size_t num_results) { + std::vector result_positions; + result_positions.reserve(num_results); + + if (interleaved_output_indices.empty()) { + for (size_t i = 0; i < num_results; ++i) { + result_positions.push_back(num_operands + i); + } + return result_positions; + } + + if (interleaved_output_indices.size() != num_results) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected an interleaving index for each of the ", num_results, + " results, got ", interleaved_output_indices.size(), " indices")); + } + + const size_t num_arguments = num_operands + num_results; + int64_t previous_index = -1; + for (const int32_t index : interleaved_output_indices) { + if (index < 0 || static_cast(index) >= num_arguments) { + return absl::InvalidArgumentError("Output index out of bounds"); + } + if (index <= previous_index) { + return absl::InvalidArgumentError( + absl::StrCat("Output indices must be strictly increasing, got ", + index, " after ", previous_index)); + } + previous_index = index; + result_positions.push_back(index); + } + return result_positions; +} + +} // namespace + +absl::StatusOr KernelArguments::CreateInternal( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, - absl::Span unmanaged_arguments) { - std::vector kernel_arguments; + absl::Span unmanaged_arguments, + absl::Span interleaved_output_indices) { absl::flat_hash_set no_invariant_operands = NonInvariantOperands(*hlo_instruction); + std::vector operand_arguments; + operand_arguments.reserve(hlo_instruction->operand_count()); for (auto [op_idx, operand] : llvm::enumerate(hlo_instruction->operands())) { ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(operand, {})); + slice_provider(*operand, {})); KernelArgument arg(operand->shape(), slice); if (no_invariant_operands.contains(op_idx)) { arg.set_invariant(false); } - kernel_arguments.emplace_back(std::move(arg)); + operand_arguments.push_back(std::move(arg)); } ABSL_ASSIGN_OR_RETURN(OutputArguments output_result, - ExtractOutputArguments(buffer_assignment, hlo_instruction)); + ExtractOutputArguments(slice_provider, hlo_instruction)); + + ABSL_ASSIGN_OR_RETURN( + std::vector result_positions, + ResultPositions(interleaved_output_indices, operand_arguments.size(), + output_result.output_arguments.size())); + + const size_t num_managed_arguments = + operand_arguments.size() + output_result.output_arguments.size(); + + std::vector kernel_arguments; + kernel_arguments.reserve(num_managed_arguments + unmanaged_arguments.size()); + ArgumentPositions positions; + positions.operands.reserve(operand_arguments.size()); + positions.results.reserve(output_result.output_arguments.size()); + + // `result_positions` is strictly increasing and holds exactly one position + // per result, so the remaining positions hold exactly the operands, in + // operand order. + size_t next_operand = 0; + size_t next_result = 0; + for (size_t position = 0; position < num_managed_arguments; ++position) { + if (next_result < result_positions.size() && + result_positions[next_result] == static_cast(position)) { + positions.results.emplace_back( + output_result.output_shape_indices[next_result], position); + kernel_arguments.push_back( + std::move(output_result.output_arguments[next_result])); + ++next_result; + continue; + } + positions.operands.push_back(position); + kernel_arguments.push_back(std::move(operand_arguments[next_operand])); + ++next_operand; + } - absl::c_move(output_result.output_arguments, - std::back_inserter(kernel_arguments)); for (const Shape& unmanaged_argument : unmanaged_arguments) { kernel_arguments.emplace_back(unmanaged_argument); } FillKernelArgumentAttributes(kernel_arguments, buffer_alignment, output_result.buffers_written); - return KernelArguments(std::move(kernel_arguments)); + return KernelArguments(std::move(kernel_arguments), std::move(positions)); } -} // namespace +absl::StatusOr KernelArguments::Create( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments) { + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, + unmanaged_arguments, + /*interleaved_output_indices=*/{}); +} absl::StatusOr KernelArguments::Create( const BufferAssignment& buffer_assignment, const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction) { - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, {}); + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, {}, + /*interleaved_output_indices=*/{}); } absl::StatusOr KernelArguments::Create( @@ -196,8 +285,13 @@ absl::StatusOr KernelArguments::Create( const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, absl::Span unmanaged_arguments) { - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, unmanaged_arguments); + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, + unmanaged_arguments, + /*interleaved_output_indices=*/{}); } absl::StatusOr KernelArguments::Create( @@ -205,65 +299,43 @@ absl::StatusOr KernelArguments::Create( const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, absl::Span interleaved_output_indices) { - if (interleaved_output_indices.empty()) { - // Fall back to regular Create method when no interleaving is requested - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, {}); - } - - const auto& operands = hlo_instruction->operands(); - - ABSL_ASSIGN_OR_RETURN(OutputArguments output_result, - ExtractOutputArguments(buffer_assignment, hlo_instruction)); - auto& [output_arguments, buffers_written] = output_result; + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, {}, + interleaved_output_indices); +} - // Check bounds: all output indices must be valid positions - size_t total_positions = operands.size() + output_arguments.size(); - for (int32_t idx : interleaved_output_indices) { - if (idx < 0 || static_cast(idx) >= total_positions) { - return absl::InvalidArgumentError("Output index out of bounds"); - } +absl::StatusOr KernelArguments::OperandIndex( + int64_t operand_index) const { + if (!positions_.has_value()) { + return absl::FailedPreconditionError( + "KernelArguments were not created from an HLO instruction, so operand " + "positions are unknown"); } - - std::vector kernel_arguments; - kernel_arguments.reserve(total_positions); - - // Interleave the inputs and outputs according to the indices - size_t arg_idx = 0; - size_t output_pos = 0; - - for (size_t i = 0; i < total_positions; ++i) { - if (output_pos < interleaved_output_indices.size() && - interleaved_output_indices[output_pos] == static_cast(i)) { - // Place output at this position - if (output_pos >= output_arguments.size()) { - return absl::InvalidArgumentError("Invalid output position index"); - } - kernel_arguments.emplace_back(output_arguments[output_pos]); - ++output_pos; - } else { - // Place input at this position - if (arg_idx >= operands.size()) { - return absl::InvalidArgumentError( - "Not enough inputs for remaining positions"); - } - ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(operands[arg_idx], {})); - kernel_arguments.emplace_back( - KernelArgument(operands[arg_idx]->shape(), slice)); - ++arg_idx; - } + if (operand_index < 0 || + operand_index >= static_cast(positions_->operands.size())) { + return absl::OutOfRangeError( + absl::StrCat("No such operand: ", operand_index)); } + return positions_->operands[operand_index]; +} - // Verify we used all inputs and outputs - if (arg_idx != operands.size() || output_pos != output_arguments.size()) { - return absl::InvalidArgumentError("Did not use all inputs/outputs"); +absl::StatusOr KernelArguments::ResultIndex( + const ShapeIndex& shape_index) const { + if (!positions_.has_value()) { + return absl::FailedPreconditionError( + "KernelArguments were not created from an HLO instruction, so result " + "positions are unknown"); } - - FillKernelArgumentAttributes(kernel_arguments, buffer_alignment, - buffers_written); - - return KernelArguments(std::move(kernel_arguments)); + for (const auto& [index, position] : positions_->results) { + if (index == shape_index) { + return position; + } + } + return absl::OutOfRangeError( + absl::StrCat("No such result: ", shape_index.ToString())); } } // namespace xla::emitters diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.h b/third_party/xla/xla/codegen/emitters/kernel_arguments.h index 9dc17aabeb2310..f8b131160854e2 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.h +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.h @@ -16,15 +16,18 @@ limitations under the License. #define XLA_CODEGEN_EMITTERS_KERNEL_ARGUMENTS_H_ #include +#include #include #include +#include "absl/functional/function_ref.h" #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/service/buffer_assignment.h" #include "xla/service/shaped_slice.h" #include "xla/shape.h" +#include "xla/shape_util.h" namespace xla::emitters { @@ -101,6 +104,28 @@ class KernelArguments { int64_t constant_buffer_align_bytes; }; + // Resolves the buffer slice backing the subshape of `instruction` at + // `index`. + // + // Emitters that can redirect an instruction's buffers (for example + // `ThunkEmitter`, which lets specialized emitters override an allocation) + // must pass their own resolver instead of handing over a `BufferAssignment`, + // otherwise those overrides are silently ignored and the kernel would be + // launched on the wrong buffers. + using SliceProvider = + absl::FunctionRef( + const HloInstruction& instruction, const ShapeIndex& index)>; + + // Creates a KernelArguments object for the given HLO instruction, resolving + // buffer slices through `slice_provider`. Argument order is the same as for + // the `BufferAssignment` overloads below: operands in operand order, then + // the array leaves of the result shape in shape-index order, then + // `unmanaged_arguments`. + static absl::StatusOr Create( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments = {}); + // Creates a KernelArguments object for the given HLO instruction. // The unmanaged_arguments are added to the end of the list of input/output // arguments. @@ -122,6 +147,10 @@ class KernelArguments { // Example: If hlo_instruction->operands() has 3 elements and hlo_instruction // shape yields 2 output arguments, and interleaved_output_indices = {1, 4}: // - Final argument order will be: input0, output0, input1, input2, output1 + // + // `interleaved_output_indices` must be strictly increasing and must name a + // position for every output argument; an empty span requests no + // interleaving. Returns `InvalidArgument` otherwise. static absl::StatusOr Create( const BufferAssignment& buffer_assignment, const BufferAlignment& buffer_alignment, @@ -133,6 +162,25 @@ class KernelArguments { const std::vector& args() const { return args_; } + // Position of the argument that carries operand `operand_index` of the + // instruction these arguments were created for. + // + // Use this instead of hardcoding argument positions when building a + // `stream_executor::KernelArgsPackingSpec`, whose relocation indices refer + // to exactly this list. + // + // Returns `FailedPrecondition` if the arguments were not built by one of the + // `Create` factories, and `OutOfRange` if there is no such operand. + absl::StatusOr OperandIndex(int64_t operand_index) const; + + // Position of the argument that carries the result subshape at + // `shape_index`. For an instruction with a non-tuple result, that is + // `ResultIndex({})`; for a one-element tuple result, `ResultIndex({0})`. + // + // Returns `FailedPrecondition` if the arguments were not built by one of the + // `Create` factories, and `OutOfRange` if there is no such result. + absl::StatusOr ResultIndex(const ShapeIndex& shape_index) const; + std::vector GetArgumentShapedSlices() const { std::vector arg_slices; arg_slices.reserve(args_.size()); @@ -179,7 +227,29 @@ class KernelArguments { } private: + // Where each operand and each result of the originating instruction ended up + // in `args_`. Only populated by the `Create` factories. + struct ArgumentPositions { + // Position in `args_` of each operand, indexed by operand number. + std::vector operands; + // Position in `args_` of each result array leaf, paired with its shape + // index. Result shapes have few leaves, so a linear scan is fine. + std::vector> results; + }; + + KernelArguments(std::vector args, ArgumentPositions positions) + : args_(std::move(args)), positions_(std::move(positions)) {} + + // Shared implementation of all `Create` overloads. See the interleaving + // `Create` overload above for the meaning of `interleaved_output_indices`. + static absl::StatusOr CreateInternal( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments, + absl::Span interleaved_output_indices); + std::vector args_; + std::optional positions_; }; } // namespace xla::emitters diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc index bc9ee1d0ae6a38..423fe75fd1bebc 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc @@ -43,6 +43,7 @@ limitations under the License. namespace xla::emitters { namespace { +using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::HasSubstr; @@ -417,5 +418,260 @@ TEST_F(KernelArgumentsTest, ElementsAre(false, true, false, true, true, true)); } +TEST_F(KernelArgumentsTest, + NoInvariantFrontendAnnotationPropagatesToInterleavedArguments) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + e { + a = s4[5] parameter(0) + b = s4[6] parameter(1) + c = s4[7] parameter(2) + t = tuple(c, b, a), + frontend_attributes={xla.no_invariant_operands="0,2"} + } + )")); + AliasInfo alias_info; + BufferAssigner::Options opts; + opts.allocate_buffers_for_constants = true; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + std::move(opts))); + + // Argument order: operand0, result0, operand1, result1, operand2, result2. + const std::vector interleaved_output_indices = {1, 3, 5}; + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + module->entry_computation()->root_instruction(), + interleaved_output_indices)); + + std::vector invariant_flags; + invariant_flags.reserve(kernel_arguments.args().size()); + for (const KernelArgument& arg : kernel_arguments.args()) { + invariant_flags.push_back(arg.invariant()); + } + // Operands 0 and 2 are annotated as non-invariant, wherever they end up. + EXPECT_THAT(invariant_flags, + ElementsAre(false, true, true, true, false, true)); +} + +TEST_F(KernelArgumentsTest, + InterleavedOutputIndicesMustCoverAllOutputsInOrder) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + e { + a = s4[5] parameter(0) + b = s4[6] parameter(1) + c = s4[7] parameter(2) + t = tuple(c, b, a) + } + )")); + AliasInfo alias_info; + BufferAssigner::Options opts; + opts.allocate_buffers_for_constants = true; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + std::move(opts))); + const HloInstruction* root = module->entry_computation()->root_instruction(); + + const std::vector too_few_indices = {1, 3}; + EXPECT_THAT( + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, too_few_indices), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Expected an interleaving index for each of the 3 " + "results"))); + + const std::vector unordered_indices = {1, 5, 3}; + EXPECT_THAT( + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, unordered_indices), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must be strictly increasing"))); +} + +TEST_F(KernelArgumentsTest, SliceProviderCanOverrideBufferAssignment) { + constexpr absl::string_view kHloString = R"( + HloModule module + + ENTRY entry { + param.0 = f32[1,2,3]{2,1,0} parameter(0) + param.1 = f32[1,2,3]{2,1,0} parameter(1) + ROOT add = f32[1,2,3]{2,1,0} add(param.0, param.1) + } + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloString)); + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + const HloInstruction* root = module->entry_computation()->root_instruction(); + const HloInstruction* param0 = root->operand(0); + ASSERT_OK_AND_ASSIGN(BufferAllocation::Slice param1_slice, + assignment->GetUniqueSlice(root->operand(1), {})); + + // Redirect operand 0 onto operand 1's buffer, the way an emitter's + // allocation overrides do. A `BufferAssignment` alone cannot express this. + auto slice_provider = + [&](const HloInstruction& instruction, + const ShapeIndex& index) -> absl::StatusOr { + if (&instruction == param0) { + return param1_slice; + } + return assignment->GetUniqueSlice(&instruction, index); + }; + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(slice_provider, gpu::GetDefaultBufferAlignment(), + root)); + + ASSERT_THAT(kernel_arguments.args(), SizeIs(3)); + EXPECT_EQ(kernel_arguments.args()[0].slice(), param1_slice); + EXPECT_EQ(kernel_arguments.args()[1].slice(), param1_slice); +} + +TEST_F(KernelArgumentsTest, SliceProviderErrorsArePropagated) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY entry { + p = f32[4] parameter(0) + ROOT add = f32[4] add(p, p) + } + )")); + + auto slice_provider = + [](const HloInstruction& instruction, + const ShapeIndex& index) -> absl::StatusOr { + return absl::NotFoundError("no slice for you"); + }; + + EXPECT_THAT( + KernelArguments::Create(slice_provider, gpu::GetDefaultBufferAlignment(), + module->entry_computation()->root_instruction()), + StatusIs(absl::StatusCode::kNotFound, HasSubstr("no slice for you"))); +} + +TEST_F(KernelArgumentsTest, OperandAndResultIndices) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + param1 = f32[20] parameter(1) + ROOT tuple_result = (f32[10], f32[20]) tuple(param0, param1) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root)); + + // Operands come first, then the array leaves of the result shape. + EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); + + EXPECT_THAT(kernel_arguments.OperandIndex(2), + StatusIs(absl::StatusCode::kOutOfRange)); + EXPECT_THAT(kernel_arguments.OperandIndex(-1), + StatusIs(absl::StatusCode::kOutOfRange)); + // The root is a tuple, so there is no array leaf at the top level. + EXPECT_THAT(kernel_arguments.ResultIndex({}), + StatusIs(absl::StatusCode::kOutOfRange)); +} + +TEST_F(KernelArgumentsTest, OperandAndResultIndicesFollowInterleaving) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + param1 = f32[20] parameter(1) + ROOT tuple_result = (f32[10], f32[20]) tuple(param0, param1) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + // Argument order becomes: input0, output0, input1, output1. + std::vector interleaved_indices = {1, 3}; + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, interleaved_indices)); + + EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); +} + +TEST_F(KernelArgumentsTest, ResultIndexOfNonTupleResult) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + ROOT result = f32[10] add(param0, param0) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root)); + + EXPECT_THAT(kernel_arguments.ResultIndex({}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), + StatusIs(absl::StatusCode::kOutOfRange)); +} + +TEST_F(KernelArgumentsTest, HandBuiltArgumentsHaveNoKnownPositions) { + KernelArguments kernel_arguments(std::vector{KernelArgument( + ShapeUtil::MakeShape(F32, {4}), BufferAllocation::Slice())}); + + EXPECT_THAT(kernel_arguments.OperandIndex(0), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_THAT(kernel_arguments.ResultIndex({}), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + } // namespace } // namespace xla::emitters From b7f4e14b57e9b77d456268baa607db9a626c989b Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Fri, 18 Sep 2026 06:47:36 -0700 Subject: [PATCH 22/32] PR #48531: [ROCm ]Enable convolution_1d_test and convolution_1d_gpu_alternative_layout_test on ROCm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48531 📝 Summary of Changes Drop cuda-only from both convolution_1d_test and convolution_1d_gpu_alternative_layout_test to enable them on ROCm CI. 🎯 Justification Add more unit test coverage on ROCm platform. 🚀 Kind of Contribution 🧪 Tests Copybara import of the project: -- c703f27f257495c05af87b7edcb971abadc9626a by Lin Chen1 : Enable convolution_1d_test and convolution_1d_gpu_alternative_layout_test on ROCm. Signed-off-by: Lin Chen1 Merging this change closes #48531 PiperOrigin-RevId: 983836641 --- third_party/xla/xla/tests/BUILD | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 717d98c173580c..d0e406542294a3 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -1522,10 +1522,7 @@ xla_test( # Turn on logging so that VLOG statements don't appear uncovered to zapfhahn. args = if_google(["--vmodule=convolution_emitter=7"]), shard_count = 40, - tags = [ - "cuda-only", - "optonly", - ], + tags = ["optonly"], deps = [ ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", @@ -1649,7 +1646,6 @@ xla_test( backend_args = {"gpu": ["--xla_backend_extra_options=xla_gpu_experimental_conv_disable_layout_heuristic"]}, backends = ["gpu"], shard_count = 25, - tags = ["cuda-only"], deps = [ ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", From 9093975301d5593f36fcb571d63b78943196ee6d Mon Sep 17 00:00:00 2001 From: Ashiq Imran Date: Fri, 18 Sep 2026 07:10:32 -0700 Subject: [PATCH 23/32] PR #48416: [XLA:GPU][Bugfix][ONEAPI] Pjrt c api test fix Imported from GitHub PR https://github.com/openxla/xla/pull/48416 This PR temporarily updates the PJRT GPU plugin name to SYCL to pass unit tests, we are reverting back to naming the PLUGIN NAME as SYCL to be consistent with the stream executor name of SYCL. Unit Test Fixes: //xla/pjrt/c/c:pjrt_c_api_gpu_test_intelgpu_any Copybara import of the project: -- a07c441910f7fdb60c7b4b6b994b8b9162147a66 by Ashiq Imran : fix for PJRT C API -- d51264f8e3d72ef99d48ae8017a4e806f3a9c1c5 by Ashiq Imran : minor changes Merging this change closes #48416 PiperOrigin-RevId: 983845645 --- third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 a09787ab3e5790..39133603faafb7 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 @@ -84,7 +84,9 @@ namespace gpu_plugin { #if TENSORFLOW_USE_ROCM #define PJRT_GPU_PLUGIN_PLATFORM_NAME "ROCM" #elif TENSORFLOW_USE_SYCL -#define PJRT_GPU_PLUGIN_PLATFORM_NAME "ONEAPI" +// TODO(Intel-tf) this will be changed to ONEAPI +// when the SYCL backend has been renamed to ONEAPI. +#define PJRT_GPU_PLUGIN_PLATFORM_NAME "SYCL" #else #define PJRT_GPU_PLUGIN_PLATFORM_NAME "CUDA" #endif @@ -337,7 +339,7 @@ PJRT_Error* PJRT_GpuDeviceTopology_Create( if (plugin_platform == "ROCM") { platform_id = xla::RocmId(); platform_name = xla::RocmName(); - } else if (plugin_platform == "ONEAPI") { + } else if (plugin_platform == "SYCL") { platform_id = xla::OneapiId(); platform_name = xla::OneapiName(); } else { From 754a7d25b68454fd9f9920fe637bc7cc5f2c887b Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Fri, 18 Sep 2026 07:26:15 -0700 Subject: [PATCH 24/32] Support CUDA 12 and CUDA 13 installation extras with dynamic DSO fallback: 1. Support CUDA 12 and CUDA 13 pip extras in setup.py (tensorflow[cuda12], tensorflow[cuda13]). 2. Dynamically resolve and fallback to alternative CUDA component DSO versions at runtime. 3. Configure Kokoro release jobs for CUDA 13 wheels. PiperOrigin-RevId: 983851549 --- tensorflow/tools/pip_package/setup.py.tpl | 137 +++++++++++++----- .../py/setup_py_nvidia_dependencies_util.py | 42 ++++-- .../xla/tsl/platform/default/dso_loader.cc | 56 +++++-- 3 files changed, 170 insertions(+), 65 deletions(-) diff --git a/tensorflow/tools/pip_package/setup.py.tpl b/tensorflow/tools/pip_package/setup.py.tpl index c506c29134c6bf..5bfb14eabe498a 100644 --- a/tensorflow/tools/pip_package/setup.py.tpl +++ b/tensorflow/tools/pip_package/setup.py.tpl @@ -58,19 +58,44 @@ cuda_version = 0 # placeholder cuda_major_version = '12' # placeholder cuda_wheel_suffix = '' # placeholder -nvidia_cublas_version = '' # placeholder -nvidia_cuda_cupti_version = '' # placeholder -nvidia_cuda_nvcc_version = '' # placeholder -nvidia_cuda_runtime_version = '' # placeholder -nvidia_cudnn_version = '' # placeholder -nvidia_cufft_version = '' # placeholder -nvidia_cusolver_version = '' # placeholder -nvidia_cusparse_version = '' # placeholder -nvidia_nccl_version = '' # placeholder -nvidia_nvjitlink_version = '' # placeholder -nvidia_cuda_nvrtc_version = '' # placeholder -nvidia_curand_version = '' # placeholder -nvidia_nvshmem_version = '' # placeholder +cuda12_nvidia_cublas_version = '' # placeholder +cuda12_nvidia_cuda_cupti_version = '' # placeholder +cuda12_nvidia_cuda_nvcc_version = '' # placeholder +cuda12_nvidia_cuda_runtime_version = '' # placeholder +cuda12_nvidia_cudnn_version = '' # placeholder +cuda12_nvidia_cufft_version = '' # placeholder +cuda12_nvidia_cusolver_version = '' # placeholder +cuda12_nvidia_cusparse_version = '' # placeholder +cuda12_nvidia_nccl_version = '' # placeholder +cuda12_nvidia_nvjitlink_version = '' # placeholder +cuda12_nvidia_cuda_nvrtc_version = '' # placeholder +cuda12_nvidia_curand_version = '' # placeholder +cuda12_nvidia_nvshmem_version = '' # placeholder + +cuda13_nvidia_cublas_version = '' # placeholder +cuda13_nvidia_cuda_cupti_version = '' # placeholder +cuda13_nvidia_cuda_nvcc_version = '' # placeholder +cuda13_nvidia_cuda_runtime_version = '' # placeholder +cuda13_nvidia_cudnn_version = '' # placeholder +cuda13_nvidia_cufft_version = '' # placeholder +cuda13_nvidia_cusolver_version = '' # placeholder +cuda13_nvidia_cusparse_version = '' # placeholder +cuda13_nvidia_nccl_version = '' # placeholder +cuda13_nvidia_nvjitlink_version = '' # placeholder +cuda13_nvidia_cuda_nvrtc_version = '' # placeholder +cuda13_nvidia_curand_version = '' # placeholder +cuda13_nvidia_nvshmem_version = '' # placeholder +cuda13_nvidia_nvvm_version = '' # placeholder +cuda13_nvidia_cuda_crt_version = '' # placeholder + +if str(cuda_major_version) == '13': + nvidia_cuda_nvcc_version = cuda13_nvidia_cuda_nvcc_version + nvidia_cuda_nvrtc_version = cuda13_nvidia_cuda_nvrtc_version + nvidia_cuda_runtime_version = cuda13_nvidia_cuda_runtime_version +else: + nvidia_cuda_nvcc_version = cuda12_nvidia_cuda_nvcc_version + nvidia_cuda_nvrtc_version = cuda12_nvidia_cuda_nvrtc_version + nvidia_cuda_runtime_version = cuda12_nvidia_cuda_runtime_version # We use the same setup.py for all tensorflow_* packages and for the nightly # equivalents (tf_nightly_*). The package is controlled from the argument line @@ -156,33 +181,67 @@ if collaborator_build: + ';platform_system=="Windows"', ] +NVIDIA_CUDA12_PACKAGES = [ + f'nvidia-cublas-cu12{cuda12_nvidia_cublas_version}', + f'nvidia-cuda-cupti-cu12{cuda12_nvidia_cuda_cupti_version}', + f'nvidia-cuda-nvcc-cu12{cuda12_nvidia_cuda_nvcc_version}', + f'nvidia-cuda-nvrtc-cu12{cuda12_nvidia_cuda_nvrtc_version}', + f'nvidia-cuda-runtime-cu12{cuda12_nvidia_cuda_runtime_version}', + f'nvidia-cudnn-cu12{cuda12_nvidia_cudnn_version}', + f'nvidia-cufft-cu12{cuda12_nvidia_cufft_version}', + f'nvidia-curand-cu12{cuda12_nvidia_curand_version}', + f'nvidia-cusolver-cu12{cuda12_nvidia_cusolver_version}', + f'nvidia-cusparse-cu12{cuda12_nvidia_cusparse_version}', + f'nvidia-nccl-cu12{cuda12_nvidia_nccl_version}', + f'nvidia-nvjitlink-cu12{cuda12_nvidia_nvjitlink_version}', + f'nvidia-nvshmem-cu12{cuda12_nvidia_nvshmem_version}', +] + +NVIDIA_CUDA13_PACKAGES = [ + f'nvidia-cublas{cuda13_nvidia_cublas_version}', + f'nvidia-cuda-cupti{cuda13_nvidia_cuda_cupti_version}', + f'nvidia-cuda-nvcc{cuda13_nvidia_cuda_nvcc_version}', + f'nvidia-cuda-nvrtc{cuda13_nvidia_cuda_nvrtc_version}', + f'nvidia-cuda-runtime{cuda13_nvidia_cuda_runtime_version}', + f'nvidia-cudnn-cu13{cuda13_nvidia_cudnn_version}', + f'nvidia-cufft{cuda13_nvidia_cufft_version}', + f'nvidia-curand{cuda13_nvidia_curand_version}', + f'nvidia-cusolver{cuda13_nvidia_cusolver_version}', + f'nvidia-cusparse{cuda13_nvidia_cusparse_version}', + f'nvidia-nccl-cu13{cuda13_nvidia_nccl_version}', + f'nvidia-nvjitlink{cuda13_nvidia_nvjitlink_version}', + f'nvidia-nvshmem-cu13{cuda13_nvidia_nvshmem_version}', + f'nvidia-nvvm{cuda13_nvidia_nvvm_version}', + f'nvidia-cuda-crt{cuda13_nvidia_cuda_crt_version}', +] + # Set up extra packages, which are optional sets of other Python package deps. -# E.g. "pip install tensorflow[and-cuda]" below installs the normal TF deps, -# plus the CUDA libraries listed. -EXTRA_PACKAGES = { - 'and-cuda': [ - # TODO(nluehr): set nvidia-* versions based on build components. - f'nvidia-cublas{cuda_wheel_suffix}{nvidia_cublas_version}', - f'nvidia-cuda-cupti{cuda_wheel_suffix}{nvidia_cuda_cupti_version}', - f'nvidia-cuda-nvcc{cuda_wheel_suffix}{nvidia_cuda_nvcc_version}', - f'nvidia-cuda-nvrtc{cuda_wheel_suffix}{nvidia_cuda_nvrtc_version}', - f'nvidia-cuda-runtime{cuda_wheel_suffix}{nvidia_cuda_runtime_version}', - f'nvidia-cudnn-cu13{nvidia_cudnn_version}' if cuda_major_version == '13' else f'nvidia-cudnn-cu12{nvidia_cudnn_version}', - f'nvidia-cufft{cuda_wheel_suffix}{nvidia_cufft_version}', - f'nvidia-curand{cuda_wheel_suffix}{nvidia_curand_version}', - f'nvidia-cusolver{cuda_wheel_suffix}{nvidia_cusolver_version}', - f'nvidia-cusparse{cuda_wheel_suffix}{nvidia_cusparse_version}', - f'nvidia-nccl-cu13{nvidia_nccl_version}' if cuda_major_version == '13' else f'nvidia-nccl-cu12{nvidia_nccl_version}', - f'nvidia-nvjitlink{cuda_wheel_suffix}{nvidia_nvjitlink_version}', - f'nvidia-nvshmem-cu13{nvidia_nvshmem_version}' if cuda_major_version == '13' else f'nvidia-nvshmem-cu12{nvidia_nvshmem_version}', - ], - 'gcs-filesystem': [ - ('tensorflow-io-gcs-filesystem>=0.23.1; ' - 'sys_platform!="win32" and python_version<"3.13"'), - ('tensorflow-io-gcs-filesystem>=0.23.1; ' - 'sys_platform=="win32" and python_version<"3.12"'), - ] -} +# E.g. "pip install tensorflow[and-cuda]" or "pip install tensorflow[cuda12]" +if str(cuda_major_version) == '13': + EXTRA_PACKAGES = { + 'and-cuda': NVIDIA_CUDA13_PACKAGES, + 'cuda': NVIDIA_CUDA13_PACKAGES, + 'cuda13': NVIDIA_CUDA13_PACKAGES, + 'cuda13-local': [], + } +else: + EXTRA_PACKAGES = { + 'and-cuda': NVIDIA_CUDA12_PACKAGES, + 'cuda': NVIDIA_CUDA12_PACKAGES, + 'cuda12': NVIDIA_CUDA12_PACKAGES, + 'cuda12-local': [], + } + +EXTRA_PACKAGES['gcs-filesystem'] = [ + ( + 'tensorflow-io-gcs-filesystem>=0.23.1; ' + 'sys_platform!="win32" and python_version<"3.13"' + ), + ( + 'tensorflow-io-gcs-filesystem>=0.23.1; ' + 'sys_platform=="win32" and python_version<"3.12"' + ), +] DOCLINES = __doc__.split('\n') diff --git a/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py b/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py index 904cdcf4ecae21..b0d135a34a7e12 100644 --- a/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py +++ b/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py @@ -48,35 +48,55 @@ def get_setup_py_content_with_nvidia_wheel_versions( ) -> str: nvidia_wheel_versions = {"12": {}, "13": {}} for line in nvidia_wheel_versions_data.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue match = NVIDIA_WHEEL_VERSIONS_PATTERN.match(line) - if match: + if match and match.group(1).startswith("nvidia"): wheel_name = match.group(1).replace("-", "_") for suffix, version in {"_cu12": "12", "_cu13": "13", "": "13"}.items(): if not wheel_name.endswith(suffix): continue - wheel_name = wheel_name.replace(suffix, "") + "_version" + if suffix: + wheel_name = wheel_name[:-len(suffix)] + wheel_name = wheel_name + "_version" nvidia_wheel_versions[version][wheel_name] = match.group(2).strip() break - if not cuda_version: - cuda_version = "12" + cuda_version_str = str(cuda_version) if cuda_version else "12" setup_py_content = setup_py_content.replace( - "cuda_version = 0 # placeholder", f"cuda_version = {cuda_version}" + "cuda_version = 0 # placeholder", f"cuda_version = {cuda_version_str}" ) setup_py_content = setup_py_content.replace( "cuda_wheel_suffix = '' # placeholder", - "cuda_wheel_suffix = '-cu12'" if cuda_version == "12" else "cuda_wheel_suffix = ''", + ( + "cuda_wheel_suffix = '-cu12'" + if cuda_version_str == "12" + else "cuda_wheel_suffix = ''" + ), ) setup_py_content = setup_py_content.replace( "cuda_major_version = '12' # placeholder", - f"cuda_major_version = '{cuda_version}'", + f"cuda_major_version = '{cuda_version_str}'", ) - for version_name, version_value in nvidia_wheel_versions[ - str(cuda_version) - ].items(): + for version_name, version_value in nvidia_wheel_versions.get("12", {}).items(): setup_py_content = setup_py_content.replace( - f"{version_name} = '' # placeholder", + f"cuda12_{version_name} = '' # placeholder", + f"cuda12_{version_name} = '{version_value}'", + ) + for version_name, version_value in nvidia_wheel_versions.get("13", {}).items(): + setup_py_content = setup_py_content.replace( + f"cuda13_{version_name} = '' # placeholder", + f"cuda13_{version_name} = '{version_value}'", + ) + for version_name, version_value in nvidia_wheel_versions.get( + cuda_version_str, {} + ).items(): + setup_py_content = re.sub( + rf"^{re.escape(version_name)} = '' # placeholder$", f"{version_name} = '{version_value}'", + setup_py_content, + flags=re.MULTILINE, ) return setup_py_content diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.cc b/third_party/xla/xla/tsl/platform/default/dso_loader.cc index e79365fd562c67..ed9ecb415e5322 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.cc +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.cc @@ -73,6 +73,36 @@ absl::StatusOr GetDsoHandle(const std::string& name, VLOG(1) << message; return absl::Status(absl::StatusCode::kFailedPrecondition, message); } +absl::StatusOr GetCudaDsoHandle(const std::string& name, + absl::string_view version) { + auto handle = GetDsoHandle(name, version); + if (handle.ok()) return handle; + auto primary_status = handle.status(); + + // Determine alternative major version across CUDA 12 and CUDA 13 ecosystems. + // Note: cuDNN does not support cross-major loading (cuDNN 8 lacks v9 entry + // points). + absl::string_view alt_version = ""; + if (name == "cufft" || name == "cusolver") { + alt_version = (version == "12") ? "11" : (version == "11" ? "12" : ""); + } else if (name == "cublas" || name == "cublasLt" || name == "cudart" || + name == "nvrtc" || name == "cusparse" || name == "cupti") { + alt_version = (version == "13") ? "12" : (version == "12" ? "13" : ""); + } + + if (!alt_version.empty()) { + handle = GetDsoHandle(name, alt_version); + if (handle.ok()) return handle; + } + + // Only CUPTI supports unversioned library fallback. + if (name == "cupti") { + handle = GetDsoHandle(name, ""); + if (handle.ok()) return handle; + } + + return primary_status; +} } // namespace namespace DsoLoader { @@ -95,51 +125,47 @@ absl::StatusOr GetNvmlDsoHandle() { } absl::StatusOr GetNvrtcDsoHandle() { - return GetDsoHandle("nvrtc", GetCudaRtVersion()); + return GetCudaDsoHandle("nvrtc", GetCudaRtVersion()); } absl::StatusOr GetCudaRuntimeDsoHandle() { - return GetDsoHandle("cudart", GetCudaRtVersion()); + return GetCudaDsoHandle("cudart", GetCudaRtVersion()); } absl::StatusOr GetCublasDsoHandle() { - return GetDsoHandle("cublas", GetCublasVersion()); + return GetCudaDsoHandle("cublas", GetCublasVersion()); } absl::StatusOr GetCublasLtDsoHandle() { - return GetDsoHandle("cublasLt", GetCublasVersion()); + return GetCudaDsoHandle("cublasLt", GetCublasVersion()); } absl::StatusOr GetCufftDsoHandle() { - return GetDsoHandle("cufft", GetCufftVersion()); + return GetCudaDsoHandle("cufft", GetCufftVersion()); } absl::StatusOr GetCusolverDsoHandle() { - return GetDsoHandle("cusolver", GetCusolverVersion()); + return GetCudaDsoHandle("cusolver", GetCusolverVersion()); } absl::StatusOr GetCusparseDsoHandle() { - return GetDsoHandle("cusparse", GetCusparseVersion()); + return GetCudaDsoHandle("cusparse", GetCusparseVersion()); } absl::StatusOr GetCuptiDsoHandle() { - // Load specific version of CUPTI this is built. - auto status_or_handle = GetDsoHandle("cupti", GetCuptiVersion()); - if (status_or_handle.ok()) return status_or_handle; - // Load whatever libcupti.so user specified. - return GetDsoHandle("cupti", ""); + return GetCudaDsoHandle("cupti", GetCuptiVersion()); } absl::StatusOr GetCudnnDsoHandle() { - return GetDsoHandle("cudnn", GetCudnnVersion()); + return GetCudaDsoHandle("cudnn", GetCudnnVersion()); } absl::StatusOr GetNcclDsoHandle() { - return GetDsoHandle("nccl", GetNcclVersion()); + return GetCudaDsoHandle("nccl", GetNcclVersion()); } absl::StatusOr GetNvshmemDsoHandle() { - return GetDsoHandle("nvshmem_host", GetNvshmemVersion()); + return GetCudaDsoHandle("nvshmem_host", GetNvshmemVersion()); } absl::StatusOr GetNvInferDsoHandle() { From 44c9d71ce14876c741b51e1ce7d30497bc6549bb Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Fri, 18 Sep 2026 07:53:31 -0700 Subject: [PATCH 25/32] Refactor CompositeRewriter precondition checks to use xla::Decision. Splits CompositeRewriter::RewriteComputation into modular precondition helper functions returning xla::Decision instead of boolean flags, and extracts instruction rewrite logic into RewriteScaledDot. PiperOrigin-RevId: 983862544 --- .../xla/xla/backends/gpu/transforms/BUILD | 2 + .../gpu/transforms/composite_rewriter.cc | 295 ++++++++++-------- 2 files changed, 169 insertions(+), 128 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index dc11f116a7cd84..f8107004477ca0 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3018,6 +3018,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/service:decision", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", @@ -3025,6 +3026,7 @@ cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@llvm-project//mlir:AsmParser", "@llvm-project//mlir:IR", diff --git a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc index aef59a56f37af0..855fe918a7285e 100644 --- a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "mlir/AsmParser/AsmParser.h" #include "mlir/IR/Attributes.h" @@ -36,6 +37,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" +#include "xla/service/decision.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/util.h" @@ -99,151 +101,188 @@ absl::StatusOr ParseDimensionNumbers( return dnums; } +Decision IsScaledDotComposite(const HloInstruction* instruction) { + if (instruction->opcode() != HloOpcode::kCall) { + return Decision::Forbid("instruction is not a call"); + } + auto call = Cast(instruction); + if (!call->is_composite()) { + return Decision::Forbid( + absl::StrCat(instruction->name(), " is not composite")); + } + if (!call->has_frontend_attributes()) { + return Decision::Forbid("No frontend attributes"); + } + const auto& frontend_attrs = call->frontend_attributes().map(); + auto it = frontend_attrs.find("composite.name"); + if (it == frontend_attrs.end() || it->second != "xla.scaled_dot") { + return Decision::Forbid( + it != frontend_attrs.end() + ? absl::StrCat("composite.name is not xla.scaled_dot: ", it->second) + : "composite.name attribute is missing"); + } + return Decision::Allow(); +} + +Decision HasSupportedDimensionNumbers( + const DotDimensionNumbers& dot_dimension_numbers) { + if (dot_dimension_numbers.lhs_contracting_dimensions_size() != 1 || + dot_dimension_numbers.rhs_contracting_dimensions_size() != 1) { + return Decision::Forbid( + absl::StrCat("Expected 1 contracting dimension on LHS and RHS, got: ", + dot_dimension_numbers.DebugString())); + } + if (dot_dimension_numbers.lhs_batch_dimensions_size() > 1 || + dot_dimension_numbers.rhs_batch_dimensions_size() > 1) { + return Decision::Forbid( + absl::StrCat("Expected at most 1 batch dimension on LHS and RHS, got: ", + dot_dimension_numbers.DebugString())); + } + return Decision::Allow(); +} + +Decision IsValidBf16Scale(const HloInstruction* operand, + const HloInstruction* scale) { + if (scale->shape().dimensions().size() != + operand->shape().dimensions().size()) { + return Decision::Forbid("scale and operand rank mismatch for BF16"); + } + for (int64_t dim : scale->shape().dimensions()) { + if (dim != 1) { + return Decision::Forbid("scale dim != 1 for BF16"); + } + } + if (scale->opcode() != HloOpcode::kConstant) { + return Decision::Forbid("scale is not constant for BF16"); + } + if (!scale->literal().IsAllFloat(1.0)) { + return Decision::Forbid("scale is not 1.0 for BF16"); + } + return Decision::Allow(); +} + +Decision IsSupportedBf16(const HloInstruction* lhs, const HloInstruction* rhs, + const HloInstruction* lhs_scale, + const HloInstruction* rhs_scale) { + if (lhs_scale->shape().element_type() != BF16 || + rhs_scale->shape().element_type() != BF16) { + return Decision::Forbid("BF16 operands require BF16 scales"); + } + if (Decision d = IsValidBf16Scale(lhs, lhs_scale); !d) { + return d; + } + return IsValidBf16Scale(rhs, rhs_scale); +} + +Decision IsSupportedScaleAndOperand(const HloInstruction* operand, + const HloInstruction* scale, + int64_t contracting_dim) { + PrimitiveType op_type = operand->shape().element_type(); + PrimitiveType scale_type = scale->shape().element_type(); + + if (scale_type != F8E8M0FNU) { + return Decision::Forbid(absl::StrCat("Unsupported scale type: ", + PrimitiveType_Name(scale_type))); + } + + if (contracting_dim >= scale->shape().dimensions().size()) { + return Decision::Forbid("contracting_dim out of bounds for scale"); + } + int64_t operand_dim_size = operand->shape().dimensions(contracting_dim); + int64_t scale_dim_size = scale->shape().dimensions(contracting_dim); + + if (scale_dim_size == 0 || operand_dim_size % scale_dim_size != 0) { + return Decision::Forbid("operand_dim_size not divisible by scale_dim_size"); + } + int64_t scale_factor = operand_dim_size / scale_dim_size; + + if (op_type == F8E4M3FN || op_type == F8E5M2 || op_type == F4E2M1FN) { + if (scale_factor % 16 != 0) { + return Decision::Forbid( + absl::StrCat("scale_factor % 16 != 0: ", scale_factor)); + } + return Decision::Allow(); + } + + return Decision::Forbid( + absl::StrCat("Unsupported operand type: ", PrimitiveType_Name(op_type))); +} + +Decision IsSupportedScaledDot( + const HloCallInstruction* call, + const DotDimensionNumbers& dot_dimension_numbers) { + if (call->operand_count() < 4) { + return Decision::Forbid("Scaled dot call has fewer than 4 operands"); + } + const HloInstruction* lhs = call->operand(0); + const HloInstruction* rhs = call->operand(1); + const HloInstruction* lhs_scale = call->operand(2); + const HloInstruction* rhs_scale = call->operand(3); + + if (lhs->shape().element_type() == BF16 && + rhs->shape().element_type() == BF16) { + return IsSupportedBf16(lhs, rhs, lhs_scale, rhs_scale); + } + + int64_t lhs_contracting_dim = + dot_dimension_numbers.lhs_contracting_dimensions(0); + int64_t rhs_contracting_dim = + dot_dimension_numbers.rhs_contracting_dimensions(0); + + if (Decision d = + IsSupportedScaleAndOperand(lhs, lhs_scale, lhs_contracting_dim); + !d) { + return d; + } + return IsSupportedScaleAndOperand(rhs, rhs_scale, rhs_contracting_dim); +} + +absl::Status RewriteScaledDot( + HloComputation* computation, HloCallInstruction* call, + const DotDimensionNumbers& dot_dimension_numbers) { + PrecisionConfig precision{}; + precision.mutable_operand_precision()->Resize(2, PrecisionConfig::DEFAULT); + auto* scaled_dot = + computation->AddInstruction(HloInstruction::CreateScaledDot( + call->shape(), call->mutable_operand(0), call->mutable_operand(1), + call->mutable_operand(2), call->mutable_operand(3), + dot_dimension_numbers, precision)); + ABSL_RETURN_IF_ERROR(call->ReplaceAllUsesWith(scaled_dot)); + ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call)); + return absl::OkStatus(); +} + } // namespace absl::StatusOr CompositeRewriter::RewriteComputation( HloComputation* computation) { bool changed = false; for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { - if (instruction->opcode() != HloOpcode::kCall) { - continue; - } - VLOG(3) << "Found call instruction: " << instruction->name(); - auto call = Cast(instruction); - if (!call->is_composite()) { - VLOG(3) << instruction->name() << " is not composite"; - continue; - } - if (!call->has_frontend_attributes()) { - VLOG(3) << "No frontend attributes"; - continue; - } - auto frontend_attrs = call->frontend_attributes().map(); - auto key = "composite.name"; - if (!frontend_attrs.contains(key) || - frontend_attrs.at(key) != "xla.scaled_dot") { - VLOG(3) << key << " is not xla.scaled_dot: " << frontend_attrs.at(key); + if (Decision d = IsScaledDotComposite(instruction); !d) { + VLOG(3) << d.Explain(); continue; } - if (!frontend_attrs.contains("composite.attributes")) { + auto* call = Cast(instruction); + const auto& frontend_attrs = call->frontend_attributes().map(); + auto it = frontend_attrs.find("composite.attributes"); + if (it == frontend_attrs.end()) { return absl::InvalidArgumentError( "composite.attributes is not set for xla.scaled_dot"); } - ABSL_ASSIGN_OR_RETURN( - DotDimensionNumbers dot_dimension_numbers, - ParseDimensionNumbers(frontend_attrs.at("composite.attributes"))); - - if (dot_dimension_numbers.lhs_contracting_dimensions_size() != 1 || - dot_dimension_numbers.rhs_contracting_dimensions_size() != 1 || - dot_dimension_numbers.lhs_batch_dimensions_size() > 1 || - dot_dimension_numbers.rhs_batch_dimensions_size() > 1) { - LOG(ERROR) << "Unsupported dimension numbers: " - << dot_dimension_numbers.DebugString(); - continue; - } + ABSL_ASSIGN_OR_RETURN(DotDimensionNumbers dot_dimension_numbers, + ParseDimensionNumbers(it->second)); - const HloInstruction* lhs = call->operand(0); - const HloInstruction* rhs = call->operand(1); - const HloInstruction* lhs_scale = call->operand(2); - const HloInstruction* rhs_scale = call->operand(3); - - int64_t lhs_contracting_dim = - dot_dimension_numbers.lhs_contracting_dimensions(0); - int64_t rhs_contracting_dim = - dot_dimension_numbers.rhs_contracting_dimensions(0); - - auto is_supported = [&](const HloInstruction* operand, - const HloInstruction* scale, - int64_t contracting_dim) { - auto op_type = operand->shape().element_type(); - auto scale_type = scale->shape().element_type(); - auto is_fp8 = [](PrimitiveType type) { - return type == F8E4M3FN || type == F8E5M2; - }; - - if (op_type == BF16 && scale_type == BF16) { - if (scale->shape().dimensions().size() != - operand->shape().dimensions().size()) { - VLOG(3) << "scale and operand rank mismatch for BF16"; - return false; - } - for (int64_t dim : scale->shape().dimensions()) { - if (dim != 1) { - VLOG(3) << "scale dim != 1 for BF16"; - return false; - } - } - if (scale->opcode() != HloOpcode::kConstant) { - VLOG(3) << "scale is not constant for BF16"; - return false; - } - bool supported = scale->literal().IsAllFloat(1.0); - if (!supported) { - VLOG(3) << "scale is not 1.0 for BF16"; - } - return supported; - } - - auto is_supported_scale_type = [](PrimitiveType type) { - return type == F8E8M0FNU; - }; - - if (!is_supported_scale_type(scale_type)) { - VLOG(3) << "Unsupported scale type: " << PrimitiveType_Name(scale_type); - return false; - } - - if (contracting_dim >= scale->shape().dimensions().size()) { - VLOG(3) << "contracting_dim out of bounds for scale"; - return false; - } - int64_t operand_dim_size = operand->shape().dimensions(contracting_dim); - int64_t scale_dim_size = scale->shape().dimensions(contracting_dim); - - if (scale_dim_size == 0 || operand_dim_size % scale_dim_size != 0) { - VLOG(3) << "operand_dim_size not divisible by scale_dim_size"; - return false; - } - int64_t scale_factor = operand_dim_size / scale_dim_size; - - if (is_fp8(op_type)) { - bool supported = scale_factor % 16 == 0; - if (!supported) { - VLOG(3) << "scale_factor % 16 != 0: " << scale_factor; - } - return supported; - } - - if (op_type == F4E2M1FN) { - bool supported = scale_factor % 16 == 0; - if (!supported) { - VLOG(3) << "scale_factor % 16 != 0: " << scale_factor; - } - return supported; - } - - VLOG(3) << "Unsupported operand type: " << PrimitiveType_Name(op_type); - return false; - }; - - if (!is_supported(lhs, lhs_scale, lhs_contracting_dim)) { - VLOG(3) << "LHS not supported"; + if (Decision d = HasSupportedDimensionNumbers(dot_dimension_numbers); !d) { + LOG(ERROR) << d.Explain(); continue; } - if (!is_supported(rhs, rhs_scale, rhs_contracting_dim)) { - VLOG(3) << "RHS not supported"; + + if (Decision d = IsSupportedScaledDot(call, dot_dimension_numbers); !d) { + VLOG(3) << "Scaled dot composite operands not supported: " << d.Explain(); continue; } - PrecisionConfig precision{}; - precision.mutable_operand_precision()->Resize(2, PrecisionConfig::DEFAULT); - auto* scaled_dot = - computation->AddInstruction(HloInstruction::CreateScaledDot( - call->shape(), call->mutable_operand(0), call->mutable_operand(1), - call->mutable_operand(2), call->mutable_operand(3), - dot_dimension_numbers, precision)); - ABSL_RETURN_IF_ERROR(call->ReplaceAllUsesWith(scaled_dot)); - ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call)); + ABSL_RETURN_IF_ERROR(RewriteScaledDot(computation, call, dot_dimension_numbers)); changed = true; } return changed; From 93355d867dc3e01bfb3c97fbdd7e779e874271a9 Mon Sep 17 00:00:00 2001 From: Saicharan Ramineni <84414237+GodlyDonuts@users.noreply.github.com> Date: Fri, 18 Sep 2026 08:07:27 -0700 Subject: [PATCH 26/32] PR #48533: Prevent overflow when rounding BFC allocation sizes Imported from GitHub PR https://github.com/openxla/xla/pull/48533 ## Summary of Changes Reject BFC allocation requests that cannot be rounded up to the minimum allocation size without overflowing size_t. Add a focused allocator regression test for the overflow boundary. ## Justification A request near SIZE_MAX currently wraps to zero during rounding. The allocator can then return an undersized chunk and later fail while validating corrupted chunk metadata. This was observed through TensorFlow in tensorflow/tensorflow#108904, but the affected allocator is maintained in OpenXLA. ## Kind of Contribution Bug fix and tests. ## Benchmark Not applicable; this only adds a cold-path bounds check before allocation. ## Unit Tests The new BFCAllocatorTest.RejectsAllocationWhoseRoundedSizeWouldOverflow test uses the fake suballocator and does not reserve a large buffer. Local validation: git diff --check. ## Execution Tests Not applicable; the regression is covered at the allocator boundary. Copybara import of the project: -- 59caf3bfacc5771f65333d5e0a0701519bc54e4f by Sai Ramineni : Prevent overflow when rounding BFC allocation sizes Merging this change closes #48533 PiperOrigin-RevId: 983868621 --- third_party/xla/xla/tsl/framework/bfc_allocator.cc | 6 ++++++ .../xla/xla/tsl/framework/bfc_allocator_test.cc | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator.cc b/third_party/xla/xla/tsl/framework/bfc_allocator.cc index 51de93e13aa239..4ebdb5d0e42681 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -526,6 +527,11 @@ void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes, VLOG(2) << "tried to allocate 0 bytes"; return nullptr; } + if (ABSL_PREDICT_FALSE(num_bytes > std::numeric_limits::max() - + (kMinAllocationSize - 1))) { + VLOG(2) << "allocation size cannot be safely rounded: " << num_bytes; + return nullptr; + } // First, always allocate memory of at least kMinAllocationSize // bytes, and always allocate multiples of kMinAllocationSize bytes // so all memory addresses are nicely byte aligned. diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc index 51284bb6e38055..9e202ab1909308 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc @@ -153,6 +153,17 @@ TEST(BFCAllocatorTest, DefaultAlignment) { alloc.DeallocateRaw(ptr); } +TEST(BFCAllocatorTest, RejectsAllocationWhoseRoundedSizeWouldOverflow) { + BFCAllocator::Options opts; + opts.allow_retry_on_failure = false; + BFCAllocator alloc(std::make_unique(), + /*total_memory=*/1 << 20, /*name=*/"test", opts); + + EXPECT_EQ( + alloc.AllocateRaw(kAlignment, std::numeric_limits::max() - 1), + nullptr); +} + TEST(BFCAllocatorTest, OomLogsAllocationAnnotations) { BFCAllocator::Options opts; opts.allow_growth = false; From babe64999a7379da7f5cf24748c25fda20123951 Mon Sep 17 00:00:00 2001 From: Ankit Vishwakarma Date: Fri, 18 Sep 2026 08:19:34 -0700 Subject: [PATCH 27/32] PR #47777: Fix ragged all to all 46982 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47777 📝 **Summary of Changes** Modified the `RaggedAllToAll` GPU collective in `xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc` to explicitly skip `LaunchSend`, `LaunchRecv`, and `LaunchPut` calls for zero-size peer slices in the NCCL fallback path. In the Put+Signal path, it now tracks the actual number of puts issued per peer and passes this correct `op_cnt` to `WaitSignal`. 🎯 **Justification** This change resolves issue #46982 by eliminating an O(n_devices) cost scaling bottleneck during ragged-all-to-all operations with sparse neighbor exchange patterns. Previously, NCCL enqueued work for zero-count operations, resulting in ~12 µs overhead for every zero-size slot. By skipping these slots, this change significantly improves performance for sparse communication on a large number of GPUs (e.g., eliminating the +625 µs scaling overhead from 16 to 64 devices on a fixed degree-10 neighbor graph). 🚀 **Kind of Contribution** 🐛 Bug Fix, ⚡️ Performance Improvement 📊 **Benchmark (for Performance Improvements)** Eliminates the ~12 µs per zero-size slot NCCL overhead reported in #46982 (cost is now constant O(actual_neighbors) rather than linear O(n_devices)). 🧪 **Unit Tests:** Existing tests in `ragged_all_to_all_thunk_test.cc` continue to validate thunk initialization and command buffer logic. 🧪 **Execution Tests:** Added `RaggedAllToAll_2GPUs_SparseNeighbors` end-to-end test to `xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc`. This 2-GPU test verifies the collective's correctness under a sparse pattern where one replica sends zero rows in one direction, ensuring that skipping zero-size slots works correctly without deadlocks and produces the expected data. Copybara import of the project: -- f90c7f3f9567c81587085eaac28ab5f34182a0e3 by ankit : Fix O(n_devices) overhead in GPU ragged-all-to-all Fixes #46982 by skipping zero-size peer slices in the NCCL fallback path. -- bc5ec206d7881d118da64513c4cfd03d23705543 by Ankit Vishwakarma : Fix zero-size ragged all-to-all operations Merging this change closes #47777 PiperOrigin-RevId: 983874052 --- .../gpu/runtime/ragged_all_to_all_thunk.cc | 62 +++++++++++++------ .../gpu/tests/ragged_all_to_all_e2e_test.cc | 47 ++++++++++++++ 2 files changed, 90 insertions(+), 19 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc index 81855ce4aecfbd..2ddd40b68ff763 100644 --- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc @@ -437,6 +437,7 @@ absl::Status RunNcclFallbackRaggedAllToAll( const int64_t* input_offsets = ragged_metadata_allocs[0]; const int64_t* send_sizes = ragged_metadata_allocs[1]; const int64_t* output_offsets = ragged_metadata_allocs[2]; + const int64_t* recv_sizes = ragged_metadata_allocs[3]; if (use_put_path) { XLA_VLOG_DEVICE(3, device_ordinal) @@ -445,12 +446,26 @@ absl::Status RunNcclFallbackRaggedAllToAll( int64_t element_byte_width = primitive_util::ByteWidth(element_type); se::DeviceAddressBase input_buffer = buffers[0].source_buffer; + // WaitSignal observes puts sent by `peer` to this rank, so its op_cnt must + // be the number of nonzero receives expected from that peer. + std::vector puts_from_peer(num_ranks, 0); + for (int peer = 0; peer < num_ranks; ++peer) { + for (int64_t i = 0; i < num_updates_per_replica; ++i) { + const int64_t idx = peer * num_updates_per_replica + i; + puts_from_peer[peer] += recv_sizes[idx] != 0; + } + } + Future<> future = gpu_comm->GroupExecute([&]() -> absl::Status { for (int peer = 0; peer < num_ranks; ++peer) { for (int64_t i = 0; i < num_updates_per_replica; ++i) { const int64_t idx = peer * num_updates_per_replica + i; const int64_t send_count = send_sizes[idx] * ragged_row_element_size; + // Skip zero-size sends: NCCL still pays ~12 us per LaunchPut even + // at count=0, causing O(n_devices) cost for sparse neighbor patterns. + if (send_count == 0) continue; + const se::DeviceAddressBase send_slice = GpuCollectives::Slice( input_buffer, element_type, input_offsets[idx] * ragged_row_element_size, send_count); @@ -472,8 +487,9 @@ absl::Status RunNcclFallbackRaggedAllToAll( GpuSignalDesc signal_desc(/*sig_idx=*/0, /*ctx=*/0); for (int peer = 0; peer < num_ranks; ++peer) { + if (puts_from_peer[peer] == 0) continue; ABSL_RETURN_IF_ERROR(comm.WaitSignal(RankId(peer), - /*op_cnt=*/num_updates_per_replica, + /*op_cnt=*/puts_from_peer[peer], signal_desc, GpuCollectives::On(stream)) .Await()); } @@ -483,7 +499,6 @@ absl::Status RunNcclFallbackRaggedAllToAll( XLA_VLOG_DEVICE(3, device_ordinal) << "RunFallbackNcclRaggedAllToAll: using Send/Recv path"; - const int64_t* recv_sizes = ragged_metadata_allocs[3]; Future<> future = gpu_comm->GroupExecute([&]() -> absl::Status { PrimitiveType element_type = buffers[0].element_type; @@ -494,23 +509,32 @@ absl::Status RunNcclFallbackRaggedAllToAll( for (int64_t i = 0; i < num_updates_per_replica; ++i) { for (int peer = 0; peer < num_ranks; ++peer) { int64_t idx = peer * num_updates_per_replica + i; - se::DeviceAddressBase send_slice = - GpuCollectives::Slice(input_buffer, element_type, - input_offsets[idx] * ragged_row_element_size, - send_sizes[idx] * ragged_row_element_size); - - se::DeviceAddressBase recv_slice = - GpuCollectives::Slice(output_buffer, element_type, - output_offsets[idx] * ragged_row_element_size, - recv_sizes[idx] * ragged_row_element_size); - - ABSL_RETURN_IF_ERROR(gpu_comm->LaunchSend( - send_slice, element_type, send_sizes[idx] * ragged_row_element_size, - RankId(peer), GpuCollectives::On(stream))); - - ABSL_RETURN_IF_ERROR(gpu_comm->LaunchRecv( - recv_slice, element_type, recv_sizes[idx] * ragged_row_element_size, - RankId(peer), GpuCollectives::On(stream))); + + // Send and receive refer to different slots. Skip each zero-size + // operation independently: NCCL pays ~12 us per LaunchSend/LaunchRecv + // even at count=0, causing O(n_devices) cost for sparse neighbor + // patterns (issue #46982). + if (send_sizes[idx] != 0) { + se::DeviceAddressBase send_slice = GpuCollectives::Slice( + input_buffer, element_type, + input_offsets[idx] * ragged_row_element_size, + send_sizes[idx] * ragged_row_element_size); + ABSL_RETURN_IF_ERROR( + gpu_comm->LaunchSend(send_slice, element_type, + send_sizes[idx] * ragged_row_element_size, + RankId(peer), GpuCollectives::On(stream))); + } + + if (recv_sizes[idx] != 0) { + se::DeviceAddressBase recv_slice = GpuCollectives::Slice( + output_buffer, element_type, + output_offsets[idx] * ragged_row_element_size, + recv_sizes[idx] * ragged_row_element_size); + ABSL_RETURN_IF_ERROR( + gpu_comm->LaunchRecv(recv_slice, element_type, + recv_sizes[idx] * ragged_row_element_size, + RankId(peer), GpuCollectives::On(stream))); + } } } diff --git a/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc index 60e09a30e872d1..1dec6a15f81264 100644 --- a/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc @@ -891,6 +891,53 @@ TEST_P(RaggedAllToAllTest, RaggedAllToAll_2GPUs_PreservesInitialData) { EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[1], results[1])); } +// Regression test for issue #46982: zero-size peer slices must be skipped so +// that the collective's cost is O(actual_neighbors) not O(n_devices). +// +// Pattern: replica_0 sends 3 rows to replica_1; replica_1 sends 0 rows to +// replica_0. This creates one direction with all-zero sizes per peer, which +// previously caused ~12 us of unnecessary NCCL overhead for every zero-size +// slot in the Send/Recv or Put+Signal paths. +TEST_P(RaggedAllToAllTest, RaggedAllToAll_2GPUs_SparseNeighbors) { + absl::string_view kModuleReplicatedStr = R"( + HloModule module, num_partitions=1 + + ENTRY entry { + input = f32[4] parameter(0) + output = f32[4] parameter(1) + input_offsets = s32[2] parameter(2) + send_sizes = s32[2] parameter(3) + output_offsets = s32[2] parameter(4) + recv_sizes = s32[2] parameter(5) + ROOT ra2a = f32[4] ragged-all-to-all(input, output, input_offsets, + send_sizes, output_offsets, recv_sizes), replica_groups={{0,1}} + })"; + + const int64_t kNumReplicas = 2; + ASSERT_GE(device_count(), kNumReplicas) + << "Test requires at least " << kNumReplicas << " devices (" + << device_count() << " available)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + kModuleReplicatedStr, kNumReplicas)); + + // Sparse pattern: replica_0 sends 3 rows to replica_1, and 1 row to itself. + // replica_1 sends 0 rows to replica_0 (zero-size slot), and 1 row to itself. + // The zero-size slot for replica_1->replica_0 must be skipped without hangs. + ASSERT_OK(CreateRandomTestData(module.get(), + /*input_sizes=*/{/*replica_0=*/{1, 3}, + /*replica_1=*/{0, 1}})); + + ASSERT_OK_AND_ASSIGN( + ExecutionResult execution_result, + ExecuteReplicated(std::move(module), GetInputLiteralPtrs())); + + const std::vector& results = execution_result.results; + ASSERT_EQ(results.size(), kNumReplicas); + EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[0], results[0])); + EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[1], results[1])); +} + TEST_P(RaggedAllToAllTest, RaggedAllToAll_8GPUs) { absl::string_view kModuleReplicatedStr = R"( HloModule module, num_partitions=1 From a1fd6dda463189a83b5b621c70db2d7357fae5fd Mon Sep 17 00:00:00 2001 From: Lizzie Liang Date: Fri, 18 Sep 2026 08:34:05 -0700 Subject: [PATCH 28/32] Fix StableHLO verifier crash - Added positive and negative type inference tests to ops_chlo.mlir and a test to chlo_legalize_to_stablehlo.mlir. - Made some changes in ChloLegalizeToStablehlo.cpp so that the bounds are preserved for dynamic tensors. PiperOrigin-RevId: 983880829 --- .../xla/third_party/stablehlo/temporary.patch | 308 ++++++++++++++++++ 1 file changed, 308 insertions(+) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index 7cc02539c261a5..b6c9e2d8cb1f8a 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -2017,6 +2017,104 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial + #endif // STABLEHLO_DIALECT_CHLO_ENUMS +diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/dialect/ChloOps.cpp +--- stablehlo/stablehlo/dialect/ChloOps.cpp ++++ stablehlo/stablehlo/dialect/ChloOps.cpp +@@ -827,6 +827,10 @@ + if (adaptor.getScanDimSize().has_value()) { + dimSize = adaptor.getScanDimSize().value(); + } ++ int64_t dimBound = ShapedType::kDynamic; ++ // Prototype encoding attribute used to construct dialect type extensions ++ // if an output does not already have an encoding. ++ Attribute prototypeEncoding; + for (auto [i, type] : llvm::enumerate(adaptor.getInputs().getTypes())) { + auto inputType = dyn_cast(type); + if (!inputType) { +@@ -839,7 +843,26 @@ + if (dimSize == ShapedType::kDynamic) { + dimSize = inputType.getDimSize(dim); + } +- } ++ if (inputType.getEncoding()) { ++ prototypeEncoding = inputType.getEncoding(); ++ } ++ ArrayRef bounds = hlo::encodingToBounds(inputType.getEncoding()); ++ if (!bounds.empty() && dim < static_cast(bounds.size())) { ++ int64_t currentBound = bounds[dim]; ++ if (currentBound != ShapedType::kDynamic) { ++ if (dimBound == ShapedType::kDynamic) { ++ dimBound = currentBound; ++ } else { ++ dimBound = std::min(dimBound, currentBound); ++ } ++ } ++ } ++ } ++ ++ // If the scan dimension size is dynamic, propagate any bound inferred from ++ // the inputs; otherwise static dimensions do not need a bound. ++ const int64_t insertBound = ++ (dimSize == ShapedType::kDynamic) ? dimBound : ShapedType::kDynamic; + + for (auto [i, type] : llvm::enumerate(terminator->getOperands().getTypes())) { + auto resultType = dyn_cast(type); +@@ -849,10 +872,34 @@ + } + SmallVector shape(resultType.getShape().begin(), + resultType.getShape().end()); ++ Attribute encoding = resultType.getEncoding(); ++ // The first `numOutputs` terminator operands correspond to scan outputs, ++ // which have the scan dimension re-inserted. The remaining operands are ++ // carries whose shapes and encodings are preserved as-is. + if (i < numOutputs) { + shape.insert(std::next(shape.begin(), dim), dimSize); +- } +- inferredReturnShapes.emplace_back(shape, resultType.getElementType()); ++ SmallVector resultBounds = ++ llvm::to_vector(hlo::encodingToBounds(encoding)); ++ // Re-insert bounds along the scan dimension if the result has bounds or ++ // if the scan dimension has a bound. ++ if (!resultBounds.empty() || insertBound != ShapedType::kDynamic) { ++ // If the body result did not have bounds, initialize all dimensions to ++ // dynamic and borrow the prototype encoding to create dialect type ++ // extensions. ++ if (resultBounds.empty()) { ++ resultBounds.assign(resultType.getRank(), ShapedType::kDynamic); ++ if (!encoding) { ++ encoding = prototypeEncoding; ++ } ++ } ++ resultBounds.insert(std::next(resultBounds.begin(), dim), insertBound); ++ if (encoding) { ++ encoding = hlo::boundsToEncoding(encoding, resultBounds); ++ } ++ } ++ } ++ inferredReturnShapes.emplace_back(shape, resultType.getElementType(), ++ encoding); + } + + return success(); +@@ -902,7 +949,17 @@ + if (i < getInputs().size()) { + auto argShape = llvm::to_vector(argType.getShape()); + argShape.erase(std::next(argShape.begin(), dim)); +- argType = argType.clone(argShape); ++ Attribute encoding = argType.getEncoding(); ++ SmallVector argBounds = ++ llvm::to_vector(hlo::encodingToBounds(encoding)); ++ if (!argBounds.empty()) { ++ argBounds.erase(std::next(argBounds.begin(), dim)); ++ encoding = argBounds.empty() ++ ? Attribute() ++ : hlo::boundsToEncoding(encoding, argBounds); ++ } ++ argType = ++ RankedTensorType::get(argShape, argType.getElementType(), encoding); + } + if (!hlo::isCompatibleForHloTypeInference( + argType, bodyBlock.getArgument(i).getType())) { diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp --- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp @@ -2974,6 +3072,65 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B + "//stablehlo/tests:stablehlo_tests", + ], +) +diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +--- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir ++++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +@@ -5424,6 +5424,55 @@ + + // ----- + ++// CHECK-LABEL: func.func @scan_bounded( ++// CHECK-SAME: %[[ARG0:.*]]: tensor>, ++// CHECK-SAME: %[[ARG1:.*]]: tensor>) -> (tensor>, tensor>) { ++// CHECK: %[[GET_DIM_SIZE_0:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 0 : (tensor>) -> tensor ++// CHECK: %[[CONVERT_0:.*]] = stablehlo.convert %[[GET_DIM_SIZE_0]] : (tensor) -> tensor ++// CHECK: %[[C0:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[C0_I32:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG0]] : tensor> -> tensor<2xindex> ++// CHECK: %[[INIT_OUT:.*]] = stablehlo.dynamic_broadcast_in_dim %[[C0_I32]], %[[SHAPE]], dims = [] : (tensor, tensor<2xindex>) -> tensor> ++// CHECK: %[[WHILE:.*]]:3 = stablehlo.while(%[[ITER_I:.*]] = %[[C0]], %[[ITER_CARRY:.*]] = %[[ARG1]], %[[ITER_OUT:.*]] = %[[INIT_OUT]]) : tensor, tensor>, tensor> ++// CHECK: cond { ++// CHECK: %[[CMP:.*]] = stablehlo.compare LT, %[[ITER_I]], %[[CONVERT_0]] : (tensor, tensor) -> tensor ++// CHECK: stablehlo.return %[[CMP]] : tensor ++// CHECK: } do { ++// CHECK: %[[C0_1:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[I_RESHAPED:.*]] = stablehlo.reshape %[[ITER_I]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[C0_RESHAPED:.*]] = stablehlo.reshape %[[C0_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[START_INDICES:.*]] = stablehlo.concatenate %[[I_RESHAPED]], %[[C0_RESHAPED]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64> ++// CHECK: %[[C1:.*]] = stablehlo.constant dense<1> : tensor ++// CHECK: %[[I_PLUS_1:.*]] = stablehlo.add %[[ITER_I]], %[[C1]] : tensor ++// CHECK: %[[GET_DIM_SIZE_1:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 1 : (tensor>) -> tensor ++// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[GET_DIM_SIZE_1]] : (tensor) -> tensor ++// CHECK: %[[LIMIT_0:.*]] = stablehlo.reshape %[[I_PLUS_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[LIMIT_1:.*]] = stablehlo.reshape %[[CONVERT_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[LIMIT_INDICES:.*]] = stablehlo.concatenate %[[LIMIT_0]], %[[LIMIT_1]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64> ++// CHECK: %[[STRIDES:.*]] = stablehlo.constant dense<1> : tensor<2xi64> ++// CHECK: %[[SLICE:.*]] = stablehlo.real_dynamic_slice %[[ARG0]], %[[START_INDICES]], %[[LIMIT_INDICES]], %[[STRIDES]] : (tensor>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x?xi32, #stablehlo.bounds> ++// CHECK: %[[INPUT_ELEM:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1x?xi32, #stablehlo.bounds>) -> tensor> ++// CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT_ELEM]], %[[ITER_CARRY]] : tensor> ++// CHECK: %[[OUT_ELEM:.*]] = stablehlo.reshape %[[ADD]] : (tensor>) -> tensor<1x?xi32, #stablehlo.bounds> ++// CHECK: %[[C0_2:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[UPDATED_OUT:.*]] = stablehlo.dynamic_update_slice %[[ITER_OUT]], %[[OUT_ELEM]], %[[ITER_I]], %[[C0_2]] : (tensor>, tensor<1x?xi32, #stablehlo.bounds>, tensor, tensor) -> tensor> ++// CHECK: %[[C1_2:.*]] = stablehlo.constant dense<1> : tensor ++// CHECK: %[[NEXT_I:.*]] = stablehlo.add %[[ITER_I]], %[[C1_2]] : tensor ++// CHECK: stablehlo.return %[[NEXT_I]], %[[ADD]], %[[UPDATED_OUT]] : tensor, tensor>, tensor> ++// CHECK: } ++// CHECK: return %[[WHILE]]#2, %[[WHILE]]#1 : tensor>, tensor> ++// CHECK: } ++func.func @scan_bounded(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ %0:2 = chlo.scan(%arg0) inits(%arg1) dimension=0 { ++ ^bb0(%scan_arg0: tensor>, %scan_arg1: tensor>): ++ %1 = stablehlo.add %scan_arg0, %scan_arg1 : tensor> ++ stablehlo.return %1, %1 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0#0, %0#1 : tensor>, tensor> ++} ++ ++// ----- ++ + // CHECK-LABEL: func.func @mulhi_i32( + // CHECK-SAME: %[[ARG0:.*]]: tensor<4xi32>, + // CHECK-SAME: %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> { diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir --- stablehlo/stablehlo/tests/interpret/all_gather.mlir +++ stablehlo/stablehlo/tests/interpret/all_gather.mlir @@ -3164,6 +3321,91 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo + func.return + } +} +diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir +--- stablehlo/stablehlo/tests/ops_chlo.mlir ++++ stablehlo/stablehlo/tests/ops_chlo.mlir +@@ -485,6 +485,50 @@ + stablehlo.return %1, %2, %1, %2 : tensor<3xf32>, tensor<3xi32>, tensor<3xf32>, tensor<3xi32> + } : (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) -> (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) + func.return %0#0, %0#1 : tensor<2x3xf32>, tensor<2x3xi32> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism ++func.func @scan_bounded_dynamism(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>) -> (tensor>, tensor>) ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor>, %carry0: tensor>): ++ %2 = stablehlo.add %input0, %carry0 : tensor> ++ stablehlo.return %2, %2 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0, %1 : tensor>, tensor> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs ++func.func @scan_bounded_dynamism_multiple_inputs(%arg0: tensor>, %arg1: tensor>, %arg2: tensor>, %arg3: tensor>) -> (tensor>, tensor>, tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 { ++ ^bb0(%input0: tensor>, %input1: tensor>, %carry0: tensor>, %carry1: tensor>): ++ %1 = stablehlo.add %input0, %carry0 : tensor> ++ %2 = stablehlo.add %input1, %carry1 : tensor> ++ stablehlo.return %1, %2, %1, %2 : tensor>, tensor>, tensor>, tensor> ++ } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ func.return %0#0, %0#1, %0#2, %0#3 : tensor>, tensor>, tensor>, tensor> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs_same_rank ++func.func @scan_bounded_dynamism_multiple_inputs_same_rank(%arg0: tensor>, %arg1: tensor>, %arg2: tensor>, %arg3: tensor>) -> (tensor>, tensor>, tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 { ++ ^bb0(%input0: tensor>, %input1: tensor>, %carry0: tensor>, %carry1: tensor>): ++ %1 = stablehlo.add %input0, %carry0 : tensor> ++ %2 = stablehlo.add %input1, %carry1 : tensor> ++ stablehlo.return %1, %2, %1, %2 : tensor>, tensor>, tensor>, tensor> ++ } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ func.return %0#0, %0#1, %0#2, %0#3 : tensor>, tensor>, tensor>, tensor> + } + + // ----- +@@ -561,6 +605,30 @@ + + // ----- + ++func.func @scan_bounded_dynamism_mismatch(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ // expected-error @+1 {{'chlo.scan' op operand and body argument 0 are incompatible}} ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor<4xf32>, %carry0: tensor>): ++ stablehlo.return %carry0, %carry0 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0, %1 : tensor>, tensor> ++} ++ ++// ----- ++ ++func.func @scan_bounded_dynamism_result_rank_mismatch(%arg0: tensor>, %arg1: tensor>) -> (tensor, tensor>) { ++ // expected-error @+2 {{op failed to infer returned types}} ++ // expected-error @+1 {{'chlo.scan' op inferred type(s) 'tensor>', 'tensor>' are incompatible with return type(s) of operation 'tensor', 'tensor>'}} ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor>, %carry0: tensor>): ++ %2 = stablehlo.add %input0, %carry0 : tensor> ++ stablehlo.return %2, %2 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor, tensor>) ++ func.return %0, %1 : tensor, tensor> ++} ++ ++// ----- ++ + // CHECK-LABEL: func @mulhi_i32 + func.func @mulhi_i32(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> { + %0 = "chlo.mulhi"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir --- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir @@ -3319,4 +3561,70 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b %2 = stablehlo.dynamic_iota %1, dim = 0 : (tensor<1xi64>) -> tensor return %2 : tensor } +diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +--- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp ++++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +@@ -2930,8 +2930,19 @@ + SmallVector sliceShape(inputType.getShape().begin(), + inputType.getShape().end()); + sliceShape[dim] = 1; ++ Attribute sliceEncoding; ++ ArrayRef inputBounds = ++ hlo::encodingToBounds(inputType.getEncoding()); ++ if (!inputBounds.empty()) { ++ SmallVector sliceBounds(inputBounds.begin(), ++ inputBounds.end()); ++ sliceBounds[dim] = ShapedType::kDynamic; ++ sliceEncoding = ++ hlo::boundsToEncoding(inputType.getEncoding(), sliceBounds); ++ } + auto sliceType = +- RankedTensorType::get(sliceShape, inputType.getElementType()); ++ RankedTensorType::get(sliceShape, inputType.getElementType(), ++ sliceEncoding); + + Value slice; + if (inputType.hasStaticShape()) { +@@ -2980,9 +2991,18 @@ + for (int64_t d = 0; d < rank; ++d) { + if (d != dim) resultShape.push_back(sliceShape[d]); + } ++ Attribute resultEncoding; ++ if (!inputBounds.empty()) { ++ SmallVector resultBounds(inputBounds.begin(), ++ inputBounds.end()); ++ resultBounds.erase(std::next(resultBounds.begin(), dim)); ++ resultEncoding = ++ hlo::boundsToEncoding(inputType.getEncoding(), resultBounds); ++ } + Value reshaped = ReshapeOp::create( + rewriter, loc, +- RankedTensorType::get(resultShape, inputType.getElementType()), ++ RankedTensorType::get(resultShape, inputType.getElementType(), ++ resultEncoding), + slice); + slicedInputs.push_back(reshaped); + } +@@ -3004,9 +3024,20 @@ + SmallVector elementShape(outputType.getShape().begin(), + outputType.getShape().end()); + elementShape[dim] = 1; ++ Attribute elementEncoding; ++ ArrayRef outputBounds = ++ hlo::encodingToBounds(outputType.getEncoding()); ++ if (!outputBounds.empty()) { ++ SmallVector elementBounds(outputBounds.begin(), ++ outputBounds.end()); ++ elementBounds[dim] = ShapedType::kDynamic; ++ elementEncoding = ++ hlo::boundsToEncoding(outputType.getEncoding(), elementBounds); ++ } + Value reshapedElement = ReshapeOp::create( + rewriter, loc, +- RankedTensorType::get(elementShape, outputType.getElementType()), ++ RankedTensorType::get(elementShape, outputType.getElementType(), ++ elementEncoding), + newElement); + + int64_t outRank = outputType.getRank(); From 89fd8b3e1dc2da2b1c757cf9048291a9bff7e74b Mon Sep 17 00:00:00 2001 From: Zviki Nozadze Date: Fri, 18 Sep 2026 09:08:04 -0700 Subject: [PATCH 29/32] [XLA] Refactor separate_compilation to parallel. PiperOrigin-RevId: 983896161 --- .../{separate_compilation => parallel}/BUILD | 8 -- .../hlo_linking_manifest.h | 10 +- .../hlo_module_linking.cc | 32 +++-- .../hlo_module_linking.h | 12 +- .../hlo_module_linking_test.cc | 136 +++++++++--------- .../hlo_module_splitting.cc | 20 ++- .../hlo_module_splitting.h | 12 +- .../hlo_module_splitting_test.cc | 99 ++++++------- 8 files changed, 162 insertions(+), 167 deletions(-) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/BUILD (93%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_linking_manifest.h (87%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_linking.cc (91%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_linking.h (81%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_linking_test.cc (65%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_splitting.cc (96%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_splitting.h (93%) rename third_party/xla/xla/hlo/{separate_compilation => parallel}/hlo_module_splitting_test.cc (73%) diff --git a/third_party/xla/xla/hlo/separate_compilation/BUILD b/third_party/xla/xla/hlo/parallel/BUILD similarity index 93% rename from third_party/xla/xla/hlo/separate_compilation/BUILD rename to third_party/xla/xla/hlo/parallel/BUILD index 776e872a486520..c6f830de3fb2a3 100644 --- a/third_party/xla/xla/hlo/separate_compilation/BUILD +++ b/third_party/xla/xla/hlo/parallel/BUILD @@ -48,8 +48,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/service:compilation_environments", "//xla/service:hlo_module_config", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/hash", @@ -74,8 +72,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/hlo/transforms/simplifiers:hlo_dce", "//xla/service:compilation_environments", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -103,8 +99,6 @@ xla_cc_test( "//xla/service:cpu_plugin", # fixdeps: keep "//xla/service:platform_util", "//xla/stream_executor:platform", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", @@ -127,8 +121,6 @@ xla_cc_test( "//xla/service:hlo_verifier", "//xla/service:platform_util", "//xla/stream_executor:platform", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h b/third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h similarity index 87% rename from third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h rename to third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h index 1558a2f6c47788..408d2991badb60 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h +++ b/third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ +#ifndef XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ +#define XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ #include @@ -24,7 +24,7 @@ limitations under the License. #include "xla/service/compilation_environments.h" #include "xla/service/hlo_module_config.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Metadata to guide linking process of HLO modules. // @@ -46,5 +46,5 @@ struct HloLinkingManifest { std::unique_ptr compilation_environment; }; -} // namespace xla::separate_compilation -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ +} // namespace xla::parallel +#endif // XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc b/third_party/xla/xla/hlo/parallel/hlo_module_linking.cc similarity index 91% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_linking.cc index f0432ea964d55b..a0c9eb23db2917 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_linking.h" +#include "xla/hlo/parallel/hlo_module_linking.h" #include #include @@ -35,13 +35,13 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/service/compilation_environments.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { struct TraversalState { @@ -123,17 +123,17 @@ class HloLinker { } state.entered = true; - PushDependencies(state); + PushDependencies(state.principal); return absl::OkStatus(); } - // Pushes dependencies of `state.principal` onto stack_ if they are stubs + // Pushes dependencies of `principal` onto stack_ if they are stubs // that need to be linked. - void PushDependencies(const TraversalState& state) { - std::vector post_order_instrs = - state.principal->MakeInstructionPostOrder(); + void PushDependencies(const HloComputation* principal) { VLOG(6) << "Processing callees:"; - for (HloInstruction* caller : post_order_instrs) { + std::vector post_order_instrs = + principal->MakeInstructionPostOrder(); + for (const HloInstruction* caller : post_order_instrs) { if (caller->opcode() != HloOpcode::kCall) { continue; } @@ -146,8 +146,8 @@ class HloLinker { it != linking_manifest_.stub_links.end()) { // Only push if its principal hasn't been mapped yet. if (clone_context_.FindComputation(callee) == nullptr) { - const HloComputation* principal = it->second; - stack_.push({principal, callee, false}); + const HloComputation* callee_principal = it->second; + stack_.push({callee_principal, callee, false}); } } } @@ -184,7 +184,7 @@ class HloLinker { absl::flat_hash_set being_linked_; absl::flat_hash_map finished_principals_; - std::stack> stack_; + std::stack stack_; HloCloneContext clone_context_; int linked_count_ = 0; @@ -222,4 +222,6 @@ absl::StatusOr> LinkComputation( return linked_module; } -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h b/third_party/xla/xla/hlo/parallel/hlo_module_linking.h similarity index 81% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h rename to third_party/xla/xla/hlo/parallel/hlo_module_linking.h index a6a68aae1d6e2c..4d55ac196fdc1f 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ +#ifndef XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ +#define XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ #include @@ -22,9 +22,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Link all the callees together with the `root_computation` in the provided // module. Returns the pointer to the linked version of `root_computation` @@ -38,6 +38,6 @@ absl::StatusOr> LinkComputation( const HloLinkingManifest& linking_manifest, const HloComputation* absl_nonnull root_computation); -} // namespace xla::separate_compilation +} // namespace xla::parallel -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ +#endif // XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc b/third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc similarity index 65% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc index 62d611e9e73905..a5a3507af7e33b 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc @@ -13,12 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_linking.h" +#include "xla/hlo/parallel/hlo_module_linking.h" #include #include #include +#include #include #include "absl/container/flat_hash_set.h" #include "absl/log/log.h" @@ -27,17 +28,17 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/compiler.h" #include "xla/service/hlo_verifier.h" #include "xla/service/platform_util.h" #include "xla/stream_executor/platform.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { // Function to normalize an HloModule by removing/replacing names. @@ -108,12 +109,13 @@ TEST_F(LinkingTest, SingleCallLinking) { ROOT %result = f32[] call(%p), to_apply=%comp } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); // LOG the split group. + // NOLINTBEGIN(*-custom-deterministic-iteration-order) for (const auto& split : module_split_group->module_splits) { LOG(INFO) << "Split: " << split->submodule->name(); LOG(INFO) << "Split module: " << split->submodule->ToString(); @@ -143,32 +145,32 @@ TEST_F(LinkingTest, SingleCallLinking) { module_split_group->linking_manifest.stub_links) { LOG(INFO) << " " << stub->name() << " ==>> " << comp->name(); } + // NOLINTEND(*-custom-deterministic-iteration-order) const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); VLOG(6) << linked_module->ToString(); - TF_ASSERT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } TEST_F(LinkingTest, ChainGraphLinking) { @@ -200,36 +202,35 @@ TEST_F(LinkingTest, ChainGraphLinking) { ROOT %result = f32[] add(%call_res, %p) } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); VLOG(6) << linked_module->ToString(); - TF_EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } TEST_F(LinkingTest, DiamondGraphLinking) { @@ -273,38 +274,39 @@ TEST_F(LinkingTest, DiamondGraphLinking) { ROOT %result = f32[] add(%call_a, %call_b) } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); VLOG(6) << linked_module->ToString(); - TF_EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } // namespace -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc similarity index 96% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc index 02254a3b5fba2c..ae315bed1e0ec6 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include #include @@ -39,14 +39,14 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" #include "xla/service/compilation_environments.h" #include "xla/service/hlo_module_config.h" #include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { constexpr absl::string_view kEntryName = "entry"; @@ -109,6 +109,7 @@ absl::StatusOr> ComposeMaps( const absl::flat_hash_map& first, const absl::flat_hash_map& second) { absl::flat_hash_map result; + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto [k, kv] : first) { if (auto it = second.find(kv); it != second.end()) { result.insert({k, it->second}); @@ -123,6 +124,7 @@ template absl::Status MergeMapInto(absl::flat_hash_map& into, const absl::flat_hash_map& from, bool error_on_duplicate_key = true) { + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [k, v] : from) { if (!into.insert({k, v}).second) { if (error_on_duplicate_key) { @@ -265,7 +267,7 @@ absl::StatusOr> CreateHloModuleSplit( if (replacement->parent() != submodule.get()) { replacement = clone_context.GetComputation(replacement); } - mapped_call_instruction->set_to_apply(callee_replacements[caller]); + mapped_call_instruction->set_to_apply(replacement); } entry_computation->SetAndSanitizeName(kEntryName); @@ -296,6 +298,7 @@ absl::StatusOr> CreateHloModuleSplitGroup( for (const auto& split : splits) { ABSL_ASSIGN_OR_RETURN(auto module_split, CreateHloModuleSplit(module, split)); module_splits.push_back(std::move(module_split)); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto* original_comp : split) { computation_address_book.insert( {original_comp, module_splits.back().get()}); @@ -311,6 +314,7 @@ absl::StatusOr> CreateHloModuleSplitGroup( for (const auto& split : module_splits) { VLOG(5) << "Split: " << split->submodule->name(); VLOG(5) << " Stub links:"; + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [stub, comp] : split->stub_map) { VLOG(5) << " " << stub->name() << " ==>> " << comp->name() << "(" @@ -354,4 +358,6 @@ absl::StatusOr HloModuleSplitGroup::GetClonedComputation( return it2->second; } -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.h similarity index 93% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting.h index 97a0031035511c..dbff1f44cefaec 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ +#ifndef XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ +#define XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ #include #include @@ -26,9 +26,9 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Returns a list of sets of computations that can be split into separate // modules. Adjacent computations in the same set can be compiled together. @@ -108,5 +108,5 @@ struct HloModuleSplitGroup { absl::StatusOr> CreateHloModuleSplitGroup( const HloModule& module); -} // namespace xla::separate_compilation -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ +} // namespace xla::parallel +#endif // XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc b/third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc similarity index 73% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc index 0eb36ac22118f9..373c3a4509b737 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include @@ -26,10 +26,8 @@ limitations under the License. #include "xla/service/compiler.h" #include "xla/service/platform_util.h" #include "xla/stream_executor/platform.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +namespace xla::parallel { namespace { using ::testing::UnorderedElementsAreArray; @@ -85,9 +83,8 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto splits, GroupComputationsForSplitting(*module)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto splits, GroupComputationsForSplitting(*module)); absl::flat_hash_set all_computations_in_splits; for (const auto& bucket : splits) { @@ -146,12 +143,11 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); auto* main = FindComputation(module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN(auto module_split, - CreateHloModuleSplit(*module, {main})); + ASSERT_OK_AND_ASSIGN(auto module_split, + CreateHloModuleSplit(*module, {main})); const int kMainModuleComputationCount = 5; // the main + 4 stubs // const int kMainModuleClonedComputations = 1; @@ -217,21 +213,20 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.fu23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); - - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); + + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } @@ -272,10 +267,9 @@ TEST_F(SplittingTest, SplitDiamondGraphModule) { } )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); // Expect all computations to be assigned somewhere. for (auto original_comp : module->computations()) { @@ -285,16 +279,16 @@ TEST_F(SplittingTest, SplitDiamondGraphModule) { EXPECT_EQ(module_split_group->module_splits.size(), module->computation_count()); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } @@ -354,23 +348,22 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.f) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); - - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); + + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } } // namespace -} // namespace xla::separate_compilation +} // namespace xla::parallel From 742d7183e49eb62d921c5a4307b40d8c96ef042d Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 18 Sep 2026 09:19:10 -0700 Subject: [PATCH 30/32] Optimize GraphDef and FunctionLibraryDefinition memory in TFRT serving runtimes. Eliminates redundant in-memory copies of GraphDef and FunctionLibraryDefinition during model initialization and signature compilation in TensorFlow serving runtimes. PiperOrigin-RevId: 983901136 --- .../common_runtime/graph_execution_state.cc | 37 ++++++++-- .../common_runtime/graph_execution_state.h | 11 +++ tensorflow/core/tfrt/fallback/BUILD | 1 + .../core/tfrt/fallback/fallback_state.cc | 3 + tensorflow/core/tfrt/utils/BUILD | 1 + .../tfrt/utils/tfrt_graph_execution_state.cc | 23 +++++- .../utils/tfrt_graph_execution_state_test.cc | 71 +++++++++++++++++++ 7 files changed, 140 insertions(+), 7 deletions(-) diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc index 2e15542512a0fa..e5f7f28a79a5d0 100644 --- a/tensorflow/core/common_runtime/graph_execution_state.cc +++ b/tensorflow/core/common_runtime/graph_execution_state.cc @@ -53,6 +53,7 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/util/device_name_utils.h" +#include "tensorflow/core/util/env_var.h" #include "tensorflow/core/util/util.h" #ifndef IS_MOBILE_PLATFORM @@ -63,12 +64,31 @@ limitations under the License. namespace tensorflow { +bool ProtoMemoryOptimizationsEnabled() { + bool val = false; + (void)ReadBoolFromEnvVar("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS", false, &val); + return val; +} + namespace { bool IsCollectiveV2(const std::string& op) { return op == "CollectiveReduceV2" || op == "CollectiveGatherV2" || op == "CollectiveBcastRecvV2" || op == "CollectiveBcastSendV2" || op == "ColectiveReduceScatterV2" || op == "ColectiveAllToAllV2"; } + +// Returns true if `flib_def` is non-null and contains definitions for all +// functions defined in `graph_def.library()`. +bool LibraryContainsAllFunctions(const FunctionLibraryDefinition* flib_def, + const GraphDef& graph_def) { + if (flib_def == nullptr) return false; + for (const auto& fdef : graph_def.library().function()) { + if (!flib_def->Contains(fdef.signature().name())) { + return false; + } + } + return true; +} } // namespace GraphExecutionState::GraphExecutionState( @@ -96,14 +116,23 @@ GraphExecutionState::~GraphExecutionState() { VLOG(4) << "Graph proto is \n" << graph_def.DebugString(); #endif // __ANDROID__ - auto flib_def = std::make_unique( - OpRegistry::Global(), graph_def); + std::unique_ptr flib_def; + if (ProtoMemoryOptimizationsEnabled() && + LibraryContainsAllFunctions(options.flib_def, graph_def)) { + flib_def = std::make_unique(*options.flib_def); + } else { + flib_def = std::make_unique(OpRegistry::Global(), + graph_def); + } TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(&graph_def, *flib_def, 0)); - if (options.session_options->config.graph_options().place_pruned_graph() || + const bool keep_original_graph_def = + options.session_options->config.graph_options().place_pruned_graph() || options.session_options->config.experimental() - .disable_optimize_for_static_graph()) { + .disable_optimize_for_static_graph(); + + if (keep_original_graph_def) { auto ret = absl::WrapUnique(new GraphExecutionState( std::make_unique(std::move(graph_def)), std::move(flib_def), options)); diff --git a/tensorflow/core/common_runtime/graph_execution_state.h b/tensorflow/core/common_runtime/graph_execution_state.h index a718b57063f10d..9722a2c6b6aab3 100644 --- a/tensorflow/core/common_runtime/graph_execution_state.h +++ b/tensorflow/core/common_runtime/graph_execution_state.h @@ -33,6 +33,11 @@ limitations under the License. #include "tensorflow/core/platform/types.h" namespace tensorflow { + +// Returns true if proto memory optimizations are enabled (default false, can be +// enabled via environment variable TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS=1). +bool ProtoMemoryOptimizationsEnabled(); + struct SessionOptions; namespace subgraph { @@ -54,6 +59,12 @@ struct GraphExecutionStateOptions { // work for almost all models. Non default values should only applied to // selective models. bool enable_tf2xla_mlir_bridge = true; + + // Optional pre-constructed FunctionLibraryDefinition. If provided, + // GraphExecutionState will copy this library definition (which performs + // an O(1) refcount increment on underlying FunctionRecords) rather than + // re-parsing or duplicating the library proto in GraphDef. + const FunctionLibraryDefinition* flib_def = nullptr; }; // A ClientGraph is simply a sub-graph of the full graph as induced by diff --git a/tensorflow/core/tfrt/fallback/BUILD b/tensorflow/core/tfrt/fallback/BUILD index 90c432bf59d01a..9ec1b4826c447f 100644 --- a/tensorflow/core/tfrt/fallback/BUILD +++ b/tensorflow/core/tfrt/fallback/BUILD @@ -67,6 +67,7 @@ cc_library( "//tensorflow/core/framework:graph_proto_cc", "//tensorflow/core/tpu:virtual_device", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/tensorflow/core/tfrt/fallback/fallback_state.cc b/tensorflow/core/tfrt/fallback/fallback_state.cc index 694fae73246469..67c6953da871e9 100644 --- a/tensorflow/core/tfrt/fallback/fallback_state.cc +++ b/tensorflow/core/tfrt/fallback/fallback_state.cc @@ -171,6 +171,9 @@ FallbackState::CreateGraphExecutionState(GraphDef graph_def, bool run_placer, options.session_handle = "tfrt_fallback_handle"; options.run_placer = run_placer; options.enable_tf2xla_mlir_bridge = enable_tf2xla_mlir_bridge; + if (tensorflow::ProtoMemoryOptimizationsEnabled()) { + options.flib_def = &func_lib_def_; + } std::unique_ptr execution_state; TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph( diff --git a/tensorflow/core/tfrt/utils/BUILD b/tensorflow/core/tfrt/utils/BUILD index 84377678e385e9..1329f53461aa1f 100644 --- a/tensorflow/core/tfrt/utils/BUILD +++ b/tensorflow/core/tfrt/utils/BUILD @@ -271,6 +271,7 @@ tf_cc_test( "//tensorflow/core/framework:graph_proto_cc", "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/framework:types_proto_cc", + "//tensorflow/core/framework:versions_proto_cc", "//tensorflow/core/grappler/utils:grappler_test", "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/platform:status", diff --git a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc index 3ab8a2c0c6536f..6fbfb9422ebfd1 100644 --- a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc +++ b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/function_body.h" #include "tensorflow/core/common_runtime/function_def_utils.h" #include "tensorflow/core/common_runtime/graph_constructor.h" +#include "tensorflow/core/common_runtime/graph_execution_state.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/function.pb.h" @@ -190,8 +191,13 @@ tensorflow::GraphDef CreateGraphDefFromGraphAndFlibDef( const tensorflow::Graph& graph, const tensorflow::FunctionLibraryDefinition& flib_def) { tensorflow::GraphDef graph_def; - graph.ToGraphDef(&graph_def); - *graph_def.mutable_library() = flib_def.ToProto(); + graph.ToGraphDef(&graph_def, /*include_flib_def=*/false); + if (tensorflow::ProtoMemoryOptimizationsEnabled()) { + *graph_def.mutable_library() = + flib_def.ReachableDefinitions(graph).ToProto(); + } else { + *graph_def.mutable_library() = flib_def.ToProto(); + } return graph_def; } @@ -693,6 +699,16 @@ TfrtGraphExecutionState::OptimizeGraph( &optimized_flib)); } + if (!options_.run_placer_grappler_on_functions && + tensorflow::ProtoMemoryOptimizationsEnabled()) { + // Fast path: Directly merge the optimized FunctionLibraryDefinition in C++ + // without serializing to a FunctionDefLibrary proto, avoiding expensive + // heap allocations and copy overhead during model initialization. + TF_RETURN_IF_ERROR(optimized_graph->mutable_flib_def()->AddLibrary( + std::move(*optimized_flib))); + return optimized_graph; + } + FunctionDefLibrary optimized_flib_proto = optimized_flib->ToProto(); if (options_.run_placer_grappler_on_functions) { TF_RETURN_IF_ERROR(OptimizeFunctions(optimized_flib_proto, *optimized_flib, @@ -704,7 +720,8 @@ TfrtGraphExecutionState::OptimizeGraph( optimized_graph->mutable_flib_def()->Clear(); } - TF_RETURN_IF_ERROR(optimized_graph->AddFunctionLibrary(optimized_flib_proto)); + TF_RETURN_IF_ERROR( + optimized_graph->AddFunctionLibrary(std::move(optimized_flib_proto))); return optimized_graph; } diff --git a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc index 2d370c0c01c644..5fa6a164ffe82c 100644 --- a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc +++ b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/core/tfrt/utils/tfrt_graph_execution_state.h" #include +#include #include #include #include @@ -42,6 +43,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/grappler/utils/grappler_test.h" #include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/platform/status.h" @@ -437,6 +439,75 @@ TEST_F(OptimizeGraphTest, OptimizeFunctions) { optimized_graph_def.library().function(0)); } +TEST_F(OptimizeGraphTest, MinConsumerBump) { + setenv("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS", "true", 1); + GraphDef graphdef; + tensorflow::FunctionDefLibrary fdef_lib; + { + auto scope = tensorflow::Scope::NewRootScope().WithDevice( + "/job:localhost/replica:0/task:0/device:CPU:0"); + + const Tensor kThree = test::AsScalar(3.0); + auto fdef = tensorflow::FunctionDefHelper::Create( + "Pow3", {"x: float"}, {"y: float"}, {}, + {{{"three"}, "Const", {}, {{"dtype", DT_FLOAT}, {"value", kThree}}}, + {{"pow3"}, "Pow", {"x", "three:output:0"}, {{"T", DT_FLOAT}}}}, + {{"y", "pow3:z:0"}}); + + tensorflow::FunctionDefLibrary fdef_lib; + *fdef_lib.add_function() = fdef; + TF_ASSERT_OK(scope.graph()->AddFunctionLibrary(fdef_lib)); + + Output a = ops::Const(scope.WithOpName("a"), 2.0, {1, 1}); + + std::vector inputs = {a}; + std::vector output_dtypes = { + fdef.signature().output_arg(0).type()}; + tensorflow::NameAttrList func_attr; + func_attr.set_name(fdef.signature().name()); + auto pcall = ops::PartitionedCall(scope, inputs, output_dtypes, func_attr); + Output b = pcall.output.front(); + + Output c = ops::Identity(scope.WithOpName("c"), b); + + TF_ASSERT_OK(scope.ToGraphDef(&graphdef)); + } + + // Ensure initial min_consumer is < 12 for testing. + graphdef.mutable_versions()->set_min_consumer(0); + + TF_ASSERT_OK_AND_ASSIGN( + auto fallback_state, + tensorflow::tfrt_stub::FallbackState::Create({}, fdef_lib)); + + TfrtGraphExecutionState::Options options; + // This triggers the early return in OptimizeGraph if + // ProtoMemoryOptimizationsEnabled() is true. + options.run_placer_grappler_on_functions = false; + TF_ASSERT_OK_AND_ASSIGN( + auto graph_execution_state, + TfrtGraphExecutionState::Create(options, graphdef, *fallback_state)); + + tensorflow::GraphImportConfig graph_import_config; + graph_import_config.prune_unused_nodes = true; + graph_import_config.enable_shape_inference = false; + tensorflow::ArrayInfo array_info; + array_info.imported_dtype = DT_FLOAT; + array_info.shape.set_unknown_rank(true); + graph_import_config.inputs["a"] = array_info; + graph_import_config.outputs = {"c"}; + + TF_ASSERT_OK_AND_ASSIGN( + auto optimized_graph, + graph_execution_state->CreateOptimizedGraph(graph_import_config)); + + // If we bypassed AddFunctionLibrary, this might still be 0 if Grappler didn't + // bump it. We want to verify if it is bumped to 12. CURRENT BEHAVIOR: It is + // NOT bumped because we bypass AddFunctionLibrary via move assignment. + EXPECT_LT(optimized_graph.graph->versions().min_consumer(), 12); + unsetenv("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS"); +} + TEST_F(OptimizeGraphTest, OptimizeFunctionsUsedByFunctionNodes) { GraphDef graphdef; tensorflow::FunctionDefLibrary fdef_lib; From 6e10ee8482f7cf44ac6be98c2cd70433906a6570 Mon Sep 17 00:00:00 2001 From: r00pd3v Date: Fri, 18 Sep 2026 09:49:07 -0700 Subject: [PATCH 31/32] PR #48513: Fix uint64 accumulation overflow in IFRT ArrayMemRegion byte-stride span Imported from GitHub PR https://github.com/openxla/xla/pull/48513 #ai-generated I found this while reviewing the recently landed #46604 byte-stride overflow fix and verified the remaining accumulation wrap against the current XLA tree with an in-tree Bazel regression test. `FromZerothElementPointer` already rejects a `stride * (dim - 1)` product that overflows `int64`. Adding those products into a `uint64_t` without an overflow check still wraps: ```text shape = [2, 2, 2] strides = [INT64_MAX, INT64_MAX, 2] dtype = u8 accepted region size = 1 ``` A second case, `strides = [INT64_MAX, INT64_MAX, 1]`, accumulates to `UINT64_MAX` and only overflows on the final `+ dtype-size` add. This change checks `__builtin_add_overflow` on the running sum and on the final add. Tests: `ByteStrideAccumOverflowUint64`, `ByteStrideFinalSizeOverflowUint64`, plus compact `u8` `[2,2,2]` success cases. Compact-layout `byte_size * shape.num_elements()` is not changed (#48286). ```text bazel test //xla/python/ifrt_proxy/common:array_util_test ``` **PASSED**, 42/42, including `ByteStrideAccumOverflowUint64` and `ByteStrideFinalSizeOverflowUint64`. Copybara import of the project: -- 4116a8f9c0eda07ba8b6d1ffe36ba22bc9541045 by roopd3v : Fix uint64 accumulation overflow in IFRT ArrayMemRegion byte-stride span. The per-term stride*(dim-1) multiply is already checked. Adding those products into a uint64_t without an overflow check still wraps, so a client-supplied stride vector such as [INT64_MAX, INT64_MAX, 2] with shape [2, 2, 2] is accepted as a 1-byte region. Check the running sum and the final + dtype-size add. Tests cover both the accumulation wrap and the final add (INT64_MAX + INT64_MAX + 1 == UINT64_MAX). Merging this change closes #48513 PiperOrigin-RevId: 983915659 --- .../python/ifrt_proxy/common/array_util.cc | 21 +++++++++++++++++-- .../ifrt_proxy/common/array_util_test.cc | 21 +++++++++++++++++-- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc index 1bae9d8549fe88..b84b969d07a52a 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc @@ -131,10 +131,27 @@ absl::StatusOr ArrayMemRegion::FromZerothElementPointer( "byte_stride[", i, "] * (dim[", i, "] - 1) overflows: stride=", stride, " dim=", shape.dims()[i])); } - last_element_byte_offset += static_cast(product); + // Each product is independently in int64 range (checked above), but the + // uint64 accumulation of those products can still wrap. Example: + // shape=[2,2,2], strides=[INT64_MAX, INT64_MAX, 2] -> each mul succeeds + // and the running sum is 2^64 == 0. + const uint64_t term = static_cast(product); + if (__builtin_add_overflow(last_element_byte_offset, term, + &last_element_byte_offset)) { + return absl::InvalidArgumentError(absl::StrCat( + "byte-stride span overflows uint64 while accumulating byte_stride[", + i, "] * (dim[", i, "] - 1): stride=", stride, + " dim=", shape.dims()[i])); + } } } - return ArrayMemRegion(mem_region_start, last_element_byte_offset + byte_size); + uint64_t total_size; + if (__builtin_add_overflow(last_element_byte_offset, + static_cast(byte_size), &total_size)) { + return absl::InvalidArgumentError( + "byte-stride span overflows uint64: last-element offset + dtype size"); + } + return ArrayMemRegion(mem_region_start, total_size); } absl::StatusOr ArrayMemRegion::FromMinimalMemRegion( diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc index 011991663a4e90..94ad8a658b95cc 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc @@ -41,6 +41,7 @@ using ::testing::TestWithParam; constexpr DType::Kind kF64 = DType::Kind::kF64; constexpr DType::Kind kS32 = DType::Kind::kS32; +constexpr DType::Kind kU8 = DType::Kind::kU8; constexpr DType::Kind kString = DType::Kind::kString; using Strides = std::vector; @@ -97,7 +98,10 @@ INSTANTIATE_TEST_SUITE_P( // Non-compact strides. TC{"NonCompactSingleDimension", kS32, {5}, Strides({16}), 68}, TC{"NonCompactDim0", kS32, {4, 3, 5}, Strides({120, 20, 4}), 420}, - TC{"PaddedElements", kS32, {4, 3, 5}, Strides({120, 40, 8}), 476}), + TC{"PaddedElements", kS32, {4, 3, 5}, Strides({120, 40, 8}), 476}, + // Products individually fit in int64; the accumulated span also fits. + TC{"U8ThreeDimsDefault", kU8, {2, 2, 2}, std::nullopt}, + TC{"U8ThreeDimsCompactStrides", kU8, {2, 2, 2}, Strides({4, 2, 1})}), testing::PrintToStringParamName()); TEST_P(ArrayMemRegionSuccess, TestCase) { const TC tc = GetParam(); @@ -148,7 +152,20 @@ INSTANTIATE_TEST_SUITE_P( TC{"ByteStrideOverflowMultiDim", kS32, {5, 3}, - Strides({INT64_C(4611686018427387904), 4})}), + Strides({INT64_C(4611686018427387904), 4})}, + // Each product stride*(dim-1) fits in int64 (the #46604 mul check + // succeeds) but the uint64 accumulation wraps: + // (INT64_MAX) + (INT64_MAX) + 2 == 2^64 == 0. + TC{"ByteStrideAccumOverflowUint64", + kU8, + {2, 2, 2}, + Strides({INT64_MAX, INT64_MAX, 2})}, + // Accumulation itself fits (INT64_MAX + INT64_MAX + 1 == UINT64_MAX); + // the final last-element-offset + dtype-size add is what overflows. + TC{"ByteStrideFinalSizeOverflowUint64", + kU8, + {2, 2, 2}, + Strides({INT64_MAX, INT64_MAX, 1})}), testing::PrintToStringParamName()); TEST_P(ArrayMemRegionFailure, TestCase) { const TC tc = GetParam(); From b131c1c386eb1fcde9e1738a20de59bd129abf67 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 18 Sep 2026 10:19:19 -0700 Subject: [PATCH 32/32] During `BufferAssignment::ToProto`, `BufferAllocation::ToProto` previously constructed `BufferAllocationProto::Assigned` protobuf objects in hash-map order and then sorted the repeated protobuf field by moving/swapping protobuf messages. In addition, `BufferAssignment::ToProto` constructed temporary stack `LogicalBufferProto` and `LogicalBufferProto::Location` messages for every logical buffer and alias before swapping them into the parent message. PiperOrigin-RevId: 983932289 --- .../xla/xla/service/buffer_assignment.cc | 44 ++++++++++++++----- .../xla/xla/service/buffer_assignment_test.cc | 41 +++++++++++++++++ 2 files changed, 73 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc index f9ee23be13925c..b8997931708ff2 100644 --- a/third_party/xla/xla/service/buffer_assignment.cc +++ b/third_party/xla/xla/service/buffer_assignment.cc @@ -680,7 +680,19 @@ BufferAllocationProto BufferAllocation::ToProto() const { } proto.set_is_constant(is_constant_); proto.set_maybe_live_out(maybe_live_out_); - for (const auto& buffer_offset_size : assigned_buffers_) { + using Entry = std::pair; + std::vector sorted; + sorted.reserve(assigned_buffers_.size()); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) + for (const auto& kv : assigned_buffers_) { + sorted.push_back(&kv); + } + absl::c_sort(sorted, [](const Entry* a, const Entry* b) { + return a->first->id() < b->first->id(); + }); + proto.mutable_assigned()->Reserve(sorted.size()); + for (const Entry* ptr : sorted) { + const auto& buffer_offset_size = *ptr; BufferAllocationProto::Assigned* proto_assigned = proto.add_assigned(); proto_assigned->set_logical_buffer_id(buffer_offset_size.first->id()); proto_assigned->set_offset(buffer_offset_size.second.offset); @@ -688,12 +700,6 @@ BufferAllocationProto BufferAllocation::ToProto() const { proto_assigned->set_element_type( buffer_offset_size.first->shape().element_type()); } - absl::c_sort(*proto.mutable_assigned(), - [](const BufferAllocationProto::Assigned& assign1, - const BufferAllocationProto::Assigned& assign2) { - return assign1.logical_buffer_id() < - assign2.logical_buffer_id(); - }); return proto; } @@ -1592,11 +1598,21 @@ void BufferAssignment::ToProto(BufferAssignmentProto* proto) const { // because we need to do the HasAllocation check for each buffer. Otherwise // the buffer_size_ call might fail for some backends. const HloDataflowAnalysis& dataflow = this->dataflow_analysis(); + proto->mutable_logical_buffers()->Reserve(dataflow.values().size()); for (BufferValue::Id id = 0; id < dataflow.values().size(); id++) { auto& value = dataflow.values().at(id); if (HasAllocation(*value)) { - LogicalBufferProto proto_buffer = value->ToProto(buffer_size_); - proto->add_logical_buffers()->Swap(&proto_buffer); + LogicalBufferProto* pb = proto->add_logical_buffers(); + pb->set_id(value->id()); + pb->set_size(buffer_size_(*value)); + LogicalBufferProto::Location* loc = pb->mutable_defined_at(); + loc->set_instruction_id(value->instruction()->unique_id()); + for (int64_t i : value->index()) { + loc->add_shape_index(i); + } + if (value->has_color()) { + pb->set_color(value->color()); + } // Fill buffer aliases. for (const HloValue* alias : @@ -1607,13 +1623,17 @@ void BufferAssignment::ToProto(BufferAssignmentProto* proto) const { } BufferAssignmentProto::BufferAlias* proto_alias = proto->add_buffer_aliases(); - LogicalBufferProto::Location proto_alias_location = - BufferValue::ToLocationProto(*alias->instruction(), alias->index()); proto_alias->set_source_buffer_id(value->id()); - proto_alias->mutable_location()->Swap(&proto_alias_location); + LogicalBufferProto::Location* aloc = proto_alias->mutable_location(); + aloc->set_instruction_id(alias->instruction()->unique_id()); + for (int64_t i : alias->index()) { + aloc->add_shape_index(i); + } } } } + proto->mutable_buffer_allocations()->Reserve(Allocations().size()); + proto->mutable_peak_buffers()->Reserve(Allocations().size()); for (const BufferAllocation& allocation : Allocations()) { BufferAllocationProto proto_allocation = allocation.ToProto(); proto->add_buffer_allocations()->Swap(&proto_allocation); diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc index 9a822813659134..68b1f8b58934c4 100644 --- a/third_party/xla/xla/service/buffer_assignment_test.cc +++ b/third_party/xla/xla/service/buffer_assignment_test.cc @@ -5689,5 +5689,46 @@ TEST_F(BufferAssignmentTest, FromProtoRejectsNegativeSize) { EXPECT_THAT(result.status().message(), ::testing::HasSubstr("negative")); } +TEST_F(BufferAssignmentTest, + ToProtoSortsAssignedBuffersAndPopulatesNestedTupleAliases) { + const char* const hlo_text = R"( + HloModule test + cond { + p = (f32[4]{0}, f32[8]{0}) parameter(0) + ROOT cond_res = pred[] constant(true) + } + body { + p = (f32[4]{0}, f32[8]{0}) parameter(0) + gte0 = f32[4]{0} get-tuple-element(p), index=0 + gte1 = f32[8]{0} get-tuple-element(p), index=1 + neg0 = f32[4]{0} negate(gte0) + ROOT out = (f32[4]{0}, f32[8]{0}) tuple(neg0, gte1) + } + ENTRY e { + p0 = (f32[4]{0}, f32[8]{0}) parameter(0) + ROOT loop = (f32[4]{0}, f32[8]{0}) while(p0), condition=cond, body=body + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + std::unique_ptr buffers = RunBufferAssignment(module.get()); + BufferAssignmentProto proto = buffers->ToProto(); + + EXPECT_FALSE(proto.logical_buffers().empty()); + EXPECT_FALSE(proto.buffer_aliases().empty()); + for (const BufferAllocationProto& alloc : proto.buffer_allocations()) { + for (int i = 1; i < alloc.assigned_size(); ++i) { + EXPECT_LT(alloc.assigned(i - 1).logical_buffer_id(), + alloc.assigned(i).logical_buffer_id()); + } + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr roundtrip, + BufferAssignment::FromProto(proto, module.get(), &BufferSizeBytes, + &alias_info_)); + EXPECT_EQ(roundtrip->Allocations().size(), buffers->Allocations().size()); +} + } // namespace } // namespace xla