From cd9c7f3938854af67ed647b846736c2978b78793 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 24 Aug 2026 14:28:30 -0700 Subject: [PATCH 01/27] Reject a negative num_elements in BarrierTakeMany --- tensorflow/core/kernels/barrier_ops.cc | 5 +++++ .../python/kernel_tests/data_structures/BUILD | 1 + .../data_structures/barrier_ops_test.py | 20 +++++++++++++++++++ 3 files changed, 26 insertions(+) diff --git a/tensorflow/core/kernels/barrier_ops.cc b/tensorflow/core/kernels/barrier_ops.cc index 001a227e781952..c0f6294b281208 100644 --- a/tensorflow/core/kernels/barrier_ops.cc +++ b/tensorflow/core/kernels/barrier_ops.cc @@ -587,6 +587,11 @@ class TakeManyOp : public BarrierOpKernel { ctx, TensorShapeUtils::IsScalar(Tnum_elements->shape()), absl::InvalidArgumentError("num_elements must be a scalar."), callback); const int32_t num_elements = Tnum_elements->scalar()(); + OP_REQUIRES_ASYNC( + ctx, num_elements >= 0, + absl::InvalidArgumentError(absl::StrCat( + "BarrierTakeMany requested ", num_elements, " < 0 elements")), + callback); DataTypeVector expected_inputs = {DT_STRING_REF, DT_INT32}; // The first output is the insertion index, the second output is the key. diff --git a/tensorflow/python/kernel_tests/data_structures/BUILD b/tensorflow/python/kernel_tests/data_structures/BUILD index 2afd6d35d8ec58..24a7a584d550c4 100644 --- a/tensorflow/python/kernel_tests/data_structures/BUILD +++ b/tensorflow/python/kernel_tests/data_structures/BUILD @@ -37,6 +37,7 @@ tf_py_strict_test( "//tensorflow/python/framework:tensor", "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:data_flow_ops", + "//tensorflow/python/ops:data_flow_ops_gen", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", ], diff --git a/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py b/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py index ec10af3f818ae9..452555c116e8e7 100644 --- a/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py +++ b/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import tensor from tensorflow.python.framework import test_util from tensorflow.python.ops import data_flow_ops +from tensorflow.python.ops import gen_data_flow_ops from tensorflow.python.platform import test @@ -102,6 +103,25 @@ def testInsertManyEmptyTensorUnknown(self): ".*Tensors with no elements are not supported.*"): insert_0_op.run() + @test_util.run_deprecated_v1 + def testTakeManyNegativeNumElements(self): + # Test case for GitHub issue 112742. A negative num_elements reached the + # underlying queue, which built an output shape with a negative dimension + # and aborted the process. The raw op is used here because Barrier.take_many + # rejects a negative count while building the graph, before the kernel runs. + with self.cached_session() as sess: + b = data_flow_ops.Barrier( + (dtypes.float32, dtypes.float32), shapes=((), ()), name="B") + keys = [b"a", b"b"] + insert_0_op = b.insert_many(0, keys, [10.0, 20.0]) + insert_1_op = b.insert_many(1, keys, [100.0, 200.0]) + sess.run([insert_0_op, insert_1_op]) + take_t = gen_data_flow_ops.barrier_take_many( + b.barrier_ref, -1, (dtypes.float32, dtypes.float32), + allow_small_batch=False) + with self.assertRaisesOpError("requested -1 < 0 elements"): + sess.run(take_t) + @test_util.run_deprecated_v1 def testTakeMany(self): with self.cached_session() as sess: From 2fa1860ebcddfbfadcbd6a5abae7e883352b45b9 Mon Sep 17 00:00:00 2001 From: Aditya Gaurav Date: Tue, 15 Sep 2026 21:46:26 +0530 Subject: [PATCH 02/27] fix(kernels): reject INT32_MAX index in TensorListScatter resize TensorListScatter resized with signed `highest_index + 1`, which overflows when the index is INT32_MAX. Apply the same guard already used by TensorListScatterIntoExistingList and size the backing vector with size_t. --- tensorflow/core/kernels/list_kernels.h | 20 +++++++++++++++++-- .../data_structures/list_ops_test.py | 15 ++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/list_kernels.h b/tensorflow/core/kernels/list_kernels.h index 4ccbc9e514fa08..78517dfb6ff9fc 100644 --- a/tensorflow/core/kernels/list_kernels.h +++ b/tensorflow/core/kernels/list_kernels.h @@ -1012,8 +1012,24 @@ class TensorListScatter : public OpKernel { highest_index = i; } } - output_list.tensors().resize(std::max(highest_index + 1, num_elements), - Tensor(DT_INVALID)); + // TensorListLength returns int32, so highest_index + 1 must be + // representable as int32 before the list is resized. Same guard as + // TensorListScatterIntoExistingList (#list_kernels.h). + OP_REQUIRES( + c, highest_index < std::numeric_limits::max(), + absl::InvalidArgumentError(absl::StrCat( + "TensorListScatter: index ", highest_index, + " would produce a list length that is not representable as " + "int32"))); + // Compute the required size as size_t: highest_index == INT32_MAX would + // overflow the signed `highest_index + 1` (undefined behavior). + const size_t needed_from_indices = + highest_index < 0 ? 0 : static_cast(highest_index) + 1; + const size_t needed = + num_elements > 0 + ? std::max(needed_from_indices, static_cast(num_elements)) + : needed_from_indices; + output_list.tensors().resize(needed, Tensor(DT_INVALID)); } OP_REQUIRES_OK(c, diff --git a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py index 212b0e3b0b75ae..448d8674deb36b 100644 --- a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py +++ b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py @@ -609,6 +609,21 @@ def testScatterIntoExistingListRejectsUnrepresentableLength(self): ) ) + def testScatterRejectsUnrepresentableLength(self): + # Mirror TensorListScatterIntoExistingList: index INT32_MAX must not reach + # the signed `highest_index + 1` resize (undefined behavior). + c0 = constant_op.constant([1.0]) + with self.assertRaisesRegex( + errors.InvalidArgumentError, + "list length that is not representable as int32", + ): + l = list_ops.tensor_list_scatter( + c0, + constant_op.constant([np.iinfo(np.int32).max], dtype=dtypes.int32), + element_shape=[], + ) + self.evaluate(l) + def testScatterFailsWithInvalidNumElements(self): c0 = constant_op.constant([1.0, 2.0]) with self.assertRaisesRegex( From 8a5df004ce2a431bf78e3476e360591bcf78c48b Mon Sep 17 00:00:00 2001 From: Riccardo Squarcialupi <44367261+RiccardoSquarcialupi@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:27:21 +0200 Subject: [PATCH 03/27] Fix ImageProjectiveTransform gradient to be a true adjoint The registered gradient for ImageProjectiveTransformV2/V3 inverted the transform matrix and re-ran the forward resampling op on the incoming gradient. That is only a correct adjoint when the forward coordinate map is a bijection on the pixel grid -- never true once fill_mode clamps or folds an out-of-range coordinate, or the transform downscales. In both cases it silently duplicates or drops gradient mass at the affected pixels. Replace it with a real scatter-based adjoint: for each output pixel, scatter-add its incoming gradient onto the same input pixel(s) the forward kernel actually read, with the same interpolation weights. The coordinate mapping (_map_coordinate) mirrors MapCoordinate in tensorflow/core/kernels/image/image_ops.h exactly, including its use of C++ truncation-toward-zero and round-half-away-from-zero, which TensorFlow's own division/round otherwise disagree with. Fixes #127241, where tf.keras.ops.image.affine_transform with interpolation="nearest" returned gradient 30.0 for a case where the correct value is 17.0. Verified against a real build (TF 2.21.0, matching the report): - The issue's exact reproducer now returns 17.0. - Adjoint identity == holds (rel. error <= 1.9e-6) across all 4 fill_modes x 2 interpolations x 6 transform types x batch/broadcast x float32/float64, 384 configs. - _map_coordinate matches the C++ kernel's actual pixel selection on 4420 probed pixels across fill_modes and transforms. - gradient_checker_v2 agrees with finite differences to <8e-5. - pylint 10.00/10 under the repo's pylintrc (Python 3.9, its CI config). Adds ImageProjectiveTransformOpTestBase covering all of the above as regression tests; every one of them fails against the old gradient and passes against the new one. --- tensorflow/python/ops/image_grad_test.py | 1 + tensorflow/python/ops/image_grad_test_base.py | 107 ++++++++ tensorflow/python/ops/image_ops.py | 258 ++++++++++++++---- 3 files changed, 309 insertions(+), 57 deletions(-) diff --git a/tensorflow/python/ops/image_grad_test.py b/tensorflow/python/ops/image_grad_test.py index 28fd605e628455..2485a09e7aae3e 100644 --- a/tensorflow/python/ops/image_grad_test.py +++ b/tensorflow/python/ops/image_grad_test.py @@ -24,6 +24,7 @@ CropAndResizeOpTest = test_base.CropAndResizeOpTestBase RGBToHSVOpTest = test_base.RGBToHSVOpTestBase AdjustContrastOpTest = test_base.AdjustContrastOpTestBase +ImageProjectiveTransformOpTest = test_base.ImageProjectiveTransformOpTestBase if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py index b412b3a9f9a5fc..d7d1229c750115 100644 --- a/tensorflow/python/ops/image_grad_test_base.py +++ b/tensorflow/python/ops/image_grad_test_base.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops from tensorflow.python.ops import array_ops_stack from tensorflow.python.ops import gen_image_ops from tensorflow.python.ops import gradient_checker_v2 @@ -677,5 +678,111 @@ def f(image_tensor, factor_tensor): self.assertLess(max_error, 1e-4) +_TRANSFORM_FILL_MODES = ['CONSTANT', 'NEAREST', 'REFLECT', 'WRAP'] +_TRANSFORM_INTERPOLATIONS = ['NEAREST', 'BILINEAR'] +_TRANSFORM_FILL_MODE_X_INTERPOLATION = [ + (fill_mode, interpolation) + for fill_mode in _TRANSFORM_FILL_MODES + for interpolation in _TRANSFORM_INTERPOLATIONS +] + + +class ImageProjectiveTransformOpTestBase(test.TestCase, parameterized.TestCase): + """Tests the gradient of ImageProjectiveTransformV2/V3. + + The forward op is linear in `images`, so for any `g` shaped like its + output, ` == ` must hold exactly + (up to float roundoff). That identity is what actually pins down that the + registered gradient is a true adjoint of the resampling, rather than + another resample through the inverted transform -- the latter silently + duplicates or drops gradient mass wherever `fill_mode` clamps or folds an + out-of-range coordinate, or the transform downscales, while still passing + a looser numerical gradient check on inputs that don't hit those cases. + """ + + def _transform(self, images, transforms, output_shape, interpolation, + fill_mode): + return gen_image_ops.image_projective_transform_v3( + images=images, + transforms=transforms, + output_shape=output_shape, + interpolation=interpolation, + fill_mode=fill_mode, + fill_value=0.0) + + @parameterized.parameters(*_TRANSFORM_FILL_MODE_X_INTERPOLATION) + def testAdjointIdentity(self, fill_mode, interpolation): + rng = np.random.default_rng(0) + batch, in_h, in_w, out_h, out_w, channels = 2, 5, 6, 4, 5, 2 + images_np = rng.standard_normal((batch, in_h, in_w, channels)) + grad_out_np = rng.standard_normal((batch, out_h, out_w, channels)) + # A representative non-trivial transform: rotation, translation, shear + # and a touch of perspective, so every coefficient participates. + transforms = constant_op.constant( + [[0.9, -0.2, 1.0, 0.2, 0.9, -0.5, 0.002, -0.001]], + dtype=dtypes.float32) + + with self.cached_session(): + images = constant_op.constant(images_np, dtype=dtypes.float64) + grad_out = constant_op.constant(grad_out_np, dtype=dtypes.float64) + with backprop.GradientTape() as tape: + tape.watch(images) + out = self._transform(images, transforms, [out_h, out_w], + interpolation, fill_mode) + lhs = math_ops.reduce_sum(out * grad_out) + grad_images = tape.gradient(lhs, images) + rhs = math_ops.reduce_sum(images * grad_images) + self.assertAllClose( + self.evaluate(lhs), self.evaluate(rhs), rtol=1e-6, atol=1e-6) + + @parameterized.parameters(*_TRANSFORM_FILL_MODE_X_INTERPOLATION) + def testGradientChecker(self, fill_mode, interpolation): + x = np.linspace(-1.0, 1.0, num=2 * 4 * 5 * 2).reshape([2, 4, 5, 2]) + x = x.astype(np.float64) + transforms = constant_op.constant( + [[1, 0.05, 0.5, -0.05, 1, -0.3, 0.001, -0.001]], dtype=dtypes.float32) + + def f(images): + return self._transform(images, transforms, [4, 5], interpolation, + fill_mode) + + with self.cached_session(): + err = gradient_checker_v2.max_error( + *gradient_checker_v2.compute_gradient(f, [constant_op.constant(x)])) + self.assertLess(err, 1e-4) + + def testNearestFillModeRegression(self): + """github.com/tensorflow/tensorflow/issues/127241. + + With `interpolation="nearest"` and a fixed transform, the selected + scalar output is linear in the image values by construction, so its + gradient has a closed form this test checks directly rather than via + finite differences: the transform maps output column 2 to input column + 3, out of range, clamped by `fill_mode="nearest"` back to column 2 -- + so `d(out[0,0,0])/dt = d(image[0,1,0])/dt` and + `d(out[1,0,0])/dt = d(image[1,1,0])/dt`, giving 1*2 + 3*5 = 17. The + previous (incorrect) gradient resampled through the inverted + transform and, hitting the same clamp in reverse, double-counted + column 0 to get 30 instead. + """ + def target(t): + image = array_ops.reshape( + array_ops_stack.stack([(i + 1) * t - i for i in range(6)]), + [1, 2, 3, 1]) + transforms = constant_op.constant( + [[1, 0, 1, 0, 1, 0, 0, 0]], dtype=dtypes.float32) + y = self._transform(image, transforms, [2, 3], 'NEAREST', 'NEAREST') + return y[0, 0, 0, 0] + 3 * y[0, 1, 0, 0] + + x = constant_op.constant(-50.0, dtype=dtypes.float64) + with self.cached_session(): + with backprop.GradientTape() as tape: + tape.watch(x) + y = target(x) + actual = tape.gradient(y, x) + self.assertAllClose(self.evaluate(y), -863.0) + self.assertAllClose(self.evaluate(actual), 17.0) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/ops/image_ops.py b/tensorflow/python/ops/image_ops.py index d5c938f4ad44af..8420645388f744 100644 --- a/tensorflow/python/ops/image_ops.py +++ b/tensorflow/python/ops/image_ops.py @@ -151,8 +151,10 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gen_image_ops -from tensorflow.python.ops import linalg_ops +from tensorflow.python.ops import array_ops_stack +from tensorflow.python.ops import clip_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.ops import math_ops # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.python.ops.gen_image_ops import * @@ -236,70 +238,212 @@ def matrices_to_flat_transforms(transform_matrices): return transforms[:, :8] -@ops.RegisterGradient("ImageProjectiveTransformV2") -def _image_projective_transform_grad(op, grad): - """Computes the gradient for ImageProjectiveTransform.""" - images = op.inputs[0] - transforms = op.inputs[1] - interpolation = op.get_attr("interpolation") - fill_mode = op.get_attr("fill_mode") +def _trunc_div(a, b): + """Integer division that truncates toward zero, like C++'s `/` on ints. + + `a / b` followed by a cast matches `static_cast(a / b)` in + `MapCoordinate` (tensorflow/core/kernels/image/image_ops.h); Python and + TensorFlow's own integer division instead floor, which disagrees for + negative inputs. + """ + return math_ops.cast(math_ops.cast(a / b, dtypes.int32), dtypes.float32) - image_or_images = ops.convert_to_tensor(images, name="images") - transform_or_transforms = ops.convert_to_tensor( - transforms, name="transforms", dtype=dtypes.float32) - if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES: - raise TypeError("Invalid dtype %s." % image_or_images.dtype) - if len(transform_or_transforms.get_shape()) == 1: - transforms = transform_or_transforms[None] - elif len(transform_or_transforms.get_shape()) == 2: - transforms = transform_or_transforms +def _map_coordinate(coord, length, fill_mode): + """Maps an output-space coordinate to an input-space one, per `fill_mode`. + + Mirrors `MapCoordinate` in + tensorflow/core/kernels/image/image_ops.h exactly, so that the adjoint + below reads from (in the forward op, sense) or writes to (here) the same + pixel the forward kernel does. + + Args: + coord: float32 tensor, arbitrary shape. + length: scalar int32 tensor, the input extent along this axis. + fill_mode: Python str, one of "CONSTANT", "NEAREST", "REFLECT", "WRAP". + + Returns: + float32 tensor, same shape as `coord`. + """ + if fill_mode == "CONSTANT": + return coord + + length_f = math_ops.cast(length, dtypes.float32) + last = length_f - 1.0 + + if fill_mode == "NEAREST": + return clip_ops.clip_by_value(coord, 0.0, last) + + if fill_mode == "WRAP": + below = coord + length_f * (_trunc_div(-coord, last) + 1.0) + above = coord - length_f * _trunc_div(coord, last) + elif fill_mode == "REFLECT": + sz2 = 2.0 * length_f + folded_low = sz2 * _trunc_div(-coord, sz2) + coord + below = array_ops.where_v2(folded_low < -length_f, folded_low + sz2, + -folded_low - 1.0) + folded_high = coord - sz2 * _trunc_div(coord, sz2) + above = array_ops.where_v2(folded_high >= length_f, + sz2 - folded_high - 1.0, folded_high) else: - raise TypeError("Transforms should have rank 1 or 2.") + raise ValueError("Unknown fill_mode %r" % (fill_mode,)) - # Invert transformations - transforms = flat_transforms_to_matrices(transforms=transforms) - inverse = linalg_ops.matrix_inverse(transforms) - transforms = matrices_to_flat_transforms(inverse) - output = gen_image_ops.image_projective_transform_v2( - images=grad, - transforms=transforms, - output_shape=array_ops.shape(image_or_images)[1:3], - interpolation=interpolation, - fill_mode=fill_mode) - return [output, None, None] + mapped = array_ops.where_v2(coord < 0.0, below, + array_ops.where_v2(coord > last, above, coord)) + # `length <= 1` collapses every coordinate to 0, matching the kernel. + mapped = array_ops.where_v2(length <= 1, array_ops.zeros_like(coord), mapped) + return clip_ops.clip_by_value(mapped, 0.0, last) -@ops.RegisterGradient("ImageProjectiveTransformV3") -def _image_projective_transform_v3_grad(op, grad): - """Computes the gradient for ImageProjectiveTransform.""" - images = op.inputs[0] - transforms = op.inputs[1] - interpolation = op.get_attr("interpolation") - fill_mode = op.get_attr("fill_mode") +def _round_half_away_from_zero(x): + """Matches C++ `std::round`, used by the forward kernel's NEAREST tap. + + TensorFlow's own `round` breaks ties to even instead, which disagrees with + the kernel at exact `.5` boundaries. + """ + return math_ops.sign(x) * math_ops.floor(math_ops.abs(x) + 0.5) + + +def _image_projective_transform_grad_impl(images, transforms, grad, + interpolation, fill_mode): + """Shared adjoint for ImageProjectiveTransformV2/V3. - image_or_images = ops.convert_to_tensor(images, name="images") - transform_or_transforms = ops.convert_to_tensor( + The forward op resamples `images` at a coordinate that `transforms` maps + each output pixel to, optionally blending up to 4 input pixels + (BILINEAR) or reading exactly 1 (NEAREST). Its true gradient scatter-adds + each output pixel's incoming gradient onto the same input pixel(s), with + the same weights. (The previous implementation instead re-ran the forward + *resampling* op on `grad` through the inverted transform, which is only + correct when the forward coordinate map is a bijection on the pixel grid + -- never true once `fill_mode` clamps or folds an out-of-range coordinate, + or the transform downscales.) + + Args: + images: the forward op's `images` input; only its shape/dtype are used. + transforms: the forward op's `transforms` input, rank 1 `[8]` or rank 2 + `[N, 8]` with `N` equal to 1 (broadcast) or to the batch size. + grad: gradient w.r.t. the forward op's output, `[B, out_H, out_W, C]`. + interpolation: Python str, "NEAREST" or "BILINEAR". + fill_mode: Python str, "CONSTANT", "NEAREST", "REFLECT", or "WRAP". + + Returns: + Gradient w.r.t. `images`, shape `[B, in_H, in_W, C]`, dtype `grad.dtype`. + """ + if images.dtype.base_dtype not in _IMAGE_DTYPES: + raise TypeError("Invalid dtype %s." % images.dtype) + + transforms = ops.convert_to_tensor( transforms, name="transforms", dtype=dtypes.float32) + if transforms.shape.ndims == 1: + transforms = transforms[None] + elif transforms.shape.ndims != 2: + raise TypeError("Transforms should have rank 1 or 2.") - if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES: - raise TypeError("Invalid dtype %s." % image_or_images.dtype) - if len(transform_or_transforms.get_shape()) == 1: - transforms = transform_or_transforms[None] - elif len(transform_or_transforms.get_shape()) == 2: - transforms = transform_or_transforms + image_shape = array_ops.shape(images) + batch, in_h, in_w, channels = (image_shape[0], image_shape[1], + image_shape[2], image_shape[3]) + grad_shape = array_ops.shape(grad) + out_h, out_w = grad_shape[1], grad_shape[2] + + # `transforms` may carry one row (broadcast to the whole batch) or one row + # per batch element, decided by its actual leading dimension at runtime -- + # gather a `[batch, 8]` view either way rather than branching on shape. + num_transforms = array_ops.shape(transforms)[0] + transform_row = array_ops.where_v2( + math_ops.equal(num_transforms, 1), + array_ops.zeros([batch], dtype=dtypes.int32), math_ops.range(batch)) + t = array_ops.gather(transforms, transform_row, axis=0) # [batch, 8] + + oy, ox = array_ops.meshgrid( + math_ops.range(out_h), math_ops.range(out_w), indexing="ij") + oy = math_ops.cast(oy, dtypes.float32) # [out_H, out_W] + ox = math_ops.cast(ox, dtypes.float32) + + def _coef(i): + return t[:, i, None, None] # [batch, 1, 1], broadcasts against ox/oy. + + proj = _coef(6) * ox + _coef(7) * oy + 1.0 # [batch, out_H, out_W] + valid_proj = math_ops.not_equal(proj, 0.0) + safe_proj = array_ops.where_v2(valid_proj, proj, array_ops.ones_like(proj)) + input_x = (_coef(0) * ox + _coef(1) * oy + _coef(2)) / safe_proj + input_y = (_coef(3) * ox + _coef(4) * oy + _coef(5)) / safe_proj + + x = _map_coordinate(input_x, in_w, fill_mode) + y = _map_coordinate(input_y, in_h, fill_mode) + + batch_idx = math_ops.cast( + array_ops.reshape(math_ops.range(batch), [-1, 1, 1]), dtypes.float32) + batch_idx += array_ops.zeros_like(x) # -> [batch, out_H, out_W] + + grad = ops.convert_to_tensor(grad) + compute_dtype = grad.dtype.base_dtype # accumulate at the tape's precision + grad_flat = array_ops.reshape(grad, [-1, channels]) + + def _scatter_tap(iy, ix, weight): + """One interpolation tap's indices/updates, out-of-range taps zeroed.""" + in_bounds = math_ops.logical_and( + math_ops.logical_and(iy >= 0.0, iy < math_ops.cast(in_h, iy.dtype)), + math_ops.logical_and(ix >= 0.0, ix < math_ops.cast(in_w, ix.dtype))) + keep = math_ops.logical_and(in_bounds, valid_proj) + zero_i = array_ops.zeros_like(iy) + safe_iy = math_ops.cast(array_ops.where_v2(keep, iy, zero_i), dtypes.int32) + safe_ix = math_ops.cast(array_ops.where_v2(keep, ix, zero_i), dtypes.int32) + indices = array_ops_stack.stack( + [math_ops.cast(batch_idx, dtypes.int32), safe_iy, safe_ix], axis=-1) + weight = array_ops.where_v2(keep, weight, array_ops.zeros_like(weight)) + weight = math_ops.cast(array_ops.reshape(weight, [-1, 1]), compute_dtype) + return array_ops.reshape(indices, [-1, 3]), weight * grad_flat + + if interpolation == "NEAREST": + iy = _round_half_away_from_zero(y) + ix = _round_half_away_from_zero(x) + taps = [(iy, ix, array_ops.ones_like(x))] + elif interpolation == "BILINEAR": + y_floor, x_floor = math_ops.floor(y), math_ops.floor(x) + y_ceil, x_ceil = y_floor + 1.0, x_floor + 1.0 + wy_floor, wy_ceil = y_ceil - y, y - y_floor + wx_floor, wx_ceil = x_ceil - x, x - x_floor + taps = [ + (y_floor, x_floor, wy_floor * wx_floor), + (y_floor, x_ceil, wy_floor * wx_ceil), + (y_ceil, x_floor, wy_ceil * wx_floor), + (y_ceil, x_ceil, wy_ceil * wx_ceil), + ] else: - raise TypeError("Transforms should have rank 1 or 2.") + raise ValueError("Unknown interpolation %r" % (interpolation,)) + + all_indices, all_updates = [], [] + for iy, ix, weight in taps: + indices, updates = _scatter_tap(iy, ix, weight) + all_indices.append(indices) + all_updates.append(updates) + + zeros = array_ops.zeros( + array_ops_stack.stack([batch, in_h, in_w, channels]), dtype=compute_dtype) + return gen_array_ops.tensor_scatter_add( + zeros, array_ops.concat(all_indices, axis=0), + array_ops.concat(all_updates, axis=0)) - # Invert transformations - transforms = flat_transforms_to_matrices(transforms=transforms) - inverse = linalg_ops.matrix_inverse(transforms) - transforms = matrices_to_flat_transforms(inverse) - output = gen_image_ops.image_projective_transform_v3( - images=grad, - transforms=transforms, - output_shape=array_ops.shape(image_or_images)[1:3], - interpolation=interpolation, - fill_mode=fill_mode, - fill_value=0.0) + +@ops.RegisterGradient("ImageProjectiveTransformV2") +def _image_projective_transform_grad(op, grad): + """Computes the gradient for ImageProjectiveTransform.""" + output = _image_projective_transform_grad_impl( + images=op.inputs[0], + transforms=op.inputs[1], + grad=grad, + interpolation=op.get_attr("interpolation").decode(), + fill_mode=op.get_attr("fill_mode").decode()) + return [output, None, None] + + +@ops.RegisterGradient("ImageProjectiveTransformV3") +def _image_projective_transform_v3_grad(op, grad): + """Computes the gradient for ImageProjectiveTransform.""" + output = _image_projective_transform_grad_impl( + images=op.inputs[0], + transforms=op.inputs[1], + grad=grad, + interpolation=op.get_attr("interpolation").decode(), + fill_mode=op.get_attr("fill_mode").decode()) return [output, None, None, None] From 14c2b8f52bd546f8006a75fe6f1ee3e069a40ffe Mon Sep 17 00:00:00 2001 From: Riccardo Squarcialupi <44367261+RiccardoSquarcialupi@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:46:03 +0200 Subject: [PATCH 04/27] Address gemini-code-assist review - Guard _map_coordinate's WRAP/REFLECT division against length == 1 (a 1-pixel-wide/tall axis), which otherwise divides by zero. The masked-out result was already correct either way, but the intermediate inf/nan getting cast to int32 is platform-dependent undefined behavior, worth avoiding regardless of whether its value is later discarded. - Build batch_idx via array_ops.broadcast_to instead of a cast-to-float + zeros_like(x) + add round trip, and keep it int32 throughout, dropping the now-redundant cast where it's used. - Use array_ops.scatter_nd directly instead of zeros + tensor_scatter_add; scatter_nd sums duplicate indices the same way, verified directly against tf.scatter_nd, so this is behavior-preserving for the overlapping-tap case this PR exists to fix. Drops the now-unused gen_array_ops import. - Add testOnePixelWideInputIsFinite covering the length == 1 case the first fix addresses. Re-verified after applying: the issue's reproducer, the 384-config adjoint-identity sweep, the 4420-pixel MapCoordinate parity probe, and gradient_checker_v2 all still pass with identical results; pylint still 10.00/10; full ImageProjectiveTransformOpTest suite (29 tests, 12 new) passes. --- tensorflow/python/ops/image_grad_test_base.py | 27 +++++++++++++++ tensorflow/python/ops/image_ops.py | 33 ++++++++++--------- 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py index d7d1229c750115..257e2c87a6eaf7 100644 --- a/tensorflow/python/ops/image_grad_test_base.py +++ b/tensorflow/python/ops/image_grad_test_base.py @@ -783,6 +783,33 @@ def target(t): self.assertAllClose(self.evaluate(y), -863.0) self.assertAllClose(self.evaluate(actual), 17.0) + @parameterized.parameters(*[(fill_mode, shape) + for fill_mode in _TRANSFORM_FILL_MODES + for shape in [(1, 1, 1, 1), (1, 1, 5, 1), + (1, 5, 1, 1)]]) + def testOnePixelWideInputIsFinite(self, fill_mode, shape): + """Regression test for a `length == 1` axis dividing by zero. + + `_map_coordinate`'s WRAP and REFLECT branches divide by `length - 1`, + which is 0 for a `length == 1` axis. The masked-out result is correct + either way, but the *intermediate* division must not compute an inf/nan + that then gets cast to int32, which is platform-dependent undefined + behavior even when its result is later discarded. + """ + images = constant_op.constant( + np.random.default_rng(0).standard_normal(shape), dtype=dtypes.float32) + transforms = constant_op.constant( + [[1, 0, 0.3, 0, 1, 0.2, 0, 0]], dtype=dtypes.float32) + + with self.cached_session(): + with backprop.GradientTape() as tape: + tape.watch(images) + out = self._transform(images, transforms, [shape[1], shape[2]], + 'BILINEAR', fill_mode) + loss = math_ops.reduce_sum(out) + grad = self.evaluate(tape.gradient(loss, images)) + self.assertTrue(np.all(np.isfinite(grad))) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/ops/image_ops.py b/tensorflow/python/ops/image_ops.py index 8420645388f744..7e6d44be954fbc 100644 --- a/tensorflow/python/ops/image_ops.py +++ b/tensorflow/python/ops/image_ops.py @@ -153,7 +153,6 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import array_ops_stack from tensorflow.python.ops import clip_ops -from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import math_ops # go/tf-wildcard-import # pylint: disable=wildcard-import @@ -274,15 +273,21 @@ def _map_coordinate(coord, length, fill_mode): if fill_mode == "NEAREST": return clip_ops.clip_by_value(coord, 0.0, last) + # `length <= 1` makes `last`/`sz2` 0, which would divide by zero below; + # the result is discarded by the `length <= 1` select further down, but + # an intermediate inf/nan cast to int32 is platform-dependent undefined + # behavior, so avoid it rather than merely discard it. + safe_last = array_ops.where_v2(length <= 1, 1.0, last) if fill_mode == "WRAP": - below = coord + length_f * (_trunc_div(-coord, last) + 1.0) - above = coord - length_f * _trunc_div(coord, last) + below = coord + length_f * (_trunc_div(-coord, safe_last) + 1.0) + above = coord - length_f * _trunc_div(coord, safe_last) elif fill_mode == "REFLECT": sz2 = 2.0 * length_f - folded_low = sz2 * _trunc_div(-coord, sz2) + coord + safe_sz2 = array_ops.where_v2(length <= 1, 2.0, sz2) + folded_low = sz2 * _trunc_div(-coord, safe_sz2) + coord below = array_ops.where_v2(folded_low < -length_f, folded_low + sz2, -folded_low - 1.0) - folded_high = coord - sz2 * _trunc_div(coord, sz2) + folded_high = coord - sz2 * _trunc_div(coord, safe_sz2) above = array_ops.where_v2(folded_high >= length_f, sz2 - folded_high - 1.0, folded_high) else: @@ -371,9 +376,9 @@ def _coef(i): x = _map_coordinate(input_x, in_w, fill_mode) y = _map_coordinate(input_y, in_h, fill_mode) - batch_idx = math_ops.cast( - array_ops.reshape(math_ops.range(batch), [-1, 1, 1]), dtypes.float32) - batch_idx += array_ops.zeros_like(x) # -> [batch, out_H, out_W] + batch_idx = array_ops.broadcast_to( + array_ops.reshape(math_ops.range(batch), [-1, 1, 1]), + array_ops.shape(x)) # int32, [batch, out_H, out_W] grad = ops.convert_to_tensor(grad) compute_dtype = grad.dtype.base_dtype # accumulate at the tape's precision @@ -388,8 +393,7 @@ def _scatter_tap(iy, ix, weight): zero_i = array_ops.zeros_like(iy) safe_iy = math_ops.cast(array_ops.where_v2(keep, iy, zero_i), dtypes.int32) safe_ix = math_ops.cast(array_ops.where_v2(keep, ix, zero_i), dtypes.int32) - indices = array_ops_stack.stack( - [math_ops.cast(batch_idx, dtypes.int32), safe_iy, safe_ix], axis=-1) + indices = array_ops_stack.stack([batch_idx, safe_iy, safe_ix], axis=-1) weight = array_ops.where_v2(keep, weight, array_ops.zeros_like(weight)) weight = math_ops.cast(array_ops.reshape(weight, [-1, 1]), compute_dtype) return array_ops.reshape(indices, [-1, 3]), weight * grad_flat @@ -418,11 +422,10 @@ def _scatter_tap(iy, ix, weight): all_indices.append(indices) all_updates.append(updates) - zeros = array_ops.zeros( - array_ops_stack.stack([batch, in_h, in_w, channels]), dtype=compute_dtype) - return gen_array_ops.tensor_scatter_add( - zeros, array_ops.concat(all_indices, axis=0), - array_ops.concat(all_updates, axis=0)) + return array_ops.scatter_nd( + array_ops.concat(all_indices, axis=0), + array_ops.concat(all_updates, axis=0), + array_ops_stack.stack([batch, in_h, in_w, channels])) @ops.RegisterGradient("ImageProjectiveTransformV2") From 036aa3df68f1c820c64748f3ce3ceb5dcff8dc19 Mon Sep 17 00:00:00 2001 From: Riccardo Squarcialupi <44367261+RiccardoSquarcialupi@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:11:42 +0200 Subject: [PATCH 05/27] Update BUILD deps for image_ops and image_grad_test_base Both targets use strict_deps, so direct deps must match direct imports: - image_ops: add :array_ops_stack, :clip_ops and :math_ops (newly imported by the adjoint implementation), drop :linalg_ops (no longer used now that the gradient no longer inverts the transform matrix). - image_grad_test_base: add :array_ops. --- tensorflow/python/ops/BUILD | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index f72ab18839b33d..134a8e62c707c4 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -1803,9 +1803,11 @@ py_library( visibility = ["//visibility:public"], deps = [ ":array_ops", + ":array_ops_stack", + ":clip_ops", ":image_ops_gen", ":image_ops_impl", - ":linalg_ops", + ":math_ops", "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:ops", @@ -3494,6 +3496,7 @@ py_library( srcs = ["image_grad_test_base.py"], strict_deps = True, deps = [ + ":array_ops", ":array_ops_stack", ":gradient_checker_v2", ":image_ops", From d810389a67a16e69e8fe23fec19c29d4e998aa04 Mon Sep 17 00:00:00 2001 From: Riccardo Squarcialupi <44367261+RiccardoSquarcialupi@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:44:52 +0200 Subject: [PATCH 06/27] Fix Python style lint warnings - Add Args/Returns sections to the _trunc_div and _round_half_away_from_zero docstrings. - Fix continuation-line indentation in ImageProjectiveTransformOpTestBase._transform. No functional change. --- tensorflow/python/ops/image_grad_test_base.py | 2 +- tensorflow/python/ops/image_ops.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py index 257e2c87a6eaf7..4c3a7e1c08b638 100644 --- a/tensorflow/python/ops/image_grad_test_base.py +++ b/tensorflow/python/ops/image_grad_test_base.py @@ -701,7 +701,7 @@ class ImageProjectiveTransformOpTestBase(test.TestCase, parameterized.TestCase): """ def _transform(self, images, transforms, output_shape, interpolation, - fill_mode): + fill_mode): return gen_image_ops.image_projective_transform_v3( images=images, transforms=transforms, diff --git a/tensorflow/python/ops/image_ops.py b/tensorflow/python/ops/image_ops.py index 7e6d44be954fbc..7fd17be1176c7b 100644 --- a/tensorflow/python/ops/image_ops.py +++ b/tensorflow/python/ops/image_ops.py @@ -244,6 +244,14 @@ def _trunc_div(a, b): `MapCoordinate` (tensorflow/core/kernels/image/image_ops.h); Python and TensorFlow's own integer division instead floor, which disagrees for negative inputs. + + Args: + a: float32 tensor, the dividend. + b: float32 tensor broadcastable with `a`, the divisor. Must be nonzero + wherever the result is used. + + Returns: + float32 tensor, `a / b` truncated toward zero. """ return math_ops.cast(math_ops.cast(a / b, dtypes.int32), dtypes.float32) @@ -305,6 +313,13 @@ def _round_half_away_from_zero(x): TensorFlow's own `round` breaks ties to even instead, which disagrees with the kernel at exact `.5` boundaries. + + Args: + x: float tensor, arbitrary shape. + + Returns: + Tensor of the same shape and dtype as `x`, with each element rounded to + the nearest integer value, ties away from zero. """ return math_ops.sign(x) * math_ops.floor(math_ops.abs(x) + 0.5) From 521b65d588f1b7a539c94650bc68d6a49834fcf2 Mon Sep 17 00:00:00 2001 From: Riccardo Squarcialupi <44367261+RiccardoSquarcialupi@users.noreply.github.com> Date: Tue, 22 Sep 2026 09:58:32 +0200 Subject: [PATCH 07/27] Disable XLA for ImageProjectiveTransformOpTestBase Per review: the internal XLA GPU stack executes the forward op's coordinate mapping with enough numerical divergence from CPU/eager to trip this suite's tight tolerances (e.g. testAdjointIdentity's rtol=atol=1e-6), unrelated to the adjoint logic itself. Follows the same test_util.for_all_test_methods(test_util.disable_xla, ...) pattern already used above for ResizeNearestNeighborOpTestBase for an analogous reason. --- tensorflow/python/ops/image_grad_test_base.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py index 4c3a7e1c08b638..a3cc703559f46e 100644 --- a/tensorflow/python/ops/image_grad_test_base.py +++ b/tensorflow/python/ops/image_grad_test_base.py @@ -687,6 +687,9 @@ def f(image_tensor, factor_tensor): ] +@test_util.for_all_test_methods( + test_util.disable_xla, + 'ImageProjectiveTransform forward XLA behaves divergently') class ImageProjectiveTransformOpTestBase(test.TestCase, parameterized.TestCase): """Tests the gradient of ImageProjectiveTransformV2/V3. From 4a69adff7312b371b5d888bca99d0a0346b60b36 Mon Sep 17 00:00:00 2001 From: Matthew Johnson Date: Wed, 23 Sep 2026 16:58:50 -0700 Subject: [PATCH 08/27] [IFRT Proxy] Increase IfrtBackend compile_thread_pool_ stack size to 2 MiB Compiling GPU convolutions invokes cuDNN v9 engine heuristics (cudnnBackendGetAttribute -> traceback_api_add), which together with XLA GPU compilation frames requires ~280+ KiB of stack space and overflows the previous 240 KiB stack size (breaking jax_api_pjrt_gpu_test after cl/986842563). PiperOrigin-RevId: 987097843 --- third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc index b34ce1cdbb219f..1d43424a310538 100644 --- a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc +++ b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc @@ -472,7 +472,7 @@ IfrtBackend::IfrtBackend(IfrtProxyVersion version, uint64_t session_id, tsl::ThreadOptions options; // Use a larger stack size since XLA often requires larger stacks // for compilation. - options.stack_size = 240 * 1024; + options.stack_size = 2 * 1024 * 1024; return options; }(), "IfrtBackend", From defcfecfc875930a6d1ca85060e09e8480bf7de7 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Wed, 23 Sep 2026 17:38:17 -0700 Subject: [PATCH 09/27] Delegate CreateAsyncTrackingEvent to raw client. PiperOrigin-RevId: 987118567 --- third_party/xla/xla/pjrt/common_pjrt_client.cc | 15 +++++++++++++++ third_party/xla/xla/pjrt/common_pjrt_client.h | 4 +--- third_party/xla/xla/pjrt/raw_pjrt_client.h | 5 +++++ 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index 0a81f753f8f137..dff4ad82416700 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -4452,8 +4452,23 @@ absl::StatusOr CommonPjRtDevice::memory_space_by_kind_id( return it->second; } +std::unique_ptr +CommonPjRtDevice::CreateAsyncTrackingEvent( + absl::string_view description) const { + if (!IsAddressable()) { + return nullptr; + } + CHECK(client_ != nullptr); + return client_->raw_client()->CreateAsyncTrackingEvent(local_device_id(), + description); +} + absl::StatusOr CommonPjRtDevice::PoisonExecution(int32_t launch_id, absl::Status error) { + if (!IsAddressable()) { + return FailedPrecondition( + "PoisonExecution() is allowed only for addressable devices"); + } CHECK(client_ != nullptr); return client_->raw_client()->PoisonExecution(local_device_id(), launch_id, std::move(error)); diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h index ca7f930f4cb384..c25558ff7a8433 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.h +++ b/third_party/xla/xla/pjrt/common_pjrt_client.h @@ -1260,9 +1260,7 @@ class CommonPjRtDevice : public PjRtDevice { absl::StatusOr memory_space_by_kind_id(int id) const; std::unique_ptr CreateAsyncTrackingEvent( - absl::string_view description) const override { - return nullptr; - } + absl::string_view description) const override; absl::StatusOr PoisonExecution(int32_t launch_id, absl::Status error) override; diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.h b/third_party/xla/xla/pjrt/raw_pjrt_client.h index 238a80bb41c526..18f893868741d2 100644 --- a/third_party/xla/xla/pjrt/raw_pjrt_client.h +++ b/third_party/xla/xla/pjrt/raw_pjrt_client.h @@ -269,6 +269,11 @@ class PjRtRawClient { LOG(FATAL) << "Implement MakeLoadState()"; } + virtual std::unique_ptr CreateAsyncTrackingEvent( + LocalDeviceId local_device_id, absl::string_view description) const { + return nullptr; + } + virtual absl::StatusOr PoisonExecution(LocalDeviceId local_device_id, int32_t launch_id, absl::Status error) { From f917fed06212917a6ae2dfde191a0cb6a3269b17 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Wed, 23 Sep 2026 18:01:43 -0700 Subject: [PATCH 10/27] Increase direct_session_test size to large and shard into 8 shards to prevent timeout. PiperOrigin-RevId: 987130327 --- tensorflow/core/common_runtime/BUILD | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD index 04b075eb16b82d..ed2680bb2aa79d 100644 --- a/tensorflow/core/common_runtime/BUILD +++ b/tensorflow/core/common_runtime/BUILD @@ -3115,11 +3115,11 @@ tf_cc_test( tf_cuda_cc_test( name = "direct_session_test", - size = "medium", + size = "large", srcs = ["direct_session_test.cc"], args = [] + if_cuda(["--heap_check="]), # The GPU tracer leaks memory features = ["-layering_check"], - shard_count = 4, + shard_count = 8, deps = [ ":core_cpu", ":core_cpu_internal", From 4920e5a7cec6ede425c269490ef7da054975ab65 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Wed, 23 Sep 2026 18:03:17 -0700 Subject: [PATCH 11/27] Assert that `kCombineDimension` is never used on PjRt buffer layouts PiperOrigin-RevId: 987131214 --- third_party/xla/xla/pjrt/BUILD | 1 + third_party/xla/xla/pjrt/pjrt_layout.h | 11 +++++++++++ 2 files changed, 12 insertions(+) diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index bc2859c3567a6b..2fe8df4824609d 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -750,6 +750,7 @@ cc_library( deps = [ "//xla:shape_util", "//xla/hlo/parser:hlo_parser", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/pjrt/pjrt_layout.h b/third_party/xla/xla/pjrt/pjrt_layout.h index 71a9ebad9bcf28..e0367f8387f2c6 100644 --- a/third_party/xla/xla/pjrt/pjrt_layout.h +++ b/third_party/xla/xla/pjrt/pjrt_layout.h @@ -16,11 +16,13 @@ limitations under the License. #ifndef XLA_PJRT_PJRT_LAYOUT_H_ #define XLA_PJRT_PJRT_LAYOUT_H_ +#include #include #include #include #include +#include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -39,6 +41,15 @@ class PjRtLayout { // Strip memory space and set it to the default. PJRT tracks memory space // separately from layout. xla_layout_.set_memory_space(xla::Layout::kDefaultMemorySpace); + + for (const auto& tile : xla_layout_.tiles()) { + CHECK(absl::c_none_of( + tile.dimensions(), + [](int64_t d) { return d == xla::Tile::kCombineDimension; })) + << "Runtime buffers must not have layouts with `kCombineDimension` " + "in tile dimensions, but got " + << xla_layout_; + } } PjRtLayout(PjRtLayout& other) = delete; From c96cc1156c84061bc273774b7c51f93351d39a85 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Wed, 23 Sep 2026 18:08:25 -0700 Subject: [PATCH 12/27] Clamp `num_lower`/`num_upper` that exceed the matrix dimensions in `tf.linalg.band_part` instead of raising `InvalidArgumentError`, so that eager execution matches the documented definition of the op and the XLA lowering. Fixes #110316 PiperOrigin-RevId: 987133747 --- RELEASE.md | 6 +- .../kernels/linalg/matrix_band_part_op.cc | 25 ++++---- .../array_ops/matrix_band_part_op_test.py | 59 ++++++++++++++++++- 3 files changed, 76 insertions(+), 14 deletions(-) diff --git a/RELEASE.md b/RELEASE.md index d115367bafc47e..38afa6c8b11cbe 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -18,7 +18,11 @@ ### Bug Fixes and Other Changes -* +* `tf.linalg.band_part`: `num_lower` and `num_upper` values that exceed the + number of rows or columns of the input are now clamped to the matrix + dimensions instead of raising an `InvalidArgumentError`. This matches the + documented definition of the op and the behavior of the XLA lowering used by + `tf.function(jit_compile=True)`. ## Thanks to our Contributors diff --git a/tensorflow/core/kernels/linalg/matrix_band_part_op.cc b/tensorflow/core/kernels/linalg/matrix_band_part_op.cc index 9033993861c34d..20155659451aed 100644 --- a/tensorflow/core/kernels/linalg/matrix_band_part_op.cc +++ b/tensorflow/core/kernels/linalg/matrix_band_part_op.cc @@ -74,24 +74,25 @@ class MatrixBandPartOp : public OpKernel { return tensor.scalar()(); } }; - const int64_t num_lower = as_int64_scalar(num_lower_in); - OP_REQUIRES( - context, num_lower <= input_reshaped.dimension(1), - errors::InvalidArgument( - "num_lower must be negative or less or equal to number of rows (", - input_reshaped.dimension(1), ") got: ", num_lower)); + // Band limits that are out of range are clamped rather than rejected. The + // op is defined by + // in_band(m, n) = (num_lower < 0 || (m - n) <= num_lower) && + // (num_upper < 0 || (n - m) <= num_upper), + // which is well defined for any limit, and this is what the XLA lowering + // already computes. Clamping to [-1, dimension] additionally keeps the + // limits exactly representable in the `int` parameters of + // MatrixBandPartFunctor below; every negative limit is equivalent to -1, + // i.e. "keep the entire triangle". + const int64_t num_lower = std::clamp( + as_int64_scalar(num_lower_in), -1, input_reshaped.dimension(1)); const Tensor& num_upper_in = context->input(2); OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_upper_in.shape()), absl::InvalidArgumentError( absl::StrCat("num_upper must be scalar, got shape ", num_upper_in.shape().DebugString()))); - const int64_t num_upper = as_int64_scalar(num_upper_in); - OP_REQUIRES(context, num_upper <= input_reshaped.dimension(2), - errors::InvalidArgument("num_upper must be negative or less or " - "equal to number of columns (", - input_reshaped.dimension(2), - ") got: ", num_upper)); + const int64_t num_upper = std::clamp( + as_int64_scalar(num_upper_in), -1, input_reshaped.dimension(2)); if (input.NumElements() == 0 || ((num_lower < 0 || num_lower == input_reshaped.dimension(1)) && diff --git a/tensorflow/python/kernel_tests/array_ops/matrix_band_part_op_test.py b/tensorflow/python/kernel_tests/array_ops/matrix_band_part_op_test.py index 80e038c3174276..20079611704615 100644 --- a/tensorflow/python/kernel_tests/array_ops/matrix_band_part_op_test.py +++ b/tensorflow/python/kernel_tests/array_ops/matrix_band_part_op_test.py @@ -36,7 +36,64 @@ def _AddTest(test, op_name, testcase_name, fn): class MatrixBandPartTest(test_lib.TestCase): - pass # Filled in below + + def testNumLowerAndNumUpperExceedingMatrixDimensions(self) -> None: + """Out-of-range band limits are clamped instead of rejected.""" + matrix = constant_op.constant( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=dtypes_lib.int32 + ) + # Both limits exceed the matrix, so every element is kept. + self.assertAllEqual( + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + self.evaluate(array_ops.matrix_band_part(matrix, 6, 6)), + ) + # Only `num_lower` is out of range: the lower triangle is kept. + self.assertAllEqual( + [[1, 0, 0], [4, 5, 0], [7, 8, 9]], + self.evaluate(array_ops.matrix_band_part(matrix, 6, 0)), + ) + # Only `num_upper` is out of range: the upper triangle is kept. + self.assertAllEqual( + [[1, 2, 3], [0, 5, 6], [0, 0, 9]], + self.evaluate(array_ops.matrix_band_part(matrix, 0, 6)), + ) + + def testNumLowerAndNumUpperClampedToInnermostDimensions(self) -> None: + """Batched, non-square inputs clamp against the two innermost dimensions.""" + # Two 2x3 matrices, so the row and column counts differ. + batch = constant_op.constant( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + dtype=dtypes_lib.int32, + ) + self.assertAllEqual( + [[[1, 0, 0], [4, 5, 0]], [[7, 0, 0], [10, 11, 0]]], + self.evaluate(array_ops.matrix_band_part(batch, 100, 0)), + ) + self.assertAllEqual( + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + self.evaluate(array_ops.matrix_band_part(batch, 100, 100)), + ) + + def testNumLowerAndNumUpperOutsideInt32RangeDoNotWrapAround(self) -> None: + """Limits far outside the `int32` range are clamped, not truncated.""" + matrix = constant_op.constant([[1, 2], [3, 4]], dtype=dtypes_lib.int32) + huge = constant_op.constant(2**40, dtype=dtypes_lib.int64) + tiny = constant_op.constant(-(2**40), dtype=dtypes_lib.int64) + zero = constant_op.constant(0, dtype=dtypes_lib.int64) + self.assertAllEqual( + [[1, 2], [3, 4]], + self.evaluate(array_ops.matrix_band_part(matrix, huge, huge)), + ) + # Every negative limit keeps the entire triangle, so this is the lower + # triangle rather than just the diagonal. + self.assertAllEqual( + [[1, 0], [3, 4]], + self.evaluate(array_ops.matrix_band_part(matrix, tiny, zero)), + ) + self.assertAllEqual( + [[1, 2], [0, 4]], + self.evaluate(array_ops.matrix_band_part(matrix, zero, huge)), + ) def _GetMatrixBandPartTest(dtype_, batch_shape_, shape_): From 2a3525476e35c62dcd1652be92985b42e9523dfc Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 18:24:34 -0700 Subject: [PATCH 13/27] [Mosaic TPU] Add some fold rules for the transpose op. Reverts 4470f2f1dd91ac973e4270bc47150a97d9644ff6 PiperOrigin-RevId: 987140670 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 22 ------------------- .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 1 - 2 files changed, 23 deletions(-) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 09432b11e5d080..6cf17ae96335fa 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -955,28 +955,6 @@ LogicalResult TransposeOp::verify() { return success(); } -OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) { - if (llvm::is_sorted(getPermutation())) { - return getVector(); - } - if (auto cst = dyn_cast_if_present(adaptor.getVector())) { - return cst.reshape(getType()); - } - if (auto prev_op = getVector().getDefiningOp()) { - ArrayRef prev_perm = prev_op.getPermutation(); - ArrayRef curr_perm = getPermutation(); - CHECK_EQ(prev_perm.size(), curr_perm.size()); - SmallVector composed_perm(curr_perm.size()); - for (size_t i = 0; i < curr_perm.size(); ++i) { - composed_perm[i] = prev_perm[curr_perm[i]]; - } - if (llvm::is_sorted(composed_perm)) { - return prev_op.getVector(); - } - } - return nullptr; -} - LogicalResult MemRefBitcastOp::verify() { auto src_ty = getInput().getType(); auto tgt_ty = getType(); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index 1130919108db8c..e23c2715ba515f 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -1921,7 +1921,6 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> { } }]; let hasVerifier = 1; - let hasFolder = 1; } def TPU_LogOp : TPU_Op<"log"> { From f797e7a7546cab7c1b421c097652ba5c7a452e5c Mon Sep 17 00:00:00 2001 From: Jun Jiang Date: Wed, 23 Sep 2026 18:41:45 -0700 Subject: [PATCH 14/27] Integrate the TFLite Swift APIs into the LiteRT Swift package. PiperOrigin-RevId: 987147690 --- tensorflow/compiler/mlir/lite/core/c/BUILD | 10 +++++++++- tensorflow/lite/core/async/c/BUILD | 7 ++++++- tensorflow/lite/core/c/BUILD | 7 ++++++- tensorflow/lite/core/c/special_rules.bzl | 6 +++--- tensorflow/lite/profiling/telemetry/c/BUILD | 7 ++++++- 5 files changed, 30 insertions(+), 7 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/core/c/BUILD b/tensorflow/compiler/mlir/lite/core/c/BUILD index 4b3745f421c3e1..885bea84af4458 100644 --- a/tensorflow/compiler/mlir/lite/core/c/BUILD +++ b/tensorflow/compiler/mlir/lite/core/c/BUILD @@ -49,6 +49,7 @@ exports_files( ], visibility = [ # copybara:uncomment_begin(google-only) + # "//third_party/odml/litert/litert/swift:__subpackages__", # "//tensorflow/lite:__subpackages__", # copybara:uncomment_end_and_comment_begin "//visibility:public", @@ -63,7 +64,14 @@ cc_library( ], compatible_with = get_compatible_with_portable(), copts = tflite_copts(), - visibility = ["//visibility:public"], + visibility = [ + # copybara:uncomment_begin(google-only) + # "//third_party/odml/litert/litert/swift:__subpackages__", + # "//visibility:public", + # copybara:uncomment_end_and_comment_begin + "//visibility:public", + # copybara:comment_end + ], ) # LINT.IfChange(common) diff --git a/tensorflow/lite/core/async/c/BUILD b/tensorflow/lite/core/async/c/BUILD index 4f665dca5b2a81..cfa2eee22c9461 100644 --- a/tensorflow/lite/core/async/c/BUILD +++ b/tensorflow/lite/core/async/c/BUILD @@ -37,7 +37,12 @@ exports_files( "types.h", ], visibility = [ - "//tensorflow/lite:__subpackages__", + # copybara:uncomment_begin(google-only) + # "//third_party/odml/litert/litert/swift:__subpackages__", + # "//tensorflow/lite:__subpackages__", + # copybara:uncomment_end_and_comment_begin + "//visibility:public", + # copybara:comment_end ], ) diff --git a/tensorflow/lite/core/c/BUILD b/tensorflow/lite/core/c/BUILD index 821c48429dc932..4b1de6f583bd4b 100644 --- a/tensorflow/lite/core/c/BUILD +++ b/tensorflow/lite/core/c/BUILD @@ -50,7 +50,12 @@ exports_files( "operator.h", ], visibility = [ - "//tensorflow/lite:__subpackages__", + # copybara:uncomment_begin(google-only) + # "//third_party/odml/litert/litert/swift:__subpackages__", + # "//tensorflow/lite:__subpackages__", + # copybara:uncomment_end_and_comment_begin + "//visibility:public", + # copybara:comment_end ], ) diff --git a/tensorflow/lite/core/c/special_rules.bzl b/tensorflow/lite/core/c/special_rules.bzl index 091e885b46f3fa..6fe05c37d4c37b 100644 --- a/tensorflow/lite/core/c/special_rules.bzl +++ b/tensorflow/lite/core/c/special_rules.bzl @@ -16,15 +16,15 @@ def c_api_visibility_allowlist(): """Returns a list of packages that can depend on c_api.""" - return [] + return ["//litert/swift:__subpackages__"] def c_api_experimental_visibility_allowlist(): """Returns a list of packages that can depend on third_party/tensorflow/lite/core/c/c_api_experimental.h.""" - return [] + return ["//litert/swift:__subpackages__"] def common_header_visibility_allowlist(): """Returns a list of packages that can depend on tensorflow/lite/core/c/common.h.""" - return [] + return ["//litert/swift:__subpackages__"] def experimental_acceleration_api_allowlist(): """Returns a list of packages that can depend on tensorflow/lite/core/experimental/acceleration/ subpackages..""" diff --git a/tensorflow/lite/profiling/telemetry/c/BUILD b/tensorflow/lite/profiling/telemetry/c/BUILD index bd0d1f3c39d5e7..1526470fc00133 100644 --- a/tensorflow/lite/profiling/telemetry/c/BUILD +++ b/tensorflow/lite/profiling/telemetry/c/BUILD @@ -33,7 +33,12 @@ exports_files( "telemetry_setting.h", ], visibility = [ - "//tensorflow/lite:__subpackages__", + # copybara:uncomment_begin(google-only) + # "//third_party/odml/litert/litert/swift:__subpackages__", + # "//tensorflow/lite:__subpackages__", + # copybara:uncomment_end_and_comment_begin + "//visibility:public", + # copybara:comment_end ], ) From 17c45530bdf172faa8b4e11f9df1990657346907 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 20:11:07 -0700 Subject: [PATCH 15/27] Automated Code Change PiperOrigin-RevId: 987183957 --- third_party/xla/xla/codegen/xtile/codegen/BUILD | 1 + third_party/xla/xla/codegen/xtile/codegen/dot_algorithms_test.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/third_party/xla/xla/codegen/xtile/codegen/BUILD b/third_party/xla/xla/codegen/xtile/codegen/BUILD index b3264481aefb8b..cd80e0c8f22bc9 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/BUILD +++ b/third_party/xla/xla/codegen/xtile/codegen/BUILD @@ -233,6 +233,7 @@ xla_cc_test( "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service/llvm_ir:llvm_util", "//xla/tsl/platform:statusor", + "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", "@llvm-project//mlir:ArithDialect", "@llvm-project//mlir:IR", diff --git a/third_party/xla/xla/codegen/xtile/codegen/dot_algorithms_test.cc b/third_party/xla/xla/codegen/xtile/codegen/dot_algorithms_test.cc index dcbd9f73a57071..9dfda098963e23 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/dot_algorithms_test.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/dot_algorithms_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" From 0b3ce1a8261ca085fbaf766d251d6fd04279bf2e Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 20:22:49 -0700 Subject: [PATCH 16/27] Automated Code Change PiperOrigin-RevId: 987188596 --- tensorflow/core/kernels/spectrogram_op.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/kernels/spectrogram_op.cc b/tensorflow/core/kernels/spectrogram_op.cc index fe9316dd369314..77032bb0c383ac 100644 --- a/tensorflow/core/kernels/spectrogram_op.cc +++ b/tensorflow/core/kernels/spectrogram_op.cc @@ -34,11 +34,12 @@ class AudioSpectrogramOp : public OpKernel { OP_REQUIRES_OK(context, context->GetAttr("stride", &stride_)); OP_REQUIRES_OK(context, context->GetAttr("magnitude_squared", &magnitude_squared_)); - OP_REQUIRES( - context, window_size_ > 1, - errors::InvalidArgument("window_size must be > 1, got ", window_size_)); + OP_REQUIRES(context, window_size_ > 1, + absl::InvalidArgumentError(absl::StrCat( + "window_size must be > 1, got ", window_size_))); OP_REQUIRES(context, stride_ > 0, - errors::InvalidArgument("stride must be > 0, got ", stride_)); + absl::InvalidArgumentError( + absl::StrCat("stride must be > 0, got ", stride_))); } void Compute(OpKernelContext* context) override { From 4fca608e690f9841296a42598bf031f7c1f52d0d Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 20:38:46 -0700 Subject: [PATCH 17/27] Automated Code Change PiperOrigin-RevId: 987195419 --- .../gradients/tape/tape_operation.cc | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tensorflow/c/experimental/gradients/tape/tape_operation.cc b/tensorflow/c/experimental/gradients/tape/tape_operation.cc index d582235b0ba833..836baa776366fc 100644 --- a/tensorflow/c/experimental/gradients/tape/tape_operation.cc +++ b/tensorflow/c/experimental/gradients/tape/tape_operation.cc @@ -109,10 +109,10 @@ absl::Status TapeOperation::SetAttrShape(const char* attr_name, const int64_t* dims, const int num_dims) { if (num_dims > TensorShape::MaxDimensions()) { - return errors::InvalidArgument("Value specified for `", attr_name, "` has ", - num_dims, - " dimensions which is over the limit of ", - TensorShape::MaxDimensions(), "."); + return absl::InvalidArgumentError( + absl::StrCat("Value specified for `", attr_name, "` has ", num_dims, + " dimensions which is over the limit of ", + TensorShape::MaxDimensions(), ".")); } TensorShapeProto proto; if (num_dims < 0) { @@ -128,19 +128,19 @@ absl::Status TapeOperation::SetAttrShape(const char* attr_name, } absl::Status TapeOperation::SetAttrFunction(const char* attr_name, const AbstractOperation* value) { - return tensorflow::errors::Unimplemented( + return absl::UnimplementedError( "SetAttrFunction has not been implemented yet."); } absl::Status TapeOperation::SetAttrFunctionName(const char* attr_name, const char* value, size_t length) { - return tensorflow::errors::Unimplemented( + return absl::UnimplementedError( "SetAttrFunctionName has not been implemented " "yet."); } absl::Status TapeOperation::SetAttrTensor(const char* attr_name, AbstractTensorInterface* tensor) { - return tensorflow::errors::Unimplemented( + return absl::UnimplementedError( "SetAttrTensor has not been implemented yet."); } absl::Status TapeOperation::SetAttrStringList(const char* attr_name, @@ -196,7 +196,7 @@ absl::Status TapeOperation::SetAttrShapeList(const char* attr_name, const auto num_dims_i = num_dims[i]; if (num_dims_i > TensorShape::MaxDimensions()) { - return errors::InvalidArgument( + return absl::InvalidArgumentError( strings::StrCat("Value specified for `", attr_name, "` has ", num_dims_i, " dimensions which is over the limit of ", TensorShape::MaxDimensions(), ".")); @@ -217,7 +217,7 @@ absl::Status TapeOperation::SetAttrShapeList(const char* attr_name, } absl::Status TapeOperation::SetAttrFunctionList( const char* attr_name, absl::Span values) { - return tensorflow::errors::Unimplemented( + return absl::UnimplementedError( "SetAttrFunctionList has not been " "implemented yet."); } From ea22af286cf928ebeea9ef6bea35b58c6713a4dc Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Wed, 23 Sep 2026 22:54:53 -0700 Subject: [PATCH 18/27] Fix lifetime issues in `ExecutionWatchdogScope` async callbacks. When `block_host_until_done` is false, `~ExecutionWatchdogScope()` releases the `HangWatchdog::Guard` asynchronously via `stream_->DoHostCallback(...)`, so the guard and its `on_timeout` and `pre_abort` callbacks outlive `ExecutionWatchdogScope` and `GpuExecutable::ExecuteThunks()`. - Copy `ExecutionTimeoutHandler` by value in `ExecutionWatchdogScope::Arm()` instead of capturing a raw `const GpuExecutableRunOptions*` pointer that can dangle when callers allocate `GpuExecutableRunOptions` on the stack. - Separate `ThunkExecutor::ProgressTracker` from `ThunkExecutor::ScopedProgressTracker` via `std::shared_ptr` so `pre_abort` in `GpuExecutable::ExecuteThunksImpl()` captures `tracker->tracker()` by value while `ScopedProgressTracker` remains a stack-local move-only RAII guard that uninstalls the thread-local pointer when leaving the dispatch scope. PiperOrigin-RevId: 987251708 --- .../backends/gpu/runtime/thunk_executor.cc | 65 ++++++++-------- .../xla/backends/gpu/runtime/thunk_executor.h | 76 ++++++++++++------- .../xla/xla/service/gpu/execution_watchdog.cc | 9 ++- .../xla/xla/service/gpu/gpu_executable.cc | 2 +- 4 files changed, 86 insertions(+), 66 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_executor.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_executor.cc index 1d969a186d903b..5331950468d919 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk_executor.cc +++ b/third_party/xla/xla/backends/gpu/runtime/thunk_executor.cc @@ -100,7 +100,7 @@ absl::Status ThunkExecutor::ExecuteOnStream( // If progress tracker is installed for current thread, verify that a // thunk indexing record exists for the given `thunk`. if (progress_tracker) { - if (!progress_tracker->indexing.contains(thunk.get())) { + if (!progress_tracker->indexing_.contains(thunk.get())) { return Internal( "[thunk=%d/%d] Progress tracker is missing a record for thunk `%s`", i, thunks_.size(), thunk->profile_annotation()); @@ -138,12 +138,12 @@ absl::Status ThunkExecutor::ExecuteOnStream( if (progress_tracker) { // Borrow an event from the pool and record it on the execution stream. ABSL_ASSIGN_OR_RETURN(auto event, - progress_tracker->event_pool->GetOrCreateEvent()); + progress_tracker->event_pool_->GetOrCreateEvent()); ABSL_RETURN_IF_ERROR(params.stream->RecordEvent(event->get())); - absl::MutexLock lock(progress_tracker->mu); - progress_tracker->events.emplace_back(thunk.get(), std::move(event), - loop_nest.nest); + absl::MutexLock lock(progress_tracker->mu_); + progress_tracker->events_.emplace_back(thunk.get(), std::move(event), + loop_nest.nest); } XLA_VLOG_DEVICE(1, device_ordinal) << absl::StreamFormat( @@ -245,12 +245,12 @@ absl::StatusOr InstallDefinitionTracker( // Tracking Thunk execution progress. //===----------------------------------------------------------------------===// -using ThunkExecution = ThunkExecutor::ScopedProgressTracker::ThunkExecution; +using ThunkExecution = ThunkExecutor::ProgressTracker::ThunkExecution; -thread_local ThunkExecutor::ScopedProgressTracker::ProgressTracker* +thread_local ThunkExecutor::ProgressTracker* ThunkExecutor::ScopedProgressTracker::installed = nullptr; -ThunkExecutor::ScopedProgressTracker::ThunkExecutionEvent::ThunkExecutionEvent( +ThunkExecutor::ProgressTracker::ThunkExecutionEvent::ThunkExecutionEvent( const Thunk* thunk, EventPool::Event event, absl::Span loop_nest) : thunk(thunk), @@ -261,7 +261,7 @@ ThunkExecutor::ScopedProgressTracker::ThunkExecutionEvent::ThunkExecutionEvent( ThunkExecutor::ScopedProgressTracker::ScopedProgressTracker( EventPool* event_pool, ThunkIndexing indexing) : tracker_( - std::make_unique(std::move(indexing), event_pool)) { + std::make_shared(std::move(indexing), event_pool)) { CHECK_EQ(installed, nullptr) // Crash OK << "Tried to install multiple progress trackers"; installed = tracker_.get(); @@ -275,31 +275,28 @@ ThunkExecutor::ScopedProgressTracker::~ScopedProgressTracker() { } } -size_t ThunkExecutor::ScopedProgressTracker::num_executions() const { - absl::MutexLock lock(tracker_->mu); - return tracker_->events.size(); +size_t ThunkExecutor::ProgressTracker::num_executions() const { + absl::MutexLock lock(mu_); + return events_.size(); } -size_t ThunkExecutor::ScopedProgressTracker::NumPendingThunks() { - absl::MutexLock lock(tracker_->mu); - return absl::c_count_if(tracker_->events, [](const auto& event) { +size_t ThunkExecutor::ProgressTracker::NumPendingThunks() { + absl::MutexLock lock(mu_); + return absl::c_count_if(events_, [](const auto& event) { return event.event->get()->PollForStatus() == se::Event::Status::kPending; }); } -size_t ThunkExecutor::ScopedProgressTracker::NumCompletedThunks() { - absl::MutexLock lock(tracker_->mu); - return absl::c_count_if(tracker_->events, [](const auto& event) { +size_t ThunkExecutor::ProgressTracker::NumCompletedThunks() { + absl::MutexLock lock(mu_); + return absl::c_count_if(events_, [](const auto& event) { return event.event->get()->PollForStatus() == se::Event::Status::kComplete; }); } -std::vector ThunkExecutor::ScopedProgressTracker::CollectThunks( +std::vector ThunkExecutor::ProgressTracker::CollectThunks( se::Event::Status status, bool most_recent_first, size_t n) { - absl::MutexLock lock(tracker_->mu); - - ThunkIndexing& indexing = tracker_->indexing; - absl::Span events = tracker_->events; + absl::MutexLock lock(mu_); // Events are naturally in chronological order (oldest first). Iterate forward // for oldest-first or backward for most-recent-first. @@ -307,45 +304,45 @@ std::vector ThunkExecutor::ScopedProgressTracker::CollectThunks( auto collect = [&](size_t exec_idx, const ThunkExecutionEvent& event) { if (event.event->get()->PollForStatus() == status) { - result.push_back({exec_idx, indexing.at(event.thunk), event.executed, + result.push_back({exec_idx, indexing_.at(event.thunk), event.executed, event.thunk->kind(), event.thunk->profile_annotation(), event.loop_nest}); } }; if (most_recent_first) { - for (size_t i = events.size(); i > 0; --i) { + for (size_t i = events_.size(); i > 0; --i) { if (result.size() >= n) { break; } - collect(i - 1, events[i - 1]); + collect(i - 1, events_[i - 1]); } } else { - for (size_t i = 0; i < events.size(); ++i) { + for (size_t i = 0; i < events_.size(); ++i) { if (result.size() >= n) { break; } - collect(i, events[i]); + collect(i, events_[i]); } } return result; } -std::vector -ThunkExecutor::ScopedProgressTracker::LastCompletedThunks(size_t n) { +std::vector ThunkExecutor::ProgressTracker::LastCompletedThunks( + size_t n) { return CollectThunks(se::Event::Status::kComplete, /*most_recent_first=*/true, n); } -std::vector -ThunkExecutor::ScopedProgressTracker::FirstPendingThunks(size_t n) { +std::vector ThunkExecutor::ProgressTracker::FirstPendingThunks( + size_t n) { return CollectThunks(se::Event::Status::kPending, /*most_recent_first=*/false, n); } -std::vector -ThunkExecutor::ScopedProgressTracker::LastPendingThunks(size_t n) { +std::vector ThunkExecutor::ProgressTracker::LastPendingThunks( + size_t n) { return CollectThunks(se::Event::Status::kPending, /*most_recent_first=*/true, n); } diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_executor.h b/third_party/xla/xla/backends/gpu/runtime/thunk_executor.h index 066a336afb7f9d..aa52eb151f2d3a 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk_executor.h +++ b/third_party/xla/xla/backends/gpu/runtime/thunk_executor.h @@ -72,6 +72,7 @@ class ThunkExecutor { // Forward declarations. See definitions below. class ScopedDefinitionTracker; + class ProgressTracker; class ScopedProgressTracker; explicit ThunkExecutor(ThunkSequence thunks); @@ -151,7 +152,7 @@ absl::StatusOr InstallDefinitionTracker( // which is installed for the current thread using a thread-local mechanism, and // it is guaranteed to work, because thunk sequence recording is // single-threaded. -class ThunkExecutor::ScopedProgressTracker { +class ThunkExecutor::ProgressTracker { public: // We use global indexing across all nested thunks in a sequence, not only the // top-level thunks executed by `ThunkExecutor`. The index is assigned by DFS @@ -179,14 +180,12 @@ class ThunkExecutor::ScopedProgressTracker { std::vector loop_nest; }; - ~ScopedProgressTracker(); - - ScopedProgressTracker(ScopedProgressTracker&&) = default; - ScopedProgressTracker& operator=(ScopedProgressTracker&&) = default; + ProgressTracker(ThunkIndexing indexing, EventPool* event_pool) + : indexing_(std::move(indexing)), event_pool_(event_pool) {} // Returns the number of unique thunks in the tracked thunk sequence. This // includes all thunks nested inside others. - size_t num_thunks() const { return tracker_->indexing.size(); } + size_t num_thunks() const { return indexing_.size(); } // Returns the total number of thunk executions (launches). This can exceed // num_thunks() when thunks are executed multiple times inside while loops. @@ -213,8 +212,7 @@ class ThunkExecutor::ScopedProgressTracker { private: friend class ThunkExecutor; - friend absl::StatusOr InstallProgressTracker( - se::StreamExecutor*, ThunkExecutor&); + friend class ScopedProgressTracker; // An event recorded for every thunk execution. The same thunk can be executed // 0 to N times during XLA program execution: it can be in a not-taken @@ -231,19 +229,50 @@ class ThunkExecutor::ScopedProgressTracker { std::vector loop_nest; // enclosing loop state snapshot }; - // Mutable state shared between the ScopedProgressTracker and the - // ThunkExecutor via thread-local pointer. Holds thunk indexing (immutable), - // the event pool, and a growing log of execution events. - struct ProgressTracker { - ProgressTracker(ThunkIndexing indexing, EventPool* event_pool) - : indexing(std::move(indexing)), event_pool(event_pool) {} + // Collects up to `n` thunks matching `status`, sorted by executed time. + // If `most_recent_first` is true, returns most recently executed thunks + // first (descending order); otherwise returns earliest executed thunks first + // (ascending order). Lazily checks event status to avoid polling all thunks. + std::vector CollectThunks(se::Event::Status status, + bool most_recent_first, size_t n); + + ThunkIndexing indexing_; + EventPool* event_pool_; - ThunkIndexing indexing; - EventPool* event_pool; + mutable absl::Mutex mu_; + std::vector events_ ABSL_GUARDED_BY(mu_); +}; - absl::Mutex mu; - std::vector events ABSL_GUARDED_BY(mu); - }; +class ThunkExecutor::ScopedProgressTracker { + public: + using ThunkIndexing = ProgressTracker::ThunkIndexing; + using ThunkExecution = ProgressTracker::ThunkExecution; + + ~ScopedProgressTracker(); + + ScopedProgressTracker(ScopedProgressTracker&&) = default; + ScopedProgressTracker& operator=(ScopedProgressTracker&&) = default; + + std::shared_ptr tracker() const { return tracker_; } + + size_t num_thunks() const { return tracker_->num_thunks(); } + size_t num_executions() const { return tracker_->num_executions(); } + size_t NumCompletedThunks() { return tracker_->NumCompletedThunks(); } + size_t NumPendingThunks() { return tracker_->NumPendingThunks(); } + std::vector LastCompletedThunks(size_t n) { + return tracker_->LastCompletedThunks(n); + } + std::vector FirstPendingThunks(size_t n) { + return tracker_->FirstPendingThunks(n); + } + std::vector LastPendingThunks(size_t n) { + return tracker_->LastPendingThunks(n); + } + + private: + friend class ThunkExecutor; + friend absl::StatusOr InstallProgressTracker( + se::StreamExecutor*, ThunkExecutor&); // Each thread can have at most one progress tracker installed and it is // automatically removed when the scoped progress tracker goes out of scope. @@ -251,14 +280,7 @@ class ThunkExecutor::ScopedProgressTracker { explicit ScopedProgressTracker(EventPool* event_pool, ThunkIndexing indexing); - // Collects up to `n` thunks matching `status`, sorted by executed time. - // If `most_recent_first` is true, returns most recently executed thunks - // first (descending order); otherwise returns earliest executed thunks first - // (ascending order). Lazily checks event status to avoid polling all thunks. - std::vector CollectThunks(se::Event::Status status, - bool most_recent_first, size_t n); - - std::unique_ptr tracker_; + std::shared_ptr tracker_; }; // Installs a progress tracker for the given sequential thunk in the current diff --git a/third_party/xla/xla/service/gpu/execution_watchdog.cc b/third_party/xla/xla/service/gpu/execution_watchdog.cc index cba7e0c86e1464..6401574eb32dbf 100644 --- a/third_party/xla/xla/service/gpu/execution_watchdog.cc +++ b/third_party/xla/xla/service/gpu/execution_watchdog.cc @@ -93,9 +93,11 @@ void ExecutionWatchdogScope::Arm(HangWatchdog::CancelCallback pre_abort) { std::weak_ptr weak_guard_holder = guard_holder_; auto watchdog_name = watchdog_name_; auto watchdog_timeout = watchdog_timeout_; - auto* gpu_run_options = gpu_run_options_; + ExecutionTimeoutHandler timeout_handler = + gpu_run_options_->execution_timeout_handler(); on_timeout = [watchdog_name, watchdog_timeout, - pre_abort = std::move(pre_abort), gpu_run_options, + pre_abort = std::move(pre_abort), + timeout_handler = std::move(timeout_handler), weak_guard_holder]() mutable { if (pre_abort) { std::move(pre_abort)(); @@ -109,8 +111,7 @@ void ExecutionWatchdogScope::Arm(HangWatchdog::CancelCallback pre_abort) { HangWatchdog::Abort("post-abort ...", absl::Minutes(1))); } - gpu_run_options->execution_timeout_handler()(watchdog_name, - watchdog_timeout); + timeout_handler(watchdog_name, watchdog_timeout); }; } else { on_timeout = HangWatchdog::Abort(watchdog_name_, watchdog_timeout_, diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc index 841f177ab8cfd5..3f628c38ea9483 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable.cc +++ b/third_party/xla/xla/service/gpu/gpu_executable.cc @@ -615,7 +615,7 @@ absl::Status GpuExecutable::ExecuteThunksImpl( if (execution_watchdog != nullptr) { HangWatchdog::CancelCallback pre_abort; if (tracker.has_value()) { - pre_abort = [&tracker, progress_tracking_n, + pre_abort = [tracker = tracker->tracker(), progress_tracking_n, device_ordinal = executor->device_ordinal()] { auto log_progress = [&](auto label, auto thunks) { LOG(ERROR) << absl::StreamFormat("[%d] %s: size=%d", device_ordinal, From 17c17ea4f0e157d76a3ee12a551aa5dc0cc9eb43 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Wed, 23 Sep 2026 23:04:39 -0700 Subject: [PATCH 19/27] Handle zero-sized inputs in tf.nn.lrn and its gradient op to prevent GPU crash. PiperOrigin-RevId: 987255629 --- tensorflow/core/kernels/lrn_op.cc | 12 ++++ tensorflow/core/kernels/lrn_op_test.cc | 38 ++++++++++++ tensorflow/python/kernel_tests/nn_ops/BUILD | 1 + .../python/kernel_tests/nn_ops/lrn_op_test.py | 58 +++++++++++++++++++ 4 files changed, 109 insertions(+) diff --git a/tensorflow/core/kernels/lrn_op.cc b/tensorflow/core/kernels/lrn_op.cc index 164bc2f9e6d4a5..46ee2d5e2e9bbe 100644 --- a/tensorflow/core/kernels/lrn_op.cc +++ b/tensorflow/core/kernels/lrn_op.cc @@ -352,6 +352,12 @@ class LRNOp : public OpKernel { context->allocate_output( 0, TensorShape({batch, rows, cols, depth}), &output)); + // NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal + // cuDNN crashes on GPU. + if (in.NumElements() == 0) { + return; + } + LaunchLRN launcher(depth_radius_, bias_, alpha_, beta_); launcher.launch(context, this, in, output); } @@ -696,6 +702,12 @@ class LRNGradOp : public OpKernel { context->allocate_output( 0, TensorShape({batch, rows, cols, depth}), &output)); + // NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal + // cuDNN crashes on GPU. + if (in_grads.NumElements() == 0) { + return; + } + LaunchLRNGrad launcher(depth_radius_, bias_, alpha_, beta_); launcher.launch(context, this, in_grads, in_image, out_image, output); } diff --git a/tensorflow/core/kernels/lrn_op_test.cc b/tensorflow/core/kernels/lrn_op_test.cc index 3c8515d522501b..98a6e418b6d23b 100644 --- a/tensorflow/core/kernels/lrn_op_test.cc +++ b/tensorflow/core/kernels/lrn_op_test.cc @@ -159,6 +159,44 @@ TEST_F(LRNFloatTest, Depth16) { EXPECT_TRUE(Compare()); } +TEST_F(LRNFloatTest, ZeroSizeInput) { + TF_ASSERT_OK(NodeDefBuilder("lrn_op", "LRN") + .Input(FakeInput()) + .Attr("depth_radius", 2) + .Attr("bias", 1.0f) + .Attr("alpha", 1.0f) + .Attr("beta", 0.5f) + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + AddInput(TensorShape({0, 0, 0, 0}), + [](int i) -> float { return 0.0f; }); + TF_ASSERT_OK(RunOpKernel()); + EXPECT_EQ(GetOutput(0)->NumElements(), 0); + EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0})); +} + +TEST_F(LRNFloatTest, ZeroSizeInputGrad) { + TF_ASSERT_OK(NodeDefBuilder("lrn_grad_op", "LRNGrad") + .Input(FakeInput()) + .Input(FakeInput()) + .Input(FakeInput()) + .Attr("depth_radius", 2) + .Attr("bias", 1.0f) + .Attr("alpha", 1.0f) + .Attr("beta", 0.5f) + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + AddInput(TensorShape({0, 0, 0, 0}), + [](int i) -> float { return 0.0f; }); + AddInput(TensorShape({0, 0, 0, 0}), + [](int i) -> float { return 0.0f; }); + AddInput(TensorShape({0, 0, 0, 0}), + [](int i) -> float { return 0.0f; }); + TF_ASSERT_OK(RunOpKernel()); + EXPECT_EQ(GetOutput(0)->NumElements(), 0); + EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0})); +} + static double RndGaussian(random::SimplePhilox* rnd) { // Box-Muller transformation. // See, for example, http://www.taygeta.com/random/gaussian.html diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD index dcbf27d841154b..ca0b76ac6bbd9a 100644 --- a/tensorflow/python/kernel_tests/nn_ops/BUILD +++ b/tensorflow/python/kernel_tests/nn_ops/BUILD @@ -596,6 +596,7 @@ cuda_py_strict_test( size = "medium", srcs = ["lrn_op_test.py"], deps = [ + "//tensorflow/python/eager:backprop", "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:for_generated_wrappers", diff --git a/tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py b/tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py index f44c7316845b21..5865048233293a 100644 --- a/tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py @@ -18,6 +18,7 @@ import numpy as np +from tensorflow.python.eager import backprop from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl @@ -191,6 +192,63 @@ def testGradients(self): if not test.is_gpu_available(): self._RunAndVerifyGradients(dtypes.float16) + @test_util.run_in_graph_and_eager_modes + def testZeroSizeInput(self) -> None: + for shape in [ + [0, 0, 0, 0], + [2, 0, 4, 3], + [0, 3, 3, 2], + [1, 2, 0, 3], + [1, 2, 3, 0], + ]: + x = array_ops.zeros(shape, dtype=dtypes.float32) + y = nn.local_response_normalization( + x, depth_radius=2, bias=1.0, alpha=1.0, beta=0.5 + ) + result = self.evaluate(y) + self.assertEqual(result.shape, tuple(shape)) + + @test_util.run_in_graph_and_eager_modes + def testZeroSizeGradInput(self) -> None: + for shape in [ + [0, 0, 0, 0], + [2, 0, 4, 3], + [0, 3, 3, 2], + [1, 2, 0, 3], + [1, 2, 3, 0], + ]: + x = array_ops.zeros(shape, dtype=dtypes.float32) + y = nn.lrn_grad( + input_grads=x, + input_image=x, + output_image=x, + depth_radius=2, + bias=1.0, + alpha=1.0, + beta=0.5, + ) + result = self.evaluate(y) + self.assertEqual(result.shape, tuple(shape)) + + @test_util.run_in_graph_and_eager_modes + def testZeroSizeAutodiff(self) -> None: + for shape in [ + [0, 0, 0, 0], + [2, 0, 4, 3], + [0, 3, 3, 2], + [1, 2, 0, 3], + [1, 2, 3, 0], + ]: + x = array_ops.zeros(shape, dtype=dtypes.float32) + with backprop.GradientTape() as tape: + tape.watch(x) + y = nn.local_response_normalization( + x, depth_radius=2, bias=1.0, alpha=1.0, beta=0.5 + ) + grad = tape.gradient(y, x) + result = self.evaluate(grad) + self.assertEqual(result.shape, tuple(shape)) + if __name__ == "__main__": test.main() From 3059318f8878e9dbfc6d36b9bfe7fc4ff47bff60 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 23:50:06 -0700 Subject: [PATCH 20/27] Automated Code Change PiperOrigin-RevId: 987275326 --- third_party/xla/xla/tests/bitcast_convert_test.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/third_party/xla/xla/tests/bitcast_convert_test.cc b/third_party/xla/xla/tests/bitcast_convert_test.cc index 1899c086204d48..facc0b03f363bc 100644 --- a/third_party/xla/xla/tests/bitcast_convert_test.cc +++ b/third_party/xla/xla/tests/bitcast_convert_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include #include -#include #include #include "absl/strings/string_view.h" From 019a0569e9542906b1e34abc950e9753af314348 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 23:52:58 -0700 Subject: [PATCH 21/27] Automated Code Change PiperOrigin-RevId: 987276596 --- .../tensorflow/debugging/dump_tensor_op.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc b/tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc index ba574415b5eb09..519b957696590c 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc +++ b/tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc @@ -63,10 +63,10 @@ REGISTER_OP("DumpTensor") class DumpTensorOp : public OpKernel { public: explicit DumpTensorOp(OpKernelConstruction* ctx) : OpKernel(ctx) { - string log_dir_path; - string file_name; - string func_name; - string node_name; + std::string log_dir_path; + std::string file_name; + std::string func_name; + std::string node_name; OP_REQUIRES_OK(ctx, ctx->GetAttr("log_dir_path", &log_dir_path)); OP_REQUIRES_OK(ctx, ctx->GetAttr("enabled", &enabled_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("file_name", &file_name)); @@ -92,7 +92,7 @@ class DumpTensorOp : public OpKernel { quant_unit_proto.set_func_name(func_name); quant_unit_proto.set_node_name(node_name); - string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb"); + std::string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb"); OP_REQUIRES_OK( ctx, SaveSerializedProtoToFile(quant_unit_proto.SerializeAsString(), quant_unit_path, ctx->env())); From 435103ead5f3218e9d1c5af51db5c1c914c3321b Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Wed, 23 Sep 2026 23:56:40 -0700 Subject: [PATCH 22/27] [TF/XLA] Remove unused `gpu_executable_run_options` dependency from `compiler/jit/kernels` Remove `"@xla//xla/service/gpu:gpu_executable_run_options"` from `XLA_OPS_DEPS` (`xla_ops` and `xla_ops_no_jit_rewrite_registration`) in `tensorflow/compiler/jit/kernels/BUILD`, as `GpuExecutableRunOptions` is not included or used by `xla_ops.{h,cc}`. PiperOrigin-RevId: 987278014 --- tensorflow/compiler/jit/kernels/BUILD | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 1c6c550c9d30c1..4bf1b48d4657c6 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -48,7 +48,6 @@ XLA_OPS_DEPS = [ "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/service:compiler", - "@xla//xla/service/gpu:gpu_executable_run_options", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", From 420ff4c2777d808dfc12254c56c103e06fa2ad0d Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Thu, 24 Sep 2026 00:09:40 -0700 Subject: [PATCH 23/27] [XLA] Remove dead `propagate_bitcast_layouts_to_backend_config` option and GPU `backend_configs_cc` dep from `mlir_hlo_to_hlo` `MlirToHloConversionOptions::propagate_bitcast_layouts_to_backend_config` defaults to `false` and is never set or referenced anywhere in the codebase. It was added as a temporary workaround for the legacy MHLO-based XLA:GPU elemental IR emitters (which have since been removed), and was the sole reason the hardware-independent `//xla/hlo/translate/mhlo_to_hlo:mlir_hlo_to_hlo` library included `xla/service/gpu/backend_configs.proto.h` (`xla::gpu::BitcastBackendConfig`) and depended on `//xla/service/gpu:backend_configs_cc`. Remove the dead option, its unreachable `if` branch in `ExportXlaOperator(mhlo::BitcastOp)`, and the `//xla/service/gpu:backend_configs_cc` dependency. PiperOrigin-RevId: 987284482 --- .../xla/xla/hlo/translate/mhlo_to_hlo/BUILD | 1 - .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc | 24 ------------------- .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.h | 5 ---- 3 files changed, 30 deletions(-) diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD index 9d83f381c4bab8..bc28776ba78b69 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD @@ -253,7 +253,6 @@ cc_library( "//xla/service:hlo_module_config", "//xla/service:hlo_proto_cc", "//xla/service:source_target_pairs", - "//xla/service/gpu:backend_configs_cc", "//xla/service/spmd/shardy:constants", "//xla/service/spmd/shardy:utils", "//xla/service/spmd/shardy/stablehlo_round_trip:export_shardings", diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc index 8d225e18d5711f..1e0a8c0c0b34f0 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc @@ -105,7 +105,6 @@ limitations under the License. #include "xla/mlir_hlo/mhlo/transforms/passes.h" #include "xla/mlir_hlo/stablehlo_ext/transforms/passes.h" #include "xla/mlir_hlo/utils/unregistered_attributes.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_module_config.h" #include "xla/service/source_target_pairs.h" @@ -4815,29 +4814,6 @@ LogicalResult ExportXlaOp(BitcastOp op, OpLoweringContext ctx) { xla::XlaOp bitcast = xla::internal::XlaBuilderFriend::BuildBitcast( ctx.builder, operand, xla::TypeToShape(op.getType())); value_map[op] = bitcast; - if (ctx.converter->GetOptions().propagate_bitcast_layouts_to_backend_config) { - // Encode the source and result layout of the bitcast into the XLA HLO - // backend config as a protobuf. Note that this is a temporary solution - // which will go away once XLA:GPU stops falling back to XLA HLO Elemental - // IR emitters. - xla::HloInstructionProto* bitcast_proto = - xla::internal::XlaBuilderFriend::GetInstruction(bitcast); - xla::HloInstructionProto* operand_proto = - xla::internal::XlaBuilderFriend::GetInstruction(operand); - xla::LayoutProto result_layout = - ExtractLayout(op, bitcast_proto->shape().dimensions_size(), - xla::kBitcastResultLayout) - .ToProto(); - xla::LayoutProto source_layout = - ExtractLayout(op, operand_proto->shape().dimensions_size(), - xla::kBitcastSourceLayout) - .ToProto(); - xla::gpu::BitcastBackendConfig bitcast_config; - *bitcast_config.mutable_source_layout() = source_layout; - *bitcast_config.mutable_result_layout() = result_layout; - *bitcast_proto->mutable_backend_config() = - bitcast_config.SerializeAsString(); - } return success(); } diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h index 9977af5eb45a73..a459a6e72f1c4f 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h @@ -44,11 +44,6 @@ struct MlirToHloConversionOptions { // TODO(timshen): Investigate the necessity of having layouts in MHLO. bool propagate_layouts = false; - // Propagate the source and result layouts from mhlo bitcast op into the - // backend config for the bitcast. This is required for XLA:GPU backend to - // use elemental IR emitters for fused bitcasts without propagating layouts. - bool propagate_bitcast_layouts_to_backend_config = false; - LayoutPreferenceFn layout_preference_fn; ShapeRepresentationFn shape_representation_fn; From 9bb279da4c18565a0e6240bfe6b6674bd93e9d03 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Thu, 24 Sep 2026 00:09:47 -0700 Subject: [PATCH 24/27] [XLA:Codegen] Remove unused GPU `backend_configs_cc` dependencies from `codegen/tiling` and `codegen/xtile` Remove unused `#include "xla/service/gpu/backend_configs.proto.h"` includes and `//xla/service/gpu:backend_configs_cc` build dependencies from `//xla/codegen/tiling:symbolic_tile_analysis`, `//xla/codegen/tiling/experimental:tiled_hlo`, `//xla/codegen/xtile:block_level_parameters`, and `//xla/codegen/xtile:tiling_from_block_parameters` (left over after `BlockLevelFusionConfig` and `Tile` were migrated to `xtile_config.proto`). PiperOrigin-RevId: 987284532 --- third_party/xla/xla/codegen/tiling/BUILD | 1 - third_party/xla/xla/codegen/tiling/experimental/BUILD | 1 - third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc | 1 - third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc | 1 - third_party/xla/xla/codegen/xtile/BUILD | 3 --- third_party/xla/xla/codegen/xtile/block_level_parameters.h | 1 - .../xla/xla/codegen/xtile/block_level_parameters_test.cc | 1 - .../xla/xla/codegen/xtile/tiling_from_block_parameters.cc | 1 - .../xla/xla/codegen/xtile/tiling_from_block_parameters.h | 1 - 9 files changed, 11 deletions(-) diff --git a/third_party/xla/xla/codegen/tiling/BUILD b/third_party/xla/xla/codegen/tiling/BUILD index 4a239fa7f36d90..84e8eebe150390 100644 --- a/third_party/xla/xla/codegen/tiling/BUILD +++ b/third_party/xla/xla/codegen/tiling/BUILD @@ -367,7 +367,6 @@ cc_library( "//xla/hlo/utils:hlo_traversal", "//xla/service:instruction_fusion", "//xla/service:name_uniquer", - "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:hlo_fusion_analysis", "//xla/service/gpu:ir_emission_utils", "//xla/tsl/platform:errors", diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD index c6b59cdd56e942..1c3c1cefb0ab2a 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/BUILD +++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD @@ -264,7 +264,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_traversal", "//xla/service:name_uniquer", - "//xla/service/gpu:backend_configs_cc", "//xla/tsl/lib/gtl:iterator_range", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc index 5965c108c529a7..c1b983f1be1274 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc @@ -52,7 +52,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/utils/hlo_traversal.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/name_uniquer.h" #include "xla/util.h" diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc index 28239cef761313..38e915a268b5b0 100644 --- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc +++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc @@ -70,7 +70,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/utils/hlo_traversal.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/hlo_fusion_analysis.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/instruction_fusion.h" diff --git a/third_party/xla/xla/codegen/xtile/BUILD b/third_party/xla/xla/codegen/xtile/BUILD index fe521fde604bed..d7d99ee69b83e6 100644 --- a/third_party/xla/xla/codegen/xtile/BUILD +++ b/third_party/xla/xla/codegen/xtile/BUILD @@ -54,7 +54,6 @@ cc_library( "//xla/codegen/tiling/experimental:tile", "//xla/codegen/tiling/experimental:tiling_space_utils", "//xla/hlo/ir:hlo", - "//xla/service/gpu:backend_configs_cc", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status", @@ -103,7 +102,6 @@ cc_library( compatible_with = get_compatible_with_portable(), deps = [ ":xtile_config_proto_cc", - "//xla/service/gpu:backend_configs_cc", ], ) @@ -113,7 +111,6 @@ xla_cc_test( deps = [ ":block_level_parameters", ":xtile_config_proto_cc", - "//xla/service/gpu:backend_configs_cc", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/codegen/xtile/block_level_parameters.h b/third_party/xla/xla/codegen/xtile/block_level_parameters.h index 24abcdbc007279..a50cad4f6eb956 100644 --- a/third_party/xla/xla/codegen/xtile/block_level_parameters.h +++ b/third_party/xla/xla/codegen/xtile/block_level_parameters.h @@ -22,7 +22,6 @@ limitations under the License. #include #include "xla/codegen/xtile/xtile_config.pb.h" -#include "xla/service/gpu/backend_configs.pb.h" namespace xla::xtile { diff --git a/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc b/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc index 8844dac161d0a0..ffff13b1b169d3 100644 --- a/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc +++ b/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include #include #include "xla/codegen/xtile/xtile_config.pb.h" -#include "xla/service/gpu/backend_configs.pb.h" namespace xla::xtile { namespace { diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc index d1751130dd66e3..094f217b2d8c68 100644 --- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc +++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc @@ -36,7 +36,6 @@ limitations under the License. #include "xla/codegen/xtile/xtile_config.pb.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/status_macros.h" #include "xla/util.h" diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.h b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.h index 3db770cafb5a92..494a3881e2832b 100644 --- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.h +++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.h @@ -25,7 +25,6 @@ limitations under the License. #include "xla/codegen/tiling/tiling_specification.h" #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/codegen/xtile/xtile_config.pb.h" -#include "xla/service/gpu/backend_configs.pb.h" namespace xla::xtile { From 65c21117a15b2fe717d863665cf63fba17997abe Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Thu, 24 Sep 2026 00:15:03 -0700 Subject: [PATCH 25/27] [XLA:GPU] Move Triton scaled dot device tests to dedicated test file. Extracts scaled dot tests from fusion_emitter_device_test into scaled_dot_device_test to isolate scaled dot test cases and reduce test compilation overhead. PiperOrigin-RevId: 987286750 --- .../backends/gpu/codegen/triton/tests/BUILD | 53 + .../tests/fusion_emitter_device_test.cc | 1247 -------------- .../triton/tests/scaled_dot_device_test.cc | 1426 +++++++++++++++++ 3 files changed, 1479 insertions(+), 1247 deletions(-) create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD index 7459d9e0139743..1532db909a751e 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD @@ -226,6 +226,59 @@ xla_test( ], ) +xla_test( + name = "scaled_dot_device_test", + srcs = ["scaled_dot_device_test.cc"], + backends = [ + "a100", + "h100", + "b200", + "amdgpu_any", + ], + shard_count = 20, + tags = ["no_mac"], + deps = [ + "//xla:comparison_util", + "//xla:error_spec", + "//xla:literal", + "//xla:literal_util", + "//xla:types", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/backends/gpu/codegen/triton:test_utils", + "//xla/backends/gpu/codegen/triton:xtile_compiler", + "//xla/backends/gpu/codegen/triton:xtile_test_base", + "//xla/backends/gpu/tests:gpu_pjrt_codegen_test", + "//xla/backends/gpu/transforms:composite_rewriter", + "//xla/codegen/xtile:block_level_parameters", + "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", + "//xla/hlo/testlib:verified_hlo_module", + "//xla/service/gpu:backend_configs_cc", + "//xla/service/gpu:gpu_compiler", + "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/stream_executor:device_description", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tests:hlo_interpreter_reference_mixin", + "//xla/tests:test_utils", + "//xla/tests:xla_internal_test_main", # fixdeps: keep + "//xla/tsl/platform:statusor", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/random:distributions", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + ], +) + xla_test( name = "fusion_emitter_large_test", srcs = ["fusion_emitter_large_test.cc"], diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc index dbf9f96b027714..f0d5a10809ef9c 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc @@ -2359,1253 +2359,6 @@ TEST_P(TritonEmitterTestWithTilingParam, RocmWavesPerEuZeroOmitsAttribute) { << "waves_per_eu=0 should not set amdgpu-waves-per-eu attribute"; } -struct ScaleDotTestParams { - std::string lhs_type; - std::string rhs_type; - std::string lhs_scale_type; - std::string rhs_scale_type; - std::string output_type; - std::string expected_triton_type; - - std::string PrepareHloText(absl::string_view hlo_template) const { - return absl::StrReplaceAll(hlo_template, - {{"$lhs_type", lhs_type}, - {"$rhs_type", rhs_type}, - {"$lhs_scale_type", lhs_scale_type}, - {"$rhs_scale_type", rhs_scale_type}, - {"$output_type", output_type}}); - } - static std::string ToString( - const ::testing::TestParamInfo& info) { - const ScaleDotTestParams& params = info.param; - auto name = absl::StrCat(params.lhs_type, "_", params.rhs_type, "_", - params.lhs_scale_type, "_", params.rhs_scale_type, - "_", params.output_type); - absl::StrReplaceAll({{"[", "_"}, {"]", "_"}, {",", "x"}}, &name); - return name; - } -}; - -std::ostream& operator<<(std::ostream& stream, const ScaleDotTestParams& tc) { - return stream << "{\n\tlhs_type:" << tc.lhs_type - << ",\n\trhs_type:" << tc.rhs_type - << ",\n\tlhs_scale_type:" << tc.lhs_scale_type - << ",\n\trhs_scale_type:" << tc.rhs_scale_type - << ",\n\toutput_type:" << tc.output_type << "\n}"; -} - -class TritonScaledDotGemmTest : public TritonEmitterTest, - public ::testing::WithParamInterface< - std::tuple> { - public: - bool EnableTilingPropagation() const override { - return std::get<1>(GetParam()); - } - - public: - DebugOptions GetDebugOptionsForTest() const override { - DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest(); - debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(true); - debug_options.set_xla_gpu_autotune_level(0); - debug_options.set_xla_gpu_cublas_fallback(false); - return debug_options; - } -}; - -TEST_P(TritonScaledDotGemmTest, - FP8ScaledDotCompilesToPtxIntrinsicsWhenAvailable) { - const ScaleDotTestParams& params = std::get<0>(GetParam()); - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule m - -triton_dot { - lhs = $lhs_type parameter(0) - rhs = $rhs_type parameter(1) - lhs_scale = $lhs_scale_type parameter(2) - rhs_scale = $rhs_scale_type parameter(3) - ROOT _ = $output_type{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={1}, - rhs_contracting_dims={0}, - backend_config={sizes:[128]} -} - -ENTRY e { - lhs = $lhs_type{1,0} parameter(0) - rhs = $rhs_type{1,0} parameter(1) - lhs_scale = $lhs_scale_type{1,0} parameter(2) - rhs_scale = $rhs_scale_type{1,0} parameter(3) - ROOT _ = $output_type{1,0} fusion(lhs, rhs, lhs_scale, rhs_scale), - kind=kCustom, - calls=triton_dot, - backend_config={ - "fusion_backend_config": { - kind: "__triton_nested_gemm_fusion", - "block_level_fusion_config":{ - "output_tiles":[{"sizes":["128", "256"]}], - "num_warps":"4", - "num_stages":"1", - "num_ctas":"1" - } - } - } -} -)hlo"; - - auto hlo_text = params.PrepareHloText(kHloTextTemplate); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); - - constexpr absl::string_view kExpectedTritonIrTmpl = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x128x$triton_type>, tensor<128x4xi8> - CHECK: tensor<128x256x$triton_type>, tensor<256x4xi8> - CHECK: -> tensor<128x256xf32> - )"; - auto expected_triton_ir = absl::StrReplaceAll( - kExpectedTritonIrTmpl, {{"$triton_type", params.expected_triton_type}}); - EXPECT_THAT( - CreateTritonIrAndFileCheckForDot( - *module->GetComputationWithName("triton_dot"), expected_triton_ir), - IsOk()); - if (GetCudaComputeCapability().IsAtLeastBlackwell()) { - CompileAndOptionallyVerifyPtx( - std::move(module), R"(CHECK: mxf8f6f4.block_scale.scale_vec::1X)"); - } -} - -TEST_P(TritonScaledDotGemmTest, FP8ScaledDotGetsFusedAndExecutesCorrectly) { - const ScaleDotTestParams& params = std::get<0>(GetParam()); - if (auto cc = GpuComputeCapability().cuda_compute_capability(); - cc && !cc->IsAtLeastBlackwell()) { - GTEST_SKIP() << "Skipping test for pre-Blackwell GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule FP8ScaledDotGetsFusedAndExecutesCorrectly - -ENTRY e { - lhs = $lhs_type parameter(0) - rhs = $rhs_type parameter(2) - lhs_scale = $lhs_scale_type parameter(1) - rhs_scale = $rhs_scale_type parameter(3) - ROOT _ = $output_type{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={1}, - rhs_contracting_dims={0} -} -)hlo"; - - auto hlo_text = params.PrepareHloText(kHloTextTemplate); - - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(hlo_text)); - EXPECT_TRUE(*RunFileCheck(optimized_module->ToString(), R"( - CHECK: fusion - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )")); - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -INSTANTIATE_TEST_SUITE_P( - TritonScaledDotGemmTest, TritonScaledDotGemmTest, - ::testing::Combine( - ::testing::Values( - ScaleDotTestParams{"f8e4m3fn[128,128]", "f8e4m3fn[128,256]", - "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]", - "bf16[128,256]", "f8E4M3FN"}, - ScaleDotTestParams{"f8e5m2[128,128]", "f8e5m2[128,256]", - "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]", - "bf16[128,256]", "f8E5M2"}), - ::testing::Bool()), - [](const ::testing::TestParamInfo>& - info) { - return absl::StrCat(ScaleDotTestParams::ToString( - ::testing::TestParamInfo( - std::get<0>(info.param), info.index)), - TilingParametersToString(std::get<1>(info.param))); - }); - -class TritonScaledDotTestBase : public TritonEmitterTest { - public: - DebugOptions GetDebugOptionsForTest() const override { - DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest(); - debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(true); - debug_options.set_xla_gpu_autotune_level(0); - debug_options.set_xla_gpu_cublas_fallback(false); - return debug_options; - } - - HloComputation* GetFirstComputationWithInstruction(const HloModule& module, - HloOpcode opcode) const { - for (const auto& computation : module.computations()) { - for (const auto& instruction : computation->instructions()) { - if (instruction->opcode() == opcode) { - return computation; - } - } - } - return nullptr; - } - - absl::Status PopulateScale(Literal* scale, std::minstd_rand0* engine) { - switch (scale->shape().element_type()) { - case F8E8M0FNU: { - absl::uniform_int_distribution exponent_distribution(-4, 1); - return scale->Populate( - [&](absl::Span /*indices*/) { - return float8_e8m0fnu( - std::ldexp(1.0f, exponent_distribution(*engine))); - }); - } - case F8E4M3FN: - for (float8_e4m3fn& value : scale->data()) { - value = float8_e4m3fn(std::abs(static_cast(value))); - } - return absl::OkStatus(); - case F8E5M2: - for (float8_e5m2& value : scale->data()) { - value = float8_e5m2(std::abs(static_cast(value))); - } - return absl::OkStatus(); - default: - return absl::InvalidArgumentError( - absl::StrCat("Unsupported scale type: ", - PrimitiveType_Name(scale->shape().element_type()))); - } - } - - absl::StatusOr> MakeScaledDotArguments( - const HloModule* module) { - std::minstd_rand0 engine; - FakeArgumentsOptions options; - options.engine = &engine; - ABSL_ASSIGN_OR_RETURN(std::vector arguments, - MakeFakeArguments(module, options)); - if (arguments.size() != 4) { - return absl::InternalError(absl::StrCat( - "Expected 4 scaled-dot arguments, got ", arguments.size())); - } - ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[2], &engine)); - ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[3], &engine)); - return arguments; - } -}; - -class TritonScaledDotTest : public TritonScaledDotTestBase, - public ::testing::WithParamInterface { - public: - bool EnableTilingPropagation() const override { return GetParam(); } -}; - -struct Fp4ScaledDotTypeCase { - PrimitiveType lhs_type; - PrimitiveType rhs_type; - PrimitiveType scale_type; - int block_size = 32; -}; - -using Fp4ScaledDotTestParam = std::tuple; - -std::string Fp4ScaledDotTestParamToString( - const ::testing::TestParamInfo& info) { - const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] = - info.param; - return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_", - PrimitiveType_Name(type_case.rhs_type), "_", - PrimitiveType_Name(type_case.scale_type), "_Block", - type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs", - rhs_k_minor, "_", - TilingParametersToString(tiling_enabled)); -} - -class TritonFp4ScaledDotTest - : public TritonScaledDotTestBase, - public ::testing::WithParamInterface { - public: - bool EnableTilingPropagation() const override { - return std::get<3>(GetParam()); - } - - void RunFp4ScaledDotExecutionTest(PrimitiveType lhs_type, - PrimitiveType rhs_type, - PrimitiveType scale_type, int block_size, - bool lhs_k_minor, bool rhs_k_minor) { - constexpr absl::string_view kHloTemplate = R"hlo( -HloModule m - -ENTRY e { - lhs = $lhs_type[$lhs_shape] parameter(0) - rhs = $rhs_type[$rhs_shape] parameter(1) - lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) - rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) - ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={$lhs_contracting_dim}, - rhs_contracting_dims={$rhs_contracting_dim} -} -)hlo"; - - ASSERT_TRUE(scale_type == F8E8M0FNU || scale_type == F8E4M3FN); - - constexpr int64_t m = 128; - constexpr int64_t n = 128; - constexpr int64_t k = 256; - const int64_t scale_k = k / block_size; - - const std::string lhs_shape = - lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); - const std::string rhs_shape = - rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); - const std::string lhs_scale_shape = lhs_k_minor - ? absl::StrCat(m, ",", scale_k) - : absl::StrCat(scale_k, ",", m); - const std::string rhs_scale_shape = rhs_k_minor - ? absl::StrCat(n, ",", scale_k) - : absl::StrCat(scale_k, ",", n); - - LOG(INFO) << "TritonFp4ScaledDotTest Params:" - << "\n LHS Type: " - << primitive_util::LowercasePrimitiveTypeName(lhs_type) - << "\n RHS Type: " - << primitive_util::LowercasePrimitiveTypeName(rhs_type) - << "\n Scale Type: " - << primitive_util::LowercasePrimitiveTypeName(scale_type) - << "\n Block Size: " << block_size - << "\n LHS Contracting Dim: " << (lhs_k_minor ? "1" : "0") - << "\n RHS Contracting Dim: " << (rhs_k_minor ? "1" : "0"); - - std::string hlo = absl::StrReplaceAll( - kHloTemplate, - {{"$lhs_type", primitive_util::LowercasePrimitiveTypeName(lhs_type)}, - {"$rhs_type", primitive_util::LowercasePrimitiveTypeName(rhs_type)}, - {"$scale_type", - primitive_util::LowercasePrimitiveTypeName(scale_type)}, - {"$lhs_shape", lhs_shape}, - {"$rhs_shape", rhs_shape}, - {"$lhs_scale_shape", lhs_scale_shape}, - {"$rhs_scale_shape", rhs_scale_shape}, - {"$output_shape", absl::StrCat(m, ",", n)}, - {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, - {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); - if (scale_type == F8E8M0FNU && block_size == 16 && - GetCudaComputeCapability().IsAtLeastBlackwell() && !lhs_k_minor) { -#ifndef NDEBUG - EXPECT_DEATH( - { (void)GetOptimizedModule(hlo); }, - "MMAv5 with kind=mxf4nvf4 does not support transpose"); - return; -#endif - } - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(hlo)); - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - "CHECK: tt.dot_scaled"), - IsOk()); - ASSERT_OK_AND_ASSIGN(std::vector arguments, - MakeScaledDotArguments(optimized_module.get())); - if (block_size == 16 && GetCudaComputeCapability().IsAtLeastBlackwell()) { - EXPECT_EXIT( - { - EXPECT_THAT(RunAndCompareNoHloPasses( - std::move(optimized_module), - LiteralUtil::MakePointers(arguments), - ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}) - .message(), - HasSubstr("CUDA_ERROR_INVALID_INSTRUCTION")); - std::_Exit(1); - }, - ::testing::ExitedWithCode(1), ""); - return; - } - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), LiteralUtil::MakePointers(arguments), - ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); - } -}; - -INSTANTIATE_TEST_SUITE_P(TritonScaledDotTestSuite, TritonScaledDotTest, - ::testing::Bool(), - [](const ::testing::TestParamInfo& info) { - return TilingParametersToString(info.param); - }); - -TEST_P(TritonFp4ScaledDotTest, Executes) { - auto cc = GpuComputeCapability().cuda_compute_capability(); - if (!cc || !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - const Fp4ScaledDotTypeCase& type_case = std::get<0>(GetParam()); - RunFp4ScaledDotExecutionTest( - type_case.lhs_type, type_case.rhs_type, type_case.scale_type, - type_case.block_size, std::get<1>(GetParam()), std::get<2>(GetParam())); -} - -INSTANTIATE_TEST_SUITE_P( - TritonFp4ScaledDotTestSuite, TritonFp4ScaledDotTest, - ::testing::Combine( - ::testing::Values( - Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, - 32}, // MXFP4 x MXFP4 (block 32) - Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, - 16}, // MXFP4 x MXFP4 (block 16) - Fp4ScaledDotTypeCase{F4E2M1FN, F8E4M3FN, F8E8M0FNU, - 32}, // MXFP4 x MXFP8 (MMAv5 mxf8f6f4) - Fp4ScaledDotTypeCase{F8E4M3FN, F4E2M1FN, F8E8M0FNU, - 32}, // MXFP8 x MXFP4 (MMAv5 mxf8f6f4) - Fp4ScaledDotTypeCase{F4E2M1FN, F8E5M2, F8E8M0FNU, - 32}, // MXFP4 x MXE5M2 - Fp4ScaledDotTypeCase{F8E5M2, F4E2M1FN, F8E8M0FNU, - 32}, // MXE5M2 x MXFP4 - Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E4M3FN, - 16}), // NVFP4 x NVFP4 (block 16) - ::testing::Bool(), ::testing::Bool(), ::testing::Bool()), - Fp4ScaledDotTestParamToString); - -TEST_P(TritonScaledDotTest, - ScaledDotWithOmmittedLhsScaleGetFusedAndExecutedCorrectly) { - if (auto cc = GpuComputeCapability().cuda_compute_capability(); - cc && !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule ScaledDotWithOmmittedLhsScaleGetFusedAndExecutedCorrectly - -ENTRY e { - lhs = bf16[3,128,128] parameter(0) - rhs = f8e4m3fn[3,128,128] parameter(1) - constant = bf16[1,1,1] constant(1.0) - rhs_scale = f8e8m0fnu[3,128,4] parameter(2) - ROOT _ = bf16[3,128,128] scaled-dot(lhs, rhs, constant, rhs_scale), - lhs_batch_dims={0}, - rhs_batch_dims={0}, - lhs_contracting_dims={2}, - rhs_contracting_dims={2} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: fusion - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - true); - for (const auto& computation : optimized_module->computations()) { - for (const auto& instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kScaledDot) { - LOG(INFO) << "Instruction: " << instruction->name(); - } - } - } - - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x128xbf16> - CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> - CHECK: -> tensor<128x16xf32> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, FP8ScaledDotLhsKNotMinorDim) { - if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { - GTEST_SKIP() << "FP8 scaled dot requires Blackwell+"; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule FP8ScaledDotLhsKNotMinorDim - -ENTRY e { - lhs = f8e4m3fn[128,64] parameter(0) - lhs_scale = f8e8m0fnu[4,64] parameter(1) - rhs = f8e4m3fn[128,256] parameter(2) - rhs_scale = f8e8m0fnu[4,256] parameter(3) - ROOT _ = bf16[64,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={0}, - rhs_contracting_dims={0} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: fusion - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - true); - - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, ScaledDotWithBatchGetFusedAndExecutedCorrectly) { - if (auto cc = GpuComputeCapability().cuda_compute_capability(); - cc && !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule ScaledDotWithBatchGetFusedAndExecutedCorrectly - -ENTRY e { - lhs = f8e4m3fn[3,128,128] parameter(0) - rhs = f8e4m3fn[3,128,128] parameter(1) - lhs_scale = f8e8m0fnu[3,128,4] parameter(2) - rhs_scale = f8e8m0fnu[3,128,4 ] parameter(3) - ROOT _ = bf16[3,128,128] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_batch_dims={0}, - rhs_batch_dims={0}, - lhs_contracting_dims={2}, - rhs_contracting_dims={2} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: fusion - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - true); - - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> - CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> - CHECK: -> tensor<128x16xf32> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, BroadcastAndReshapeGetFused) { - if (auto cc = GpuComputeCapability().cuda_compute_capability(); - cc && !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule ScaledDotWithBatchGetFusedAndExecutedCorrectly - -ENTRY e { - lhs = f8e4m3fn[3,128,128] parameter(0) - rhs = f8e4m3fn[3,128,128] parameter(1) - lhs_scale = f8e8m0fnu[3,128,1] parameter(2) - lhs_scale_broadcasted = f8e8m0fnu[3,128,1,4] broadcast(lhs_scale), - dimensions={0,1,2} - lhs_scale_reshaped = f8e8m0fnu[3,128,4] reshape(lhs_scale_broadcasted) - rhs_scale = f8e8m0fnu[3,128,1] parameter(3) - rhs_scale_broadcasted = f8e8m0fnu[3,128,1,4] broadcast(rhs_scale), - dimensions={0,1,2} - rhs_scale_reshaped = f8e8m0fnu[3,128,4] reshape(rhs_scale_broadcasted) - ROOT _ = bf16[3,128,128] scaled-dot( - lhs, - rhs, - lhs_scale_reshaped, - rhs_scale_reshaped), - lhs_batch_dims={0}, - rhs_batch_dims={0}, - lhs_contracting_dims={2}, - rhs_contracting_dims={2} -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: %fusion - CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1} - CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1} - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - true); - - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> - CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> - CHECK: -> tensor<128x16xf32> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, Mxfp8ScaledDotSmallBlockKAndNExecutes) { - if (!GetCudaComputeCapability().IsAtLeastHopper()) { - GTEST_SKIP() << "Requires Hopper+."; - } - constexpr absl::string_view kHloText = R"hlo( -HloModule m - -fusion__ { - parameter_0 = f8e4m3fn[128,256]{1,0} parameter(0) - parameter_1 = f8e4m3fn[256,128]{1,0} parameter(1) - parameter_2 = f8e8m0fnu[128,8]{1,0} parameter(2) - parameter_3 = f8e8m0fnu[8,128]{1,0} parameter(3) - ROOT _.1 = bf16[128,128]{1,0} scaled-dot( - parameter_0, parameter_1, parameter_2, parameter_3), - lhs_contracting_dims={1}, rhs_contracting_dims={0}, - backend_config={"sizes":["64"]} -} - -ENTRY e { - lhs = f8e4m3fn[128,256]{1,0} parameter(0) - rhs = f8e4m3fn[256,128]{1,0} parameter(1) - lhs_scale = f8e8m0fnu[128,8]{1,0} parameter(2) - rhs_scale = f8e8m0fnu[8,128]{1,0} parameter(3) - ROOT fusion = bf16[128,128]{1,0} fusion( - lhs, rhs, lhs_scale, rhs_scale), - kind=kCustom, calls=fusion__, - backend_config={"fusion_backend_config":{ - "kind":"__triton_nested_gemm_fusion", - "block_level_fusion_config":{ - "output_tiles":[{"sizes":["128","16"]}], - "num_warps":"4","num_ctas":"1","num_stages":"1"}}} -} -)hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText)); - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, Fp4Succeeds) { - if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { - GTEST_SKIP() << "Scaled dot with FP4 isn't supported by Triton for " - "pre-Blackwell GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( - HloModule jit_scaled_dot_fn - - ENTRY %main.2 { - %lhs = f4e2m1fn[1,1024,256]{2,1,0} parameter(0) - %rhs = f4e2m1fn[1,256,256]{2,1,0} parameter(1) - %lhs_scale = f8e8m0fnu[1,1024,8]{2,1,0} parameter(2) - %rhs_scale = f8e8m0fnu[1,8,256]{2,1,0} parameter(3) - ROOT %scaled-dot = bf16[1,1024,256]{2,1,0} scaled-dot(%lhs, %rhs, %lhs_scale, %rhs_scale), - lhs_batch_dims={0}, - lhs_contracting_dims={2}, - rhs_batch_dims={0}, - rhs_contracting_dims={1} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - ASSERT_OK_AND_ASSIGN(GpuBackendConfig gpu_config, - scaled_dot_computation->FusionInstruction() - ->backend_config()); - const BlockLevelFusionConfig& block_level_fusion_config = - gpu_config.fusion_backend_config().block_level_fusion_config(); - ASSERT_EQ(block_level_fusion_config.output_tiles_size(), 1); - const auto& output_tile_sizes = - block_level_fusion_config.output_tiles(0).sizes(); - ASSERT_EQ(output_tile_sizes.size(), 3); - const int64_t output_m = output_tile_sizes.Get(1); - const int64_t output_n = output_tile_sizes.Get(2); - ASSERT_EQ(output_n % 2, 0); - const std::string expected_triton_ir = - absl::Substitute(R"( - CHECK: tt.dot_scaled - CHECK: tensor<$0x64xi8>, tensor<$0x4xi8> - CHECK: tensor<128x$1xi8>, tensor<$2x4xi8> - CHECK: -> tensor<$0x$2xf32> - )", - output_m, output_n / 2, output_n); - - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - expected_triton_ir), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, Mxfp4KPackedRhsBlackwellLowersWithTranspose) { - if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { - GTEST_SKIP() << "Requires Blackwell+."; - } - constexpr absl::string_view kHloText = R"hlo( -HloModule m -fusion__ { - parameter_0 = f4e2m1fn[128,256]{1,0:E(4)} parameter(0) - parameter_1 = f4e2m1fn[128,256]{1,0:E(4)} parameter(1) - parameter_2 = f8e8m0fnu[128,8]{1,0} parameter(2) - parameter_3 = f8e8m0fnu[128,8]{1,0} parameter(3) - ROOT _.1 = bf16[128,128] scaled-dot(parameter_0, parameter_1, parameter_2, parameter_3), - lhs_contracting_dims={1}, rhs_contracting_dims={1}, - backend_config={"sizes":["64"]} -} -ENTRY e { - lhs = f4e2m1fn[128,256]{1,0:E(4)} parameter(0) - rhs = f4e2m1fn[128,256]{1,0:E(4)} parameter(1) - lhs_scale = f8e8m0fnu[128,8]{1,0} parameter(2) - rhs_scale = f8e8m0fnu[128,8]{1,0} parameter(3) - ROOT fusion = bf16[128,128] fusion(lhs, rhs, lhs_scale, rhs_scale), - kind=kCustom, calls=fusion__, - backend_config={"fusion_backend_config":{"kind":"__triton_nested_gemm_fusion", - "block_level_fusion_config":{ - "output_tiles":[{"sizes":["128","128"]}], - "num_warps":"4","num_ctas":"1","num_stages":"1"}}} -} -)hlo"; - TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText)); - HloComputation* scaled_dot_computation = - GetFirstComputationWithInstruction(*module, HloOpcode::kScaledDot); - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, R"( - CHECK: tt.trans - CHECK: tensor<128x32xi8> -> tensor<32x128xi8> - CHECK-NOT: unrealized_conversion_cast - CHECK: tt.dot_scaled - CHECK: tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> - )"), - IsOk()) - << "We expect to see tt.trans and the follow up dot_scaled"; - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, GlobalScalerSucceeds) { - if (!GetCudaComputeCapability().IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - constexpr absl::string_view kHloTextTemplate = R"hlo( -HloModule ScaledDotWithGlobalScaler - -ENTRY e { - lhs = f8e4m3fn[3,128,128] parameter(0) - rhs = f8e4m3fn[3,128,128] parameter(1) - lhs_scale = f8e8m0fnu[3,128,4] parameter(2) - rhs_scale = f8e8m0fnu[3,128,4] parameter(3) - scaled_dot = bf16[3,128,128] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_batch_dims={0}, - rhs_batch_dims={0}, - lhs_contracting_dims={2}, - rhs_contracting_dims={2} - global_scaler = bf16[] constant(1.42) - global_scaler_broadcasted = bf16[3,128,128] broadcast(global_scaler), - dimensions={} - ROOT _ = bf16[3,128,128] multiply(scaled_dot, global_scaler_broadcasted) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, - GetOptimizedModule(kHloTextTemplate)); - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: %[[fusion_name:.*]] (parameter - CHECK: %[[scaled_dot:.*]] = bf16[3,128,128]{2,1,0} scaled-dot - CHECK: %[[global_scaler:.*]] = bf16[3,128,128]{2,1,0} broadcast - CHECK: ROOT %{{.*}} = bf16[3,128,128]{2,1,0} multiply(%[[scaled_dot]], %[[global_scaler]]) - CHECK: ENTRY - CHECK: ROOT {{.*}} fusion({{.*}}), kind=kCustom, calls=%[[fusion_name]] - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - true); - - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> - CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> - CHECK: -> tensor<128x16xf32> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, ScaledDotWithE8m0Scale) { - if (!GetCudaComputeCapability().IsAtLeastHopper()) { - GTEST_SKIP() - << "ScaledDot with Triton requires Hopper or newer architecture."; - } - constexpr absl::string_view kHloText = R"hlo( -HloModule E8m0ScaledDot - -ENTRY e { - lhs = f8e4m3fn[128,128] parameter(0) - rhs = f8e4m3fn[128,256] parameter(2) - lhs_scale = f8e8m0fnu[128,4] parameter(1) - rhs_scale = f8e8m0fnu[4,256] parameter(3) - ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={1}, - rhs_contracting_dims={0} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); - - // Verify HLO fusion - constexpr absl::string_view kExpectedOptimizedHLO = R"( - CHECK: fusion - CHECK: ROOT {{.*}} scaled-dot - CHECK: ENTRY - CHECK: __triton_nested_gemm_fusion - )"; - EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), - absl_testing::IsOkAndHolds(true)); - - // Verify Triton IR - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x4xi8> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - // Execute on GPU hardware and compare with reference - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); -} - -TEST_P(TritonScaledDotTest, ScaledDotWithE4m3Scale) { - if (!GetCudaComputeCapability().IsAtLeastHopper()) { - GTEST_SKIP() - << "ScaledDot with Triton requires Hopper or newer architecture."; - } - constexpr absl::string_view kHloText = R"hlo( -HloModule E4m3ScaledDot - -ENTRY e { - lhs = f8e4m3fn[128,128] parameter(0) - rhs = f8e4m3fn[128,256] parameter(2) - lhs_scale = f8e4m3fn[128,4] parameter(1) - rhs_scale = f8e4m3fn[4,256] parameter(3) - ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={1}, - rhs_contracting_dims={0} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); - - // Verify Triton IR contains tt.dot_scaled with f8e4m3fn scale tensors - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x4xf8E4M3FN> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - // Execute on GPU hardware and compare with reference (Hopper only; Blackwell - // TMEM scale expects E8M0 scale). - if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); - } -} - -TEST_P(TritonScaledDotTest, ScaledDotWithE5m2Scale) { - if (!GetCudaComputeCapability().IsAtLeastHopper()) { - GTEST_SKIP() - << "ScaledDot with Triton requires Hopper or newer architecture."; - } - constexpr absl::string_view kHloText = R"hlo( -HloModule E5m2ScaledDot - -ENTRY e { - lhs = f8e4m3fn[128,128] parameter(0) - rhs = f8e4m3fn[128,256] parameter(2) - lhs_scale = f8e5m2[128,4] parameter(1) - rhs_scale = f8e5m2[4,256] parameter(3) - ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={1}, - rhs_contracting_dims={0} -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); - - // Verify Triton IR contains tt.dot_scaled with f8e5m2 scale tensors - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); - constexpr absl::string_view kExpectedTritonIr = R"( - CHECK: tt.dot_scaled - CHECK: tensor<128x4xf8E5M2> - )"; - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - kExpectedTritonIr), - IsOk()); - - // Execute on GPU hardware and compare with reference (Hopper only; Blackwell - // TMEM scale expects E8M0 scale). - if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { - EXPECT_TRUE(RunAndCompareNoHloPasses( - std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); - } -} - -struct ScaledDotCoverageTestCase { - PrimitiveType lhs_type; - PrimitiveType rhs_type; - PrimitiveType scale_type; - int block_size; -}; - -class TritonScaledDotCoverageTest : public TritonScaledDotTestBase, - public ::testing::WithParamInterface< - std::tuple> { - public: - bool EnableTilingPropagation() const override { - return std::get<3>(GetParam()); - } -}; - -std::vector GetCoverageTestCases() { - std::vector input_types = {F8E4M3FN, F8E5M2, F4E2M1FN}; - std::vector cases; - - // 1. All input combinations with E8M0 scale and block size 32 (9 cases) - for (auto lhs : input_types) { - for (auto rhs : input_types) { - cases.push_back({lhs, rhs, F8E8M0FNU, 32}); - } - } - - // 2. Block size 16 cases - cases.push_back({F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16}); - cases.push_back({F4E2M1FN, F4E2M1FN, F8E4M3FN, 16}); - - return cases; -} - -std::string ScaledDotCoverageTestParamToString( - const ::testing::TestParamInfo< - std::tuple>& info) { - const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] = - info.param; - return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_", - PrimitiveType_Name(type_case.rhs_type), "_", - PrimitiveType_Name(type_case.scale_type), "_Block", - type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs", - rhs_k_minor, "_", - TilingParametersToString(tiling_enabled)); -} - -TEST_P(TritonScaledDotCoverageTest, Executes) { - const auto& [param, lhs_k_minor, rhs_k_minor, tiling_enabled] = GetParam(); - - auto cc = GpuComputeCapability().cuda_compute_capability(); - std::string device_name = "Unknown"; - if (cc) { - device_name = cc->IsAtLeastBlackwell() - ? "B200" - : (cc->IsAtLeastHopper() ? "H100" : "Pre-Hopper"); - } - - std::string lhs_name = - primitive_util::LowercasePrimitiveTypeName(param.lhs_type); - std::string rhs_name = - primitive_util::LowercasePrimitiveTypeName(param.rhs_type); - std::string scale_name = - primitive_util::LowercasePrimitiveTypeName(param.scale_type); - std::string lhs_display = - lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T"); - std::string rhs_display = - rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name; - std::string tiling_name = tiling_enabled ? "Symbolic" : "Experimental"; - - LOG(ERROR) << "Report Device: " << device_name; - LOG(ERROR) << "Report LHS: " << lhs_display; - LOG(ERROR) << "Report RHS: " << rhs_display; - LOG(ERROR) << "Report Scale: " << scale_name; - LOG(ERROR) << "Report BlockSize: " << param.block_size; - LOG(ERROR) << "Report Tiling: " << tiling_name; - - if (!cc || !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - - constexpr int64_t m = 128; - constexpr int64_t n = 128; - constexpr int64_t k = 256; - const int64_t scale_k = k / param.block_size; - - const std::string lhs_shape = - lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); - const std::string rhs_shape = - rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); - const std::string lhs_scale_shape = lhs_k_minor - ? absl::StrCat(m, ",", scale_k) - : absl::StrCat(scale_k, ",", m); - const std::string rhs_scale_shape = rhs_k_minor - ? absl::StrCat(n, ",", scale_k) - : absl::StrCat(scale_k, ",", n); - - // 1. Stage: After Composite - std::string after_composite = "N/A"; - constexpr absl::string_view kCompositeHloTemplate = R"hlo( -HloModule test_composite - -%xla.scaled_dot.1 { - %p0 = $lhs_type[$lhs_shape]{1,0} parameter(0) - %p1 = $rhs_type[$rhs_shape]{1,0} parameter(1) - %p2 = $scale_type[$lhs_scale_shape]{1,0} parameter(2) - %p3 = $scale_type[$rhs_scale_shape]{1,0} parameter(3) - ROOT %dummy = bf16[$output_shape]{1,0} constant({...}) -} - -ENTRY %main { - %lhs = $lhs_type[$lhs_shape]{1,0} parameter(0) - %rhs = $rhs_type[$rhs_shape]{1,0} parameter(1) - %lhs_scale = $scale_type[$lhs_scale_shape]{1,0} parameter(2) - %rhs_scale = $scale_type[$rhs_scale_shape]{1,0} parameter(3) - ROOT %call = bf16[$output_shape]{1,0} call(%lhs, %rhs, %lhs_scale, %rhs_scale), - to_apply=%xla.scaled_dot.1, - is_composite=true, - frontend_attributes={ - composite.attributes="{dimension_numbers=[[[$lhs_contracting_dim],[$rhs_contracting_dim]],[[],[]]]}", - composite.name="xla.scaled_dot", - composite.version="1" - } -} -)hlo"; - - std::string composite_hlo = - absl::StrReplaceAll(kCompositeHloTemplate, - {{"$lhs_type", lhs_name}, - {"$rhs_type", rhs_name}, - {"$scale_type", scale_name}, - {"$lhs_shape", lhs_shape}, - {"$rhs_shape", rhs_shape}, - {"$lhs_scale_shape", lhs_scale_shape}, - {"$rhs_scale_shape", rhs_scale_shape}, - {"$output_shape", absl::StrCat(m, ",", n)}, - {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, - {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); - - if (auto comp_module = ParseAndReturnUnverifiedModule(composite_hlo); - comp_module.ok()) { - CompositeRewriter rewriter; - if (auto rewrite_status = rewriter.Run(comp_module->get()); - rewrite_status.ok()) { - const auto* root = - (*comp_module)->entry_computation()->root_instruction(); - after_composite = HloOpcodeString(root->opcode()); - } - } - LOG(ERROR) << "Report AfterComposite: " << after_composite; - - // 2. Stage: Optimized HLO - constexpr absl::string_view kHloTemplate = R"hlo( -HloModule m - -ENTRY e { - lhs = $lhs_type[$lhs_shape] parameter(0) - rhs = $rhs_type[$rhs_shape] parameter(1) - lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) - rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) - ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={$lhs_contracting_dim}, - rhs_contracting_dims={$rhs_contracting_dim} -} -)hlo"; - - std::string hlo = absl::StrReplaceAll( - kHloTemplate, {{"$lhs_type", lhs_name}, - {"$rhs_type", rhs_name}, - {"$scale_type", scale_name}, - {"$lhs_shape", lhs_shape}, - {"$rhs_shape", rhs_shape}, - {"$lhs_scale_shape", lhs_scale_shape}, - {"$rhs_scale_shape", rhs_scale_shape}, - {"$output_shape", absl::StrCat(m, ",", n)}, - {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, - {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); - - if (param.scale_type == F8E8M0FNU && param.block_size == 16 && - GetCudaComputeCapability().IsAtLeastBlackwell() && !lhs_k_minor) { -#ifndef NDEBUG - EXPECT_DEATH( - { (void)GetOptimizedModule(hlo); }, - "MMAv5 with kind=mxf4nvf4 does not support transpose"); - return; -#endif - } - - std::string optimized_hlo = "N/A"; - auto optimized_module_or = GetOptimizedModule(hlo); - HloComputation* scaled_dot_computation = nullptr; - if (optimized_module_or.ok()) { - scaled_dot_computation = GetFirstComputationWithInstruction( - **optimized_module_or, HloOpcode::kScaledDot); - if (scaled_dot_computation) { - optimized_hlo = "scaled-dot"; - } else if (GetFirstComputationWithInstruction(**optimized_module_or, - HloOpcode::kDot)) { - optimized_hlo = "dot"; - } - } - LOG(ERROR) << "Report OptimizedHLO: " << optimized_hlo; - - // 3. Stage: Triton Dot Lowering - std::string triton_dot = "N/A"; - if (scaled_dot_computation && scaled_dot_computation->FusionInstruction()) { - mlir::MLIRContext mlir_context; - BlockLevelParameters block_level_parameters; - if (auto gpu_config = scaled_dot_computation->FusionInstruction() - ->backend_config(); - gpu_config.ok() && gpu_config->has_fusion_backend_config() && - gpu_config->fusion_backend_config().has_block_level_fusion_config()) { - block_level_parameters = BlockLevelParameters::FromBlockLevelFusionConfig( - gpu_config->fusion_backend_config().block_level_fusion_config()); - } - auto triton_source_or = - CreateTritonModule("triton_fn", - *Cast( - scaled_dot_computation->FusionInstruction()), - TestGpuDeviceInfo::RTXA6000DeviceInfo(*cc), - block_level_parameters, mlir_context); - if (triton_source_or.ok()) { - std::string mlir_str; - llvm::raw_string_ostream os(mlir_str); - triton_source_or->module()->print(os); - if (absl::StrContains(mlir_str, "tt.dot_scaled")) { - triton_dot = "tt.dot_scaled"; - } else if (absl::StrContains(mlir_str, "tt.dot")) { - triton_dot = "tt.dot"; - } else if (absl::StrContains(mlir_str, "xtile.dot_scaled")) { - triton_dot = "xtile.dot_scaled"; - } - } - } - LOG(ERROR) << "Report TritonDot: " << triton_dot; - - // 4. Stage: PTX Generation - std::vector ptx_instructions; - GpuCompiler* gpu_compiler = dynamic_cast(compiler()); - bool compilation_succeeded = false; - if (gpu_compiler && optimized_module_or.ok()) { - gpu_compiler->SetAsmHook([&](absl::string_view ptx) { - std::string ptx_str(ptx); - std::string line; - size_t pos = 0; - while ((pos = ptx_str.find('\n')) != std::string::npos) { - line = ptx_str.substr(0, pos); - size_t first_non_space = line.find_first_not_of(" \t"); - if (first_non_space != std::string::npos) { - line = line.substr(first_non_space); - } - std::vector tokens = - absl::StrSplit(line, absl::ByAnyChar(" \t;"), absl::SkipEmpty()); - if (!tokens.empty()) { - std::string inst = tokens[0]; - if (absl::StartsWith(inst, "@") && tokens.size() > 1) { - inst = tokens[1]; - } - if (absl::StrContains(inst, "mma") && - (absl::StartsWith(inst, "wgmma.mma") || - absl::StartsWith(inst, "tcgen05.mma") || - absl::StartsWith(inst, "mma.sync"))) { - ptx_instructions.push_back(inst); - } - } - ptx_str.erase(0, pos + 1); - } - }); - - auto cloned_module = (*optimized_module_or)->Clone(); - auto executable_or = CompileToExecutable(std::move(cloned_module), - /*run_optimization_passes=*/true); - compilation_succeeded = executable_or.ok(); - gpu_compiler->RemoveAsmHook(); - } - - std::vector unique_ptx; - for (const auto& inst : ptx_instructions) { - if (absl::c_find(unique_ptx, inst) == unique_ptx.end()) { - unique_ptx.push_back(inst); - } - } - std::string ptx_str = - unique_ptx.empty() ? "N/A" : absl::StrJoin(unique_ptx, ", "); - LOG(ERROR) << "Report PTX: " << ptx_str; - - // 5. Stage: Execution Status - std::string status = "FAILURE"; - if (compilation_succeeded) { - auto arguments_or = MakeScaledDotArguments((*optimized_module_or).get()); - if (arguments_or.ok()) { - auto run_result = - RunAndCompareNoHloPasses(std::move(*optimized_module_or), - LiteralUtil::MakePointers(*arguments_or), - ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}); - status = run_result ? "SUCCESS" : "FAILURE"; - if (!run_result) { - if (absl::StrContains(run_result.message(), "ILLEGAL_INSTRUCTION")) { - status = "ILLEGAL_INSTRUCTION"; - } - } - } - } - LOG(ERROR) << "Report Status: " << status; - if (GetCudaComputeCapability().IsAtLeastBlackwell() && - status == "ILLEGAL_INSTRUCTION") { - EXPECT_DEATH( - { CHECK_NE(status, "ILLEGAL_INSTRUCTION"); }, "ILLEGAL_INSTRUCTION"); - } -} - -INSTANTIATE_TEST_SUITE_P( - TritonScaledDotCoverageTestSuite, TritonScaledDotCoverageTest, - ::testing::Combine(::testing::ValuesIn(GetCoverageTestCases()), - ::testing::Bool(), ::testing::Bool(), ::testing::Bool()), - ScaledDotCoverageTestParamToString); - // Scans are only emitted with the experimental tiling propagation enabled. class ScanTritonEmitterTest : public TritonEmitterTest { public: diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc new file mode 100644 index 00000000000000..d8a21a19a5f94a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc @@ -0,0 +1,1426 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/algorithm/container.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/random/uniform_int_distribution.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/status_matchers.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "absl/strings/str_replace.h" +#include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" +#include "absl/strings/substitute.h" +#include "absl/types/span.h" +#include "llvm/Support/raw_ostream.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "xla/backends/gpu/codegen/triton/test_utils.h" +#include "xla/backends/gpu/codegen/triton/xtile_compiler.h" +#include "xla/backends/gpu/codegen/triton/xtile_test_base.h" +#include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h" +#include "xla/backends/gpu/transforms/composite_rewriter.h" +#include "xla/codegen/xtile/block_level_parameters.h" +#include "xla/error_spec.h" +#include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" +#include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/primitive_util.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu/gpu_compiler.h" +#include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/tests/hlo_interpreter_reference_mixin.h" +#include "xla/tests/test_utils.h" +#include "xla/tsl/platform/statusor.h" +#include "xla/types.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace gpu { +namespace { + +using ::absl_testing::IsOk; +using ::absl_testing::IsOkAndHolds; +using ::testing::HasSubstr; +using ::xla::xtile::BlockLevelFusionConfig; +using ::xla::xtile::BlockLevelParameters; + +std::string TilingParametersToString(bool tiling_propagation_enabled) { + return tiling_propagation_enabled ? "SymbolicTiling" : "ExperimentalTiling"; +} + +class TritonEmitterTest + : public HloInterpreterReferenceMixin, + public XTileTestBase { + public: + virtual bool EnableTilingPropagation() const = 0; + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloInterpreterReferenceMixin< + GpuPjRtCodegenTest>::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_unsupported_enable_triton_multi_output_fusion( + true); + debug_options.set_xla_gpu_experimental_disable_binary_libraries(true); + debug_options.set_xla_gpu_experimental_enable_tiling_propagation( + EnableTilingPropagation()); + return debug_options; + } + + const stream_executor::GpuComputeCapability& GpuComputeCapability() { + return device_description().gpu_compute_capability(); + } + stream_executor::CudaComputeCapability GetCudaComputeCapability() { + return device_description().cuda_compute_capability(); + } + absl::StatusOr< + std::pair, std::unique_ptr>> + CreateXTileIrAndFileCheck(absl::string_view hlo_text, + absl::string_view triton_fusion_name, + absl::string_view filecheck_pattern) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + return XTileTestBase::CreateXTileIrAndFileCheck( + std::move(module), triton_fusion_name, filecheck_pattern); + } + absl::Status CreateTritonIrFromHloTextAndFileCheck( + absl::string_view hlo_text, absl::string_view triton_fusion_name, + absl::string_view filecheck_pattern) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + return CreateTritonIrAndFileCheck(module.get(), triton_fusion_name, + filecheck_pattern); + } + absl::Status CreateTritonIrFromHloTextAndFileCheckForDot( + absl::string_view hlo_text, absl::string_view triton_fusion_name, + absl::string_view filecheck_pattern) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + return CreateTritonIrAndFileCheckForDot(module.get(), triton_fusion_name, + filecheck_pattern); + } +}; + +class TritonEmitterTestWithTilingParam + : public TritonEmitterTest, + public ::testing::WithParamInterface { + public: + bool EnableTilingPropagation() const override { return GetParam(); } +}; + +TEST_P(TritonEmitterTestWithTilingParam, + ScaledDotIsSupportedByReferencePlatform) { + constexpr absl::string_view kHloText = R"( + HloModule ScaledDotIsSupportedByReferencePlatform + + ENTRY entry { + lhs = bf16[16,128] parameter(0) + rhs = bf16[128,16] parameter(1) + lhs_scale = bf16[1,4] parameter(2) + rhs_scale = bf16[4,1] parameter(3) + ROOT dot = bf16[16,16] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} + } + )"; + + EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +INSTANTIATE_TEST_SUITE_P(TritonEmitterTestWithTilingParamTestSuite, + TritonEmitterTestWithTilingParam, ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return TilingParametersToString(info.param); + }); + +struct ScaleDotTestParams { + std::string lhs_type; + std::string rhs_type; + std::string lhs_scale_type; + std::string rhs_scale_type; + std::string output_type; + std::string expected_triton_type; + + std::string PrepareHloText(absl::string_view hlo_template) const { + return absl::StrReplaceAll(hlo_template, + {{"$lhs_type", lhs_type}, + {"$rhs_type", rhs_type}, + {"$lhs_scale_type", lhs_scale_type}, + {"$rhs_scale_type", rhs_scale_type}, + {"$output_type", output_type}}); + } + static std::string ToString( + const ::testing::TestParamInfo& info) { + const ScaleDotTestParams& params = info.param; + auto name = absl::StrCat(params.lhs_type, "_", params.rhs_type, "_", + params.lhs_scale_type, "_", params.rhs_scale_type, + "_", params.output_type); + absl::StrReplaceAll({{"[", "_"}, {"]", "_"}, {",", "x"}}, &name); + return name; + } +}; + +std::ostream& operator<<(std::ostream& stream, const ScaleDotTestParams& tc) { + return stream << "{\n\tlhs_type:" << tc.lhs_type + << ",\n\trhs_type:" << tc.rhs_type + << ",\n\tlhs_scale_type:" << tc.lhs_scale_type + << ",\n\trhs_scale_type:" << tc.rhs_scale_type + << ",\n\toutput_type:" << tc.output_type << "\n}"; +} + +class TritonScaledDotGemmTest : public TritonEmitterTest, + public ::testing::WithParamInterface< + std::tuple> { + public: + bool EnableTilingPropagation() const override { + return std::get<1>(GetParam()); + } + + public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(true); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.set_xla_gpu_cublas_fallback(false); + return debug_options; + } +}; + +TEST_P(TritonScaledDotGemmTest, + FP8ScaledDotCompilesToPtxIntrinsicsWhenAvailable) { + const ScaleDotTestParams& params = std::get<0>(GetParam()); + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule m + +triton_dot { + lhs = $lhs_type parameter(0) + rhs = $rhs_type parameter(1) + lhs_scale = $lhs_scale_type parameter(2) + rhs_scale = $rhs_scale_type parameter(3) + ROOT _ = $output_type{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0}, + backend_config={sizes:[128]} +} + +ENTRY e { + lhs = $lhs_type{1,0} parameter(0) + rhs = $rhs_type{1,0} parameter(1) + lhs_scale = $lhs_scale_type{1,0} parameter(2) + rhs_scale = $rhs_scale_type{1,0} parameter(3) + ROOT _ = $output_type{1,0} fusion(lhs, rhs, lhs_scale, rhs_scale), + kind=kCustom, + calls=triton_dot, + backend_config={ + "fusion_backend_config": { + kind: "__triton_nested_gemm_fusion", + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["128", "256"]}], + "num_warps":"4", + "num_stages":"1", + "num_ctas":"1" + } + } + } +} +)hlo"; + + auto hlo_text = params.PrepareHloText(kHloTextTemplate); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + + constexpr absl::string_view kExpectedTritonIrTmpl = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x128x$triton_type>, tensor<128x4xi8> + CHECK: tensor<128x256x$triton_type>, tensor<256x4xi8> + CHECK: -> tensor<128x256xf32> + )"; + auto expected_triton_ir = absl::StrReplaceAll( + kExpectedTritonIrTmpl, {{"$triton_type", params.expected_triton_type}}); + EXPECT_THAT( + CreateTritonIrAndFileCheckForDot( + *module->GetComputationWithName("triton_dot"), expected_triton_ir), + IsOk()); + if (GetCudaComputeCapability().IsAtLeastBlackwell()) { + CompileAndOptionallyVerifyPtx( + std::move(module), R"(CHECK: mxf8f6f4.block_scale.scale_vec::1X)"); + } +} + +TEST_P(TritonScaledDotGemmTest, FP8ScaledDotGetsFusedAndExecutesCorrectly) { + const ScaleDotTestParams& params = std::get<0>(GetParam()); + if (auto cc = GpuComputeCapability().cuda_compute_capability(); + cc && !cc->IsAtLeastBlackwell()) { + GTEST_SKIP() << "Skipping test for pre-Blackwell GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule FP8ScaledDotGetsFusedAndExecutesCorrectly + +ENTRY e { + lhs = $lhs_type parameter(0) + rhs = $rhs_type parameter(2) + lhs_scale = $lhs_scale_type parameter(1) + rhs_scale = $rhs_scale_type parameter(3) + ROOT _ = $output_type{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} +} +)hlo"; + + auto hlo_text = params.PrepareHloText(kHloTextTemplate); + + ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(hlo_text)); + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), R"( + CHECK: fusion + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"), + IsOkAndHolds(true)); + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +INSTANTIATE_TEST_SUITE_P( + TritonScaledDotGemmTest, TritonScaledDotGemmTest, + ::testing::Combine( + ::testing::Values( + ScaleDotTestParams{"f8e4m3fn[128,128]", "f8e4m3fn[128,256]", + "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]", + "bf16[128,256]", "f8E4M3FN"}, + ScaleDotTestParams{"f8e5m2[128,128]", "f8e5m2[128,256]", + "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]", + "bf16[128,256]", "f8E5M2"}), + ::testing::Bool()), + [](const ::testing::TestParamInfo>& + info) { + return absl::StrCat(ScaleDotTestParams::ToString( + ::testing::TestParamInfo( + std::get<0>(info.param), info.index)), + TilingParametersToString(std::get<1>(info.param))); + }); + +class TritonScaledDotTestBase : public TritonEmitterTest { + public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(true); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.set_xla_gpu_cublas_fallback(false); + return debug_options; + } + + HloComputation* GetFirstComputationWithInstruction(const HloModule& module, + HloOpcode opcode) const { + for (const auto& computation : module.computations()) { + for (const auto& instruction : computation->instructions()) { + if (instruction->opcode() == opcode) { + return computation; + } + } + } + return nullptr; + } + + absl::Status PopulateScale(Literal* scale, std::minstd_rand0* engine) { + switch (scale->shape().element_type()) { + case F8E8M0FNU: { + absl::uniform_int_distribution exponent_distribution(-4, 1); + return scale->Populate( + [&](absl::Span /*indices*/) { + return float8_e8m0fnu( + std::ldexp(1.0f, exponent_distribution(*engine))); + }); + } + case F8E4M3FN: + for (float8_e4m3fn& value : scale->data()) { + value = float8_e4m3fn(std::abs(static_cast(value))); + } + return absl::OkStatus(); + case F8E5M2: + for (float8_e5m2& value : scale->data()) { + value = float8_e5m2(std::abs(static_cast(value))); + } + return absl::OkStatus(); + default: + return absl::InvalidArgumentError( + absl::StrCat("Unsupported scale type: ", + PrimitiveType_Name(scale->shape().element_type()))); + } + } + + absl::StatusOr> MakeScaledDotArguments( + const HloModule* module) { + std::minstd_rand0 engine; + FakeArgumentsOptions options; + options.engine = &engine; + ABSL_ASSIGN_OR_RETURN(std::vector arguments, + MakeFakeArguments(module, options)); + if (arguments.size() != 4) { + return absl::InternalError(absl::StrCat( + "Expected 4 scaled-dot arguments, got ", arguments.size())); + } + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[2], &engine)); + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[3], &engine)); + return arguments; + } +}; + +class TritonScaledDotTest : public TritonScaledDotTestBase, + public ::testing::WithParamInterface { + public: + bool EnableTilingPropagation() const override { return GetParam(); } +}; + +struct Fp4ScaledDotTypeCase { + PrimitiveType lhs_type; + PrimitiveType rhs_type; + PrimitiveType scale_type; + int block_size = 32; +}; + +using Fp4ScaledDotTestParam = std::tuple; + +std::string Fp4ScaledDotTestParamToString( + const ::testing::TestParamInfo& info) { + const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] = + info.param; + return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_", + PrimitiveType_Name(type_case.rhs_type), "_", + PrimitiveType_Name(type_case.scale_type), "_Block", + type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs", + rhs_k_minor, "_", + TilingParametersToString(tiling_enabled)); +} + +class TritonFp4ScaledDotTest + : public TritonScaledDotTestBase, + public ::testing::WithParamInterface { + public: + bool EnableTilingPropagation() const override { + return std::get<3>(GetParam()); + } + + void RunFp4ScaledDotExecutionTest(PrimitiveType lhs_type, + PrimitiveType rhs_type, + PrimitiveType scale_type, int block_size, + bool lhs_k_minor, bool rhs_k_minor) { + constexpr absl::string_view kHloTemplate = R"hlo( +HloModule m + +ENTRY e { + lhs = $lhs_type[$lhs_shape] parameter(0) + rhs = $rhs_type[$rhs_shape] parameter(1) + lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) + rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) + ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={$lhs_contracting_dim}, + rhs_contracting_dims={$rhs_contracting_dim} +} +)hlo"; + + ASSERT_TRUE(scale_type == F8E8M0FNU || scale_type == F8E4M3FN); + + constexpr int64_t m = 128; + constexpr int64_t n = 128; + constexpr int64_t k = 256; + const int64_t scale_k = k / block_size; + + const std::string lhs_shape = + lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); + const std::string rhs_shape = + rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); + const std::string lhs_scale_shape = lhs_k_minor + ? absl::StrCat(m, ",", scale_k) + : absl::StrCat(scale_k, ",", m); + const std::string rhs_scale_shape = rhs_k_minor + ? absl::StrCat(n, ",", scale_k) + : absl::StrCat(scale_k, ",", n); + + LOG(INFO) << "TritonFp4ScaledDotTest Params:" + << "\n LHS Type: " + << primitive_util::LowercasePrimitiveTypeName(lhs_type) + << "\n RHS Type: " + << primitive_util::LowercasePrimitiveTypeName(rhs_type) + << "\n Scale Type: " + << primitive_util::LowercasePrimitiveTypeName(scale_type) + << "\n Block Size: " << block_size + << "\n LHS Contracting Dim: " << (lhs_k_minor ? "1" : "0") + << "\n RHS Contracting Dim: " << (rhs_k_minor ? "1" : "0"); + + std::string hlo = absl::StrReplaceAll( + kHloTemplate, + {{"$lhs_type", primitive_util::LowercasePrimitiveTypeName(lhs_type)}, + {"$rhs_type", primitive_util::LowercasePrimitiveTypeName(rhs_type)}, + {"$scale_type", + primitive_util::LowercasePrimitiveTypeName(scale_type)}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + if (scale_type == F8E8M0FNU && block_size == 16 && + GetCudaComputeCapability().IsAtLeastBlackwell() && !lhs_k_minor) { +#ifndef NDEBUG + EXPECT_DEATH( + { (void)GetOptimizedModule(hlo); }, + "MMAv5 with kind=mxf4nvf4 does not support transpose"); + return; +#endif + } + ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(hlo)); + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + "CHECK: tt.dot_scaled"), + IsOk()); + ASSERT_OK_AND_ASSIGN(std::vector arguments, + MakeScaledDotArguments(optimized_module.get())); + if (block_size == 16 && GetCudaComputeCapability().IsAtLeastBlackwell()) { + EXPECT_EXIT( + { + EXPECT_THAT(RunAndCompareNoHloPasses( + std::move(optimized_module), + LiteralUtil::MakePointers(arguments), + ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}) + .message(), + HasSubstr("CUDA_ERROR_INVALID_INSTRUCTION")); + std::_Exit(1); + }, + ::testing::ExitedWithCode(1), ""); + return; + } + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), LiteralUtil::MakePointers(arguments), + ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); + } +}; + +INSTANTIATE_TEST_SUITE_P(TritonScaledDotTestSuite, TritonScaledDotTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return TilingParametersToString(info.param); + }); + +TEST_P(TritonFp4ScaledDotTest, Executes) { + auto cc = GpuComputeCapability().cuda_compute_capability(); + if (!cc || !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + const Fp4ScaledDotTypeCase& type_case = std::get<0>(GetParam()); + RunFp4ScaledDotExecutionTest( + type_case.lhs_type, type_case.rhs_type, type_case.scale_type, + type_case.block_size, std::get<1>(GetParam()), std::get<2>(GetParam())); +} + +INSTANTIATE_TEST_SUITE_P( + TritonFp4ScaledDotTestSuite, TritonFp4ScaledDotTest, + ::testing::Combine( + ::testing::Values( + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, + 32}, // MXFP4 x MXFP4 (block 32) + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, + 16}, // MXFP4 x MXFP4 (block 16) + Fp4ScaledDotTypeCase{F4E2M1FN, F8E4M3FN, F8E8M0FNU, + 32}, // MXFP4 x MXFP8 (MMAv5 mxf8f6f4) + Fp4ScaledDotTypeCase{F8E4M3FN, F4E2M1FN, F8E8M0FNU, + 32}, // MXFP8 x MXFP4 (MMAv5 mxf8f6f4) + Fp4ScaledDotTypeCase{F4E2M1FN, F8E5M2, F8E8M0FNU, + 32}, // MXFP4 x MXE5M2 + Fp4ScaledDotTypeCase{F8E5M2, F4E2M1FN, F8E8M0FNU, + 32}, // MXE5M2 x MXFP4 + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E4M3FN, + 16}), // NVFP4 x NVFP4 (block 16) + ::testing::Bool(), ::testing::Bool(), ::testing::Bool()), + Fp4ScaledDotTestParamToString); + +TEST_P(TritonScaledDotTest, + ScaledDotWithOmmittedLhsScaleGetFusedAndExecutedCorrectly) { + if (auto cc = GpuComputeCapability().cuda_compute_capability(); + cc && !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule ScaledDotWithOmmittedLhsScaleGetFusedAndExecutedCorrectly + +ENTRY e { + lhs = bf16[3,128,128] parameter(0) + rhs = f8e4m3fn[3,128,128] parameter(1) + constant = bf16[1,1,1] constant(1.0) + rhs_scale = f8e8m0fnu[3,128,4] parameter(2) + ROOT _ = bf16[3,128,128] scaled-dot(lhs, rhs, constant, rhs_scale), + lhs_batch_dims={0}, + rhs_batch_dims={0}, + lhs_contracting_dims={2}, + rhs_contracting_dims={2} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: fusion + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + true); + for (const auto& computation : optimized_module->computations()) { + for (const auto& instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kScaledDot) { + LOG(INFO) << "Instruction: " << instruction->name(); + } + } + } + + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x128xbf16> + CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> + CHECK: -> tensor<128x16xf32> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, FP8ScaledDotLhsKNotMinorDim) { + if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { + GTEST_SKIP() << "FP8 scaled dot requires Blackwell+"; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule FP8ScaledDotLhsKNotMinorDim + +ENTRY e { + lhs = f8e4m3fn[128,64] parameter(0) + lhs_scale = f8e8m0fnu[4,64] parameter(1) + rhs = f8e4m3fn[128,256] parameter(2) + rhs_scale = f8e8m0fnu[4,256] parameter(3) + ROOT _ = bf16[64,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={0}, + rhs_contracting_dims={0} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: fusion + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + true); + + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, ScaledDotWithBatchGetFusedAndExecutedCorrectly) { + if (auto cc = GpuComputeCapability().cuda_compute_capability(); + cc && !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule ScaledDotWithBatchGetFusedAndExecutedCorrectly + +ENTRY e { + lhs = f8e4m3fn[3,128,128] parameter(0) + rhs = f8e4m3fn[3,128,128] parameter(1) + lhs_scale = f8e8m0fnu[3,128,4] parameter(2) + rhs_scale = f8e8m0fnu[3,128,4 ] parameter(3) + ROOT _ = bf16[3,128,128] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_batch_dims={0}, + rhs_batch_dims={0}, + lhs_contracting_dims={2}, + rhs_contracting_dims={2} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: fusion + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + true); + + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> + CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> + CHECK: -> tensor<128x16xf32> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, BroadcastAndReshapeGetFused) { + if (auto cc = GpuComputeCapability().cuda_compute_capability(); + cc && !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule ScaledDotWithBatchGetFusedAndExecutedCorrectly + +ENTRY e { + lhs = f8e4m3fn[3,128,128] parameter(0) + rhs = f8e4m3fn[3,128,128] parameter(1) + lhs_scale = f8e8m0fnu[3,128,1] parameter(2) + lhs_scale_broadcasted = f8e8m0fnu[3,128,1,4] broadcast(lhs_scale), + dimensions={0,1,2} + lhs_scale_reshaped = f8e8m0fnu[3,128,4] reshape(lhs_scale_broadcasted) + rhs_scale = f8e8m0fnu[3,128,1] parameter(3) + rhs_scale_broadcasted = f8e8m0fnu[3,128,1,4] broadcast(rhs_scale), + dimensions={0,1,2} + rhs_scale_reshaped = f8e8m0fnu[3,128,4] reshape(rhs_scale_broadcasted) + ROOT _ = bf16[3,128,128] scaled-dot( + lhs, + rhs, + lhs_scale_reshaped, + rhs_scale_reshaped), + lhs_batch_dims={0}, + rhs_batch_dims={0}, + lhs_contracting_dims={2}, + rhs_contracting_dims={2} +} + )hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: %fusion + CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1} + CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1} + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + true); + + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> + CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> + CHECK: -> tensor<128x16xf32> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, Mxfp8ScaledDotSmallBlockKAndNExecutes) { + if (!GetCudaComputeCapability().IsAtLeastHopper()) { + GTEST_SKIP() << "Requires Hopper+."; + } + constexpr absl::string_view kHloText = R"hlo( +HloModule m + +fusion__ { + parameter_0 = f8e4m3fn[128,256]{1,0} parameter(0) + parameter_1 = f8e4m3fn[256,128]{1,0} parameter(1) + parameter_2 = f8e8m0fnu[128,8]{1,0} parameter(2) + parameter_3 = f8e8m0fnu[8,128]{1,0} parameter(3) + ROOT _.1 = bf16[128,128]{1,0} scaled-dot( + parameter_0, parameter_1, parameter_2, parameter_3), + lhs_contracting_dims={1}, rhs_contracting_dims={0}, + backend_config={"sizes":["64"]} +} + +ENTRY e { + lhs = f8e4m3fn[128,256]{1,0} parameter(0) + rhs = f8e4m3fn[256,128]{1,0} parameter(1) + lhs_scale = f8e8m0fnu[128,8]{1,0} parameter(2) + rhs_scale = f8e8m0fnu[8,128]{1,0} parameter(3) + ROOT fusion = bf16[128,128]{1,0} fusion( + lhs, rhs, lhs_scale, rhs_scale), + kind=kCustom, calls=fusion__, + backend_config={"fusion_backend_config":{ + "kind":"__triton_nested_gemm_fusion", + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["128","16"]}], + "num_warps":"4","num_ctas":"1","num_stages":"1"}}} +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText)); + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, Fp4Succeeds) { + if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { + GTEST_SKIP() << "Scaled dot with FP4 isn't supported by Triton for " + "pre-Blackwell GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( + HloModule jit_scaled_dot_fn + + ENTRY %main.2 { + %lhs = f4e2m1fn[1,1024,256]{2,1,0} parameter(0) + %rhs = f4e2m1fn[1,256,256]{2,1,0} parameter(1) + %lhs_scale = f8e8m0fnu[1,1024,8]{2,1,0} parameter(2) + %rhs_scale = f8e8m0fnu[1,8,256]{2,1,0} parameter(3) + ROOT %scaled-dot = bf16[1,1024,256]{2,1,0} scaled-dot(%lhs, %rhs, %lhs_scale, %rhs_scale), + lhs_batch_dims={0}, + lhs_contracting_dims={2}, + rhs_batch_dims={0}, + rhs_contracting_dims={1} + } + )hlo"; + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + ASSERT_OK_AND_ASSIGN(GpuBackendConfig gpu_config, + scaled_dot_computation->FusionInstruction() + ->backend_config()); + const BlockLevelFusionConfig& block_level_fusion_config = + gpu_config.fusion_backend_config().block_level_fusion_config(); + ASSERT_EQ(block_level_fusion_config.output_tiles_size(), 1); + const auto& output_tile_sizes = + block_level_fusion_config.output_tiles(0).sizes(); + ASSERT_EQ(output_tile_sizes.size(), 3); + const int64_t output_m = output_tile_sizes.Get(1); + const int64_t output_n = output_tile_sizes.Get(2); + ASSERT_EQ(output_n % 2, 0); + const std::string expected_triton_ir = + absl::Substitute(R"( + CHECK: tt.dot_scaled + CHECK: tensor<$0x64xi8>, tensor<$0x4xi8> + CHECK: tensor<128x$1xi8>, tensor<$2x4xi8> + CHECK: -> tensor<$0x$2xf32> + )", + output_m, output_n / 2, output_n); + + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + expected_triton_ir), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, Mxfp4KPackedRhsBlackwellLowersWithTranspose) { + if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { + GTEST_SKIP() << "Requires Blackwell+."; + } + constexpr absl::string_view kHloText = R"hlo( +HloModule m +fusion__ { + parameter_0 = f4e2m1fn[128,256]{1,0:E(4)} parameter(0) + parameter_1 = f4e2m1fn[128,256]{1,0:E(4)} parameter(1) + parameter_2 = f8e8m0fnu[128,8]{1,0} parameter(2) + parameter_3 = f8e8m0fnu[128,8]{1,0} parameter(3) + ROOT _.1 = bf16[128,128] scaled-dot(parameter_0, parameter_1, parameter_2, parameter_3), + lhs_contracting_dims={1}, rhs_contracting_dims={1}, + backend_config={"sizes":["64"]} +} +ENTRY e { + lhs = f4e2m1fn[128,256]{1,0:E(4)} parameter(0) + rhs = f4e2m1fn[128,256]{1,0:E(4)} parameter(1) + lhs_scale = f8e8m0fnu[128,8]{1,0} parameter(2) + rhs_scale = f8e8m0fnu[128,8]{1,0} parameter(3) + ROOT fusion = bf16[128,128] fusion(lhs, rhs, lhs_scale, rhs_scale), + kind=kCustom, calls=fusion__, + backend_config={"fusion_backend_config":{"kind":"__triton_nested_gemm_fusion", + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["128","128"]}], + "num_warps":"4","num_ctas":"1","num_stages":"1"}}} +} +)hlo"; + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText)); + HloComputation* scaled_dot_computation = + GetFirstComputationWithInstruction(*module, HloOpcode::kScaledDot); + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, R"( + CHECK: tt.trans + CHECK: tensor<128x32xi8> -> tensor<32x128xi8> + CHECK-NOT: unrealized_conversion_cast + CHECK: tt.dot_scaled + CHECK: tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> + )"), + IsOk()) + << "We expect to see tt.trans and the follow up dot_scaled"; + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, GlobalScalerSucceeds) { + if (!GetCudaComputeCapability().IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + constexpr absl::string_view kHloTextTemplate = R"hlo( +HloModule ScaledDotWithGlobalScaler + +ENTRY e { + lhs = f8e4m3fn[3,128,128] parameter(0) + rhs = f8e4m3fn[3,128,128] parameter(1) + lhs_scale = f8e8m0fnu[3,128,4] parameter(2) + rhs_scale = f8e8m0fnu[3,128,4] parameter(3) + scaled_dot = bf16[3,128,128] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_batch_dims={0}, + rhs_batch_dims={0}, + lhs_contracting_dims={2}, + rhs_contracting_dims={2} + global_scaler = bf16[] constant(1.42) + global_scaler_broadcasted = bf16[3,128,128] broadcast(global_scaler), + dimensions={} + ROOT _ = bf16[3,128,128] multiply(scaled_dot, global_scaler_broadcasted) +} + )hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, + GetOptimizedModule(kHloTextTemplate)); + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: %[[fusion_name:.*]] (parameter + CHECK: %[[scaled_dot:.*]] = bf16[3,128,128]{2,1,0} scaled-dot + CHECK: %[[global_scaler:.*]] = bf16[3,128,128]{2,1,0} broadcast + CHECK: ROOT %{{.*}} = bf16[3,128,128]{2,1,0} multiply(%[[scaled_dot]], %[[global_scaler]]) + CHECK: ENTRY + CHECK: ROOT {{.*}} fusion({{.*}}), kind=kCustom, calls=%[[fusion_name]] + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + true); + + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x128xf8E4M3FN>, tensor<128x4xi8> + CHECK: tensor<128x16xf8E4M3FN>, tensor<16x4xi8> + CHECK: -> tensor<128x16xf32> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, ScaledDotWithE8m0Scale) { + if (!GetCudaComputeCapability().IsAtLeastHopper()) { + GTEST_SKIP() + << "ScaledDot with Triton requires Hopper or newer architecture."; + } + constexpr absl::string_view kHloText = R"hlo( +HloModule E8m0ScaledDot + +ENTRY e { + lhs = f8e4m3fn[128,128] parameter(0) + rhs = f8e4m3fn[128,256] parameter(2) + lhs_scale = f8e8m0fnu[128,4] parameter(1) + rhs_scale = f8e8m0fnu[4,256] parameter(3) + ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); + + // Verify HLO fusion + constexpr absl::string_view kExpectedOptimizedHLO = R"( + CHECK: fusion + CHECK: ROOT {{.*}} scaled-dot + CHECK: ENTRY + CHECK: __triton_nested_gemm_fusion + )"; + EXPECT_THAT(RunFileCheck(optimized_module->ToString(), kExpectedOptimizedHLO), + absl_testing::IsOkAndHolds(true)); + + // Verify Triton IR + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x4xi8> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + // Execute on GPU hardware and compare with reference + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); +} + +TEST_P(TritonScaledDotTest, ScaledDotWithE4m3Scale) { + if (!GetCudaComputeCapability().IsAtLeastHopper()) { + GTEST_SKIP() + << "ScaledDot with Triton requires Hopper or newer architecture."; + } + constexpr absl::string_view kHloText = R"hlo( +HloModule E4m3ScaledDot + +ENTRY e { + lhs = f8e4m3fn[128,128] parameter(0) + rhs = f8e4m3fn[128,256] parameter(2) + lhs_scale = f8e4m3fn[128,4] parameter(1) + rhs_scale = f8e4m3fn[4,256] parameter(3) + ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); + + // Verify Triton IR contains tt.dot_scaled with f8e4m3fn scale tensors + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x4xf8E4M3FN> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + // Execute on GPU hardware and compare with reference (Hopper only; Blackwell + // TMEM scale expects E8M0 scale). + if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); + } +} + +TEST_P(TritonScaledDotTest, ScaledDotWithE5m2Scale) { + if (!GetCudaComputeCapability().IsAtLeastHopper()) { + GTEST_SKIP() + << "ScaledDot with Triton requires Hopper or newer architecture."; + } + constexpr absl::string_view kHloText = R"hlo( +HloModule E5m2ScaledDot + +ENTRY e { + lhs = f8e4m3fn[128,128] parameter(0) + rhs = f8e4m3fn[128,256] parameter(2) + lhs_scale = f8e5m2[128,4] parameter(1) + rhs_scale = f8e5m2[4,256] parameter(3) + ROOT _ = bf16[128,256]{1,0} scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} +} +)hlo"; + + ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(kHloText)); + + // Verify Triton IR contains tt.dot_scaled with f8e5m2 scale tensors + HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( + *optimized_module, HloOpcode::kScaledDot); + constexpr absl::string_view kExpectedTritonIr = R"( + CHECK: tt.dot_scaled + CHECK: tensor<128x4xf8E5M2> + )"; + EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, + kExpectedTritonIr), + IsOk()); + + // Execute on GPU hardware and compare with reference (Hopper only; Blackwell + // TMEM scale expects E8M0 scale). + if (!GetCudaComputeCapability().IsAtLeastBlackwell()) { + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); + } +} + +struct ScaledDotCoverageTestCase { + PrimitiveType lhs_type; + PrimitiveType rhs_type; + PrimitiveType scale_type; + int block_size; +}; + +class TritonScaledDotCoverageTest : public TritonScaledDotTestBase, + public ::testing::WithParamInterface< + std::tuple> { + public: + bool EnableTilingPropagation() const override { + return std::get<3>(GetParam()); + } +}; + +std::vector GetCoverageTestCases() { + std::vector input_types = {F8E4M3FN, F8E5M2, F4E2M1FN}; + std::vector cases; + + // 1. All input combinations with E8M0 scale and block size 32 (9 cases) + for (auto lhs : input_types) { + for (auto rhs : input_types) { + cases.push_back({lhs, rhs, F8E8M0FNU, 32}); + } + } + + // 2. Block size 16 cases + cases.push_back({F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16}); + cases.push_back({F4E2M1FN, F4E2M1FN, F8E4M3FN, 16}); + + return cases; +} + +std::string ScaledDotCoverageTestParamToString( + const ::testing::TestParamInfo< + std::tuple>& info) { + const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] = + info.param; + return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_", + PrimitiveType_Name(type_case.rhs_type), "_", + PrimitiveType_Name(type_case.scale_type), "_Block", + type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs", + rhs_k_minor, "_", + TilingParametersToString(tiling_enabled)); +} + +TEST_P(TritonScaledDotCoverageTest, Executes) { + const auto& [param, lhs_k_minor, rhs_k_minor, tiling_enabled] = GetParam(); + + auto cc = GpuComputeCapability().cuda_compute_capability(); + std::string device_name = "Unknown"; + if (cc) { + device_name = cc->IsAtLeastBlackwell() + ? "B200" + : (cc->IsAtLeastHopper() ? "H100" : "Pre-Hopper"); + } + + std::string lhs_name = + primitive_util::LowercasePrimitiveTypeName(param.lhs_type); + std::string rhs_name = + primitive_util::LowercasePrimitiveTypeName(param.rhs_type); + std::string scale_name = + primitive_util::LowercasePrimitiveTypeName(param.scale_type); + std::string lhs_display = + lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T"); + std::string rhs_display = + rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name; + std::string tiling_name = tiling_enabled ? "Symbolic" : "Experimental"; + + LOG(ERROR) << "Report Device: " << device_name; + LOG(ERROR) << "Report LHS: " << lhs_display; + LOG(ERROR) << "Report RHS: " << rhs_display; + LOG(ERROR) << "Report Scale: " << scale_name; + LOG(ERROR) << "Report BlockSize: " << param.block_size; + LOG(ERROR) << "Report Tiling: " << tiling_name; + + if (!cc || !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } + + constexpr int64_t m = 128; + constexpr int64_t n = 128; + constexpr int64_t k = 256; + const int64_t scale_k = k / param.block_size; + + const std::string lhs_shape = + lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); + const std::string rhs_shape = + rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); + const std::string lhs_scale_shape = lhs_k_minor + ? absl::StrCat(m, ",", scale_k) + : absl::StrCat(scale_k, ",", m); + const std::string rhs_scale_shape = rhs_k_minor + ? absl::StrCat(n, ",", scale_k) + : absl::StrCat(scale_k, ",", n); + + // 1. Stage: After Composite + std::string after_composite = "N/A"; + constexpr absl::string_view kCompositeHloTemplate = R"hlo( +HloModule test_composite + +%xla.scaled_dot.1 { + %p0 = $lhs_type[$lhs_shape]{1,0} parameter(0) + %p1 = $rhs_type[$rhs_shape]{1,0} parameter(1) + %p2 = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %p3 = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %dummy = bf16[$output_shape]{1,0} constant({...}) +} + +ENTRY %main { + %lhs = $lhs_type[$lhs_shape]{1,0} parameter(0) + %rhs = $rhs_type[$rhs_shape]{1,0} parameter(1) + %lhs_scale = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %rhs_scale = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %call = bf16[$output_shape]{1,0} call(%lhs, %rhs, %lhs_scale, %rhs_scale), + to_apply=%xla.scaled_dot.1, + is_composite=true, + frontend_attributes={ + composite.attributes="{dimension_numbers=[[[$lhs_contracting_dim],[$rhs_contracting_dim]],[[],[]]]}", + composite.name="xla.scaled_dot", + composite.version="1" + } +} +)hlo"; + + std::string composite_hlo = + absl::StrReplaceAll(kCompositeHloTemplate, + {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (auto comp_module = ParseAndReturnUnverifiedModule(composite_hlo); + comp_module.ok()) { + CompositeRewriter rewriter; + if (auto rewrite_status = rewriter.Run(comp_module->get()); + rewrite_status.ok()) { + const auto* root = + (*comp_module)->entry_computation()->root_instruction(); + after_composite = HloOpcodeString(root->opcode()); + } + } + LOG(ERROR) << "Report AfterComposite: " << after_composite; + + // 2. Stage: Optimized HLO + constexpr absl::string_view kHloTemplate = R"hlo( +HloModule m + +ENTRY e { + lhs = $lhs_type[$lhs_shape] parameter(0) + rhs = $rhs_type[$rhs_shape] parameter(1) + lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) + rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) + ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={$lhs_contracting_dim}, + rhs_contracting_dims={$rhs_contracting_dim} +} +)hlo"; + + std::string hlo = absl::StrReplaceAll( + kHloTemplate, {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (param.scale_type == F8E8M0FNU && param.block_size == 16 && + GetCudaComputeCapability().IsAtLeastBlackwell() && !lhs_k_minor) { +#ifndef NDEBUG + EXPECT_DEATH( + { (void)GetOptimizedModule(hlo); }, + "MMAv5 with kind=mxf4nvf4 does not support transpose"); + return; +#endif + } + + std::string optimized_hlo = "N/A"; + auto optimized_module_or = GetOptimizedModule(hlo); + HloComputation* scaled_dot_computation = nullptr; + if (optimized_module_or.ok()) { + scaled_dot_computation = GetFirstComputationWithInstruction( + **optimized_module_or, HloOpcode::kScaledDot); + if (scaled_dot_computation) { + optimized_hlo = "scaled-dot"; + } else if (GetFirstComputationWithInstruction(**optimized_module_or, + HloOpcode::kDot)) { + optimized_hlo = "dot"; + } + } + LOG(ERROR) << "Report OptimizedHLO: " << optimized_hlo; + + // 3. Stage: Triton Dot Lowering + std::string triton_dot = "N/A"; + if (scaled_dot_computation && scaled_dot_computation->FusionInstruction()) { + mlir::MLIRContext mlir_context; + BlockLevelParameters block_level_parameters; + if (auto gpu_config = scaled_dot_computation->FusionInstruction() + ->backend_config(); + gpu_config.ok() && gpu_config->has_fusion_backend_config() && + gpu_config->fusion_backend_config().has_block_level_fusion_config()) { + block_level_parameters = BlockLevelParameters::FromBlockLevelFusionConfig( + gpu_config->fusion_backend_config().block_level_fusion_config()); + } + auto triton_source_or = + CreateTritonModule("triton_fn", + *Cast( + scaled_dot_computation->FusionInstruction()), + TestGpuDeviceInfo::RTXA6000DeviceInfo(*cc), + block_level_parameters, mlir_context); + if (triton_source_or.ok()) { + std::string mlir_str; + llvm::raw_string_ostream os(mlir_str); + triton_source_or->module()->print(os); + if (absl::StrContains(mlir_str, "tt.dot_scaled")) { + triton_dot = "tt.dot_scaled"; + } else if (absl::StrContains(mlir_str, "tt.dot")) { + triton_dot = "tt.dot"; + } else if (absl::StrContains(mlir_str, "xtile.dot_scaled")) { + triton_dot = "xtile.dot_scaled"; + } + } + } + LOG(ERROR) << "Report TritonDot: " << triton_dot; + + // 4. Stage: PTX Generation + std::vector ptx_instructions; + GpuCompiler* gpu_compiler = dynamic_cast(compiler()); + bool compilation_succeeded = false; + if (gpu_compiler && optimized_module_or.ok()) { + gpu_compiler->SetAsmHook([&](absl::string_view ptx) { + std::string ptx_str(ptx); + std::string line; + size_t pos = 0; + while ((pos = ptx_str.find('\n')) != std::string::npos) { + line = ptx_str.substr(0, pos); + size_t first_non_space = line.find_first_not_of(" \t"); + if (first_non_space != std::string::npos) { + line = line.substr(first_non_space); + } + std::vector tokens = + absl::StrSplit(line, absl::ByAnyChar(" \t;"), absl::SkipEmpty()); + if (!tokens.empty()) { + std::string inst = tokens[0]; + if (absl::StartsWith(inst, "@") && tokens.size() > 1) { + inst = tokens[1]; + } + if (absl::StrContains(inst, "mma") && + (absl::StartsWith(inst, "wgmma.mma") || + absl::StartsWith(inst, "tcgen05.mma") || + absl::StartsWith(inst, "mma.sync"))) { + ptx_instructions.push_back(inst); + } + } + ptx_str.erase(0, pos + 1); + } + }); + + auto cloned_module = (*optimized_module_or)->Clone(); + auto executable_or = CompileToExecutable(std::move(cloned_module), + /*run_optimization_passes=*/true); + compilation_succeeded = executable_or.ok(); + gpu_compiler->RemoveAsmHook(); + } + + std::vector unique_ptx; + for (const auto& inst : ptx_instructions) { + if (absl::c_find(unique_ptx, inst) == unique_ptx.end()) { + unique_ptx.push_back(inst); + } + } + std::string ptx_str = + unique_ptx.empty() ? "N/A" : absl::StrJoin(unique_ptx, ", "); + LOG(ERROR) << "Report PTX: " << ptx_str; + + // 5. Stage: Execution Status + std::string status = "FAILURE"; + if (compilation_succeeded) { + auto arguments_or = MakeScaledDotArguments((*optimized_module_or).get()); + if (arguments_or.ok()) { + auto run_result = + RunAndCompareNoHloPasses(std::move(*optimized_module_or), + LiteralUtil::MakePointers(*arguments_or), + ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}); + status = run_result ? "SUCCESS" : "FAILURE"; + if (!run_result) { + if (absl::StrContains(run_result.message(), "ILLEGAL_INSTRUCTION")) { + status = "ILLEGAL_INSTRUCTION"; + } + } + } + } + LOG(ERROR) << "Report Status: " << status; + if (GetCudaComputeCapability().IsAtLeastBlackwell() && + status == "ILLEGAL_INSTRUCTION") { + EXPECT_DEATH( + { CHECK_NE(status, "ILLEGAL_INSTRUCTION"); }, "ILLEGAL_INSTRUCTION"); + } +} + +INSTANTIATE_TEST_SUITE_P( + TritonScaledDotCoverageTestSuite, TritonScaledDotCoverageTest, + ::testing::Combine(::testing::ValuesIn(GetCoverageTestCases()), + ::testing::Bool(), ::testing::Bool(), ::testing::Bool()), + ScaledDotCoverageTestParamToString); + +} // namespace +} // namespace gpu +} // namespace xla From 3312c76b6cbc49860be5dd6a58542f10115459ca Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Thu, 24 Sep 2026 01:32:10 -0700 Subject: [PATCH 26/27] PR #48977: [ROCm]Add new rocm_blas_lt_test for hipBLASLt device GEMM. 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/48977 ๐Ÿ“ Summary of Changes - Add //xla/stream_executor/rocm:rocm_blas_lt_test as the StreamExecutor twin of cuda_blas_lt_test. - Keep cuda-only on the CUDA test. Use GpuPlatformName(), required F32 and S8S32 GEMM, skip S8S8F32 if hipBLASLt returns no algorithms. ๐ŸŽฏ Justification Add more unit test coverage on ROCm platform. ๐Ÿš€ Kind of Contribution ๐Ÿงช Tests Copybara import of the project: -- c9d526efeb59b64ace5a3d6b3d81dea2a38f2350 by Lin Chen1 : Add new rocm_blas_lt_test for hipBLASLt device GEMM. Signed-off-by: Lin Chen1 -- 7611d7aed1b43c18a91d4a45efbee51346b8215d by Lin Chen1 : fix clang-format error in rocm_blas_lt_test. Signed-off-by: Lin Chen1 -- 4e60098349f79f2403982075a07f624265573b28 by linchen1 : Add a trailing newline to rocm_blas_lt_test.cc. Signed-off-by: linchen1 Merging this change closes #48977 PiperOrigin-RevId: 987322640 --- .../xla/xla/stream_executor/rocm/BUILD | 24 +++ .../stream_executor/rocm/rocm_blas_lt_test.cc | 169 ++++++++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 third_party/xla/xla/stream_executor/rocm/rocm_blas_lt_test.cc diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD index f156e8b4168251..0a73c4da4aa0cc 100644 --- a/third_party/xla/xla/stream_executor/rocm/BUILD +++ b/third_party/xla/xla/stream_executor/rocm/BUILD @@ -783,6 +783,30 @@ cc_library( alwayslink = True, ) +xla_test( + name = "rocm_blas_lt_test", + srcs = ["rocm_blas_lt_test.cc"], + backends = ["gpu"], + tags = ["rocm-only"], + deps = [ + ":amdhipblaslt_plugin", # buildcleaner: keep + ":rocblas_plugin", # buildcleaner: keep + ":rocm_platform", # buildcleaner: keep + "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla/stream_executor:blas", + "//xla/stream_executor:device_address", + "//xla/stream_executor:device_address_handle", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/gpu:gpu_blas_lt", + "@com_google_absl//absl/log", + "@com_google_googletest//:gtest_main", # buildcleaner: keep + ], +) + cc_library( name = "hipblas_lt_header", hdrs = [ diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_blas_lt_test.cc b/third_party/xla/xla/stream_executor/rocm/rocm_blas_lt_test.cc new file mode 100644 index 00000000000000..bfa4c8068f3daa --- /dev/null +++ b/third_party/xla/xla/stream_executor/rocm/rocm_blas_lt_test.cc @@ -0,0 +1,169 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/log/log.h" +#include "xla/primitive_util.h" +#include "xla/stream_executor/blas.h" +#include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_address_handle.h" +#include "xla/stream_executor/gpu/gpu_blas_lt.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/xla_data.pb.h" + +namespace stream_executor::gpu { +namespace { + +class RocmBlasLtTest : public ::testing::Test { + protected: + void SetUp() override { + ASSERT_OK_AND_ASSIGN(platform_, PlatformManager::PlatformWithName("ROCM")); + ASSERT_OK_AND_ASSIGN(executor_, platform_->ExecutorForDevice(0)); + LOG(INFO) << "Device name: " << executor_->GetDeviceDescription().name(); + ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream()); + ASSERT_OK_AND_ASSIGN(blas_lt_, gpu::BlasLt::Get(executor_)); + } + + // Column-major A (m x k) and B (k x n) are ones, C is zero, alpha=1, beta=0. + // Each output element must equal k. + // skip_if_unsupported skips only an empty algorithm list (SKU/library gap). + // GetMatmulPlan failures are always test failures. + template + void RunGemm(blas::ComputationType compute_type, bool skip_if_unsupported) { + xla::PrimitiveType input_primitive_type = + xla::primitive_util::NativeToPrimitiveType(); + xla::PrimitiveType output_primitive_type = + xla::primitive_util::NativeToPrimitiveType(); + + const int64_t m = 32, n = 32, k = 32; + + gpu::MatrixLayout a_layout(input_primitive_type, m, k, + gpu::MatrixLayout::Order::kColumnMajor); + gpu::MatrixLayout b_layout(input_primitive_type, k, n, + gpu::MatrixLayout::Order::kColumnMajor); + gpu::MatrixLayout c_layout(output_primitive_type, m, n, + gpu::MatrixLayout::Order::kColumnMajor); + + gpu::GemmConfig cfg = { + a_layout, // lhs_layout + b_layout, // rhs_layout + c_layout, // c_layout + c_layout, // output_layout + {1.0, 0.0}, // alpha + 0.0, // beta + 0, // compute_precision + xla::PrecisionConfig::ALG_UNSET, // precision_algorithm + std::nullopt, // algorithm + false, // grad_x + false, // grad_y + gpu::ScaleMode::kNone, // scale_mode + compute_type // compute_type + }; + + ASSERT_OK_AND_ASSIGN(auto plan, blas_lt_->GetMatmulPlan( + cfg, gpu::BlasLt::Epilogue::kDefault)); + + const size_t workspace_size = 32 * 1024 * 1024; // 32 MB + ASSERT_OK_AND_ASSIGN(auto algorithms, + plan->GetAlgorithms(128, workspace_size)); + if (skip_if_unsupported && algorithms.empty()) { + GTEST_SKIP() << "hipBLASLt returned no algorithms for this dtype."; + } + ASSERT_FALSE(algorithms.empty()); + ASSERT_OK(plan->SetAlgorithm(algorithms[0])); + + std::vector h_a(m * k, static_cast(1)); + std::vector h_b(k * n, static_cast(1)); + + DeviceAddress d_a = executor_->AllocateArray(h_a.size()); + DeviceAddress d_b = executor_->AllocateArray(h_b.size()); + DeviceAddress d_c = executor_->AllocateArray(m * n); + DeviceAddressBase workspace = executor_->Allocate(workspace_size); + DeviceAddressHandle a_keep(executor_, d_a); + DeviceAddressHandle b_keep(executor_, d_b); + DeviceAddressHandle c_keep(executor_, d_c); + DeviceAddressHandle workspace_keep(executor_, workspace); + + ASSERT_OK(stream_->Memcpy(&d_a, h_a.data(), h_a.size() * sizeof(InputT))); + ASSERT_OK(stream_->Memcpy(&d_b, h_b.data(), h_b.size() * sizeof(InputT))); + ASSERT_OK(stream_->MemZero(&d_c, m * n * sizeof(OutputT))); + + gpu::BlasLt::MemoryArgs args{ + /*a=*/d_a, + /*b=*/d_b, + /*c=*/d_c, + /*d=*/d_c, + /*bias=*/DeviceAddressBase{}, + /*aux=*/DeviceAddressBase{}, + /*a_scale=*/DeviceAddressBase{}, + /*b_scale=*/DeviceAddressBase{}, + /*c_scale=*/DeviceAddressBase{}, + /*d_scale=*/DeviceAddressBase{}, + /*d_amax=*/{DeviceAddressBase{}}, + workspace, + /*scratch_allocator=*/nullptr, + }; + + ASSERT_OK(plan->ExecuteOnStream(stream_.get(), args, nullptr)); + + std::vector h_result(m * n); + ASSERT_OK(stream_->Memcpy(h_result.data(), d_c, + h_result.size() * sizeof(OutputT))); + ASSERT_OK(stream_->BlockHostUntilDone()); + + for (int64_t i = 0; i < m * n; ++i) { + if constexpr (std::is_floating_point_v) { + ASSERT_FLOAT_EQ(h_result[i], static_cast(k)) + << "at index " << i; + } else { + ASSERT_EQ(h_result[i], static_cast(k)) << "at index " << i; + } + } + } + + Platform* platform_ = nullptr; + StreamExecutor* executor_ = nullptr; + std::unique_ptr stream_; + gpu::BlasLt* blas_lt_ = nullptr; +}; + +TEST_F(RocmBlasLtTest, F32F32F32Gemm) { + RunGemm(blas::ComputationType::kF32, + /*skip_if_unsupported=*/false); +} + +// hipBLASLt int8 coverage is SKU/library dependent (CUDA skips S8 below sm61). +TEST_F(RocmBlasLtTest, S8S8S32Gemm) { + RunGemm(blas::ComputationType::kI32, + /*skip_if_unsupported=*/true); +} + +TEST_F(RocmBlasLtTest, S8S8F32Gemm) { + RunGemm(blas::ComputationType::kF32, + /*skip_if_unsupported=*/true); +} + +} // namespace +} // namespace stream_executor::gpu From ee1d5b73a407c8e8d47c5a8c30a7757a2b1f7214 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 24 Sep 2026 02:32:10 -0700 Subject: [PATCH 27/27] [XLA:GPU] Clamp per-dimension offsets in dynamic-slice analysis HLO dynamic-slice and dynamic-update-slice clamp each dimension's start index i to [0, dim_size[i] - slice_size[i]]. Previously, AnalyzeDynamicSlice and VerifySliceOffset computed the 1D byte offset directly from unclamped start indices, relying only on runtime 1D buffer clamping to [0, buffer_size - slice_size]. When a slice is contiguous because all dimensions major to a partially-sliced dimension have slice_size == 1, an out-of-bounds offset on a minor dimension or size-1 dimension could stay within the total buffer bounds and write or read at the wrong element offset. Clamp constant offset operands per-dimension in AnalyzeDynamicSlice and VerifySliceOffset, and for loop-dependent offsets verify that the runtime 1D buffer clamping matches the per-dimension clamped HLO byte offset on every loop iteration. Also preserve the sliced operand shape before walking through bitcasts in DynamicSliceFusion::ResolveParameter so Parameter.shape matches the slice rank and strides in ComputeSliceOffset. Fixes #49120 PiperOrigin-RevId: 987353923 --- .../runtime/dynamic_slice_fusion_v2_thunk.cc | 20 +- .../xla/backends/gpu/tests/gpu_copy_test.cc | 37 +++- .../gpu/transforms/dynamic_slice_analysis.cc | 115 +++++++--- .../transforms/dynamic_slice_analysis_test.cc | 197 ++++++++++++++++++ .../gpu/transforms/dynamic_slice_copy_test.cc | 18 +- .../gpu/transforms/dynamic_slice_fusion.cc | 3 +- .../gpu/transforms/dynamic_slice_fusion.h | 6 +- .../transforms/dynamic_slice_fusion_test.cc | 36 ++++ 8 files changed, 384 insertions(+), 48 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc index fb4eaf8bc1a819..8542acfdf287c1 100644 --- a/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc @@ -76,11 +76,18 @@ static int64_t ComputeSliceOffset(const DynamicSliceConfig& config, return config.byte_offset() + iteration * config.byte_stride(); } -// Computes the raw byte offset from actual offset expressions. Runtime scalar -// parameters are D2H-copied from device before evaluation. +// Computes the byte offset from actual offset expressions, clamping each +// dimension's start index to [0, src_dim - dst_dim] per DS/DUS semantics. +// Runtime scalar parameters are D2H-copied from device before evaluation. static absl::StatusOr ComputeSliceOffset( - const Shape& src_shape, absl::Span offsets, + const Shape& src_shape, const Shape& dst_shape, + absl::Span offsets, absl::Span> parameters) { + if (src_shape.dimensions().size() != dst_shape.dimensions().size()) { + return Internal( + "Source shape %s and destination shape %s must have the same rank", + ShapeUtil::HumanString(src_shape), ShapeUtil::HumanString(dst_shape)); + } auto byte_strides = ShapeUtil::ByteStrides(src_shape); if (!byte_strides.has_value()) { return InvalidArgument("Failed to compute byte strides for shape %s", @@ -92,6 +99,8 @@ static absl::StatusOr ComputeSliceOffset( int64_t dim = offset.dimension_number; ABSL_ASSIGN_OR_RETURN(int64_t idx, DynamicSliceFusion::Evaluate(offset.expr, parameters)); + int64_t max_idx = src_shape.dimensions(dim) - dst_shape.dimensions(dim); + idx = std::clamp(idx, int64_t{0}, max_idx); byte_offset += idx * (*byte_strides)[dim]; } return byte_offset; @@ -308,8 +317,9 @@ static absl::Status VerifySliceOffset( // Compare offsets after clamping both to [0, buffer_size - slice_size]. int64_t buffer_size = ShapeUtil::ByteSizeOf(src_shape); int64_t slice_size = ShapeUtil::ByteSizeOf(dst_shape); - ABSL_ASSIGN_OR_RETURN(int64_t offset_from_exprs, - ComputeSliceOffset(src_shape, *offsets, parameters)); + ABSL_ASSIGN_OR_RETURN( + int64_t offset_from_exprs, + ComputeSliceOffset(src_shape, dst_shape, *offsets, parameters)); int64_t actual_offset = ClampSliceOffset(offset_from_exprs, buffer_size, slice_size); int64_t annotated_offset = ClampSliceOffset( diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc index ebcde83065aa1c..8ce5f32d7a6133 100644 --- a/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc @@ -248,9 +248,9 @@ TEST_F(GpuCopyTest, UseMemcpyForDynamicUpdateSliceWithBitcasts) { constexpr char kSliceMemcpyModuleUnfused[] = R"( body { - p0 = (s32[], s32[4,8,1000000], s32[1,1,1000000]) parameter(0) + p0 = (s32[], s32[6,8,1000000], s32[1,1,1000000]) parameter(0) ivar = s32[] get-tuple-element(p0), index=0 - input = s32[4,8,1000000] get-tuple-element(p0), index=1 + input = s32[6,8,1000000] get-tuple-element(p0), index=1 ivar_copy = s32[] copy(ivar) c1 = s32[] constant(1) @@ -258,7 +258,7 @@ constexpr char kSliceMemcpyModuleUnfused[] = R"( dynamic_slice_sizes={1,1,1000000} next_ivar = s32[] add(ivar_copy, c1) - ROOT result = (s32[], s32[4,8,1000000], s32[1,1,1000000]) + ROOT result = (s32[], s32[6,8,1000000], s32[1,1,1000000]) tuple(next_ivar, input, slice) } @@ -269,18 +269,18 @@ constexpr char kSliceMemcpyModuleUnfused[] = R"( } condition { - p0 = (s32[], s32[4,8,1000000], s32[1,1,1000000]) parameter(0) + p0 = (s32[], s32[6,8,1000000], s32[1,1,1000000]) parameter(0) ivar = s32[] get-tuple-element(p0), index=0 c6 = s32[] constant(6) ROOT cmp = pred[] compare(ivar, c6), direction=LT } ENTRY main { - p0 = s32[4,8,1000000] parameter(0) + p0 = s32[6,8,1000000] parameter(0) p1 = s32[1,1,1000000] parameter(1) c0 = s32[] constant(0) - tuple = (s32[], s32[4,8,1000000], s32[1,1,1000000]) tuple(c0, p0, p1) - ROOT while = (s32[], s32[4,8,1000000], s32[1,1,1000000]) while(tuple), + tuple = (s32[], s32[6,8,1000000], s32[1,1,1000000]) tuple(c0, p0, p1) + ROOT while = (s32[], s32[6,8,1000000], s32[1,1,1000000]) while(tuple), condition=condition, body=body })"; @@ -326,5 +326,28 @@ TEST_F(GpuCopyTest, UseDynamicMemcpyIntegrationTest) { /*run_optimization_passes=*/true)); } +constexpr char kDUSOutOfBoundsConstantOffsetsModule[] = R"( + ENTRY main { + p0 = s32[3,1,33] parameter(0) + p1 = s32[1,1,8] parameter(1) + c1 = s32[] constant(1) + c31 = s32[] constant(31) + ROOT dus = s32[3,1,33] dynamic-update-slice(p0, p1, c1, c1, c31) + })"; + +TEST_F(GpuCopyTest, + UseMemcpyForDynamicUpdateSliceWithOutOfBoundsConstantOffsets) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(kDUSOutOfBoundsConstantOffsetsModule)); + + ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module), + "; CHECK-NOT: define {{.*}}@", + /*match_optimized_ir=*/false, + /*run_optimization_passes=*/true)); + EXPECT_TRUE( + RunAndCompare(kDUSOutOfBoundsConstantOffsetsModule, ErrorSpec{0, 0})); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis.cc index dfa367badb4ade..ee12cba4914ccc 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/backends/gpu/transforms/dynamic_slice_analysis.h" +#include #include #include #include @@ -81,6 +82,13 @@ static bool IsZeroOffset(const HloInstruction* slice, int32_t dim) { return GetSliceSize(slice, dim) == slice->operand(0)->shape().dimensions(dim); } +static int64_t ClampDimensionOffset(const HloInstruction* slice, int32_t dim, + int64_t offset) { + int64_t max_offset = + slice->operand(0)->shape().dimensions(dim) - GetSliceSize(slice, dim); + return std::clamp(offset, int64_t{0}, max_offset); +} + int32_t GetFirstOffsetOperandIndex(const HloInstruction* slice) { CHECK(slice->opcode() == HloOpcode::kDynamicSlice || slice->opcode() == HloOpcode::kDynamicUpdateSlice); @@ -164,9 +172,18 @@ static absl::StatusOr EvaluateFunctionalOffset( return evaluator->Evaluate(offset, {}, true, substitutions); } +namespace { +struct EvaluatedByteOffsets { + int64_t unclamped; + int64_t clamped; +}; +} // namespace + // Evaluates the total byte offset for a DS/DUS at a given induction variable -// value by substituting into HloEvaluator. -static absl::StatusOr EvaluateByteOffsetAtIteration( +// value by substituting into HloEvaluator. Constant offset operands are always +// clamped per-dimension; for dynamic operands, both the unclamped linear offset +// and the per-dimension clamped HLO offset are returned. +static absl::StatusOr EvaluateByteOffsetAtIteration( const HloInstruction* instr, absl::Span byte_strides, const HloInstruction* induction_var, const FunctionalDependencies& functional_dependencies, int64_t ivar_value) { @@ -180,7 +197,8 @@ static absl::StatusOr EvaluateByteOffsetAtIteration( substitutions[induction_var] = &ivar_literal; HloEvaluator evaluator(/*max_loop_iterations=*/0); - int64_t total_byte_offset = 0; + int64_t unclamped_byte_offset = 0; + int64_t clamped_byte_offset = 0; // Iterate over each dimension's offset operand of the DS/DUS. for (int32_t i = 0; i < rank; ++i) { @@ -199,7 +217,10 @@ static absl::StatusOr EvaluateByteOffsetAtIteration( if (!value) { return Internal("Failed to read constant offset."); } - total_byte_offset += *value * byte_strides[i]; + int64_t dim_byte_offset = + ClampDimensionOffset(instr, i, *value) * byte_strides[i]; + unclamped_byte_offset += dim_byte_offset; + clamped_byte_offset += dim_byte_offset; continue; } @@ -219,10 +240,12 @@ static absl::StatusOr EvaluateByteOffsetAtIteration( return Internal("Failed to evaluate dynamic offset."); } - total_byte_offset += *offset_value * byte_strides[i]; + unclamped_byte_offset += *offset_value * byte_strides[i]; + clamped_byte_offset += + ClampDimensionOffset(instr, i, *offset_value) * byte_strides[i]; } - return total_byte_offset; + return EvaluatedByteOffsets{unclamped_byte_offset, clamped_byte_offset}; } // Attempts to identify a staggered induction variable created by loop @@ -449,12 +472,15 @@ absl::StatusOr> AnalyzeDynamicSlice( if (IsZeroOffset(instr, i)) { continue; } + if ((*strides)[i] < 0) { + return std::nullopt; + } const HloInstruction* operand = instr->operand(i + first_offset_index); auto value = LiteralUtil::LiteralAsScalarInt64(operand->literal()); if (!value.has_value()) { return std::nullopt; } - byte_offset += *value * (*strides)[i]; + byte_offset += ClampDimensionOffset(instr, i, *value) * (*strides)[i]; } return DynamicSliceDescriptor{std::nullopt, std::nullopt, byte_offset, 0}; } @@ -504,29 +530,66 @@ absl::StatusOr> AnalyzeDynamicSlice( // Step 8: Evaluate the byte offset for every iteration by substituting the // induction variable value into HloEvaluator. - std::vector offsets(trip_count); + std::vector offsets(trip_count); for (int64_t iter = 0; iter < trip_count; ++iter) { int64_t ivar = effective_init + iter * init_step.step; ABSL_ASSIGN_OR_RETURN(offsets[iter], EvaluateByteOffsetAtIteration( instr, *strides, induction_var, functional_dependencies, ivar)); VLOG(3) << instr->name() << ": iteration " << iter << " (ivar=" << ivar - << ") -> byte_offset=" << offsets[iter]; - } - - // Step 9: Verify linearity โ€” all consecutive differences must be equal. - // This confirms the offset is an affine function of the iteration count: - // byte_address = base + byte_offset + byte_stride * iteration - int64_t byte_offset = offsets[0]; - int64_t byte_stride = (trip_count > 1) ? (offsets[1] - offsets[0]) : 0; - - for (int64_t iter = 2; iter < trip_count; ++iter) { - int64_t actual_stride = offsets[iter] - offsets[iter - 1]; - if (actual_stride != byte_stride) { - VLOG(3) << instr->name() << ": non-linear offset pattern at iteration " - << iter << ": stride " << actual_stride << " != " << byte_stride; - return std::nullopt; + << ") -> unclamped_byte_offset=" << offsets[iter].unclamped + << ", clamped_byte_offset=" << offsets[iter].clamped; + } + + // Step 9: Verify linearity of the affine progression and ensure that the + // runtime 1D buffer clamping performed by DynamicSliceFusionV2Thunk matches + // the per-dimension clamped HLO byte offset on every iteration. First check + // whether the per-dimension clamped offsets already form a linear progression + // (which covers all in-bounds slices as well as slices whose inner-dimension + // offsets clamp to a constant across all iterations). Otherwise, check + // whether the unclamped offsets form a linear progression whose 1D buffer + // clamping matches the per-dimension clamped offset on every iteration. + const Shape& slice_shape = (instr->opcode() == HloOpcode::kDynamicSlice) + ? instr->shape() + : instr->operand(1)->shape(); + int64_t buffer_size = ShapeUtil::ByteSizeOf(slice_input_shape); + int64_t slice_size = ShapeUtil::ByteSizeOf(slice_shape); + int64_t max_buffer_offset = std::max(0, buffer_size - slice_size); + + auto get_linear_stride = [&](auto get_offset) -> std::optional { + if (trip_count <= 1) { + return 0; + } + int64_t stride = get_offset(offsets[1]) - get_offset(offsets[0]); + for (int64_t iter = 2; iter < trip_count; ++iter) { + if (get_offset(offsets[iter]) - get_offset(offsets[iter - 1]) != stride) { + return std::nullopt; + } } + return stride; + }; + + int64_t byte_offset = 0; + int64_t byte_stride = 0; + if (auto stride = get_linear_stride( + [](const EvaluatedByteOffsets& o) { return o.clamped; })) { + byte_offset = offsets[0].clamped; + byte_stride = *stride; + } else if (auto stride = get_linear_stride( + [](const EvaluatedByteOffsets& o) { return o.unclamped; }); + stride.has_value() && + absl::c_all_of(offsets, [&](const EvaluatedByteOffsets& o) { + return std::clamp(o.unclamped, 0, max_buffer_offset) == + o.clamped; + })) { + byte_offset = offsets[0].unclamped; + byte_stride = *stride; + } else { + VLOG(3) << instr->name() + << ": neither unclamped nor clamped offsets form a valid linear " + "progression over " + << trip_count << " iterations"; + return std::nullopt; } VLOG(2) << instr->name() << ": linear pattern confirmed over " << trip_count @@ -680,12 +743,16 @@ std::optional IsNonOverlapping(const DynamicSliceChain& chain) { return std::nullopt; } int64_t trip_count = loop_config->known_trip_count().n(); + int64_t buffer_byte_size = + ShapeUtil::ByteSizeOf(chain.updates.front()->operand(0)->shape()); for (int64_t iter = 0; iter < trip_count; ++iter) { std::vector ranges; ranges.reserve(dus_descriptors.size()); for (const auto& [desc, byte_size] : dus_descriptors) { - int64_t offset = desc.byte_offset + desc.byte_stride * iter; + int64_t max_offset = std::max(0, buffer_byte_size - byte_size); + int64_t offset = std::clamp( + desc.byte_offset + desc.byte_stride * iter, 0, max_offset); ranges.push_back({offset, byte_size}); } diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis_test.cc index 668a5238e05fea..40046291de3880 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_analysis_test.cc @@ -1034,5 +1034,202 @@ TEST_F(DynamicSliceAnalysisTest, ClassifiesLoopCarriedVariables) { .has_value()); } +TEST_F(DynamicSliceAnalysisTest, ConstantOffsetsClampedPerDimension) { + constexpr absl::string_view kHlo = R"( + ENTRY main { + p0 = s32[3,1,33] parameter(0) + p1 = s32[1,1,8] parameter(1) + c1 = s32[] constant(1) + c31 = s32[] constant(31) + ROOT dus = s32[3,1,33] dynamic-update-slice(p0, p1, c1, c1, c31) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + auto* dus = module->entry_computation()->root_instruction(); + ASSERT_OK_AND_ASSIGN(auto desc, AnalyzeDynamicSlice(dus)); + ASSERT_TRUE(desc.has_value()); + EXPECT_EQ(desc->loop_index, std::nullopt); + // dim 0: min(1, 3 - 1) = 1 (stride 33 * 4 = 132) + // dim 1: min(1, 1 - 1) = 0 (stride 33 * 4 = 132) + // dim 2: min(31, 33 - 8) = 25 (stride 4) + // Total byte offset = 1 * 132 + 0 * 132 + 25 * 4 = 232. + EXPECT_EQ(desc->byte_offset, 232); + EXPECT_EQ(desc->byte_stride, 0); +} + +TEST_F(DynamicSliceAnalysisTest, + LoopDependentWithOutOfBoundsConstantOffsetClamped) { + constexpr absl::string_view kHlo = R"( + body { + p0 = (s32[], s32[4,1,33]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + buf = s32[4,1,33] get-tuple-element(p0), index=1 + upd = s32[1,1,8] broadcast(s32[] constant(7)), dimensions={} + c1 = s32[] constant(1) + c31 = s32[] constant(31) + dus = s32[4,1,33] dynamic-update-slice(buf, upd, ivar, c1, c31) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], s32[4,1,33]) tuple(next_ivar, dus) + } + + condition { + p0 = (s32[], s32[4,1,33]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c4 = s32[] constant(4) + ROOT cmp = pred[] compare(ivar, c4), direction=LT + } + + ENTRY main { + input = s32[4,1,33] parameter(0) + c0 = s32[] constant(0) + tuple = (s32[], s32[4,1,33]) tuple(c0, input) + ROOT while = (s32[], s32[4,1,33]) while(tuple), + condition=condition, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(auto desc, AnalyzeByName(module.get(), "dus")); + ASSERT_TRUE(desc.has_value()); + EXPECT_THAT(desc->loop_index, ::testing::Optional(0)); + // At iter 0: dim 0 = 0, dim 1 = min(1, 0) = 0, dim 2 = min(31, 25) = 25 -> + // 100 + EXPECT_EQ(desc->byte_offset, 100); + EXPECT_EQ(desc->byte_stride, 132); +} + +TEST_F(DynamicSliceAnalysisTest, + RejectsLoopOutOfBoundsWhen1DClampMismatchesPerDimensionClamp) { + constexpr absl::string_view kHlo = R"( + body { + p0 = (s32[], s32[4,8,8], s32[1,1,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + input = s32[4,8,8] get-tuple-element(p0), index=1 + c1 = s32[] constant(1) + slice = s32[1,1,8] dynamic-slice(input, ivar, c1, c1), + dynamic_slice_sizes={1,1,8} + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], s32[4,8,8], s32[1,1,8]) + tuple(next_ivar, input, slice) + } + + condition { + p0 = (s32[], s32[4,8,8], s32[1,1,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c6 = s32[] constant(6) + ROOT cmp = pred[] compare(ivar, c6), direction=LT + } + + ENTRY main { + p0 = s32[4,8,8] parameter(0) + p1 = s32[1,1,8] parameter(1) + c0 = s32[] constant(0) + tuple = (s32[], s32[4,8,8], s32[1,1,8]) tuple(c0, p0, p1) + ROOT while = (s32[], s32[4,8,8], s32[1,1,8]) while(tuple), + condition=condition, body=body, + backend_config={"known_trip_count":{"n":"6"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(auto desc, AnalyzeByName(module.get(), "slice")); + EXPECT_FALSE(desc.has_value()); +} + +TEST_F(DynamicSliceAnalysisTest, + ClampedOffsetsLinearWhenDynamicInnerDimensionAlwaysClampsToMax) { + constexpr absl::string_view kHlo = R"( + body { + p0 = (s32[], s32[4,1,33]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + buf = s32[4,1,33] get-tuple-element(p0), index=1 + upd = s32[1,1,8] broadcast(s32[] constant(7)), dimensions={} + c1 = s32[] constant(1) + c31 = s32[] constant(31) + dyn_oob = s32[] add(ivar, c31) + dus = s32[4,1,33] dynamic-update-slice(buf, upd, ivar, c1, dyn_oob) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], s32[4,1,33]) tuple(next_ivar, dus) + } + + condition { + p0 = (s32[], s32[4,1,33]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c4 = s32[] constant(4) + ROOT cmp = pred[] compare(ivar, c4), direction=LT + } + + ENTRY main { + input = s32[4,1,33] parameter(0) + c0 = s32[] constant(0) + tuple = (s32[], s32[4,1,33]) tuple(c0, input) + ROOT while = (s32[], s32[4,1,33]) while(tuple), + condition=condition, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(auto desc, AnalyzeByName(module.get(), "dus")); + ASSERT_TRUE(desc.has_value()); + EXPECT_THAT(desc->loop_index, ::testing::Optional(0)); + // dyn_oob = ivar + 31 clamps to 25 on all 4 iterations, so the clamped byte + // offsets are 100 + iter * 132. + EXPECT_EQ(desc->byte_offset, 100); + EXPECT_EQ(desc->byte_stride, 132); +} + +TEST_F(DynamicSliceAnalysisTest, OverlappingTwoDusAfterClamping) { + constexpr absl::string_view kHlo = R"( + body { + p0 = (s32[], s32[4,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + buf = s32[4,8] get-tuple-element(p0), index=1 + c0 = s32[] constant(0) + c1 = s32[] constant(1) + c2 = s32[] constant(2) + idx0 = s32[] multiply(ivar, c2) + idx1 = s32[] add(idx0, c1) + val = s32[1,8] constant({{1,2,3,4,5,6,7,8}}) + dus0 = s32[4,8] dynamic-update-slice(buf, val, idx0, c0) + dus1 = s32[4,8] dynamic-update-slice(dus0, val, idx1, c0) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], s32[4,8]) tuple(next_ivar, dus1) + } + + condition { + p0 = (s32[], s32[4,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c3 = s32[] constant(3) + ROOT cmp = pred[] compare(ivar, c3), direction=LT + } + + ENTRY main { + input = s32[4,8] parameter(0) + c0 = s32[] constant(0) + tuple = (s32[], s32[4,8]) tuple(c0, input) + ROOT while = (s32[], s32[4,8]) while(tuple), + condition=condition, body=body, + backend_config={"known_trip_count":{"n":"3"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + auto* dus0 = + module->GetComputationWithName("body")->GetInstructionWithName("dus0"); + ASSERT_OK_AND_ASSIGN(auto chain, FindDynamicSliceChain(dus0)); + EXPECT_EQ(chain.updates.size(), 2); + // On iter 2, idx0 = 4 and idx1 = 5 both clamp to row 3 of s32[4,8], so the + // two DUS updates overlap after clamping. + auto result = IsNonOverlapping(chain); + ASSERT_TRUE(result.has_value()); + EXPECT_FALSE(*result); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc index 581d30908095d4..5496865f69967b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc @@ -305,7 +305,7 @@ TEST_F(DynamicSliceCopyTest, AnalyzesStaticSliceRootCopy) { TEST_F(DynamicSliceCopyTest, DynamicVariableUsesPerVariableInitStep) { constexpr char kHlo[] = R"( dynamic_slice_comp { - p0 = s32[4,8,8] parameter(0) + p0 = s32[6,8,8] parameter(0) p1 = s32[] parameter(1) c0 = s32[] constant(0) ROOT slice = s32[1,8,8] dynamic-slice(p0, p1, c0, c0), @@ -313,9 +313,9 @@ TEST_F(DynamicSliceCopyTest, DynamicVariableUsesPerVariableInitStep) { } body { - p0 = (s32[], s32[4,8,8], s32[]) parameter(0) + p0 = (s32[], s32[6,8,8], s32[]) parameter(0) ivar = s32[] get-tuple-element(p0), index=0 - input = s32[4,8,8] get-tuple-element(p0), index=1 + input = s32[6,8,8] get-tuple-element(p0), index=1 counter = s32[] get-tuple-element(p0), index=2 sliced = s32[1,8,8] fusion(input, counter), kind=kLoop, @@ -325,23 +325,23 @@ TEST_F(DynamicSliceCopyTest, DynamicVariableUsesPerVariableInitStep) { next_ivar = s32[] add(ivar, c1) next_counter = s32[] add(counter, c1) - ROOT result = (s32[], s32[4,8,8], s32[]) + ROOT result = (s32[], s32[6,8,8], s32[]) tuple(next_ivar, input, next_counter) } condition { - p0 = (s32[], s32[4,8,8], s32[]) parameter(0) + p0 = (s32[], s32[6,8,8], s32[]) parameter(0) ivar = s32[] get-tuple-element(p0), index=0 c3 = s32[] constant(3) ROOT cmp = pred[] compare(ivar, c3), direction=LT } ENTRY main { - input = s32[4,8,8] parameter(0) + input = s32[6,8,8] parameter(0) c2 = s32[] constant(2) c3 = s32[] constant(3) - tuple = (s32[], s32[4,8,8], s32[]) tuple(c2, input, c3) - ROOT while = (s32[], s32[4,8,8], s32[]) while(tuple), + tuple = (s32[], s32[6,8,8], s32[]) tuple(c2, input, c3) + ROOT while = (s32[], s32[6,8,8], s32[]) while(tuple), condition=condition, body=body, backend_config={"known_trip_count":{"n":"3"}, "known_init_step":{"init":"2","step":"1"}, @@ -360,7 +360,7 @@ TEST_F(DynamicSliceCopyTest, DynamicVariableUsesPerVariableInitStep) { AnalyzeDynamicSliceCopyFusion(fusion)); ASSERT_TRUE(copy.has_value()); EXPECT_THAT(copy->parameters, - ElementsAre(Parameter{0, ShapeUtil::MakeShape(S32, {4, 8, 8}), + ElementsAre(Parameter{0, ShapeUtil::MakeShape(S32, {6, 8, 8}), ShapeUtil::MakeShape(S32, {1, 8, 8}), MakeConfig(0, 768, 256), Offsets{{0, Offset::Parameter(1)}, diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.cc index 063253fe400870..7fed947fb19bad 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.cc @@ -508,6 +508,7 @@ DynamicSliceFusion::ResolveParameter(const HloInstruction* operand) { source = slice->operand(0); } + Shape source_shape = source->shape(); source = WalkThroughBitcasts(source); auto* parameter = DynCast(source); if (parameter == nullptr) { @@ -519,7 +520,7 @@ DynamicSliceFusion::ResolveParameter(const HloInstruction* operand) { return DynamicSliceFusion::Parameter{ parameter->parameter_number(), - source->shape(), + std::move(source_shape), slice_shape, config, std::move(offsets), diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.h b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.h index 314787fb080fc4..fa01729467c46c 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.h +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion.h @@ -153,7 +153,8 @@ struct DynamicSliceFusion { // (the DS source or direct parameter). int64_t parameter_number; - // Shape of the fusion parameter (the full, unsliced buffer). + // Shape of the unsliced source buffer fed into the slice (after any + // bitcast/reshape on the fusion parameter, matching `slice_shape` rank). Shape parameter_shape; // Shape of the slice fed to the hero (DS output shape, or same as @@ -182,7 +183,8 @@ struct DynamicSliceFusion { // the leaves are numbered 0, 1, 2. int64_t result_number = 0; - // Shape of the DUS target buffer (the full output buffer). + // Shape of the unsliced DUS target buffer (after any bitcast/reshape on the + // fusion parameter, matching `update_shape` rank). Shape result_shape; // Shape of the DUS update (the hero output slice inserted into the diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc index 99ddc830851833..94169b4807a3bf 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc @@ -395,6 +395,42 @@ TEST_F(DynamicSliceFusionTest, ResolveParamDirectParameter) { EXPECT_EQ(params[1], (Parameter{1, f32_4x4, f32_4x4})); } +TEST_F(DynamicSliceFusionTest, ResolveParamWithBitcastBeforeDynamicSlice) { + const char* hlo = R"( + HloModule test + + %fused { + %p0 = f32[2,4,4] parameter(0) + %p1 = s32[] parameter(1) + %bc = f32[8,4] bitcast(%p0) + %zero = s32[] constant(0) + %ds = f32[1,4] dynamic-slice(%bc, %p1, %zero), dynamic_slice_sizes={1,4}, + backend_config={"dynamic_slice_config":{"loop_index":0,"byte_offset":0,"byte_stride":16}} + ROOT %custom = f32[1,4] custom-call(%ds), custom_call_target="hero" + } + + ENTRY main { + %input = f32[2,4,4] parameter(0) + %ivar = s32[] parameter(1) + ROOT %fusion = f32[1,4] fusion(%input, %ivar), kind=kCustom, calls=%fused + } + )"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + const HloComputation* body = module->GetComputationWithName("fused"); + const HloInstruction* hero = DynamicSliceFusion::FindHero(body); + ASSERT_NE(hero, nullptr); + + ASSERT_OK_AND_ASSIGN(auto params, + DynamicSliceFusion::ResolveParameters(hero)); + ASSERT_EQ(params.size(), 1); + EXPECT_EQ(params[0], + (Parameter{ + 0, ShapeUtil::MakeShape(F32, {8, 4}), + ShapeUtil::MakeShape(F32, {1, 4}), MakeConfig(0, 0, 16), + Offsets{{0, Offset::Parameter(1)}, {1, Offset::Constant(0)}}})); +} + //===----------------------------------------------------------------------===// // DynamicSliceFusion::ResolveResults tests //===----------------------------------------------------------------------===//