From ad892ef81d3ac453586bcfc3960ec538a665d815 Mon Sep 17 00:00:00 2001 From: Venkat6871 Date: Tue, 21 Jul 2026 12:06:05 +0530 Subject: [PATCH 01/42] Fix segmentation fault in tf.raw_ops.WriteScalarSummary on empty and multi-element tensors - Add OP_REQUIRES validation in WriteScalarSummaryOp::Compute (summary_kernels.cc) to verify that the value tensor contains exactly one element (NumElements() == 1) before delegating to the summary writer. - Return errors::InvalidArgument for empty or multi-element tensors, ensuring Python cleanly raises tf.errors.InvalidArgumentError instead of terminating with a segmentation fault. - Allow single-element tensors of any rank (e.g. shape (), (1,), or (1, 1)) in accordance with established summary op API contracts. - Add comprehensive graph and eager mode unit tests in summary_ops_test.py covering valid scalar and single-element tensor successes alongside empty and multi-element tensor rejections. Fixes #113060. --- tensorflow/core/kernels/summary_kernels.cc | 4 ++ .../summary_ops/summary_ops_test.py | 45 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/tensorflow/core/kernels/summary_kernels.cc b/tensorflow/core/kernels/summary_kernels.cc index 81d30e8dbee42b..deffe04cc0f05c 100644 --- a/tensorflow/core/kernels/summary_kernels.cc +++ b/tensorflow/core/kernels/summary_kernels.cc @@ -224,6 +224,10 @@ class WriteScalarSummaryOp : public OpKernel { const Tensor* t; OP_REQUIRES_OK(ctx, ctx->input("value", &t)); + OP_REQUIRES(ctx, t->NumElements() == 1, + errors::InvalidArgument( + "value must contain exactly one element, got shape ", + t->shape().DebugString())); OP_REQUIRES_OK(ctx, s->WriteScalar(step, *t, tag)); } diff --git a/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py b/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py index cdccc3da78c9ea..4665ec718d9e4d 100644 --- a/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py +++ b/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py @@ -34,6 +34,7 @@ from tensorflow.python.framework import test_util from tensorflow.python.lib.io import tf_record from tensorflow.python.module import module +from tensorflow.python.ops import gen_summary_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import summary_ops_v2 as summary_ops from tensorflow.python.ops import variables @@ -1570,6 +1571,50 @@ def summary_op_fn(): self.exec_summary_op(summary_op_fn) +class WriteScalarSummaryRawOpsTest(test_util.TensorFlowTestCase): + + @test_util.run_in_graph_and_eager_modes + def testValidScalarAndSingleElementInputs(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer_v2(logdir) + self.evaluate(writer.init()) + step = constant_op.constant(1, dtype=dtypes.int64) + tag = constant_op.constant("test_tag", dtype=dtypes.string) + + for val in [2.5, [2.5], [[2.5]]]: + t = constant_op.constant(val, dtype=dtypes.float32) + self.evaluate( + gen_summary_ops.write_scalar_summary( + writer=writer._resource, # pylint: disable=protected-access + step=step, + tag=tag, + value=t, + ) + ) + + @test_util.run_in_graph_and_eager_modes + def testInvalidEmptyAndMultiElementInputs(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer_v2(logdir) + self.evaluate(writer.init()) + step = constant_op.constant(1, dtype=dtypes.int64) + tag = constant_op.constant("test_tag", dtype=dtypes.string) + + for invalid_val in [[], [1.0, 2.0], [[1.0, 2.0]]]: + t = constant_op.constant(invalid_val, dtype=dtypes.float32) + with self.assertRaisesRegex( + errors.InvalidArgumentError, "must contain exactly one element" + ): + self.evaluate( + gen_summary_ops.write_scalar_summary( + writer=writer._resource, # pylint: disable=protected-access + step=step, + tag=tag, + value=t, + ) + ) + + def events_from_file(filepath): """Returns all events in a single event file. From 7414ec0d4f6c78a1b94d44fc10353634389c82b4 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Wed, 29 Jul 2026 16:06:02 +0530 Subject: [PATCH 02/42] Fix GPU Pack kernel failure when autopacking 0-D scalar tensors in tf.concat Fixes #124017 Problem: When tf.rank() is used inside a tf.function with a dynamically-shaped input (i.e. TensorSpec(shape=None, ...)), the static rank is unknown at trace time, so tf.rank(image) produces a 0-D (scalar) Tensor. Building a permutation for tf.transpose with: perm = tf.concat([[num_dims - 1], tf.range(num_dims - 1)], axis=0) triggers TensorFlow's auto-packing conversion on the Python list [num_dims - 1]. Auto-packing calls gen_array_ops.pack (i.e. tf.stack) on the list elements, which is the Pack kernel. On GPU, the Pack kernel rejects 0-D scalar inputs with: InvalidArgumentError: ConcatOp : Can't concatenate scalars (use tf.stack instead) [Op:Pack] name: packed The same pattern works on CPU because the CPU Pack kernel handles scalars correctly. Root cause: _autopacking_helper (in tensorflow/python/ops/array_ops.py) calls gen_array_ops.pack directly when all elements are already Tensor objects. The GPU implementation of the underlying Pack op does not support 0-D scalar inputs. Fix: Detect when all elements being packed are 0-D scalar tensors and, in that case, use reshape(scalar, [1]) + concat_v2 instead of pack. This produces an identical 1-D output tensor and works correctly on both CPU and GPU. The guard is applied in both the eager fast-path and the graph-building code path of _autopacking_helper. Tests: Added testConcatScalarRankDynamic to tensorflow/python/kernel_tests/math_ops/transpose_op_test.py. The test traces a tf.function with TensorSpec(shape=None, ...) to force dynamic rank, then constructs channel-first and channel-last permutations using the exact pattern from the issue report, and verifies a correct round-trip result. --- .../math_ops/transpose_op_test.py | 55 +++++++++++++++++++ tensorflow/python/ops/array_ops.py | 22 ++++++++ 2 files changed, 77 insertions(+) diff --git a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py index 20cf5f9a58aab4..b6aa25f141ae69 100644 --- a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import tensor_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker_v2 +from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -519,6 +520,60 @@ def testNullTensor(self): xt = array_ops.transpose(x, [0, 2, 1]) self.assertAllEqual(xt.shape, (1, 0, 4)) + def testConcatScalarRankDynamic(self): + """Regression test for GitHub issue #124017. + + When tf.rank() returns a 0-D scalar tensor at graph/function build time + (i.e. the static rank is unknown) and the user builds a permutation via + tf.concat([[num_dims - 1], tf.range(num_dims - 1)], axis=0), the + autopacking of [num_dims - 1] previously failed on GPU with: + "ConcatOp: Can't concatenate scalars (use tf.stack instead)" + because the GPU Pack kernel does not support 0-D inputs. + + The fix ensures that scalar tensors inside a Python list are reshaped to + [1] before being concatenated, making the operation device-agnostic. + """ + + @def_function.function(input_signature=[ + tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32), + ]) + def move_channel_first(image): + """Move the last axis to the front (channels-first transform).""" + # tf.rank returns a 0-D scalar when static rank is unknown. + num_dims = array_ops.rank(image) + # Building perm with a scalar in a Python list triggers autopacking. + perm = array_ops.concat( + [[num_dims - 1], math_ops.range(num_dims - 1)], + axis=0, + ) + return array_ops.transpose(image, perm) + + # 4-D input: [batch, H, W, C] -> [C, batch, H, W] + image = constant_op.constant( + [[[[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]]]]) # shape (1, 2, 2, 2) + result = move_channel_first(image) + # channel dim (size 2) should now be first + self.assertAllEqual(result.shape[-3:], (1, 2, 2)) + self.assertEqual(result.shape[0], 2) + + @def_function.function(input_signature=[ + tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32), + ]) + def move_channel_last(image): + """Move the first axis to the end (channels-last transform).""" + num_dims = array_ops.rank(image) + inverse_perm = array_ops.concat( + [math_ops.range(1, num_dims), [0]], + axis=0, + ) + return array_ops.transpose(image, inverse_perm) + + # Verify round-trip: channels-first -> channels-last gives original shape. + channels_first = move_channel_first(image) # (2, 1, 2, 2) + restored = move_channel_last(channels_first) # (1, 2, 2, 2) + self.assertAllEqual(restored, image) + def testScalar(self): with self.cached_session(): x = constant_op.constant(42, dtype=dtypes.float32, shape=[]) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 237005b783a7a5..523f8e874670aa 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -1224,6 +1224,16 @@ def _autopacking_helper(list_or_tuple, dtype, name): # NOTE: Fast path when all the items are tensors, this doesn't do any type # checking. if all(isinstance(elem, core.Tensor) for elem in list_or_tuple): + # GPU Pack kernels do not support 0-D (scalar) inputs. Reshape each + # scalar to shape [1] so that concat can be used instead, then squeeze + # the result back to the expected 1-D vector. This matches the CPU + # behaviour of pack([s0, s1, ...]) == [s0, s1, ...]. + if all(elem.shape.rank == 0 for elem in list_or_tuple): + return gen_array_ops.concat_v2( + [gen_array_ops.reshape(elem, [1]) for elem in list_or_tuple], + axis=0, + name=name, + ) return gen_array_ops.pack(list_or_tuple, name=name) must_pack = False converted_elems = [] @@ -1253,6 +1263,18 @@ def _autopacking_helper(list_or_tuple, dtype, name): # convertible-to-tensor types, such as numpy arrays. elems_as_tensors.append( constant_op.constant(elem, dtype=dtype, name=str(i))) + # GPU Pack kernels do not support 0-D (scalar) inputs. When all + # elements are scalars, use reshape+concat to build the 1-D vector + # instead so the operation is device-agnostic. + if all( + isinstance(e, core.Tensor) and e.shape.rank == 0 + for e in elems_as_tensors + ): + return gen_array_ops.concat_v2( + [gen_array_ops.reshape(e, [1]) for e in elems_as_tensors], + axis=0, + name=scope, + ) return gen_array_ops.pack(elems_as_tensors, name=scope) else: return converted_elems From 20064d4a488f0649f1d01e31af36c2329ac3a2b3 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Wed, 29 Jul 2026 16:21:09 +0530 Subject: [PATCH 03/42] Address review feedback on _autopacking_helper scalar fix - Fix misleading comment: remove 'then squeeze the result back' since concatenating [1]-shaped tensors along axis=0 already produces the correct 1-D vector of shape [N] with no squeezing needed. - Add non-empty guard (list_or_tuple and / elems_as_tensors and) in both code paths so that an empty input does not cause concat_v2 to be called with an empty list, which would raise an error. --- tensorflow/python/ops/array_ops.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 523f8e874670aa..91cd88c7f2cee2 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -1225,10 +1225,10 @@ def _autopacking_helper(list_or_tuple, dtype, name): # checking. if all(isinstance(elem, core.Tensor) for elem in list_or_tuple): # GPU Pack kernels do not support 0-D (scalar) inputs. Reshape each - # scalar to shape [1] so that concat can be used instead, then squeeze - # the result back to the expected 1-D vector. This matches the CPU - # behaviour of pack([s0, s1, ...]) == [s0, s1, ...]. - if all(elem.shape.rank == 0 for elem in list_or_tuple): + # scalar to shape [1] so that concat can be used instead to produce + # the expected 1-D vector. This matches the CPU behaviour of + # pack([s0, s1, ...]) == [s0, s1, ...]. + if list_or_tuple and all(elem.shape.rank == 0 for elem in list_or_tuple): return gen_array_ops.concat_v2( [gen_array_ops.reshape(elem, [1]) for elem in list_or_tuple], axis=0, @@ -1266,7 +1266,7 @@ def _autopacking_helper(list_or_tuple, dtype, name): # GPU Pack kernels do not support 0-D (scalar) inputs. When all # elements are scalars, use reshape+concat to build the 1-D vector # instead so the operation is device-agnostic. - if all( + if elems_as_tensors and all( isinstance(e, core.Tensor) and e.shape.rank == 0 for e in elems_as_tensors ): From 0c3bce13c034dfc83d4dad9887c88366bfcb1db6 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Fri, 31 Jul 2026 15:59:34 +0530 Subject: [PATCH 04/42] Fix ConcatV2 single-element crash in scalar autopacking path concat_v2 requires at least 2 inputs (it raises: ValueError: List argument 'values' to 'ConcatV2' Op with length 1 shorter than minimum length 2.) When the Python list passed to tf.concat contains exactly one 0-D scalar tensor (the common case: [num_dims - 1]), the previous fix called concat_v2 with a single-element list and crashed. Fix: when there is only one scalar, return reshape(scalar, [1]) directly instead of wrapping it in concat_v2. Two or more scalars continue to use reshape+concat_v2 as before. Applied to both the eager fast-path and the graph-building path of _autopacking_helper. --- tensorflow/python/ops/array_ops.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 91cd88c7f2cee2..c52d5ad29e66ac 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -1225,15 +1225,15 @@ def _autopacking_helper(list_or_tuple, dtype, name): # checking. if all(isinstance(elem, core.Tensor) for elem in list_or_tuple): # GPU Pack kernels do not support 0-D (scalar) inputs. Reshape each - # scalar to shape [1] so that concat can be used instead to produce - # the expected 1-D vector. This matches the CPU behaviour of - # pack([s0, s1, ...]) == [s0, s1, ...]. + # scalar to shape [1] so that the result is a 1-D vector, matching the + # CPU behaviour of pack([s0, s1, ...]) == [s0, s1, ...]. + # concat_v2 requires >= 2 inputs, so a single scalar is handled with a + # plain reshape. if list_or_tuple and all(elem.shape.rank == 0 for elem in list_or_tuple): - return gen_array_ops.concat_v2( - [gen_array_ops.reshape(elem, [1]) for elem in list_or_tuple], - axis=0, - name=name, - ) + reshaped = [gen_array_ops.reshape(elem, [1]) for elem in list_or_tuple] + if len(reshaped) == 1: + return reshaped[0] + return gen_array_ops.concat_v2(reshaped, axis=0, name=name) return gen_array_ops.pack(list_or_tuple, name=name) must_pack = False converted_elems = [] @@ -1264,17 +1264,17 @@ def _autopacking_helper(list_or_tuple, dtype, name): elems_as_tensors.append( constant_op.constant(elem, dtype=dtype, name=str(i))) # GPU Pack kernels do not support 0-D (scalar) inputs. When all - # elements are scalars, use reshape+concat to build the 1-D vector - # instead so the operation is device-agnostic. + # elements are scalars, use reshape to build the 1-D vector instead so + # the operation is device-agnostic. concat_v2 requires >= 2 inputs, so + # a single scalar is handled with a plain reshape. if elems_as_tensors and all( isinstance(e, core.Tensor) and e.shape.rank == 0 for e in elems_as_tensors ): - return gen_array_ops.concat_v2( - [gen_array_ops.reshape(e, [1]) for e in elems_as_tensors], - axis=0, - name=scope, - ) + reshaped = [gen_array_ops.reshape(e, [1]) for e in elems_as_tensors] + if len(reshaped) == 1: + return reshaped[0] + return gen_array_ops.concat_v2(reshaped, axis=0, name=scope) return gen_array_ops.pack(elems_as_tensors, name=scope) else: return converted_elems From 7e81a27c136daf6f0d1cb8d6ade069c227bb59d6 Mon Sep 17 00:00:00 2001 From: Artem Bolshov Date: Sat, 31 Jan 2026 10:45:19 -0600 Subject: [PATCH 05/42] validate rank of params vs batch_dims in ResourceGatherOp::Compute --- tensorflow/core/kernels/resource_variable_ops.cc | 12 +++++++----- .../variables/resource_variable_ops_test.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index bafa896eb46d6c..62e4a8c2d9d9ce 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -767,14 +767,16 @@ class ResourceGatherOp : public OpKernel { tf_shared_lock ml(*v->mu()); const Tensor& params = *v->tensor(); const Tensor& indices = c->input(1); + OP_REQUIRES( c, TensorShapeUtils::IsVectorOrHigher(params.shape()), absl::InvalidArgumentError("params must be at least 1 dimensional")); - OP_REQUIRES(c, params.shape().dims() >= batch_dims_, - absl::InvalidArgumentError( - absl::StrCat("params must have at least ", batch_dims_, - " (batch_dims) dimensions but it has shape ", - params.shape().DebugString()))); + + OP_REQUIRES( + c, batch_dims_ < params.dims(), + absl::InvalidArgumentError(absl::StrCat( + "batch_dims (", batch_dims_, ") must be less than rank(params) (", + params.dims(), ")."))); // Check that we have enough index space const int64_t N = indices.NumElements(); diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index 7c5262667f2710..4f50a1da5cf6d6 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1750,6 +1750,21 @@ def testGatherWithBatchDimsMatchesTensor(self, params_shape, indices_shape, self.assertAllEqual(output_shape, result.shape.as_list()) self.assertAllEqual(expected, result) + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsEqualTensorRank(self): + var = resource_variable_ops.ResourceVariable( + [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32) + indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) + with self.assertRaisesRegex( + errors.InvalidArgumentError, "rank"): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True + ) + @parameterized.parameters([ dict(dtype=dtypes.bool), dict(dtype=dtypes.int64), From 43408d9f35b3361ac0de84c79e989e0167b760b0 Mon Sep 17 00:00:00 2001 From: abolshov Date: Sun, 5 Jul 2026 17:26:14 -0500 Subject: [PATCH 06/42] update test case to catch ValueError --- .../kernel_tests/variables/resource_variable_ops_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index 4f50a1da5cf6d6..cf20e6d2786ce7 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1755,8 +1755,8 @@ def testGatherWithBatchDimsEqualTensorRank(self): var = resource_variable_ops.ResourceVariable( [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32) indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) - with self.assertRaisesRegex( - errors.InvalidArgumentError, "rank"): + with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), + "rank"): _ = resource_variable_ops.resource_gather( resource=var.handle, indices=indices, From 0e3729327ed2832442e171896c4c41d011efa29c Mon Sep 17 00:00:00 2001 From: abolshov Date: Sat, 18 Jul 2026 13:37:48 -0500 Subject: [PATCH 07/42] add validation of batch_dims vs rank of indices and add a corresponding test --- tensorflow/core/kernels/resource_variable_ops.cc | 6 ++++++ .../variables/resource_variable_ops_test.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 62e4a8c2d9d9ce..4f08a88a4def39 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -778,6 +778,12 @@ class ResourceGatherOp : public OpKernel { "batch_dims (", batch_dims_, ") must be less than rank(params) (", params.dims(), ")."))); + OP_REQUIRES( + c, batch_dims_ <= indices.dims(), + absl::InvalidArgumentError(absl::StrCat( + "Expected batch_dims <= rank(indices) (", indices.dims(), + "), but got ", batch_dims_))); + // Check that we have enough index space const int64_t N = indices.NumElements(); OP_REQUIRES(c, params.dim_size(0) <= std::numeric_limits::max(), diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index cf20e6d2786ce7..b94f4b1bc560bd 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1765,6 +1765,22 @@ def testGatherWithBatchDimsEqualTensorRank(self): validate_indices=True ) + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsGreaterThanIndicesRank(self): + var = resource_variable_ops.ResourceVariable( + [[10, 20], [30, 40]], name="var0", dtype=dtypes.float32) + # indices has rank 0 (scalar) + indices = constant_op.constant(1, dtype=dtypes.int32) + # batch_dims = 1 > indices.dims() = 0 + with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), + "rank"): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True) + @parameterized.parameters([ dict(dtype=dtypes.bool), dict(dtype=dtypes.int64), From 59e1f7dadf51172ebaf52ffa7ec91b48f8a29104 Mon Sep 17 00:00:00 2001 From: abolshov Date: Sun, 9 Aug 2026 10:34:36 -0500 Subject: [PATCH 08/42] update validation logic in ResourceGatherOp::Compute and update unit tests --- .../core/kernels/resource_variable_ops.cc | 9 ++--- .../variables/resource_variable_ops_test.py | 39 +++++++++++-------- 2 files changed, 26 insertions(+), 22 deletions(-) diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 4f08a88a4def39..8b54ab10ac4ccb 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -778,11 +778,10 @@ class ResourceGatherOp : public OpKernel { "batch_dims (", batch_dims_, ") must be less than rank(params) (", params.dims(), ")."))); - OP_REQUIRES( - c, batch_dims_ <= indices.dims(), - absl::InvalidArgumentError(absl::StrCat( - "Expected batch_dims <= rank(indices) (", indices.dims(), - "), but got ", batch_dims_))); + OP_REQUIRES(c, batch_dims_ <= indices.dims(), + absl::InvalidArgumentError( + absl::StrCat("Expected batch_dims <= rank(indices) (", + indices.dims(), "), but got ", batch_dims_))); // Check that we have enough index space const int64_t N = indices.NumElements(); diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index b94f4b1bc560bd..0200a36fc76673 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1750,36 +1750,41 @@ def testGatherWithBatchDimsMatchesTensor(self, params_shape, indices_shape, self.assertAllEqual(output_shape, result.shape.as_list()) self.assertAllEqual(expected, result) - @test_util.run_in_graph_and_eager_modes - def testGatherWithBatchDimsEqualTensorRank(self): - var = resource_variable_ops.ResourceVariable( - [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32) - indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) - with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), - "rank"): - _ = resource_variable_ops.resource_gather( - resource=var.handle, - indices=indices, - dtype=dtypes.float32, - batch_dims=1, - validate_indices=True + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsEqualTensorRank(self): + var = resource_variable_ops.ResourceVariable( + [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32 ) + indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), "rank" + ): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True, + ) @test_util.run_in_graph_and_eager_modes def testGatherWithBatchDimsGreaterThanIndicesRank(self): var = resource_variable_ops.ResourceVariable( - [[10, 20], [30, 40]], name="var0", dtype=dtypes.float32) + [[10, 20], [30, 40]], name="var0", dtype=dtypes.float32 + ) # indices has rank 0 (scalar) indices = constant_op.constant(1, dtype=dtypes.int32) # batch_dims = 1 > indices.dims() = 0 - with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), - "rank"): + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), "rank" + ): _ = resource_variable_ops.resource_gather( resource=var.handle, indices=indices, dtype=dtypes.float32, batch_dims=1, - validate_indices=True) + validate_indices=True, + ) @parameterized.parameters([ dict(dtype=dtypes.bool), From 360ea8a86268839770258b2785c9cde771e4e6e0 Mon Sep 17 00:00:00 2001 From: abolshov Date: Sun, 9 Aug 2026 13:10:06 -0500 Subject: [PATCH 09/42] fix indentation --- .../variables/resource_variable_ops_test.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index 0200a36fc76673..8ef38e4349539e 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1750,22 +1750,22 @@ def testGatherWithBatchDimsMatchesTensor(self, params_shape, indices_shape, self.assertAllEqual(output_shape, result.shape.as_list()) self.assertAllEqual(expected, result) - @test_util.run_in_graph_and_eager_modes - def testGatherWithBatchDimsEqualTensorRank(self): - var = resource_variable_ops.ResourceVariable( - [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32 + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsEqualTensorRank(self): + var = resource_variable_ops.ResourceVariable( + [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32 + ) + indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), "rank" + ): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True, ) - indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) - with self.assertRaisesRegex( - (errors.InvalidArgumentError, ValueError), "rank" - ): - _ = resource_variable_ops.resource_gather( - resource=var.handle, - indices=indices, - dtype=dtypes.float32, - batch_dims=1, - validate_indices=True, - ) @test_util.run_in_graph_and_eager_modes def testGatherWithBatchDimsGreaterThanIndicesRank(self): From 76da4e457c16a79d24a18641bbb1a3a51e4c79db Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 10 Aug 2026 07:32:30 -0700 Subject: [PATCH 10/42] Demote GPU CheckNumerics anomaly log from LOG(ERROR) to VLOG(1) The GPU CheckNumerics kernel emitted a LOG(ERROR) line on every NaN/Inf detection, printing a host buffer address and raw indicator flags, e.g. abnormal_detected_host @0x7fd6e3c05600 = {1, 0} immediately before setting an InvalidArgumentError that reports the same condition in readable form. The log line therefore only duplicates the error, floods stderr when anomalies are frequent, and exposes an address that is of no use outside debugging this kernel. Both siblings already report the condition without logging: the CPU CheckNumerics kernel and CheckNumericsV2 set the status only. Demote the line to VLOG(1) so it is silent by default and still reachable via TF_CPP_VMODULE for kernel debugging. Fixes #96524 --- tensorflow/core/kernels/check_numerics_op.cc | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/check_numerics_op.cc b/tensorflow/core/kernels/check_numerics_op.cc index fb34c7a4bd79ba..c53bf44b5fde8d 100644 --- a/tensorflow/core/kernels/check_numerics_op.cc +++ b/tensorflow/core/kernels/check_numerics_op.cc @@ -289,8 +289,14 @@ class CheckNumericsOp : public AsyncOpKernel { const int is_nan = abnormality_indicators(0); const int is_inf = abnormality_indicators(1); if (is_nan || is_inf) { - LOG(ERROR) << "abnormal_detected_host @" << abnormality_indicators.data() - << " = {" << is_nan << ", " << is_inf << "} " << message_; + // Logged at VLOG level rather than LOG(ERROR): the InvalidArgumentError + // set below already reports the anomaly to the caller in readable form, + // so emitting an error line per detection only duplicates it, and the + // host buffer address is useful for debugging this kernel alone. The + // CPU kernel and CheckNumericsV2 report the same condition without any + // logging. + VLOG(1) << "abnormal_detected_host @" << abnormality_indicators.data() + << " = {" << is_nan << ", " << is_inf << "} " << message_; std::string anomalies; if (is_nan && is_inf) { From 21bf1d8545b9948bc5d78e3a58336892ed941531 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 10 Aug 2026 07:41:13 -0700 Subject: [PATCH 11/42] Trim comment on the demoted CheckNumerics log Six lines of justification for a logging change was more than the code needs. Keep two lines explaining why this is intentionally not an error log, so the choice is not silently reverted, and drop the rest, which belongs in the PR rather than the source. --- tensorflow/core/kernels/check_numerics_op.cc | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/check_numerics_op.cc b/tensorflow/core/kernels/check_numerics_op.cc index c53bf44b5fde8d..5535fe048589e9 100644 --- a/tensorflow/core/kernels/check_numerics_op.cc +++ b/tensorflow/core/kernels/check_numerics_op.cc @@ -289,12 +289,8 @@ class CheckNumericsOp : public AsyncOpKernel { const int is_nan = abnormality_indicators(0); const int is_inf = abnormality_indicators(1); if (is_nan || is_inf) { - // Logged at VLOG level rather than LOG(ERROR): the InvalidArgumentError - // set below already reports the anomaly to the caller in readable form, - // so emitting an error line per detection only duplicates it, and the - // host buffer address is useful for debugging this kernel alone. The - // CPU kernel and CheckNumericsV2 report the same condition without any - // logging. + // Not LOG(ERROR): the InvalidArgumentError set below already reports + // this to the caller, so logging per detection would duplicate it. VLOG(1) << "abnormal_detected_host @" << abnormality_indicators.data() << " = {" << is_nan << ", " << is_inf << "} " << message_; From 5db7703ab7672c117e381b8cf8cfd7f684af63b1 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Tue, 11 Aug 2026 19:30:50 +0530 Subject: [PATCH 12/42] Restore accidentally deleted testScalarConjugateTranspose The testScalarConjugateTranspose test (regression test for issue #118345) was accidentally removed when inserting testConcatScalarRankDynamic above testScalar. Restore it to its original position between testScalar and _testError. --- .../python/kernel_tests/math_ops/transpose_op_test.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py index b6aa25f141ae69..3cbd0ff9fb66cf 100644 --- a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py @@ -580,6 +580,15 @@ def testScalar(self): xt = array_ops.transpose(x) self.assertAllEqual(xt, x) + def testScalarConjugateTranspose(self): + # Regression test for https://github.com/tensorflow/tensorflow/issues/118345 + # ConjugateTranspose on a rank-0 tensor previously aborted via a + # CHECK_GT in CalculateTFStrides reached through the MKL kernel. + with self.cached_session(): + x = constant_op.constant(42, dtype=dtypes.float32, shape=[]) + xt = array_ops.transpose(x, conjugate=True) + self.assertAllEqual(xt, x) + def _testError(self, x, p, err): with self.cached_session(): with self.assertRaisesOpError(err): From 80e51095105894111efb388d7911b0b96b9b094d Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Thu, 13 Aug 2026 05:23:15 -0400 Subject: [PATCH 13/42] Validate random index shuffle rounds --- tensorflow/core/kernels/random_index_shuffle_ops.cc | 6 ++++++ .../kernel_tests/random/random_index_shuffle_test.py | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index f817f69b7d4bd0..649c0daf102f9e 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -68,6 +68,12 @@ class RandomIndexShuffleOp : public OpKernel { explicit RandomIndexShuffleOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr(kRounds, &rounds_)); + OP_REQUIRES( + context, rounds_ >= 4 && rounds_ % 2 == 0, + absl::InvalidArgumentError(absl::StrFormat( + "rounds must be an even integer greater than or equal to 4, but " + "got %d", + rounds_))); } void Compute(OpKernelContext* context) override { diff --git a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py index c7ab54a722b506..be535735bde817 100644 --- a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py +++ b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py @@ -118,6 +118,18 @@ def test_index_greater_than_max_index(self): ): self.evaluate(stateless.index_shuffle(5, seed=(1, 2), max_index=4)) + @parameterized.parameters(-1, 0, 2, 3, 5) + def test_invalid_rounds(self, rounds): + with self.assertRaisesRegex( + errors.InvalidArgumentError, + 'rounds must be an even integer greater than or equal to 4', + ): + self.evaluate( + gen_random_index_shuffle_ops.random_index_shuffle( + index=0, seed=[1, 2, 3], max_index=10, rounds=rounds + ) + ) + if __name__ == '__main__': test.main() From a2ed4f30202bb604c798dc9422e5d335522f3c06 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Thu, 13 Aug 2026 05:30:36 -0400 Subject: [PATCH 14/42] Use type-safe round count formatting --- tensorflow/core/kernels/random_index_shuffle_ops.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index 649c0daf102f9e..17e047b5a5d676 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -70,9 +70,9 @@ class RandomIndexShuffleOp : public OpKernel { OP_REQUIRES_OK(context, context->GetAttr(kRounds, &rounds_)); OP_REQUIRES( context, rounds_ >= 4 && rounds_ % 2 == 0, - absl::InvalidArgumentError(absl::StrFormat( + absl::InvalidArgumentError(absl::StrCat( "rounds must be an even integer greater than or equal to 4, but " - "got %d", + "got ", rounds_))); } From 0775e47593e73526568263460b814d8a8f1c1616 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Thu, 13 Aug 2026 05:38:27 -0400 Subject: [PATCH 15/42] Bound random index shuffle rounds --- tensorflow/core/kernels/random_index_shuffle_ops.cc | 7 ++++--- .../kernel_tests/random/random_index_shuffle_test.py | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index 17e047b5a5d676..6376a4beb4023d 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -36,6 +36,7 @@ namespace tensorflow { namespace { constexpr absl::string_view kRounds = "rounds"; +constexpr int32_t kMaxRounds = 1000; template std::array CastSeedFrom(const Tensor& seed_t, const int row) { @@ -69,10 +70,10 @@ class RandomIndexShuffleOp : public OpKernel { : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr(kRounds, &rounds_)); OP_REQUIRES( - context, rounds_ >= 4 && rounds_ % 2 == 0, + context, rounds_ >= 4 && rounds_ <= kMaxRounds && rounds_ % 2 == 0, absl::InvalidArgumentError(absl::StrCat( - "rounds must be an even integer greater than or equal to 4, but " - "got ", + "rounds must be an even integer between 4 and ", kMaxRounds, + ", but got ", rounds_))); } diff --git a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py index be535735bde817..68debf5eee47d7 100644 --- a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py +++ b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py @@ -118,11 +118,11 @@ def test_index_greater_than_max_index(self): ): self.evaluate(stateless.index_shuffle(5, seed=(1, 2), max_index=4)) - @parameterized.parameters(-1, 0, 2, 3, 5) + @parameterized.parameters(-1, 0, 2, 3, 5, 1002) def test_invalid_rounds(self, rounds): with self.assertRaisesRegex( errors.InvalidArgumentError, - 'rounds must be an even integer greater than or equal to 4', + 'rounds must be an even integer between 4 and 1000', ): self.evaluate( gen_random_index_shuffle_ops.random_index_shuffle( From 018fab6f8dcc3a1d2d89d2477954446db4cd0057 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Tue, 25 Aug 2026 07:09:41 +0200 Subject: [PATCH 16/42] Fix out-of-bounds write in GPU bincount kernels with negative input BincountReduceKernel, BincountColReduceKernel, and BincountColReduceSharedKernel only checked `bin < num_bins` before writing to the output buffer. A negative input value (e.g. from corrupt or adversarial input, or overflowed indices) satisfies that check and indexes before the start of the buffer, causing an out-of-bounds write on GPU. The CPU kernel already skips out-of-range bins in both directions, so this brings the GPU kernels in line with it by also checking `bin >= 0`. Fixes #103995 --- tensorflow/core/kernels/bincount_op_gpu.cu.cc | 6 +-- .../kernel_tests/math_ops/bincount_op_test.py | 54 +++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/kernels/bincount_op_gpu.cu.cc b/tensorflow/core/kernels/bincount_op_gpu.cu.cc index 056d78e555d762..5b5fe2a0220eff 100644 --- a/tensorflow/core/kernels/bincount_op_gpu.cu.cc +++ b/tensorflow/core/kernels/bincount_op_gpu.cu.cc @@ -118,7 +118,7 @@ __global__ void BincountReduceKernel(const Tidx* in, T* out, const int nthreads, const Tidx num_bins) { GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (bin >= 0 && bin < num_bins) { out[bin] = T(1); } } @@ -149,7 +149,7 @@ __global__ void BincountColReduceKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (bin >= 0 && bin < num_bins) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { @@ -179,7 +179,7 @@ __global__ void BincountColReduceSharedKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (bin >= 0 && bin < num_bins) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py index 70d4a758cb5059..45bfaf56545216 100644 --- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py @@ -273,6 +273,60 @@ def test_bincount_all_binary_with_weights(self, dtype): gen_math_ops.dense_bincount( input=inp, weights=np_weight, size=size, binary_output=True))) + @test_util.run_gpu_only + def test_bincount_all_binary_with_negative_input_on_gpu(self): + # Regression test for https://github.com/tensorflow/tensorflow/issues/103995 + # BincountReduceKernel only checked `bin < num_bins`, so a negative + # input value indexed before the start of the output buffer on GPU + # (an out-of-bounds write) instead of being skipped like the CPU + # kernel does. + size = 7 + inp = np.array( + [-1737321568, 32584, 7, 4, 3, 7, 7, 2, 5, 4, 1, 7, 5, 1], + dtype=np.int32) + expected_out = np.zeros((size,)) + for v in inp: + if 0 <= v < size: + expected_out[v] = 1 + with ops.device("/GPU:0"): + out = self.evaluate( + gen_math_ops.dense_bincount( + input=inp, weights=[], size=size, binary_output=True)) + self.assertAllEqual(expected_out, out) + + def _test_bincount_col_binary_with_negative_input(self, num_rows, num_cols, + size, dtype): + # Regression test for the same out-of-bounds write as above, but in + # BincountColReduceKernel / BincountColReduceSharedKernel, which had + # the identical missing lower-bound check. + np.random.seed(42) + inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype) + inp[0, 0] = -1737321568 + inp[num_rows - 1, num_cols - 1] = -1 + expected_out = np.zeros((num_rows, size)) + for row in range(num_rows): + for v in inp[row]: + if 0 <= v < size: + expected_out[row, v] = 1 + with ops.device("/GPU:0"): + out = self.evaluate( + gen_math_ops.dense_bincount( + input=inp, weights=[], size=size, binary_output=True)) + self.assertAllEqual(expected_out, out) + + @test_util.run_gpu_only + def test_col_reduce_shared_memory_with_negative_input(self): + # num_rows * size small enough to select BincountColReduceSharedKernel, + # matching test_col_reduce_shared_memory below. + self._test_bincount_col_binary_with_negative_input(128, 27, 10, np.int32) + + @test_util.run_gpu_only + def test_col_reduce_global_memory_with_negative_input(self): + # num_rows * size large enough to select BincountColReduceKernel, + # matching test_col_reduce_global_memory below. + self._test_bincount_col_binary_with_negative_input( + 128, 27, 1024, np.int32) + def _test_bincount_col_count(self, num_rows, num_cols, size, dtype): np.random.seed(42) inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype) From de67fad5af2fcf06c4ffa7d7bc6062485a622d02 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Tue, 25 Aug 2026 07:18:25 +0200 Subject: [PATCH 17/42] Combine bounds check into single unsigned comparison Per review feedback: cast bin and num_bins to their unsigned counterpart so the 0 <= bin < num_bins check is a single comparison instead of two, and avoid -Wtype-limits warnings if Tidx is ever unsigned. --- tensorflow/core/kernels/bincount_op_gpu.cu.cc | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/kernels/bincount_op_gpu.cu.cc b/tensorflow/core/kernels/bincount_op_gpu.cu.cc index 5b5fe2a0220eff..98ed253ba6c46f 100644 --- a/tensorflow/core/kernels/bincount_op_gpu.cu.cc +++ b/tensorflow/core/kernels/bincount_op_gpu.cu.cc @@ -17,6 +17,8 @@ limitations under the License. #define EIGEN_USE_GPU +#include + #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" @@ -118,7 +120,8 @@ __global__ void BincountReduceKernel(const Tidx* in, T* out, const int nthreads, const Tidx num_bins) { GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin >= 0 && bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { out[bin] = T(1); } } @@ -149,7 +152,8 @@ __global__ void BincountColReduceKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin >= 0 && bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { @@ -179,7 +183,8 @@ __global__ void BincountColReduceSharedKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin >= 0 && bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { From f448aca09cd34fc19810829aab85c97550adee73 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Tue, 25 Aug 2026 09:25:24 +0200 Subject: [PATCH 18/42] Parameterize negative-input col reduce tests over int32/int64 Per review feedback, cover 64-bit index types in the BincountColReduceKernel/BincountColReduceSharedKernel negative-input regression tests, matching test_col_reduce_shared_memory / test_col_reduce_global_memory. --- .../kernel_tests/math_ops/bincount_op_test.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py index 45bfaf56545216..fdca57f949d260 100644 --- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py @@ -314,18 +314,27 @@ def _test_bincount_col_binary_with_negative_input(self, num_rows, num_cols, input=inp, weights=[], size=size, binary_output=True)) self.assertAllEqual(expected_out, out) + @parameterized.parameters([{ + "dtype": np.int32, + }, { + "dtype": np.int64, + }]) @test_util.run_gpu_only - def test_col_reduce_shared_memory_with_negative_input(self): + def test_col_reduce_shared_memory_with_negative_input(self, dtype): # num_rows * size small enough to select BincountColReduceSharedKernel, # matching test_col_reduce_shared_memory below. - self._test_bincount_col_binary_with_negative_input(128, 27, 10, np.int32) + self._test_bincount_col_binary_with_negative_input(128, 27, 10, dtype) + @parameterized.parameters([{ + "dtype": np.int32, + }, { + "dtype": np.int64, + }]) @test_util.run_gpu_only - def test_col_reduce_global_memory_with_negative_input(self): + def test_col_reduce_global_memory_with_negative_input(self, dtype): # num_rows * size large enough to select BincountColReduceKernel, # matching test_col_reduce_global_memory below. - self._test_bincount_col_binary_with_negative_input( - 128, 27, 1024, np.int32) + self._test_bincount_col_binary_with_negative_input(128, 27, 1024, dtype) def _test_bincount_col_count(self, num_rows, num_cols, size, dtype): np.random.seed(42) From 9618c6c149399a9b2e7377c12a353ff7ca08637c Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Sun, 30 Aug 2026 13:29:22 +0300 Subject: [PATCH 19/42] Zero the SparseSegment*Grad output when indices is empty `SparseSegmentGradOpBase::Compute` sizes its output from `output_dim0` (input 3, a runtime tensor) but takes its early return on `N = indices.NumElements()`: allocate_output(0, output_shape, &output); // [M, d1...] from output_dim0 if (M == 0 || N == 0) return; // output never written M and N are independent, so an empty `indices` with `output_dim0 > 0` returns a non-empty output that no code has written. `allocate_output` does not zero the buffer, so the caller receives whatever the allocator last held there. The forward op in this same file already documents and handles exactly this case, roughly 700 lines above: // Note that we do not initialize the output buffer with a default value, // so we need to explicitly set missing indices to the default value. if (num_indices == 0) { if (output_rows > 0) { output->flat_outer_dims().setConstant(default_value_); } return; } This mirrors that treatment for the gradient, where the identity is zero: a row that received no contribution has a gradient of exactly zero. Behaviour is unchanged for every defined input. When M == 0 the new guard is false and the path is identical. When N > 0 the branch is not taken at all, so the hot path is untouched. The only values that change are the ones that were previously undefined. Affects SparseSegmentSumGrad, SparseSegmentMeanGrad and SparseSegmentSqrtNGrad, for float, double, half and bfloat16. Not compile-tested (no Bazel environment) - please run CI. --- tensorflow/core/kernels/segment_reduction_ops_impl.h | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index ad7d9f8e4916b4..463d19e38c790b 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -1342,7 +1342,16 @@ class SparseSegmentGradOpBase : public OpKernel { OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, M)); Tensor* output = nullptr; OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); - if (M == 0 || N == 0) return; + if (M == 0 || N == 0) { + // `allocate_output` does not initialize the buffer, and the functor below + // only runs when N > 0. Without this, an empty `indices` combined with + // `output_dim0` > 0 returns the freshly allocated output uninitialized. + // Zero is the correct gradient for a row that received no contribution. + if (M > 0) { + output->flat_outer_dims().setZero(); + } + return; + } functor::SparseSegmentGradFunctor()( context, operation_, input_flat, indices_vec, segment_vec, output); From 057963da9d9b120f77b65cfa6a8da9a22dc8fb88 Mon Sep 17 00:00:00 2001 From: kokol16 Date: Tue, 1 Sep 2026 17:17:15 -0400 Subject: [PATCH 20/42] Validate var and accum_update shapes in ApplyAdadeltaOp::DoValidate. Fixes host-side shape checking for ApplyAdadelta and ResourceApplyAdadelta. --- tensorflow/core/kernels/training_ops.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc index 1def8b120223a6..d3e75abc9c32fc 100644 --- a/tensorflow/core/kernels/training_ops.cc +++ b/tensorflow/core/kernels/training_ops.cc @@ -1117,6 +1117,12 @@ class ApplyAdadeltaOp : public OpKernel { absl::InvalidArgumentError(absl::StrCat( "var and grad do not have the same shape", var.shape().DebugString(), " ", grad.shape().DebugString()))); + OP_REQUIRES( + ctx, var.shape().IsSameSize(accum_update.shape()), + absl::InvalidArgumentError(absl::StrCat( + "var and accum_update do not have the same shape", + var.shape().DebugString(), " ", + accum_update.shape().DebugString()))); } void DoCompute(OpKernelContext* ctx) { From d8684e3ff0cc095f6a29e96116305f57c8d9ac8c Mon Sep 17 00:00:00 2001 From: kokol16 Date: Tue, 1 Sep 2026 17:17:15 -0400 Subject: [PATCH 21/42] Add ApplyAdadelta mismatched-shape regression tests. --- .../python/training/training_ops_test.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py index 158ace06a26d54..05c5c3864247d5 100644 --- a/tensorflow/python/training/training_ops_test.py +++ b/tensorflow/python/training/training_ops_test.py @@ -198,6 +198,60 @@ def testApplyAdagrad(self): grad = np.arange(100).astype(dtype) self._testTypesForAdagrad(x, y, lr, grad, use_gpu) + @test_util.run_v1_only("ApplyAdadelta op returns a ref, so it is not " + "supported in eager mode.") + def testApplyAdadeltaMismatchedShape(self): + for dtype in [dtypes.float32]: + var = variable_v1.VariableV1([1.0, 2.0], dtype=dtype) + accum = variable_v1.VariableV1([1.0, 2.0], dtype=dtype) + accum_update = variable_v1.VariableV1([1.0], dtype=dtype) + lr = constant_op.constant(1.0, dtype=dtype) + rho = constant_op.constant(0.9, dtype=dtype) + epsilon = constant_op.constant(1e-6, dtype=dtype) + grad = constant_op.constant([0.1, 0.2], dtype=dtype) + with self.cached_session() as sess: + self.evaluate(variables.global_variables_initializer()) + with self.assertRaises((ValueError, errors.InvalidArgumentError)): + self.evaluate( + gen_training_ops.apply_adadelta( + var, + accum, + accum_update, + lr, + rho, + epsilon, + grad, + use_locking=False, + ) + ) + + @test_util.run_v1_only("ResourceApplyAdadelta op returns a ref, so it is not " + "supported in eager mode.") + def testResourceApplyAdadeltaMismatchedShape(self): + for dtype in [dtypes.float32]: + var = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + accum = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + accum_update = resource_variable_ops.ResourceVariable([1.0], dtype=dtype) + lr = constant_op.constant(1.0, dtype=dtype) + rho = constant_op.constant(0.9, dtype=dtype) + epsilon = constant_op.constant(1e-6, dtype=dtype) + grad = constant_op.constant([0.1, 0.2], dtype=dtype) + with self.cached_session() as sess: + self.evaluate(variables.global_variables_initializer()) + with self.assertRaises((ValueError, errors.InvalidArgumentError)): + self.evaluate( + gen_training_ops.resource_apply_adadelta( + var.handle, + accum.handle, + accum_update.handle, + lr, + rho, + epsilon, + grad, + use_locking=False, + ) + ) + @test_util.run_v1_only("ApplyFtrl op returns a ref, so it is not " "supported in eager mode.") def testApplyFtrl(self): From 84d2f0b47586a084a33b40e72e755c52ab03acf9 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 2 Sep 2026 12:00:34 +0300 Subject: [PATCH 22/42] Use SetZeroFunctor for the empty-indices zeroing and add a regression test Addresses both review comments: - The previous `output->flat_outer_dims().setZero()` runs on the host CPU. `SparseSegmentGradOpBase` is templated on `Device` and is registered for GPU (`segment_reduction_ops_impl_5.cc`), so on GPU that wrote through a device pointer from the host. `functor::SetZeroFunctor` dispatches to the correct device. The guard is now `output->NumElements() > 0`, which is exact when the inner dimensions are empty as well. - Adds `testGradientEmptyIndicesWithNonZeroOutputDim` covering empty `indices` with `output_dim0 > 0` for sum/mean/sqrt_n over float16, bfloat16, float32 and float64. The test fails on an unpatched build, where the returned values are uninitialized heap contents rather than zero. --- .../core/kernels/segment_reduction_ops_impl.h | 6 ++++-- .../math_ops/segment_reduction_ops_test.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 463d19e38c790b..9d2afe083154e4 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -48,6 +48,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/tensor_util.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/kernels/fill_functor.h" #include "tensorflow/core/kernels/segment_reduction_ops.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/bfloat16.h" @@ -1347,8 +1348,9 @@ class SparseSegmentGradOpBase : public OpKernel { // only runs when N > 0. Without this, an empty `indices` combined with // `output_dim0` > 0 returns the freshly allocated output uninitialized. // Zero is the correct gradient for a row that received no contribution. - if (M > 0) { - output->flat_outer_dims().setZero(); + if (output->NumElements() > 0) { + functor::SetZeroFunctor()(context->eigen_device(), + output->flat()); } return; } diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 24cfb3eec109ec..9b37ba8a1f4702 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -1182,6 +1182,27 @@ def testGradientExplicitSingleOutput(self): tf_xgrad = tf_op(tf_ygrad, indices, segment_ids, output_dim0) self.assertAllClose(tf_xgrad, np_xgrad) + def testGradientEmptyIndicesWithNonZeroOutputDim(self): + ops_list = [ + math_ops.sparse_segment_sum_grad, + math_ops.sparse_segment_mean_grad, + math_ops.sparse_segment_sqrt_n_grad, + ] + indices = [] + segment_ids = [] + output_dim0 = 4 + for dtype in [ + dtypes_lib.float16, + dtypes_lib.bfloat16, + dtypes_lib.float32, + dtypes_lib.float64, + ]: + grad = constant_op.constant([], shape=[0, 3], dtype=dtype) + for tf_op in ops_list: + res = self.evaluate(tf_op(grad, indices, segment_ids, output_dim0)) + self.assertEqual(res.shape, (4, 3)) + self.assertAllEqual(res, np.zeros((4, 3), dtype=dtype.as_numpy_dtype)) + def testGradientV2Explicit(self): # Note that the GPU implem has different paths for different inner sizes. for inner_size in (1, 2, 3, 32): From 24c6e61dadcdc599ff419ccf0ff289ea240f1ffa Mon Sep 17 00:00:00 2001 From: adi-IL Date: Thu, 3 Sep 2026 02:11:12 +0530 Subject: [PATCH 23/42] fix: support broadcast-shaped bias with leading 1s in fused MatMul and Conv2D When Grappler remapper fuses MatMul or Conv2D with Add, the bias tensor can have broadcast dimensions where leading axes are 1 (e.g. [1, N]). While mkl_matmul_op_fused allows bias_dims > 1 with leading 1s, InitBiasAddArgs in fused_eigen_output_kernels.h, GPU LaunchFusedMatMulOp, and LaunchFusedConv2DOp strictly required bias.dims() == 1. This caused InvalidArgumentError: bias must be 1-dimensional during graph execution in tf.function. Allow bias tensors with rank >= 1 whose leading dimensions are 1 in InitBiasAddArgs, GPU LaunchFusedMatMulOp, and LaunchFusedConv2DOp, aligning the Eigen and GPU kernels with the remapper and MKL implementations. Add test coverage in matmul_op_test.cc and matmul_op_test.py. Fixes #126661 --- tensorflow/core/kernels/conv_ops_fused_impl.h | 16 +++++++++---- .../core/kernels/fused_eigen_output_kernels.h | 16 ++++++++++--- tensorflow/core/kernels/matmul_op_fused.cc | 18 ++++++++++---- tensorflow/core/kernels/matmul_op_test.cc | 24 +++++++++++++++++++ .../kernel_tests/math_ops/matmul_op_test.py | 18 ++++++++++++++ 5 files changed, 81 insertions(+), 11 deletions(-) diff --git a/tensorflow/core/kernels/conv_ops_fused_impl.h b/tensorflow/core/kernels/conv_ops_fused_impl.h index 2a62c27d1f4a6e..c4dad84d4a7678 100644 --- a/tensorflow/core/kernels/conv_ops_fused_impl.h +++ b/tensorflow/core/kernels/conv_ops_fused_impl.h @@ -375,12 +375,20 @@ struct LaunchFusedConv2DOp { const int64_t out_cols = GetTensorDim(*output, params.data_format, 'W'); const int64_t out_depths = GetTensorDim(*output, params.data_format, 'C'); - // Bias of the following dimensions: [ output_depth ] + // Bias of the following dimensions: [ output_depth ] or [1, ..., 1, output_depth] const Tensor& bias = context->input(2); - OP_REQUIRES(context, bias.dims() == 1, + OP_REQUIRES(context, bias.dims() >= 1, absl::InvalidArgumentError(absl::StrCat( - "bias must be 1-dimensional", bias.shape().DebugString()))); - OP_REQUIRES(context, bias.dim_size(0) == out_depths, + "bias must be at least 1-dimensional, got: ", + bias.shape().DebugString()))); + for (int i = 0; i < bias.dims() - 1; ++i) { + OP_REQUIRES(context, bias.dim_size(i) == 1, + absl::InvalidArgumentError(absl::StrCat( + "For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString()))); + } + OP_REQUIRES(context, bias.dim_size(bias.dims() - 1) == out_depths, absl::InvalidArgumentError( absl::StrCat("bias depth must be equal to out depth", bias.shape().DebugString()))); diff --git a/tensorflow/core/kernels/fused_eigen_output_kernels.h b/tensorflow/core/kernels/fused_eigen_output_kernels.h index 7cc19a9a67ab92..a33eccc0214e4a 100644 --- a/tensorflow/core/kernels/fused_eigen_output_kernels.h +++ b/tensorflow/core/kernels/fused_eigen_output_kernels.h @@ -411,12 +411,22 @@ using WithFusedBatchNormAndLeakyRelu = FusedBatchNormOutputKernel; template absl::Status InitBiasAddArgs(OpKernelContext* context, BiasAddArgs* args, const float* leakyrelu_alpha = nullptr) { - // Bias of the following dimensions: [ output_depth ] + // Bias of the following dimensions: [ output_depth ] or [1, ..., 1, output_depth] const Tensor& bias = context->input(2); - if (bias.dims() != 1) + if (bias.dims() < 1) { return absl::InvalidArgumentError( - absl::StrCat("bias must be 1-dimensional", bias.shape().DebugString())); + absl::StrCat("bias must be at least 1-dimensional, got: ", + bias.shape().DebugString())); + } + for (int i = 0; i < bias.dims() - 1; ++i) { + if (bias.dim_size(i) != 1) { + return absl::InvalidArgumentError( + absl::StrCat("For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString())); + } + } const auto data_ptr = [](const Tensor& tensor) -> const T* { return reinterpret_cast(tensor.tensor_data().data()); diff --git a/tensorflow/core/kernels/matmul_op_fused.cc b/tensorflow/core/kernels/matmul_op_fused.cc index 25e4959aff8136..fe1bfe3e547e18 100644 --- a/tensorflow/core/kernels/matmul_op_fused.cc +++ b/tensorflow/core/kernels/matmul_op_fused.cc @@ -453,10 +453,20 @@ struct LaunchFusedMatMulOp { // + . Therefore, the bias tensor is required. const Tensor& bias = context->input(2); - if (bias.dims() != 1) { - OP_REQUIRES_OK(context, absl::InvalidArgumentError( - absl::StrCat("bias must be 1-dimensional", - bias.shape().DebugString()))); + if (bias.dims() < 1) { + OP_REQUIRES_OK(context, + absl::InvalidArgumentError(absl::StrCat( + "bias must be at least 1-dimensional, got: ", + bias.shape().DebugString()))); + } + for (int i = 0; i < bias.dims() - 1; ++i) { + if (bias.dim_size(i) != 1) { + OP_REQUIRES_OK(context, + absl::InvalidArgumentError(absl::StrCat( + "For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString()))); + } } auto a_ptr = AsDeviceMemory(a.template flat().data(), diff --git a/tensorflow/core/kernels/matmul_op_test.cc b/tensorflow/core/kernels/matmul_op_test.cc index 51b8fc71e1a1e6..cf43fffeef3c11 100644 --- a/tensorflow/core/kernels/matmul_op_test.cc +++ b/tensorflow/core/kernels/matmul_op_test.cc @@ -288,6 +288,25 @@ class FusedMatMulOpTest : public OpsTestBase { run_fused); } + void VerifyMatMulWithBroadcastBias(int m, int k, int n, bool transpose_a, + bool transpose_b) { + DataType dtype = DataTypeToEnum::v(); + Tensor lhs(dtype, {transpose_a ? k : m, transpose_a ? m : k}); + lhs.flat() = lhs.flat().setRandom(); + Tensor rhs(dtype, {transpose_b ? n : k, transpose_b ? k : n}); + rhs.flat() = rhs.flat().setRandom(); + Tensor bias(dtype, {1, n}); + bias.flat() = bias.flat().setRandom(); + + Tensor fused_matmul; + bool skipped = false; + RunFusedMatMulOp(lhs, rhs, {bias}, {"BiasAdd"}, transpose_a, transpose_b, + &fused_matmul, /*allow_gpu_device=*/true, &skipped); + if (!skipped) { + ASSERT_EQ(fused_matmul.shape(), TensorShape({m, n})); + } + } + // Verifies that computing MatMul+BiasAdd+{Activation} in a graph is identical // to FusedMatMul. void VerifyConv2DWithBiasAndActivation(int m, int k, int n, bool transpose_a, @@ -356,6 +375,10 @@ TYPED_TEST_P(FusedMatMulWithBiasOpTest, MatMul1x256x1) { this->VerifyMatMulWithBias(1, 256, 1, false, false); } +TYPED_TEST_P(FusedMatMulWithBiasOpTest, MatMulWithBroadcastBias) { + this->VerifyMatMulWithBroadcastBias(8, 16, 32, false, false); +} + static auto GetActivations(DataType dtype) { // "GeluExact", "Tanh", "Sigmoid" fusions are only supported for half-float // datatype @@ -407,6 +430,7 @@ REGISTER_TYPED_TEST_SUITE_P(FusedMatMulWithBiasOpTest, // MatMul1x256x256, // MatMul256x256x1, // MatMul1x256x1, // + MatMulWithBroadcastBias, // MatMul256x128x64WithActivation, // MatMul1x256x256WithActivation, // MatMul256x256x1WithActivation, // diff --git a/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py b/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py index ae764c8ff126b9..bb7d462e141845 100644 --- a/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py @@ -19,6 +19,7 @@ import numpy as np from tensorflow.python import tf2 +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -55,6 +56,23 @@ def testBatchMatMulV3MixedPrec(self): self.assertAllEqual([[5, 11], [11, 25]], c) +@test_util.with_eager_op_as_function +class MatMulAddFusionTest(test_lib.TestCase): + """Test broadcasting behavior when MatMul and Add are fused in graph execution.""" + + def testMatMulAddBroadcastInTfFunction(self): + @def_function.function + def f(x): + return math_ops.add( + math_ops.matmul(x, x, transpose_a=True), + math_ops.matmul(array_ops.ones([1, 1], dtype=x.dtype), x), + ) + + x = constant_op.constant(np.ones([1, 7], dtype=np.float32)) + res = f(x) + self.assertEqual(tuple(res.shape), (7, 7)) + + @test_util.with_eager_op_as_function class MatVecTest(test_lib.TestCase): """Simple test for matvec, which is sugar on top of matmul.""" From 1953bb51b4faa695a3917b6d87c7e7e7a8c876ec Mon Sep 17 00:00:00 2001 From: Som Tripathi Date: Wed, 2 Sep 2026 18:31:31 -0500 Subject: [PATCH 24/42] Reject RaggedGather value-count overflow instead of a heap-buffer-overflow write RaggedGather with int32 Tsplits could accumulate a total value count above int32, wrap it, allocate an undersized output, and then write past it (heap-buffer-overflow). Accumulate the count in int64 and reject counts that exceed the Tsplits maximum with InvalidArgumentError before allocating, matching the existing guard in ragged_range_op.cc. Adds a regression test that pins Tsplits=int32 via the raw op. Fixes #126135. --- tensorflow/core/kernels/ragged_gather_op.cc | 11 ++++++++++- tensorflow/python/ops/ragged/BUILD | 1 + .../python/ops/ragged/ragged_gather_op_test.py | 18 ++++++++++++++++++ 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc index d280c948303179..72e5c840e86ad5 100644 --- a/tensorflow/core/kernels/ragged_gather_op.cc +++ b/tensorflow/core/kernels/ragged_gather_op.cc @@ -128,6 +128,7 @@ class RaggedGatherOpBase : public OpKernel { std::vector>* value_slices, SPLITS_TYPE* num_values) { *num_values = 0; + int64_t total_values = 0; value_slices->clear(); int num_splits = indices_in.dims() - 1 + params_nested_splits_in.size(); @@ -188,9 +189,17 @@ class RaggedGatherOpBase : public OpKernel { } if (limit != start) { value_slices->emplace_back(start, limit); - *num_values += limit - start; + total_values += limit - start; } } + // The value count is accumulated in a wide type and validated before it is + // narrowed to SPLITS_TYPE, so an int32 Tsplits cannot silently wrap and + // drive an undersized output allocation in WriteValues. + if (total_values > std::numeric_limits::max()) { + return absl::InvalidArgumentError(absl::StrCat( + "Number of values ", total_values, " exceeds limits of Tsplits")); + } + *num_values = static_cast(total_values); return absl::OkStatus(); } diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD index 2f1bb55b55d132..4c8416eb018483 100644 --- a/tensorflow/python/ops/ragged/BUILD +++ b/tensorflow/python/ops/ragged/BUILD @@ -965,6 +965,7 @@ py_test( "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:gradients_impl", + "//tensorflow/python/ops:ragged_array_ops_gen", "//tensorflow/python/platform:test", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/python/ops/ragged/ragged_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_op_test.py index b7bd1f1a9aa150..cf4c28a0b1f04a 100644 --- a/tensorflow/python/ops/ragged/ragged_gather_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_gather_op_test.py @@ -26,6 +26,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_ragged_array_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.ops.ragged import ragged_gather_ops @@ -222,6 +223,23 @@ def testOutOfBoundsError(self): r'indices\[1\] = 3 is not in \[0, 2\)'): self.evaluate(ragged_gather_ops.gather(ragged_params, ragged_indices)) + def testValueCountOverflowError(self): + # The gathered value count (num_indices * gathered_row_length) overflows an + # int32 Tsplits: 65537 * 65536 = 2**32 + 65536. Before the fix this wrapped + # to 65536, drove an undersized output allocation, and the copy loop then + # wrote past it (heap-buffer-overflow). The raw op pins Tsplits=int32; the + # high-level gather defaults to int64 splits and cannot reproduce. + with self.assertRaisesRegex(errors.InvalidArgumentError, + 'exceeds limits of Tsplits'): + self.evaluate( + gen_ragged_array_ops.ragged_gather( + params_nested_splits=[ + constant_op.constant([0, 65536], dtype=dtypes.int32) + ], + params_dense_values=array_ops.zeros([65536], dtype=dtypes.uint8), + indices=array_ops.zeros([65537], dtype=dtypes.int32), + OUTPUT_RAGGED_RANK=1)) + def testUnknownIndicesRankError(self): if context.executing_eagerly(): return From 2a741fd7ead610739e857ebb0522397a7c6c4924 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Fri, 4 Sep 2026 17:48:31 +0530 Subject: [PATCH 25/42] Fix GPU Pack kernel to handle 0-D scalar inputs directly Address review feedback on #124017: The previous fix intercepted 0-D scalar tensors in the Python-level _autopacking_helper and replaced gen_array_ops.pack with reshape(scalar, [1]) + concat_v2. As pointed out in the review, this approach has several regressions: - Breaks static shape inference and constant propagation: PackOp retains known scalar values during shape inference, allowing TensorFlow to infer partial shapes (e.g. [None, 3]) for SparseTensor, RaggedTensor, and graph compilation. Replacing Pack with ConcatV2 loses this and produces [None, None] instead. - Incompatible with DT_RESOURCE and DT_VARIANT: ConcatV2 has no registered GPU kernel for these types, so tf.stack([r1, r2]) over resource handles fails with "No OpKernel was registered to support Op 'ConcatV2' with T=DT_RESOURCE". - Breaks graph pattern matching: TFLite OpHint AGGREGATE_STACK and XLA/Grappler passes pattern-match on Pack nodes; inserting Reshape + ConcatV2 prevents these transformations. Fix: Revert the Python-level _autopacking_helper changes entirely so that gen_array_ops.pack is always used, restoring full shape inference and type support. Add an explicit scalar-input branch in the C++ GPU PackOp kernel (tensorflow/core/kernels/pack_op.cc): when all inputs are 0-D scalars (first_input.dims() == 0, num > 1), write each scalar directly into the output flat vector using Eigen chip operations rather than routing through ConcatGPU, which was incompatible with 1-element-per-row configurations on some GPU drivers. The num == 1 scalar path already works via the existing CopyFrom reshape shortcut and is unchanged. Update the test docstring in testConcatScalarRankDynamic to reflect that the fix now lives in the C++ kernel rather than the Python autopacking helper. --- tensorflow/core/kernels/pack_op.cc | 26 +++++++++++++++++++ .../math_ops/transpose_op_test.py | 8 +++--- tensorflow/python/ops/array_ops.py | 22 ---------------- 3 files changed, 31 insertions(+), 25 deletions(-) diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc index 95331ce42dc481..9675595a3a8598 100644 --- a/tensorflow/core/kernels/pack_op.cc +++ b/tensorflow/core/kernels/pack_op.cc @@ -80,6 +80,32 @@ class PackOp : public OpKernel { Tensor* output; OP_REQUIRES_OK(c, c->allocate_output(0, output_shape, &output)); + // Special case: packing 0-D (scalar) inputs. + // + // The generic path below flattens each input to a {before_dim, after_dim} + // matrix and then calls ConcatGPU / ConcatCPU. For scalar inputs + // before_dim == after_dim == 1, so every matrix has shape {1, 1} and the + // output matrix has shape {1, num}. While the arithmetic is correct, the + // GPU ConcatGPU helper validates that its inputs have at least one element + // per row (dimension(1) > 0 is always true here), but some GPU drivers + // reject the launch configuration produced for 1-element rows. Avoid the + // issue entirely by treating the scalar case as a flat vector copy: + // each scalar contributes exactly one element at position i in the output. + if (first_input.dims() == 0) { + auto output_vec = output->flat(); + for (int i = 0; i < num; ++i) { + const Tensor& input = c->input(i); + OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), + absl::InvalidArgumentError(absl::StrCat( + "Shapes of all inputs must match: values[0].shape = ", + first_input.shape().DebugString(), " != values[", i, + "].shape = ", input.shape().DebugString()))); + output_vec.template chip<0>(i).device(c->eigen_device()) = + input.scalar(); + } + return; + } + int64_t before_dim = 1; for (int i = 0; i < axis; ++i) { before_dim *= output_shape.dim_size(i); diff --git a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py index 3cbd0ff9fb66cf..32a3ff005342cb 100644 --- a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py @@ -528,10 +528,12 @@ def testConcatScalarRankDynamic(self): tf.concat([[num_dims - 1], tf.range(num_dims - 1)], axis=0), the autopacking of [num_dims - 1] previously failed on GPU with: "ConcatOp: Can't concatenate scalars (use tf.stack instead)" - because the GPU Pack kernel does not support 0-D inputs. + because the GPU Pack kernel did not support 0-D scalar inputs. - The fix ensures that scalar tensors inside a Python list are reshaped to - [1] before being concatenated, making the operation device-agnostic. + The fix adds an explicit scalar-input path in the C++ GPU PackOp kernel + (tensorflow/core/kernels/pack_op.cc) that writes each scalar element + directly into the output vector via Eigen chip operations, bypassing the + ConcatGPU helper that was incompatible with 0-D inputs. """ @def_function.function(input_signature=[ diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index c52d5ad29e66ac..237005b783a7a5 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -1224,16 +1224,6 @@ def _autopacking_helper(list_or_tuple, dtype, name): # NOTE: Fast path when all the items are tensors, this doesn't do any type # checking. if all(isinstance(elem, core.Tensor) for elem in list_or_tuple): - # GPU Pack kernels do not support 0-D (scalar) inputs. Reshape each - # scalar to shape [1] so that the result is a 1-D vector, matching the - # CPU behaviour of pack([s0, s1, ...]) == [s0, s1, ...]. - # concat_v2 requires >= 2 inputs, so a single scalar is handled with a - # plain reshape. - if list_or_tuple and all(elem.shape.rank == 0 for elem in list_or_tuple): - reshaped = [gen_array_ops.reshape(elem, [1]) for elem in list_or_tuple] - if len(reshaped) == 1: - return reshaped[0] - return gen_array_ops.concat_v2(reshaped, axis=0, name=name) return gen_array_ops.pack(list_or_tuple, name=name) must_pack = False converted_elems = [] @@ -1263,18 +1253,6 @@ def _autopacking_helper(list_or_tuple, dtype, name): # convertible-to-tensor types, such as numpy arrays. elems_as_tensors.append( constant_op.constant(elem, dtype=dtype, name=str(i))) - # GPU Pack kernels do not support 0-D (scalar) inputs. When all - # elements are scalars, use reshape to build the 1-D vector instead so - # the operation is device-agnostic. concat_v2 requires >= 2 inputs, so - # a single scalar is handled with a plain reshape. - if elems_as_tensors and all( - isinstance(e, core.Tensor) and e.shape.rank == 0 - for e in elems_as_tensors - ): - reshaped = [gen_array_ops.reshape(e, [1]) for e in elems_as_tensors] - if len(reshaped) == 1: - return reshaped[0] - return gen_array_ops.concat_v2(reshaped, axis=0, name=scope) return gen_array_ops.pack(elems_as_tensors, name=scope) else: return converted_elems From 6438edc0d526154464e6cfa1f8b473541c7d6c76 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Sat, 5 Sep 2026 19:54:09 +0530 Subject: [PATCH 26/42] Fix compilation errors on non-GPU builds in GPU Pack scalar path The previous commit added a 0-D scalar shortcut in PackOp::Compute that called: output_vec.template chip<0>(i).device(c->eigen_device()) = input.scalar(); without restricting the code path to GPU-only. This caused compilation failures on Linux x86-CPU, ARM64, and Windows x86 CI targets because: 1. The expression is instantiated for PackOp (where CPUDevice is Eigen::ThreadPoolDevice). Eigen's device() assignment asserts at compile time that the executor matches the device type; for ThreadPoolDevice it fires: error: static assertion failed: Default executor instantiated with non-default device. You must #define EIGEN_USE_THREADS ... 2. Types like tensorflow::Variant and Eigen::QInt8 are only registered for the CPU PackOp and never for GPU, so they are never affected by the GPU bug; there is no reason to touch their code path at all. Fix: Wrap the entire scalar shortcut inside: #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM if constexpr (std::is_same_v) { ... } #endif The outer preprocessor guard prevents the block from being compiled at all on CPU-only builds. The inner if constexpr ensures that even in a GPU-enabled build (where both PackOp and PackOp are instantiated) the Eigen chip expression is only compiled into the GPUDevice specialisation. CPU builds and all non-scalar inputs are completely unaffected: they fall through to the unchanged ConcatCPU / ConcatGPU generic path. --- tensorflow/core/kernels/pack_op.cc | 48 ++++++++++++++++++------------ 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc index 9675595a3a8598..ccf72664e2d14f 100644 --- a/tensorflow/core/kernels/pack_op.cc +++ b/tensorflow/core/kernels/pack_op.cc @@ -80,31 +80,41 @@ class PackOp : public OpKernel { Tensor* output; OP_REQUIRES_OK(c, c->allocate_output(0, output_shape, &output)); - // Special case: packing 0-D (scalar) inputs. +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + // Special case: packing 0-D (scalar) inputs on GPU. // // The generic path below flattens each input to a {before_dim, after_dim} // matrix and then calls ConcatGPU / ConcatCPU. For scalar inputs // before_dim == after_dim == 1, so every matrix has shape {1, 1} and the - // output matrix has shape {1, num}. While the arithmetic is correct, the - // GPU ConcatGPU helper validates that its inputs have at least one element - // per row (dimension(1) > 0 is always true here), but some GPU drivers - // reject the launch configuration produced for 1-element rows. Avoid the - // issue entirely by treating the scalar case as a flat vector copy: - // each scalar contributes exactly one element at position i in the output. - if (first_input.dims() == 0) { - auto output_vec = output->flat(); - for (int i = 0; i < num; ++i) { - const Tensor& input = c->input(i); - OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), - absl::InvalidArgumentError(absl::StrCat( - "Shapes of all inputs must match: values[0].shape = ", - first_input.shape().DebugString(), " != values[", i, - "].shape = ", input.shape().DebugString()))); - output_vec.template chip<0>(i).device(c->eigen_device()) = - input.scalar(); + // output matrix has shape {1, num}. While the arithmetic is correct, some + // GPU drivers reject the launch configuration produced for 1-element rows + // inside ConcatGPU, causing an "Can't concatenate scalars" error at + // runtime. Handle scalars on GPU by writing each scalar element directly + // into the output flat vector via an Eigen chip assignment. + // + // The if constexpr guard is required: the Eigen chip + scalar expression + // only compiles for Eigen::GpuDevice. ThreadPoolDevice (the CPU device + // type used for types like Variant and QInt8) would trigger a static + // assertion in Eigen ("Default executor instantiated with non-default + // device"). The CPU ConcatCPU path already handles scalars correctly and + // is used unchanged for all non-GPU devices. + if constexpr (std::is_same_v) { + if (first_input.dims() == 0) { + auto output_vec = output->flat(); + for (int i = 0; i < num; ++i) { + const Tensor& input = c->input(i); + OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), + absl::InvalidArgumentError(absl::StrCat( + "Shapes of all inputs must match: values[0].shape = ", + first_input.shape().DebugString(), " != values[", i, + "].shape = ", input.shape().DebugString()))); + output_vec.template chip<0>(i).device(c->eigen_device()) = + input.scalar(); + } + return; } - return; } +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM int64_t before_dim = 1; for (int i = 0; i < axis; ++i) { From 0c1396c7d3a79355987bb794bfa873b3071ede7f Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Sat, 5 Sep 2026 16:45:34 -0400 Subject: [PATCH 27/42] Reject negative GPU NMS output sizes --- .../image/non_max_suppression_op.cu.cc | 3 ++ .../image/non_max_suppression_op_gpu_test.cc | 38 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc b/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc index 9b734a2cdfa39f..2b666090b77879 100644 --- a/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc +++ b/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc @@ -484,6 +484,9 @@ absl::Status CheckValidInputs(const Tensor& boxes, const Tensor& scores, max_output_size.shape().DebugString(), " (Shape must be rank 0 but is ", "rank ", max_output_size.dims(), ")")); } + if (max_output_size.scalar()() < 0) { + return absl::InvalidArgumentError("max_output_size must be non-negative"); + } if (!TensorShapeUtils::IsScalar(iou_threshold.shape())) { return absl::InvalidArgumentError(absl::StrCat( "iou_threshold must be 0-D, got shape ", diff --git a/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc b/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc index 76b2a9670d6460..08bd4bea40185f 100644 --- a/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc +++ b/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc @@ -213,6 +213,18 @@ TEST_F(NonMaxSuppressionV2GPUOpTest, TestInvalidIOUThreshold) { << s; } +TEST_F(NonMaxSuppressionV2GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + TEST_F(NonMaxSuppressionV2GPUOpTest, TestEmptyInput) { MakeOp(); AddInputFromArray(TensorShape({0, 4}), {}); @@ -446,6 +458,19 @@ TEST_F(NonMaxSuppressionV3GPUOpTest, TestInvalidIOUThreshold) { << s; } +TEST_F(NonMaxSuppressionV3GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + AddInputFromArray(TensorShape({}), {0.0f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + TEST_F(NonMaxSuppressionV3GPUOpTest, TestEmptyInput) { MakeOp(); AddInputFromArray(TensorShape({0, 4}), {}); @@ -539,6 +564,19 @@ TEST_F(NonMaxSuppressionV4GPUOpTest, DifferentInputAndThresholdTypesWorks) { test::ExpectTensorEqual(expected_num_valid, *GetOutput(1)); } +TEST_F(NonMaxSuppressionV4GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + AddInputFromArray(TensorShape({}), {0.0f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + #endif } // namespace From 0701f6e84bce93924bbb2bd9f8585765fc7af949 Mon Sep 17 00:00:00 2001 From: Hasrat Ali Date: Tue, 8 Sep 2026 18:04:55 +0530 Subject: [PATCH 28/42] Move GPU scalar-pack logic to pack_op_gpu.cu.cc to fix build errors The previous fix placed the Eigen chip+device() assignment inside pack_op.cc behind if constexpr (std::is_same_v). This still caused compilation failures on GPU builds because: * pack_op.cc is a plain .cc file; NVCC does NOT define EIGEN_USE_GPU for it, only for .cu.cc sources. * Inside #if GOOGLE_CUDA, TF registers PackOp for the GPU device with HostMemory (line 208). That means the compiler instantiates PackOp::Compute in this TU, and even the discarded if-constexpr branch for GPUDevice causes Eigen to emit its static assertion: "Default executor instantiated with non-default device. You must #define EIGEN_USE_THREADS, EIGEN_USE_GPU or EIGEN_USE_SYCL before including Eigen headers." Fix: * Remove all Eigen device() expressions from pack_op.cc. * Add pack_op_gpu.cu.cc which defines PackScalarsOnGPU: a free function that performs the scalar-element chip assignments using the Eigen::GpuDevice. This file is compiled by NVCC/hipcc (as a .cu.cc source), which defines EIGEN_USE_GPU before any Eigen headers are included, satisfying Eigen's requirement. * tf_kernel_library(name="pack_op", prefix="pack_op") in BUILD automatically adds pack_op*.cu.cc to gpu_srcs, so no BUILD change is needed. * pack_op.cc forward-declares PackScalarsOnGPU inside #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM and calls it when std::is_same::value && first_input.dims() == 0. The regular std::is_same runtime check (not if constexpr) is used here because no Eigen device expression appears in pack_op.cc any more, so there is no risk of a static assertion. * The CPU path (ConcatCPU) continues to handle all scalar inputs on CPU unchanged. * Explicit instantiations in pack_op_gpu.cu.cc cover every type registered with REGISTER_GPU: TF_CALL_GPU_ALL_TYPES, int64, int16, uint32, uint64, float8_e5m2, float8_e4m3fn. --- tensorflow/core/kernels/pack_op.cc | 58 ++++++++--------- tensorflow/core/kernels/pack_op_gpu.cu.cc | 79 +++++++++++++++++++++++ 2 files changed, 107 insertions(+), 30 deletions(-) create mode 100644 tensorflow/core/kernels/pack_op_gpu.cu.cc diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc index ccf72664e2d14f..a43723bd438cf6 100644 --- a/tensorflow/core/kernels/pack_op.cc +++ b/tensorflow/core/kernels/pack_op.cc @@ -28,6 +28,20 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM +// Forward declaration of the GPU scalar-pack helper defined in +// pack_op_gpu.cu.cc (compiled with EIGEN_USE_GPU by NVCC/hipcc). +// Calling into that translation unit avoids instantiating Eigen +// GPU-device expressions inside a regular .cc file where +// EIGEN_USE_GPU is not defined, which would trigger Eigen's +// static assertion: +// "Default executor instantiated with non-default device." +namespace tensorflow { +template +void PackScalarsOnGPU(OpKernelContext* c, int num, Tensor* output); +} // namespace tensorflow +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + namespace tensorflow { typedef Eigen::ThreadPoolDevice CPUDevice; @@ -83,36 +97,20 @@ class PackOp : public OpKernel { #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM // Special case: packing 0-D (scalar) inputs on GPU. // - // The generic path below flattens each input to a {before_dim, after_dim} - // matrix and then calls ConcatGPU / ConcatCPU. For scalar inputs - // before_dim == after_dim == 1, so every matrix has shape {1, 1} and the - // output matrix has shape {1, num}. While the arithmetic is correct, some - // GPU drivers reject the launch configuration produced for 1-element rows - // inside ConcatGPU, causing an "Can't concatenate scalars" error at - // runtime. Handle scalars on GPU by writing each scalar element directly - // into the output flat vector via an Eigen chip assignment. - // - // The if constexpr guard is required: the Eigen chip + scalar expression - // only compiles for Eigen::GpuDevice. ThreadPoolDevice (the CPU device - // type used for types like Variant and QInt8) would trigger a static - // assertion in Eigen ("Default executor instantiated with non-default - // device"). The CPU ConcatCPU path already handles scalars correctly and - // is used unchanged for all non-GPU devices. - if constexpr (std::is_same_v) { - if (first_input.dims() == 0) { - auto output_vec = output->flat(); - for (int i = 0; i < num; ++i) { - const Tensor& input = c->input(i); - OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), - absl::InvalidArgumentError(absl::StrCat( - "Shapes of all inputs must match: values[0].shape = ", - first_input.shape().DebugString(), " != values[", i, - "].shape = ", input.shape().DebugString()))); - output_vec.template chip<0>(i).device(c->eigen_device()) = - input.scalar(); - } - return; - } + // The generic path below flattens inputs to {before_dim, after_dim} + // matrices and calls ConcatGPU. For scalars before_dim == after_dim == 1, + // giving {1,1} input matrices and a {1,num} output matrix. Some GPU + // drivers reject the resulting launch configuration, causing a + // "Can't concatenate scalars" error at runtime. Dispatch to + // PackScalarsOnGPU (defined in pack_op_gpu.cu.cc, compiled with + // EIGEN_USE_GPU) which writes each scalar element directly into the + // output flat vector. The indirection through a separate .cu.cc TU is + // required: Eigen device() assignments must be compiled with + // EIGEN_USE_GPU defined (i.e. by NVCC/hipcc), otherwise Eigen fires a + // static assertion even inside discarded if-constexpr branches. + if (std::is_same::value && first_input.dims() == 0) { + PackScalarsOnGPU(c, num, output); + return; } #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/pack_op_gpu.cu.cc b/tensorflow/core/kernels/pack_op_gpu.cu.cc new file mode 100644 index 00000000000000..730930e1fb785a --- /dev/null +++ b/tensorflow/core/kernels/pack_op_gpu.cu.cc @@ -0,0 +1,79 @@ +/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// GPU-side implementation of the PackScalarsOnGPU helper for PackOp. +// +// This file is compiled by NVCC / hipcc (i.e. as a .cu.cc source) which +// defines EIGEN_USE_GPU before including Eigen headers. That is required for +// the Eigen device() assignment used below; doing the same in the plain .cc +// file (where EIGEN_USE_GPU is NOT defined) would trigger: +// +// static assertion failed: "Default executor instantiated with non-default +// device. You must #define EIGEN_USE_THREADS, EIGEN_USE_GPU or +// EIGEN_USE_SYCL before including Eigen headers." +// +// pack_op.cc forward-declares PackScalarsOnGPU and calls it for the GPU +// device path; the linker resolves the call to the explicit instantiations +// at the bottom of this file. + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + +#define EIGEN_USE_GPU + +#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_types.h" + +namespace tensorflow { + +// Packs `num` 0-D scalar input tensors (c->input(0) … c->input(num-1)) into +// a pre-allocated 1-D output tensor `output` of shape [num]. +// Each scalar element i is written to output[i] via an Eigen GPU chip +// assignment so that the operation is correctly dispatched to the GPU stream. +template +void PackScalarsOnGPU(OpKernelContext* c, int num, Tensor* output) { + const Tensor& first_input = c->input(0); + auto output_vec = output->flat(); + const Eigen::GpuDevice& d = c->eigen_device(); + for (int i = 0; i < num; ++i) { + const Tensor& input = c->input(i); + OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), + absl::InvalidArgumentError(absl::StrCat( + "Shapes of all inputs must match: values[0].shape = ", + first_input.shape().DebugString(), " != values[", i, + "].shape = ", input.shape().DebugString()))); + output_vec.template chip<0>(i).device(d) = input.scalar(); + } +} + +// Explicit instantiations for all types registered with REGISTER_GPU in +// pack_op.cc. Must match TF_CALL_GPU_ALL_TYPES + the extra integer types. +#define INSTANTIATE(T) template void PackScalarsOnGPU(OpKernelContext*, int, Tensor*); + +TF_CALL_int64(INSTANTIATE); +TF_CALL_int16(INSTANTIATE); +TF_CALL_uint32(INSTANTIATE); +TF_CALL_uint64(INSTANTIATE); +TF_CALL_GPU_ALL_TYPES(INSTANTIATE); +TF_CALL_float8_e5m2(INSTANTIATE); +TF_CALL_float8_e4m3fn(INSTANTIATE); + +#undef INSTANTIATE + +} // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM From 66eea1cea303e1db26944146f4bf824981c3fc9d Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 8 Sep 2026 20:08:10 -0400 Subject: [PATCH 29/42] Allow large valid index shuffle round counts --- tensorflow/core/kernels/random_index_shuffle_ops.cc | 8 +++----- .../kernel_tests/random/random_index_shuffle_test.py | 4 ++-- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index 6376a4beb4023d..1ea88c8626e66d 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -36,7 +36,6 @@ namespace tensorflow { namespace { constexpr absl::string_view kRounds = "rounds"; -constexpr int32_t kMaxRounds = 1000; template std::array CastSeedFrom(const Tensor& seed_t, const int row) { @@ -70,11 +69,10 @@ class RandomIndexShuffleOp : public OpKernel { : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr(kRounds, &rounds_)); OP_REQUIRES( - context, rounds_ >= 4 && rounds_ <= kMaxRounds && rounds_ % 2 == 0, + context, rounds_ >= 4 && rounds_ % 2 == 0, absl::InvalidArgumentError(absl::StrCat( - "rounds must be an even integer between 4 and ", kMaxRounds, - ", but got ", - rounds_))); + "rounds must be an even integer greater than or equal to 4, but ", + "got ", rounds_))); } void Compute(OpKernelContext* context) override { diff --git a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py index 68debf5eee47d7..251e01e12b9e5b 100644 --- a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py +++ b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py @@ -118,11 +118,11 @@ def test_index_greater_than_max_index(self): ): self.evaluate(stateless.index_shuffle(5, seed=(1, 2), max_index=4)) - @parameterized.parameters(-1, 0, 2, 3, 5, 1002) + @parameterized.parameters(-1, 0, 2, 3, 5) def test_invalid_rounds(self, rounds): with self.assertRaisesRegex( errors.InvalidArgumentError, - 'rounds must be an even integer between 4 and 1000', + 'rounds must be an even integer', ): self.evaluate( gen_random_index_shuffle_ops.random_index_shuffle( From d39bf88666275ef59072acf708fc3b6a62d33125 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 8 Sep 2026 20:44:29 -0400 Subject: [PATCH 30/42] Include Abseil string concatenation header --- tensorflow/core/kernels/random_index_shuffle_ops.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index 1ea88c8626e66d..515d08679cc350 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "tensorflow/core/framework/bounds_check.h" From 3b9e9095772d074d53a1b1f76d69dfdb111db3a4 Mon Sep 17 00:00:00 2001 From: Aditya Gaurav Date: Wed, 16 Sep 2026 00:09:14 +0530 Subject: [PATCH 31/42] fix(kernels): fix 32-bit signed integer overflow in Dilation GPU kernels Use 64-bit launch configuration with occupancy-aware GetGpuLaunchConfig64 to prevent integer overflow on large tensors in 2D and 3D dilation kernels, while staying within device resource limits. --- .../core/kernels/dilation_ops_gpu.cu.cc | 229 ++++++++++-------- 1 file changed, 122 insertions(+), 107 deletions(-) diff --git a/tensorflow/core/kernels/dilation_ops_gpu.cu.cc b/tensorflow/core/kernels/dilation_ops_gpu.cu.cc index 788ad9f5127b6d..057fc9c518c972 100644 --- a/tensorflow/core/kernels/dilation_ops_gpu.cu.cc +++ b/tensorflow/core/kernels/dilation_ops_gpu.cu.cc @@ -36,27 +36,28 @@ namespace { template __global__ void DilationKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, - const T* __restrict__ filter_ptr, int batch, int input_rows, int input_cols, - int depth, int filter_rows, int filter_cols, int output_rows, - int output_cols, int stride_rows, int stride_cols, int rate_rows, - int rate_cols, int pad_top, int pad_left, T* __restrict__ output_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + const int64_t nthreads, const T* __restrict__ input_ptr, + const T* __restrict__ filter_ptr, int64_t batch, int64_t input_rows, + int64_t input_cols, int64_t depth, int64_t filter_rows, int64_t filter_cols, + int64_t output_rows, int64_t output_cols, int stride_rows, int stride_cols, + int rate_rows, int rate_cols, int pad_top, int pad_left, + T* __restrict__ output_ptr) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -75,33 +76,33 @@ __global__ void DilationKernel( template __global__ void DilationBackpropInputKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, + const int64_t nthreads, const T* __restrict__ input_ptr, const T* __restrict__ filter_ptr, const T* __restrict__ out_backprop_ptr, - int batch, int input_rows, int input_cols, int depth, int filter_rows, - int filter_cols, int output_rows, int output_cols, int stride_rows, - int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, - T* __restrict__ in_backprop_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + int64_t batch, int64_t input_rows, int64_t input_cols, int64_t depth, + int64_t filter_rows, int64_t filter_cols, int64_t output_rows, + int64_t output_cols, int stride_rows, int stride_cols, int rate_rows, + int rate_cols, int pad_top, int pad_left, T* __restrict__ in_backprop_ptr) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - int h_in_max = (h_beg < 0) ? 0 : h_beg; - int w_in_max = (w_beg < 0) ? 0 : w_beg; + int64_t h_in_max = (h_beg < 0) ? 0 : h_beg; + int64_t w_in_max = (w_beg < 0) ? 0 : w_beg; // In the case of multiple argmax branches, we only back-propagate along the // last branch, i.e., the one with largest value of `h * filter_cols + w`, // similarly to the max-pooling backward routines. - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -125,33 +126,34 @@ __global__ void DilationBackpropInputKernel( template __global__ void DilationBackpropFilterKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, + const int64_t nthreads, const T* __restrict__ input_ptr, const T* __restrict__ filter_ptr, const T* __restrict__ out_backprop_ptr, - int batch, int input_rows, int input_cols, int depth, int filter_rows, - int filter_cols, int output_rows, int output_cols, int stride_rows, - int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, + int64_t batch, int64_t input_rows, int64_t input_cols, int64_t depth, + int64_t filter_rows, int64_t filter_cols, int64_t output_rows, + int64_t output_cols, int stride_rows, int stride_cols, int rate_rows, + int rate_cols, int pad_top, int pad_left, T* __restrict__ filter_backprop_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - int h_max = 0; - int w_max = 0; + int64_t h_max = 0; + int64_t w_max = 0; // In the case of multiple argmax branches, we only back-propagate along the // last branch, i.e., the one with largest value of `h * filter_cols + w`, // similarly to the max-pooling backward routines. - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -182,23 +184,26 @@ struct Dilation { typename TTypes::ConstTensor filter, int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor output) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); - - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); - - const int output_rows = output.dimension(1); - const int output_cols = output.dimension(2); - - const int total_count = batch * output_rows * output_cols * depth; - GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d); - + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); + + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); + + const int64_t output_rows = output.dimension(1); + const int64_t output_cols = output.dimension(2); + + const int64_t total_count = batch * output_rows * output_cols * depth; + // Occupancy-aware config: heuristic 1024-thread blocks exceed registers + // on cuda13 for these kernels (too many resources requested for launch). + absl::StatusOr config = + GetGpuLaunchConfig64(total_count, d, DilationKernel, 0, 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationKernel, config.block_count, config.thread_per_block, 0, - d.stream(), config.virtual_thread_count, input.data(), filter.data(), + DilationKernel, config->block_count, config->thread_per_block, 0, + d.stream(), config->virtual_thread_count, input.data(), filter.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, output.data())); @@ -213,33 +218,37 @@ struct DilationBackpropInput { int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor in_backprop) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); - const int output_rows = out_backprop.dimension(1); - const int output_cols = out_backprop.dimension(2); + const int64_t output_rows = out_backprop.dimension(1); + const int64_t output_cols = out_backprop.dimension(2); - int total_count; - GpuLaunchConfig config; + int64_t total_count; + absl::StatusOr config; // Initialize in_backprop with all zeros. total_count = batch * input_rows * input_cols * depth; - config = GetGpuLaunchConfig(total_count, d); - TF_CHECK_OK(GpuLaunchKernel(SetZero, config.block_count, - config.thread_per_block, 0, d.stream(), - total_count, in_backprop.data())); + config = GetGpuLaunchConfig64(total_count, d); + CHECK_OK(config.status()); // Crash OK + TF_CHECK_OK(GpuLaunchKernel(SetZero, config->block_count, + config->thread_per_block, 0, d.stream(), + config->virtual_thread_count, + in_backprop.data())); // Accumulate. total_count = batch * output_rows * output_cols * depth; - config = GetGpuLaunchConfig(total_count, d); + config = GetGpuLaunchConfig64(total_count, d, DilationBackpropInputKernel, + 0, 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationBackpropInputKernel, config.block_count, - config.thread_per_block, 0, d.stream(), config.virtual_thread_count, + DilationBackpropInputKernel, config->block_count, + config->thread_per_block, 0, d.stream(), config->virtual_thread_count, input.data(), filter.data(), out_backprop.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, @@ -255,33 +264,38 @@ struct DilationBackpropFilter { int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor filter_backprop) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); - const int output_rows = out_backprop.dimension(1); - const int output_cols = out_backprop.dimension(2); + const int64_t output_rows = out_backprop.dimension(1); + const int64_t output_cols = out_backprop.dimension(2); - int total_count; - GpuLaunchConfig config; + int64_t total_count; + absl::StatusOr config; // Initialize filter_backprop with all zeros. total_count = filter_rows * filter_cols * depth; - config = GetGpuLaunchConfig(total_count, d); - TF_CHECK_OK(GpuLaunchKernel(SetZero, config.block_count, - config.thread_per_block, 0, d.stream(), - total_count, filter_backprop.data())); + config = GetGpuLaunchConfig64(total_count, d); + CHECK_OK(config.status()); // Crash OK + TF_CHECK_OK(GpuLaunchKernel(SetZero, config->block_count, + config->thread_per_block, 0, d.stream(), + config->virtual_thread_count, + filter_backprop.data())); // Accumulate. total_count = batch * output_rows * output_cols * depth; - config = GetGpuLaunchConfig(total_count, d); + config = + GetGpuLaunchConfig64(total_count, d, DilationBackpropFilterKernel, 0, + 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationBackpropFilterKernel, config.block_count, - config.thread_per_block, 0, d.stream(), config.virtual_thread_count, + DilationBackpropFilterKernel, config->block_count, + config->thread_per_block, 0, d.stream(), config->virtual_thread_count, input.data(), filter.data(), out_backprop.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, @@ -303,3 +317,4 @@ TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPECS); } // namespace tensorflow #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + From 22867aec054eb40cf819316aa1a24aa4cd38b505 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Fri, 18 Sep 2026 10:37:32 -0700 Subject: [PATCH 32/42] Restore the Apache license header in gpu_prim.h The disclaimer paragraph in this file's header is word-scrambled and does not match the Apache 2.0 boilerplate every other file in the repository uses. It reads "To in writing unless required by applicable law or agreed, distributed on an, software distributed under the license is AS IS BASIS, WITHOUT OF ANY KIND WARRANTIES OR CONDITIONS ... the license you must see." Replace it with the standard wording and restore the blank line after the copyright line, so the header matches the neighbouring gpu_prim_helpers.h byte for byte from line 2 onward. The 2020 copyright year is kept. Searching the repository for each scrambled phrase returns only this file, and it has been this way since 25956c47d6fa created it in January 2020, so nothing else needs the same treatment. Comments and code are untouched; this changes no behavior. Fixes #127660 --- tensorflow/core/kernels/gpu_prim.h | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/gpu_prim.h b/tensorflow/core/kernels/gpu_prim.h index 7f27caa9b825f2..52954116f0218a 100644 --- a/tensorflow/core/kernels/gpu_prim.h +++ b/tensorflow/core/kernels/gpu_prim.h @@ -1,15 +1,16 @@ /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 -To in writing unless required by applicable law or agreed, -distributed on an, software distributed under the license is "AS IS" -BASIS, WITHOUT OF ANY KIND WARRANTIES OR CONDITIONS, either express -or implied. For the specific language governing permissions and -limitations under the license, the license you must see. +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. ==============================================================================*/ #ifndef TENSORFLOW_CORE_KERNELS_GPU_PRIM_H_ #define TENSORFLOW_CORE_KERNELS_GPU_PRIM_H_ From 755d5b17501161693042b8caa7a69fa13e138f16 Mon Sep 17 00:00:00 2001 From: Yehor Cherednichenko Date: Wed, 23 Sep 2026 12:26:36 -0700 Subject: [PATCH 33/42] PR #49007: [XLA:GPU] Rewrite FP8 convolutions in all computations of a module 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/49007 ### 📝 Summary of Changes `CudnnFusedConvRewriter::RunImpl` iterates over the non-fusion computations of a module and, for each of them, first runs `F8GraphConv`. If `F8GraphConv` rewrote anything, the function returned immediately, so the loop never reached the remaining computations. Any FP8 convolution in a later computation was left as a plain convolution: it got neither the `__cudnn$convForwardGraph` rewrite nor the bias/activation fusions, and the FP8 fallback later demoted it to a legacy non-FP8 `__cudnn$convForward` call. This change replaces the early `return` with `any_changed |= changed; continue;`. The `continue` keeps the existing behaviour for the rewritten computation (the remaining fusions only match the legacy conv custom calls and do not apply to graph convolutions), and the loop now processes every computation of the module. ### 🎯 Justification Hit in practice with JAX 0.10.2 on an RTX 5070 (sm_120, CUDA 12.9, cuDNN 9.24). A `jit` for MCTS self-play evaluates an FP8 convolutional network once in the entry computation and again inside a `lax.while_loop` body. Only one of the two copies was rewritten: the compiled HLO contained 20 `__cudnn$convForwardGraph` calls instead of 40, and the other 20 convolutions ran as legacy f32 convolutions. With this change all 40 are FP8 graph convolutions. Which copy loses depends on the computation order, so the workaround was to keep one copy out of FP8 entirely. Minimal JAX repro (same FP8 convolution in the entry computation and in a while body): ```python import jax, jax.numpy as jnp from jax import lax def conv(x, w): return lax.conv_general_dilated(x, w, (1, 1), "SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=jnp.float32) @jax.jit def f(x, w): y0 = conv(x, w) # entry computation def body(c): i, y = c return i + 1, y + conv(y.astype(jnp.float8_e4m3fn), w) # while body return lax.while_loop(lambda c: c[0] < 4, body, (0, y0))[1] x = jnp.ones((1, 6, 6, 128), jnp.float8_e4m3fn) w = jnp.ones((3, 3, 128, 128), jnp.float8_e4m3fn) hlo = f.lower(x, w).compile().as_text() print(hlo.count('custom_call_target="__cudnn$convForwardGraph"'), hlo.count('custom_call_target="__cudnn$convForward"')) ``` Before this change (jax-cuda12-pjrt 0.10.2): `1 1` (one graph convolution, one legacy convolution). With this change (plugin rebuilt from the same JAX release with the patched XLA): `2 0`. ### 🚀 Kind of Contribution 🐛 Bug Fix ### 🧪 Unit Tests `CudnnFusedConvRewriterTest.TestConvF8InMultipleComputations` in `cudnn_fused_conv_rewriter_test.cc`. The module has two `conditional` branches, each with an FP8 convolution on a different filter. After `ConvRewriter` and `CudnnFusedConvRewriter` (run with a Hopper compute capability, like the existing F8 tests on pre-FP8 hardware) both branch computations must contain a `__cudnn$convForwardGraph` custom call. Without the fix only the first branch visited is rewritten and the test fails on the second one. ### 🧪 Execution Tests On FP8-capable GPUs the same test additionally compiles the module through the full pipeline, checks with FileCheck that both convolutions are `__cudnn$convForwardGraph` and no legacy `__cudnn$convForward` remains, and compares the result against the interpreter (`RunAndCompare`, `ErrorSpec{0.15, 0.15}`). Tested locally on an RTX 5070 (sm_120) with the hermetic CUDA toolchain (cuDNN 9.12): `bazel test //xla/backends/gpu/transforms:cudnn_fused_conv_rewriter_test_nvgpu_any --test_filter='*TestConvF8InMultipleComputations*:*.TestConvF8:*TestConvScaledOutputF8'` passes (the new test runs the full-pipeline branch: `fp8_supported: true`). With the test but without the change in `cudnn_fused_conv_rewriter.cc` the new test fails as intended: `branch_b` keeps the legacy `__cudnn$convForward` custom call. Copybara import of the project: -- faf7a8a6ffa198a62772dcfe03ef90cfbb51eae3 by Yehor Cherednichenko : [XLA:GPU] Rewrite FP8 convolutions in all computations of a module CudnnFusedConvRewriter::RunImpl returned as soon as F8GraphConv had rewritten the FP8 convolutions of one computation. FP8 convolutions in the remaining computations of the module (e.g. a while body next to the entry computation) were left unrewritten and later demoted to unfused f32 convolutions by the FP8 fallback. Continue with the next computation instead; the remaining fusions only match the legacy conv custom calls, so they still do not apply to the rewritten computation. Add TestConvF8InMultipleComputations: both branches of a conditional contain an FP8 convolution and both must get the ForwardGraph custom call. Merging this change closes #49007 PiperOrigin-RevId: 986927090 --- .../transforms/cudnn_fused_conv_rewriter.cc | 4 +- .../cudnn_fused_conv_rewriter_test.cc | 79 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc index bb6f8587dd4448..7e0e395940053f 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc @@ -1698,9 +1698,7 @@ absl::StatusOr CudnnFusedConvRewriter::RunImpl( auto* cc = compute_capability_.cuda_compute_capability(); ABSL_ASSIGN_OR_RETURN(changed, F8GraphConv(comp, *cc, dnn_version_, toolkit_version_)); - if (changed) { - return changed; - } + any_changed |= changed; } // Fuse "inside out" starting with the operations closest to the conv. ABSL_ASSIGN_OR_RETURN(changed, FuseRemoveConvertInConv(comp)); diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc index c5205ee226492b..d55df350ed27a2 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/comparison_util.h" #include "xla/debug_options_flags.h" #include "xla/error_spec.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" @@ -907,6 +908,84 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvF8) { )"); } +// The pass used to stop after the first computation in which an FP8 +// convolution was rewritten, leaving the FP8 convolutions of the remaining +// computations of the module unrewritten. +TEST_F(CudnnFusedConvRewriterTest, TestConvF8InMultipleComputations) { + MAYBE_SKIP_TEST("F8"); + const std::string kHloString = R"( + HloModule Test + + branch_a { + p = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) parameter(0) + input_a = f8e4m3fn[1,128,6,6] get-tuple-element(p), index=0 + filter_a = f8e4m3fn[3,3,128,16] get-tuple-element(p), index=1 + ROOT conv_a = f8e4m3fn[1,16,6,6] convolution(input_a, filter_a), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, feature_group_count=1 + } + + branch_b { + p = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) parameter(0) + input_b = f8e4m3fn[1,128,6,6] get-tuple-element(p), index=0 + filter_b = f8e4m3fn[3,3,128,16] get-tuple-element(p), index=1 + ROOT conv_b = f8e4m3fn[1,16,6,6] convolution(input_b, filter_b), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, feature_group_count=1 + } + + ENTRY Test { + predicate = pred[] parameter(0) + input = f8e4m3fn[1,128,6,6] parameter(1) + filter_a = f8e4m3fn[3,3,128,16] parameter(2) + filter_b = f8e4m3fn[3,3,128,16] parameter(3) + operands_a = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) tuple(input, filter_a) + operands_b = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) tuple(input, filter_b) + ROOT result = f8e4m3fn[1,16,6,6] conditional(predicate, operands_a, operands_b), true_computation=branch_a, false_computation=branch_b + })"; + + // Both branch computations must get a ForwardGraph Custom Call. + const se::CudaComputeCapability hopper{se::CudaComputeCapability::kHopper, 0}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloString)); + ASSERT_OK_AND_ASSIGN(bool changed, + RunHloPass(ConvRewriter(hopper), module.get())); + EXPECT_TRUE(changed); + ASSERT_OK_AND_ASSIGN( + changed, RunHloPass(CudnnFusedConvRewriter(hopper, GetDnnVersion(), + GetToolkitVersion()), + module.get())); + EXPECT_TRUE(changed); + for (absl::string_view name : {"branch_a", "branch_b"}) { + HloComputation* branch = module->GetComputationWithName(name); + ASSERT_THAT(branch, NotNull()) << name; + const HloInstruction* conv = nullptr; + for (const HloInstruction* instr : branch->instructions()) { + if (instr->opcode() == HloOpcode::kCustomCall) { + conv = instr; + } + } + ASSERT_THAT(conv, NotNull()) << name; + EXPECT_EQ(conv->custom_call_target(), kCudnnConvForwardGraphCallTarget) + << name; + } + + bool fp8_supported = GetDnnVersion() >= se::dnn::VersionInfo{9, 8, 0} + ? GetCudaComputeCapability().IsAtLeastAda() + : GetCudaComputeCapability().IsAtLeastHopper(); + if (fp8_supported) { + // Through the full pipeline, neither convolution may fall back to the + // legacy (non-FP8) Custom Call. + absl::StatusOr filecheck_result = + RunFileCheck(GetOptimizedHlo(kHloString), R"( +// CHECK-NOT: custom_call_target="__cudnn$convForward" +// CHECK: custom_call_target="__cudnn$convForwardGraph" +// CHECK-NOT: custom_call_target="__cudnn$convForward" +// CHECK: custom_call_target="__cudnn$convForwardGraph" +// CHECK-NOT: custom_call_target="__cudnn$convForward" + )"); + ASSERT_TRUE(filecheck_result.ok()) << filecheck_result.status(); + EXPECT_TRUE(*filecheck_result); + EXPECT_TRUE(RunAndCompare(kHloString, ErrorSpec{0.15, 0.15})) << kHloString; + } +} + TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputF8) { MAYBE_SKIP_TEST("F8"); TestF8( From 873b71f1c434d4eb9999ea00958e17af22aa8419 Mon Sep 17 00:00:00 2001 From: Saicharan Ramineni Date: Wed, 23 Sep 2026 12:38:31 -0700 Subject: [PATCH 34/42] PR #123156: Support int32 argmax in CPU MaxPool gradients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/123156 Fixes #103454 `MaxPoolWithArgmax` already supports `output_dtype=tf.int32` on CPU, but its gradient only registered an int64 argmax kernel. This caused automatic differentiation to fail with “No OpKernel was registered” when users selected int32 output indices. Template the CPU gradient implementation on the argmax dtype, register the existing CPU gradient for `Targmax=int32`, and extend the pooling gradient test matrix to validate both int32 and int64 CPU argmax tensors. GPU remains unchanged because its forward argmax kernel only supports int64. Validation: `git diff --check`; `python3 -m py_compile tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py`; `bazel build --config=macos_arm64 //tensorflow/core/kernels:pooling_ops` is running locally. Copybara import of the project: -- 7ce54520e9faf502224067834ace16482755dac9 by Saicharan Ramineni : Support int32 argmax in CPU max pool gradients Merging this change closes #123156 COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/123156 from GodlyDonuts:fix/maxpool-int32-argmax 3ef9f723edfa347b84e33332ef6c1b032353d4c8 PiperOrigin-RevId: 986934206 --- tensorflow/core/kernels/maxpooling_op.cc | 46 ++++++++++--------- .../kernel_tests/nn_ops/pooling_ops_test.py | 16 ++++--- 2 files changed, 35 insertions(+), 27 deletions(-) diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc index 6f7d68cfdec580..e5f94a0ff9bf8f 100644 --- a/tensorflow/core/kernels/maxpooling_op.cc +++ b/tensorflow/core/kernels/maxpooling_op.cc @@ -1056,11 +1056,11 @@ class MaxPoolingWithArgmaxOp : public OpKernel { bool include_batch_in_index_; }; -template +template struct LaunchMaxPoolingGradWithArgmax; -template -struct LaunchMaxPoolingGradWithArgmax { +template +struct LaunchMaxPoolingGradWithArgmax { typedef Eigen::Map> EigenMatrixMap; @@ -1080,7 +1080,7 @@ struct LaunchMaxPoolingGradWithArgmax { { auto grad_out_flat = grad_out->flat(); - auto argmax_flat = argmax.flat(); + auto argmax_flat = argmax.flat(); auto grad_in_flat = grad_in.flat(); const int64_t output_start = start * output_size_per_batch; @@ -1116,8 +1116,7 @@ struct LaunchMaxPoolingGradWithArgmax { } }; -// TODO(b/175733711): Support int32 argmax type in MaxPoolGradWithArgmax op. -template +template class MaxPoolingGradWithArgmaxOp : public OpKernel { public: explicit MaxPoolingGradWithArgmaxOp(OpKernelConstruction* context) @@ -1195,7 +1194,7 @@ class MaxPoolingGradWithArgmaxOp : public OpKernel { if (out_shape.num_elements() == 0) return; // nothing to be done - LaunchMaxPoolingGradWithArgmax::launch( + LaunchMaxPoolingGradWithArgmax::launch( context, params, grad_in, argmax, grad_out, include_batch_in_index_); } @@ -1529,7 +1528,7 @@ struct LaunchMaxPoolingWithArgmax { }; template -struct LaunchMaxPoolingGradWithArgmax { +struct LaunchMaxPoolingGradWithArgmax { static void launch(OpKernelContext* context, const PoolParameters& params, const Tensor& grad_in, const Tensor& argmax, Tensor* grad_out, const bool include_batch_in_index) { @@ -1605,21 +1604,26 @@ struct LaunchMaxPoolingGradGradWithArgmax { .Device(DEVICE_##D) \ .TypeConstraint("T") \ .TypeConstraint("Targmax"), \ - MaxPoolingGradWithArgmaxOp); + MaxPoolingGradWithArgmaxOp); // Below kernels implemented only for CPU device. -#define REGISTER_CPU_ONLY_POOL_KERNELS(T) \ - REGISTER_KERNEL_BUILDER( \ - Name("MaxPool").Device(DEVICE_CPU).TypeConstraint("T"), \ - MaxPoolingOp); \ - REGISTER_KERNEL_BUILDER( \ - Name("MaxPoolV2").Device(DEVICE_CPU).TypeConstraint("T"), \ - MaxPoolingV2Op); \ - REGISTER_KERNEL_BUILDER(Name("MaxPoolWithArgmax") \ - .Device(DEVICE_CPU) \ - .TypeConstraint("Targmax") \ - .TypeConstraint("T"), \ - MaxPoolingWithArgmaxOp); +#define REGISTER_CPU_ONLY_POOL_KERNELS(T) \ + REGISTER_KERNEL_BUILDER( \ + Name("MaxPool").Device(DEVICE_CPU).TypeConstraint("T"), \ + MaxPoolingOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("MaxPoolV2").Device(DEVICE_CPU).TypeConstraint("T"), \ + MaxPoolingV2Op); \ + REGISTER_KERNEL_BUILDER(Name("MaxPoolWithArgmax") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("Targmax") \ + .TypeConstraint("T"), \ + MaxPoolingWithArgmaxOp); \ + REGISTER_KERNEL_BUILDER(Name("MaxPoolGradWithArgmax") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Targmax"), \ + MaxPoolingGradWithArgmaxOp); TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_ONLY_POOL_KERNELS); #undef REGISTER_CPU_ONLY_POOL_KERNELS diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py index c7f707285e1e23..78a7a851619332 100644 --- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py @@ -1091,12 +1091,15 @@ def testMaxPoolingGradWithArgmax(self): tensor_input = [11.0, 12.0, 13.0, 14.0, 21.0, 22.0, 23.0, 24.0] Config = collections.namedtuple( - "Config", ["use_gpu", "include_batch_in_index", "argmax"]) + "Config", ["use_gpu", "include_batch_in_index", "argmax", "Targmax"] + ) configs = [ - Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8]), - Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17]), - Config(True, False, [0, 1, 3, 5, 0, 2, 6, 8]), - Config(True, True, [0, 1, 3, 5, 9, 11, 15, 17]) + Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int64), + Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int64), + Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int32), + Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int32), + Config(True, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int64), + Config(True, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int64), ] for config in configs: @@ -1104,7 +1107,8 @@ def testMaxPoolingGradWithArgmax(self): orig_in = constant_op.constant(orig_input, shape=[2, 3, 3, 1]) t = constant_op.constant(tensor_input, shape=[2, 2, 2, 1]) argmax_t = constant_op.constant( - config.argmax, shape=[2, 2, 2, 1], dtype=dtypes.int64) + config.argmax, shape=[2, 2, 2, 1], dtype=config.Targmax + ) out_op = gen_nn_ops.max_pool_grad_with_argmax( orig_in, t, From 82987dc45be0c37cf12fcee4d7a8300f74df0e7e Mon Sep 17 00:00:00 2001 From: Matthias Guenther Date: Wed, 23 Sep 2026 13:10:13 -0700 Subject: [PATCH 35/42] Sync repo state before enabling per-change Copybara migrations This unblocks the initial bulk migration, which will in turn unblock iterative migrations on a per-change basis. PiperOrigin-RevId: 986954805 --- .../mlir/tf2xla/tests/legalize-tf.mlir | 2 +- .../xla/third_party/stablehlo/temporary.patch | 3094 +---------------- .../xla/third_party/stablehlo/workspace.bzl | 4 +- 3 files changed, 31 insertions(+), 3069 deletions(-) diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir index a4d4f6a6ced441..195c1fadfad69b 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir @@ -2712,7 +2712,7 @@ func.func @neg_dynamic(%arg0: tensor) -> tensor { func.func @random_shuffle_float_vector(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK: "mhlo.sort" // CHECK: ^bb0([[KEY_LHS:%.*]]: tensor, [[KEY_RHS:%.*]]: tensor, {{.*}}: tensor, {{.*}}: tensor): - // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]], TOTALORDER : (tensor, tensor) -> tensor + // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]] : (tensor, tensor) -> tensor %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32> func.return %0 : tensor<2xf32> } diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ec90dce6dc5dc6..fb55a869629448 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -18,1117 +18,11 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel --- stablehlo/BUILD.bazel +++ stablehlo/BUILD.bazel -@@ -11,24 +11,42 @@ - # 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. -+ - load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "gentbl_filegroup", "td_library") --load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test") -+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") -+# copybara:uncomment load("@rules_license//rules:license.bzl", "license") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = [":license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], +@@ -1960,3 +1960,17 @@ + name = "all_tests", + tests = [":stablehlo_ci_tests"], ) - -+# copybara:uncomment_begin -+# license( -+# name = "license", -+# package_name = "stablehlo", -+# ) -+# copybara:uncomment_end + - exports_files([ - "LICENSE", - # Export to be used by MLIR-free python binding impls - "stablehlo/integrations/python/StablehloApi.cpp", - "stablehlo/integrations/python/StablehloApi.h", -+ # Export TD files for builder code generation. -+ "stablehlo/dialect/ChloOps.td", -+ "stablehlo/dialect/StablehloOps.td", - ]) - - filegroup( - name = "stablehlo_ops_td_filegroup", - srcs = glob(["stablehlo/dialect/*.td"]), -+ compatible_with = get_compatible_with_portable(), - ) - - cc_library( -@@ -39,7 +57,8 @@ - hdrs = [ - "stablehlo/dialect/Base.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base_attr_interfaces_inc_gen", - "@llvm-project//llvm:Support", -@@ -57,7 +76,7 @@ - - gentbl_cc_library( - name = "base_attr_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/BaseAttrInterfaces.cpp.inc": ["-gen-attr-interface-defs"], - "stablehlo/dialect/BaseAttrInterfaces.h.inc": ["-gen-attr-interface-decls"], -@@ -72,6 +91,7 @@ - srcs = [ - "stablehlo/dialect/Base.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -@@ -88,7 +108,8 @@ - hdrs = [ - "stablehlo/dialect/BroadcastUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - "@llvm-project//llvm:Support", -@@ -100,7 +121,7 @@ - - gentbl_cc_library( - name = "chlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/ChloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -126,7 +147,6 @@ - name = "chlo_capi", - srcs = CHLO_CAPI_SOURCES, - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":chlo_ops", - "@llvm-project//llvm:Support", -@@ -138,7 +158,7 @@ - cc_library( - name = "chlo_capi_headers", - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -149,7 +169,6 @@ - name = "chlo_capi_objects", - srcs = CHLO_CAPI_SOURCES, - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":chlo_ops", - "@llvm-project//llvm:Support", -@@ -160,7 +179,7 @@ - - gentbl_cc_library( - name = "chlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/ChloEnums.h.inc": ["-gen-enum-decls"], -@@ -174,7 +193,7 @@ - - gentbl_cc_library( - name = "chlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/ChloOps.h.inc": ["-gen-op-decls"], -@@ -200,6 +219,7 @@ - - gentbl_filegroup( - name = "chlo_ops_py_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = [ - ( - [ -@@ -223,6 +243,7 @@ - "stablehlo/dialect/ChloEnums.td", - "stablehlo/dialect/ChloOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -249,7 +270,8 @@ - "stablehlo/dialect/ChloBytecode.h", - "stablehlo/dialect/ChloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":broadcast_utils", -@@ -278,6 +300,7 @@ - - gentbl_cc_library( - name = "chlo_rewriters_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/ChloDecompositionPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/ChloDecompositionPatterns.td", -@@ -289,6 +312,7 @@ - - gentbl_cc_library( - name = "stablehlo_aggressive_simplification_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td", -@@ -299,6 +323,7 @@ - - gentbl_cc_library( - name = "stablehlo_legalize_deprecated_ops_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td", -@@ -309,6 +334,7 @@ - - gentbl_cc_library( - name = "vhlo_rewriters_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/VhloToVersionPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/VhloToVersionPatterns.td", -@@ -319,6 +345,7 @@ - - gentbl_cc_library( - name = "stablehlo_create_compatibility_expander_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloCompatibilityExpanderPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloCompatibilityExpanderPatterns.td", -@@ -329,6 +356,7 @@ - - gentbl_cc_library( - name = "stablehlo_create_complex_math_expander_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloComplexMathExpanderPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloComplexMathExpanderPatterns.td", -@@ -345,7 +373,8 @@ - hdrs = [ - "stablehlo/reference/InterpreterOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":interpreter_ops_inc_gen", -@@ -364,7 +393,7 @@ - - gentbl_cc_library( - name = "interpreter_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/reference/InterpreterOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/reference/InterpreterOps.h.inc": ["-gen-op-decls"], -@@ -381,6 +410,7 @@ - srcs = [ - "stablehlo/reference/InterpreterOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -389,6 +419,7 @@ - - gentbl_cc_library( - name = "interpreter_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/reference/InterpreterPasses.h.inc": [ -@@ -412,7 +443,8 @@ - hdrs = [ - "stablehlo/reference/InterpreterPasses.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":interpreter_ops", - ":interpreter_pass_inc_gen", -@@ -426,6 +458,7 @@ - - gentbl_cc_library( - name = "linalg_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/conversions/linalg/transforms/Passes.h.inc": [ -@@ -461,6 +494,7 @@ - "stablehlo/conversions/linalg/transforms/Rewriters.h", - "stablehlo/conversions/linalg/transforms/TypeConversion.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - "stablehlo_type_conversions", -@@ -503,6 +537,7 @@ - hdrs = [ - "stablehlo/reference/Api.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":interpreter_ops", -@@ -534,7 +569,8 @@ - hdrs = [ - "stablehlo/reference/Axes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -550,6 +586,7 @@ - hdrs = [ - "stablehlo/reference/Configuration.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":reference_errors", -@@ -569,7 +606,8 @@ - hdrs = [ - "stablehlo/reference/Element.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":reference_types", -@@ -588,7 +626,8 @@ - hdrs = [ - "stablehlo/reference/Errors.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - ], -@@ -602,7 +641,8 @@ - hdrs = [ - "stablehlo/reference/Index.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -618,6 +658,7 @@ - hdrs = [ - "stablehlo/reference/NumPy.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":reference_tensor", -@@ -635,7 +676,8 @@ - hdrs = [ - "stablehlo/reference/Ops.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":reference_axes", -@@ -669,7 +711,8 @@ - hdrs = [ - "stablehlo/reference/Process.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_process_grid", - ":reference_tensor", -@@ -686,7 +729,8 @@ - hdrs = [ - "stablehlo/reference/ProcessGrid.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_tensor", - "@llvm-project//llvm:Support", -@@ -703,7 +747,8 @@ - hdrs = [ - "stablehlo/reference/Scope.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_tensor", - ":reference_token", -@@ -722,7 +767,8 @@ - hdrs = [ - "stablehlo/reference/Tensor.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_axes", - ":reference_element", -@@ -743,7 +789,8 @@ - hdrs = [ - "stablehlo/reference/Token.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":stablehlo_ops", -@@ -761,7 +808,8 @@ - hdrs = [ - "stablehlo/reference/Types.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", -@@ -776,7 +824,8 @@ - hdrs = [ - "stablehlo/reference/Value.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":reference_tensor", -@@ -795,7 +844,7 @@ - hdrs = [ - "stablehlo/dialect/Register.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":chlo_ops", - ":stablehlo_ops", -@@ -817,7 +866,7 @@ - hdrs = [ - "stablehlo/dialect/AssemblyFormat.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":base", - "@llvm-project//llvm:Support", -@@ -828,7 +877,7 @@ - - gentbl_cc_library( - name = "stablehlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/StablehloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -860,7 +909,6 @@ - name = "stablehlo_dialect_capi", - srcs = STABLEHLO_DIALECT_CAPI_SOURCES, - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":interpreter_ops", - ":stablehlo_ops", -@@ -878,7 +926,7 @@ - cc_library( - name = "stablehlo_dialect_capi_headers", - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -889,7 +937,6 @@ - name = "stablehlo_dialect_capi_objects", - srcs = STABLEHLO_DIALECT_CAPI_SOURCES, - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":interpreter_ops", - ":stablehlo_ops", -@@ -923,7 +970,6 @@ - name = "stablehlo_unified_capi", - srcs = STABLEHLO_UNIFIED_CAPI_SOURCES, - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":linalg_passes", - ":reference_api", -@@ -951,7 +997,7 @@ - cc_library( - name = "stablehlo_unified_capi_headers", - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -967,7 +1013,6 @@ - name = "stablehlo_unified_capi_objects", - srcs = STABLEHLO_UNIFIED_CAPI_SOURCES, - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":linalg_passes", - ":reference_api", -@@ -993,7 +1038,7 @@ - - gentbl_cc_library( - name = "stablehlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/StablehloEnums.h.inc": ["-gen-enum-decls"], -@@ -1007,7 +1052,7 @@ - - gentbl_cc_library( - name = "stablehlo_types_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloTypeDefs.cpp.inc": [ - "-gen-typedef-defs", -@@ -1027,7 +1072,7 @@ - - gentbl_cc_library( - name = "stablehlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/StablehloOps.h.inc": ["-gen-op-decls"], -@@ -1049,6 +1094,7 @@ - - gentbl_filegroup( - name = "stablehlo_ops_py_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = [ - ( - [ -@@ -1085,6 +1131,7 @@ - "stablehlo/dialect/StablehloTypes.td", - "stablehlo/transforms/ChloDecompositionPatternsMath.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -1096,6 +1143,7 @@ - - gentbl_cc_library( - name = "stablehlo_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/transforms/Passes.h.inc": [ -@@ -1118,7 +1166,8 @@ - hdrs = [ - "stablehlo/transforms/StablehloBroadcastLowering.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - "@llvm-project//llvm:Support", -@@ -1136,7 +1185,8 @@ - hdrs = [ - "stablehlo/transforms/PassUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":chlo_ops", -@@ -1145,6 +1195,7 @@ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:ComplexDialect", - "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", - "@llvm-project//mlir:Support", - ], - ) -@@ -1179,7 +1230,8 @@ - "stablehlo/transforms/Passes.h", - "stablehlo/transforms/StablehloRefineShapes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":broadcast_utils", -@@ -1229,6 +1281,7 @@ - - gentbl_cc_library( - name = "stablehlo_passes_optimization_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/transforms/optimization/Passes.h.inc": [ -@@ -1254,7 +1307,8 @@ - hdrs = [ - "stablehlo/transforms/optimization/Passes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":chlo_ops", -@@ -1286,7 +1340,8 @@ - hdrs = [ - "stablehlo/api/PortableApi.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":register", - ":stablehlo_ops", -@@ -1310,6 +1365,7 @@ - hdrs = [ - "stablehlo/dialect/Serialization.h", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":stablehlo_ops", - ":stablehlo_passes", -@@ -1333,7 +1389,8 @@ - hdrs = [ - "stablehlo/transforms/conversions/TypeConversion.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -1350,7 +1407,8 @@ - hdrs = [ - "stablehlo/dialect/TypeInference.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1371,7 +1429,8 @@ - hdrs = [ - "stablehlo/dialect/ReplicaGroupUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - "@llvm-project//llvm:Support", -@@ -1390,7 +1449,8 @@ - "stablehlo/dialect/StablehloBytecode.h", - "stablehlo/dialect/StablehloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1439,11 +1499,13 @@ - srcs = [ - "stablehlo/tools/StablehloTranslateMain.cpp", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":interpreter_ops", - ":reference_api", - ":reference_configuration", - ":reference_errors", -+ ":reference_numpy", - ":reference_ops", - ":reference_process", - ":reference_process_grid", -@@ -1478,6 +1540,7 @@ - srcs = [ - "stablehlo/tools/StablehloOptMain.cpp", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":interpreter_ops", - ":interpreter_passes", -@@ -1486,8 +1549,6 @@ - ":stablehlo_passes", - ":stablehlo_passes_optimization", - ":tosa_passes", -- "//stablehlo/tests:check_ops", -- "//stablehlo/tests:test_utils", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:AllExtensions", - "@llvm-project//mlir:AllPassesAndDialects", -@@ -1497,18 +1558,24 @@ - "@llvm-project//mlir:RegisterAllExtensions", - "@llvm-project//mlir:RegisterAllPasses", - "@llvm-project//mlir:TosaDialect", -+ # copybara:uncomment "@shardy//shardy/dialect/sdy/ir:register", -+ "//stablehlo/tests:check_ops", -+ "//stablehlo/tests:test_utils", - ], - ) - - gentbl_cc_library( - name = "tosa_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/conversions/tosa/transforms/Passes.h.inc": [ - "-gen-pass-decls", - "-name=StablehloTOSATransforms", - ], -- "stablehlo/conversions/tosa/transforms/stablehlo_tosa_passes.md": ["-gen-pass-doc"], -+ "stablehlo/conversions/tosa/transforms/stablehlo_tosa_passes.md": [ -+ "-gen-pass-doc", -+ ], - }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/conversions/tosa/transforms/Passes.td", -@@ -1526,7 +1593,8 @@ - hdrs = [ - "stablehlo/conversions/tosa/transforms/Passes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - ":tosa_pass_inc_gen", -@@ -1549,6 +1617,8 @@ - - gentbl_cc_library( - name = "tosa_pdll_inc_gen", -+ compatible_with = get_compatible_with_portable(), -+ strip_include_prefix = ".", - tbl_outs = {"stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll.h.inc": ["-x=cpp"]}, - tblgen = "@llvm-project//mlir:mlir-pdll", - td_file = "stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll", -@@ -1567,7 +1637,8 @@ - hdrs = [ - "stablehlo/dialect/Version.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -1587,7 +1658,6 @@ - name = "vhlo_capi", - srcs = VHLO_CAPI_SOURCES, - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":vhlo_ops", - "@llvm-project//mlir:CAPIIR", -@@ -1598,7 +1668,7 @@ - cc_library( - name = "vhlo_capi_headers", - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -1609,7 +1679,6 @@ - name = "vhlo_capi_objects", - srcs = VHLO_CAPI_SOURCES, - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":vhlo_ops", - "@llvm-project//mlir:CAPIIRObjects", -@@ -1627,15 +1696,13 @@ - - gentbl_filegroup( - name = "vhlo_ops_py_gen", -- tbl_outs = [ -- ( -- [ -- "-gen-python-op-bindings", -- "-bind-dialect=vhlo", -- ], -- "stablehlo/integrations/python/mlir/dialects/_vhlo_ops_gen.py", -- ), -- ], -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "stablehlo/integrations/python/mlir/dialects/_vhlo_ops_gen.py": [ -+ "-gen-python-op-bindings", -+ "-bind-dialect=vhlo", -+ ], -+ }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/integrations/python/mlir/dialects/VhloOps.td", - deps = [ -@@ -1646,7 +1713,7 @@ - - gentbl_cc_library( - name = "vhlo_attr_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloAttrInterfaces.cpp.inc": ["-gen-attr-interface-defs"], - "stablehlo/dialect/VhloAttrInterfaces.h.inc": ["-gen-attr-interface-decls"], -@@ -1660,7 +1727,7 @@ - - gentbl_cc_library( - name = "vhlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/VhloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -1674,7 +1741,7 @@ - - gentbl_cc_library( - name = "vhlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/VhloEnums.h.inc": ["-gen-enum-decls"], -@@ -1688,7 +1755,7 @@ - - gentbl_cc_library( - name = "vhlo_op_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloOpInterfaces.cpp.inc": ["-gen-op-interface-defs"], - "stablehlo/dialect/VhloOpInterfaces.h.inc": ["-gen-op-interface-decls"], -@@ -1710,7 +1777,8 @@ - "stablehlo/dialect/VhloBytecode.h", - "stablehlo/dialect/VhloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1734,7 +1802,7 @@ - - gentbl_cc_library( - name = "vhlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/VhloOps.h.inc": ["-gen-op-decls"], -@@ -1756,6 +1824,7 @@ - "stablehlo/dialect/VhloOps.td", - "stablehlo/dialect/VhloTypes.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:BuiltinDialectTdFiles", -@@ -1772,7 +1841,8 @@ - hdrs = [ - "stablehlo/dialect/VhloTypes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_assembly_format", - ":version", -@@ -1790,7 +1860,7 @@ - - gentbl_cc_library( - name = "vhlo_type_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloTypeInterfaces.cpp.inc": ["-gen-type-interface-defs"], - "stablehlo/dialect/VhloTypeInterfaces.h.inc": ["-gen-type-interface-decls"], -@@ -1804,7 +1874,7 @@ - - gentbl_cc_library( - name = "vhlo_types_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloTypeDefs.cpp.inc": ["-gen-typedef-defs"], - "stablehlo/dialect/VhloTypeDefs.h.inc": ["-gen-typedef-decls"], -@@ -1826,227 +1896,21 @@ - ], - ) - --### Builder APIs -- --##### --# TableGen --##### -- --cc_binary( -- name = "mlir_builder_tblgen", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilderTblgen.cpp"], -- deps = [ -- "@llvm-project//llvm:Support", -- "@llvm-project//llvm:TableGen", -- "@llvm-project//mlir:Support", -- "@llvm-project//mlir:TableGen", -- ], --) -- --#### --## MlirBuilder base class --#### -- --cc_library( -- name = "mlir_builder", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/MlirBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":attr_type_builder_util", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "mlir_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":chlo_builder", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//mlir:TosaDialect", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --#### --## Attr / Type Builder Helpers --#### -- --cc_library( -- name = "attr_type_builder_util", -- srcs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.h"], -- strip_include_prefix = ".", -- deps = [ -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "attr_type_builder_util_test", -- srcs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtilTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --##### --## Dialect-specific builders --#### -- --gentbl_cc_library( -- name = "chlo_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/ChloBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/ChloBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/ChloBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = ":stablehlo/dialect/ChloOps.td", -- deps = [ -- ":chlo_ops_td_files", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "chlo_builder", -- srcs = ["stablehlo/integrations/cpp/builder/ChloBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/ChloBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":chlo_builder_inc", -- ":chlo_ops", -- ":mlir_builder", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "chlo_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/ChloBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":chlo_builder", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --gentbl_cc_library( -- name = "func_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/FuncBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/FuncBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/FuncBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = "@llvm-project//mlir:include/mlir/Dialect/Func/IR/FuncOps.td", -- deps = [ -- "@llvm-project//mlir:FuncTdFiles", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "func_builder", -- srcs = ["stablehlo/integrations/cpp/builder/FuncBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/FuncBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":func_builder_inc", -- ":mlir_builder", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --gentbl_cc_library( -- name = "stablehlo_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/StablehloBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/StablehloBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/StablehloBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = ":stablehlo/dialect/StablehloOps.td", -- deps = [ -- ":base_td_files", -- ":stablehlo_ops_td_filegroup", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:ShapeOpsTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "stablehlo_builder", -- srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/StablehloBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":attr_type_builder_util", -- ":mlir_builder", -- ":stablehlo_broadcast_lowering", -- ":stablehlo_builder_inc", -- ":stablehlo_ops", -- ":stablehlo_type_inference", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:InferTypeOpInterface", -- "@llvm-project//mlir:Support", -- ], -+test_suite( -+ name = "all_tests", -+ tests = [":stablehlo_ci_tests"], - ) - --cc_test( -- name = "stablehlo_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) +# copybara:uncomment_begin +# alias( +# name = "check_ops", @@ -1142,744 +36,34 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel +# deprecation = "Use //stablehlo/tests:test_utils instead.", +# ) +# copybara:uncomment_end -diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.bazel ---- stablehlo/build_tools/BUILD.bazel -+++ stablehlo/build_tools/BUILD.bazel -@@ -0,0 +1,38 @@ -+# Copyright 2025 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+exports_files( -+ [ -+ "lit_test_suite.bzl", -+ "utils.bzl", -+ ], -+ visibility = ["//:__subpackages__"], -+) diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/lit_test_suite.bzl --- stablehlo/build_tools/lit_test_suite.bzl +++ stablehlo/build_tools/lit_test_suite.bzl -@@ -0,0 +1,71 @@ -+# Copyright 2026 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+"""Defines the `lit_test_suite` macro.""" -+ -+load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+ -+def lit_test_suite( -+ name, -+ srcs, -+ size_per_test = "small", -+ data = None, -+ tags = None, -+ args = None): -+ """Defines a `lit_test` for each file in `srcs`; wraps them in a `test_suite`. -+ -+ Args: -+ name: The name of the test suite. -+ srcs: The list/glob of `.mlir` LIT test files included in the test suite. -+ size_per_test: The size of the individual tests comprising the suite. -+ data: List of labels for the test's data dependencies (i.e. runtime deps). -+ tags: List of tags to attach to the tests. -+ args: List of extra arguments to pass to the test driver. -+ """ -+ -+ if data == None: -+ data = [] -+ if tags == None: -+ tags = [] -+ if args == None: -+ args = [] -+ -+ tests_in_suite = [] -+ -+ for src in srcs: +@@ -44,11 +44,13 @@ + tests_in_suite = [] + + for src in srcs: +- test_name = ("%s.test" % src).replace("/", "_") + test_name = ("%s.test" % src).replace("/", "_") # copybara:comment + + # copybara:uncomment test_name = "%s.test" % src -+ lit_test( -+ name = test_name, # copybara:comment -+ # copybara:uncomment name = src, + lit_test( + name = test_name, # copybara:comment + # copybara:uncomment name = src, +- srcs = [src], + srcs = [src], # copybara:comment -+ # copybara:uncomment src = src, -+ args = args, -+ data = data + native.glob( -+ ["%s.bc" % src], -+ # Allow empty since not all `.mlir` tests have `.mlir.bc` files. -+ allow_empty = True, -+ ), -+ deps = ["@rules_python//python/runfiles"], # copybara:comment + # copybara:uncomment src = src, + args = args, + data = data + native.glob( +@@ -57,7 +59,7 @@ + allow_empty = True, + ), + deps = ["@rules_python//python/runfiles"], # copybara:comment +- # copybara:uncomment driver = "//third_party/llvm/llvm-project/mlir:run_lit.sh", + # copybara:uncomment driver = "@llvm-project//mlir:run_lit.sh", -+ tags = tags, -+ size = size_per_test, -+ ) -+ tests_in_suite.append(test_name) -+ -+ native.test_suite( -+ name = name, -+ tests = tests_in_suite, -+ ) -diff --ruN a/stablehlo/build_tools/utils.bzl b/stablehlo/build_tools/utils.bzl ---- stablehlo/build_tools/utils.bzl -+++ stablehlo/build_tools/utils.bzl -@@ -24,3 +24,6 @@ - def workspace_name(): - """Return the name of the workspace.""" - return "_main" if is_bzlmod_enabled() else "stablehlo" -+ -+def get_compatible_with_portable(): -+ return [] -diff --ruN a/stablehlo/examples/CMakeLists.txt b/stablehlo/examples/CMakeLists.txt ---- stablehlo/examples/CMakeLists.txt -+++ stablehlo/examples/CMakeLists.txt -@@ -12,7 +12,7 @@ - # See the License for the specific language governing permissions and - # limitations under the License. - --add_subdirectory(c++) -+add_subdirectory(cpp) - - if(STABLEHLO_ENABLE_BINDINGS_PYTHON) - add_subdirectory(python) -diff --ruN a/stablehlo/examples/c++/BUILD.bazel b/stablehlo/examples/c++/BUILD.bazel ---- stablehlo/examples/c++/BUILD.bazel -+++ stablehlo/examples/c++/BUILD.bazel -@@ -1,53 +0,0 @@ --# Copyright 2023 The StableHLO Authors. All Rights Reserved. --# --# Licensed under the Apache License, Version 2.0 (the "License"); --# you may not use this file except in compliance with the License. --# You may obtain a copy of the License at --# --# https://www.apache.org/licenses/LICENSE-2.0 --# --# Unless required by applicable law or agreed to in writing, software --# distributed under the License is distributed on an "AS IS" BASIS, --# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. --# See the License for the specific language governing permissions and --# limitations under the License. --load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_test") -- --package( -- default_visibility = ["//visibility:public"], -- licenses = ["notice"], --) -- --cc_binary( -- name = "example_add", -- srcs = [ -- "ExampleAdd.cpp", -- ], -- deps = [ -- "//:reference_api", -- "//:reference_configuration", -- "//:stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:QuantOps", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "example_add_test", -- srcs = [ -- "ExampleAdd.cpp", -- ], -- deps = [ -- "//:reference_api", -- "//:reference_configuration", -- "//:stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:QuantOps", -- "@llvm-project//mlir:Support", -- ], --) -diff --ruN a/stablehlo/examples/c++/CMakeLists.txt b/stablehlo/examples/c++/CMakeLists.txt ---- stablehlo/examples/c++/CMakeLists.txt -+++ stablehlo/examples/c++/CMakeLists.txt -@@ -1,29 +0,0 @@ --# Copyright 2024 The StableHLO 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 --# --# https://www.apache.org/licenses/LICENSE-2.0 --# --# Unless required by applicable law or agreed to in writing, software --# distributed under the License is distributed on an "AS IS" BASIS, --# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. --# See the License for the specific language governing permissions and --# limitations under the License. -- --# TODO(fzakaria): As we think about adding more tests we should consider --# making this a function to be DRY. --add_executable(example-add ExampleAdd.cpp) --llvm_update_compile_flags(example-add) --target_link_libraries(example-add PRIVATE StablehloOps StablehloReferenceApi) -- --mlir_check_all_link_libraries(example-add) --add_custom_target(example-add-test -- COMMAND example-add -- WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -- COMMENT "Executing example-add to validate it works." -- DEPENDS -- example-add --) --add_dependencies(check-stablehlo-quick example-add-test) -diff --ruN a/stablehlo/examples/c++/ExampleAdd.cpp b/stablehlo/examples/c++/ExampleAdd.cpp ---- stablehlo/examples/c++/ExampleAdd.cpp -+++ stablehlo/examples/c++/ExampleAdd.cpp -@@ -1,93 +0,0 @@ --/* Copyright 2022 The StableHLO 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 "llvm/ADT/ArrayRef.h" --#include "llvm/ADT/SmallVector.h" --#include "llvm/Support/LogicalResult.h" --#include "mlir/Dialect/Func/IR/FuncOps.h" --#include "mlir/Dialect/Quant/IR/Quant.h" --#include "mlir/IR/Attributes.h" --#include "mlir/IR/Block.h" --#include "mlir/IR/Builders.h" --#include "mlir/IR/BuiltinAttributes.h" --#include "mlir/IR/BuiltinOps.h" --#include "mlir/IR/BuiltinTypes.h" --#include "mlir/IR/Location.h" --#include "mlir/IR/MLIRContext.h" --#include "mlir/IR/OwningOpRef.h" --#include "mlir/IR/Value.h" --#include "mlir/IR/Verifier.h" --#include "mlir/Support/LLVM.h" --#include "stablehlo/dialect/StablehloOps.h" --#include "stablehlo/reference/Api.h" --#include "stablehlo/reference/Configuration.h" -- --int main() { -- mlir::MLIRContext context; -- -- /** create module **/ -- mlir::OwningOpRef module = -- mlir::ModuleOp::create(mlir::UnknownLoc::get(&context)); -- module->getContext()->loadDialect(); -- module->getContext()->loadDialect(); -- module->getContext()->loadDialect(); -- module->setName("test_module"); -- -- /** create function **/ -- // create function argument and result types. -- auto tensorType = -- mlir::RankedTensorType::get({3, 4}, mlir::Float32Type::get(&context)); -- auto func_type = -- mlir::FunctionType::get(&context, {tensorType, tensorType}, {tensorType}); -- -- // create the function and map arguments. -- llvm::ArrayRef attrs; -- auto function = mlir::func::FuncOp::create(mlir::UnknownLoc::get(&context), -- "main", func_type, attrs); -- function.setVisibility(mlir::func::FuncOp::Visibility::Public); -- module->push_back(function); -- -- // create function block with add operations. -- mlir::Block* block = function.addEntryBlock(); -- llvm::SmallVector arguments(block->args_begin(), -- block->args_end()); -- mlir::OpBuilder block_builder = mlir::OpBuilder::atBlockEnd(block); -- mlir::Location loc = block_builder.getUnknownLoc(); -- -- llvm::SmallVector attributes; -- mlir::Operation* op = -- mlir::stablehlo::AddOp::create(block_builder, loc, arguments, attributes) -- .getOperation(); -- mlir::func::ReturnOp::create(block_builder, loc, op->getResult(0)); -- -- /** verify and dump the module **/ -- assert(mlir::succeeded(mlir::verify(module.get()))); -- -- /* interpret the function "main" with concrete inputs **/ -- auto getConstValue = [&](double val) { -- return mlir::DenseElementsAttr::get( -- tensorType, -- block_builder.getFloatAttr(tensorType.getElementType(), val)); -- }; -- -- auto inputValue1 = getConstValue(10.0); -- auto inputValue2 = getConstValue(20.0); -- auto expectedValue = getConstValue(30.0); -- -- mlir::stablehlo::InterpreterConfiguration config; -- auto results = evalModule(*module, {inputValue1, inputValue2}, config); -- return failed(results) || (*results)[0] != expectedValue; --} -diff --ruN a/stablehlo/examples/c++/README.md b/stablehlo/examples/c++/README.md ---- stablehlo/examples/c++/README.md -+++ stablehlo/examples/c++/README.md -@@ -1,45 +0,0 @@ --# Examples -- --This directory should contain a series of examples as a starting point --for how to use StableHLO. -- --Note: If you have a great example to highlight, we welcome contributions! -- --* [example-add](./ExampleAdd.cpp): A simple example that demonstrates how to -- * Use the StableHLO library to add two numbers. -- * Interpret a StableHLO program using concrete inputs. -- --```c++ -- // Assume 'module' is an MLIR module with function "main" containing StableHLO -- // operations. -- llvm::outs() << "Program:\n " << module << "\n"; -- -- // Create concrete inputs to be used for interpreting "main". -- auto inputValue1 = mlir::DenseElementsAttr::get( -- tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -- static_cast(10))); -- auto inputValue2 = mlir::DenseElementsAttr::get( -- tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -- static_cast(20))); -- llvm::outs() << "Inputs: " << inputValue1 << ", " << inputValue2 << "\n"; -- -- -- mlir::stablehlo::InterpreterConfiguration config; -- auto results = evalModule(module, {inputValue1, inputValue2}, config); -- llvm::outs() << "Output: " << (*results)[0]; --``` -- --Output: -- --```mlir --Program: --module @test_module { -- func.func @main(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<3x4xf32> { -- %0 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -- %1 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -- return %1 : tensor<3x4xf32> -- } --} --Inputs: dense<1.000000e+01> : tensor<3x4xf32>, dense<2.000000e+01> : tensor<3x4xf32> --Output: dense<3.000000e+01> : tensor<3x4xf32> --``` -diff --ruN a/stablehlo/examples/cpp/BUILD.bazel b/stablehlo/examples/cpp/BUILD.bazel ---- stablehlo/examples/cpp/BUILD.bazel -+++ stablehlo/examples/cpp/BUILD.bazel -@@ -0,0 +1,66 @@ -+# Copyright 2023 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_test") -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+cc_binary( -+ name = "example_add", -+ srcs = [ -+ "ExampleAdd.cpp", -+ ], -+ deps = [ -+ "//:reference_api", -+ "//:reference_configuration", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "example_add_test", -+ srcs = [ -+ "ExampleAdd.cpp", -+ ], -+ deps = [ -+ "//:reference_api", -+ "//:reference_configuration", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", -+ "@llvm-project//mlir:Support", -+ ], -+) -diff --ruN a/stablehlo/examples/cpp/CMakeLists.txt b/stablehlo/examples/cpp/CMakeLists.txt ---- stablehlo/examples/cpp/CMakeLists.txt -+++ stablehlo/examples/cpp/CMakeLists.txt -@@ -0,0 +1,29 @@ -+# Copyright 2024 The StableHLO 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 -+# -+# https://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+# TODO(fzakaria): As we think about adding more tests we should consider -+# making this a function to be DRY. -+add_executable(example-add ExampleAdd.cpp) -+llvm_update_compile_flags(example-add) -+target_link_libraries(example-add PRIVATE StablehloOps StablehloReferenceApi) -+ -+mlir_check_all_link_libraries(example-add) -+add_custom_target(example-add-test -+ COMMAND example-add -+ WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -+ COMMENT "Executing example-add to validate it works." -+ DEPENDS -+ example-add -+) -+add_dependencies(check-stablehlo-quick example-add-test) -diff --ruN a/stablehlo/examples/cpp/ExampleAdd.cpp b/stablehlo/examples/cpp/ExampleAdd.cpp ---- stablehlo/examples/cpp/ExampleAdd.cpp -+++ stablehlo/examples/cpp/ExampleAdd.cpp -@@ -0,0 +1,93 @@ -+/* Copyright 2022 The StableHLO 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 "llvm/ADT/ArrayRef.h" -+#include "llvm/ADT/SmallVector.h" -+#include "llvm/Support/LogicalResult.h" -+#include "mlir/Dialect/Func/IR/FuncOps.h" -+#include "mlir/Dialect/Quant/IR/Quant.h" -+#include "mlir/IR/Attributes.h" -+#include "mlir/IR/Block.h" -+#include "mlir/IR/Builders.h" -+#include "mlir/IR/BuiltinAttributes.h" -+#include "mlir/IR/BuiltinOps.h" -+#include "mlir/IR/BuiltinTypes.h" -+#include "mlir/IR/Location.h" -+#include "mlir/IR/MLIRContext.h" -+#include "mlir/IR/OwningOpRef.h" -+#include "mlir/IR/Value.h" -+#include "mlir/IR/Verifier.h" -+#include "mlir/Support/LLVM.h" -+#include "stablehlo/dialect/StablehloOps.h" -+#include "stablehlo/reference/Api.h" -+#include "stablehlo/reference/Configuration.h" -+ -+int main() { -+ mlir::MLIRContext context; -+ -+ /** create module **/ -+ mlir::OwningOpRef module = -+ mlir::ModuleOp::create(mlir::UnknownLoc::get(&context)); -+ module->getContext()->loadDialect(); -+ module->getContext()->loadDialect(); -+ module->getContext()->loadDialect(); -+ module->setName("test_module"); -+ -+ /** create function **/ -+ // create function argument and result types. -+ auto tensorType = -+ mlir::RankedTensorType::get({3, 4}, mlir::Float32Type::get(&context)); -+ auto func_type = -+ mlir::FunctionType::get(&context, {tensorType, tensorType}, {tensorType}); -+ -+ // create the function and map arguments. -+ llvm::ArrayRef attrs; -+ auto function = mlir::func::FuncOp::create(mlir::UnknownLoc::get(&context), -+ "main", func_type, attrs); -+ function.setVisibility(mlir::func::FuncOp::Visibility::Public); -+ module->push_back(function); -+ -+ // create function block with add operations. -+ mlir::Block* block = function.addEntryBlock(); -+ llvm::SmallVector arguments(block->args_begin(), -+ block->args_end()); -+ mlir::OpBuilder block_builder = mlir::OpBuilder::atBlockEnd(block); -+ mlir::Location loc = block_builder.getUnknownLoc(); -+ -+ llvm::SmallVector attributes; -+ mlir::Operation* op = -+ mlir::stablehlo::AddOp::create(block_builder, loc, arguments, attributes) -+ .getOperation(); -+ mlir::func::ReturnOp::create(block_builder, loc, op->getResult(0)); -+ -+ /** verify and dump the module **/ -+ assert(mlir::succeeded(mlir::verify(module.get()))); -+ -+ /* interpret the function "main" with concrete inputs **/ -+ auto getConstValue = [&](double val) { -+ return mlir::DenseElementsAttr::get( -+ tensorType, -+ block_builder.getFloatAttr(tensorType.getElementType(), val)); -+ }; -+ -+ auto inputValue1 = getConstValue(10.0); -+ auto inputValue2 = getConstValue(20.0); -+ auto expectedValue = getConstValue(30.0); -+ -+ mlir::stablehlo::InterpreterConfiguration config; -+ auto results = evalModule(*module, {inputValue1, inputValue2}, config); -+ return failed(results) || (*results)[0] != expectedValue; -+} -diff --ruN a/stablehlo/examples/cpp/README.md b/stablehlo/examples/cpp/README.md ---- stablehlo/examples/cpp/README.md -+++ stablehlo/examples/cpp/README.md -@@ -0,0 +1,45 @@ -+# Examples -+ -+This directory should contain a series of examples as a starting point -+for how to use StableHLO. -+ -+Note: If you have a great example to highlight, we welcome contributions! -+ -+* [example-add](./ExampleAdd.cpp): A simple example that demonstrates how to -+ * Use the StableHLO library to add two numbers. -+ * Interpret a StableHLO program using concrete inputs. -+ -+```c++ -+ // Assume 'module' is an MLIR module with function "main" containing StableHLO -+ // operations. -+ llvm::outs() << "Program:\n " << module << "\n"; -+ -+ // Create concrete inputs to be used for interpreting "main". -+ auto inputValue1 = mlir::DenseElementsAttr::get( -+ tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -+ static_cast(10))); -+ auto inputValue2 = mlir::DenseElementsAttr::get( -+ tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -+ static_cast(20))); -+ llvm::outs() << "Inputs: " << inputValue1 << ", " << inputValue2 << "\n"; -+ -+ -+ mlir::stablehlo::InterpreterConfiguration config; -+ auto results = evalModule(module, {inputValue1, inputValue2}, config); -+ llvm::outs() << "Output: " << (*results)[0]; -+``` -+ -+Output: -+ -+```mlir -+Program: -+module @test_module { -+ func.func @main(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<3x4xf32> { -+ %0 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -+ %1 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -+ return %1 : tensor<3x4xf32> -+ } -+} -+Inputs: dense<1.000000e+01> : tensor<3x4xf32>, dense<2.000000e+01> : tensor<3x4xf32> -+Output: dense<3.000000e+01> : tensor<3x4xf32> -+``` -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel ---- stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel -+++ stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,25 +52,15 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_linalg_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_linalg_tests", -+ srcs = glob(["**/*.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_linalg_tests"], - ) -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir b/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir ---- stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir -+++ stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir -@@ -52,8 +52,8 @@ - // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM0]], %[[DIM2]]) - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]]{{.*}}outs(%[[INIT]] - // CHECK: linalg.conv_1d_nwc_wcf --// CHECK-SAME: {dilations = dense<1> : tensor<1xi64> --// CHECK-SAME: someattr -+// CHECK-SAME: {someattr -+// CHECK-SAME: dilations = dense<1> : tensor<1xi64> - // CHECK-SAME: strides = dense<1> : tensor<1xi64>} - // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor, tensor<2x?x?xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor) -> tensor -@@ -360,7 +360,7 @@ - // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2x3xf32> - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> - // CHECK: %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm --// CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, someattr, strides = dense<1> : tensor<2xi64>} -+// CHECK-SAME: {someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} - // CHECK-SAME: ins(%[[IN]], %[[EXPAND]] : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> - // CHECK: %{{.+}} = tensor.collapse_shape %[[OUT]] -@@ -409,7 +409,7 @@ - // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x6x2x2xf32> - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]] : f32) outs(%[[INIT]] : tensor<2x3x6x2x2xf32>) -> tensor<2x3x6x2x2xf32> - // CHECK: %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm --// CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, someattr, strides = dense<1> : tensor<2xi64>} -+// CHECK-SAME: {someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} - // CHECK-SAME: ins(%[[PAD]], %[[EXPAND]] : tensor<2x4x7x2xf32>, tensor<2x2x2x2xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<2x3x6x2x2xf32>) -> tensor<2x3x6x2x2xf32> - // CHECK: %{{.+}} = tensor.collapse_shape %[[OUT]] -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir b/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir ---- stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir -+++ stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir -@@ -394,8 +394,8 @@ - // CHECK: %[[INIT_VAL:.+]] = tensor.extract %[[ARG1]][] : tensor - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[INIT_VAL]] : f32) outs(%[[INIT]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32> - // CHECK: %[[RES:.+]] = linalg.pooling_nhwc_min --// CHECK-SAME: {dilations = dense<1> : vector<2xi64> --// CHECK-SAME: someattr, -+// CHECK-SAME: {someattr, -+// CHECK-SAME: dilations = dense<1> : vector<2xi64> - // CHECK-SAME: strides = dense<2> : vector<2xi64>} - // CHECK-SAME: ins(%[[ARG0]], %[[WINDOW]] : tensor<1x17x17x64xf32>, tensor<3x3xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32> + tags = tags, + size = size_per_test, + ) diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h --- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h +++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h @@ -1939,76 +123,6 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU /// Extracts integer values from the attribute |elements|. SmallVector extract1DVector(DenseIntElementsAttr elements); -diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel ---- stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel -+++ stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,26 +52,16 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_tosa_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_tosa_tests", -+ srcs = glob(["**/*.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_tosa_tests"], - ) diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp --- stablehlo/stablehlo/dialect/ChloBytecode.cpp +++ stablehlo/stablehlo/dialect/ChloBytecode.cpp @@ -2175,7 +289,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial //===----------------------------------------------------------------------===// // Ragged dot op definitions. -@@ -90,9 +94,12 @@ +@@ -90,7 +94,9 @@ let cppNamespace = "::mlir::chlo"; } @@ -2186,9 +300,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial def CHLO_PrecisionConfigAttr: TypedArrayAttrBase; - -+ - #endif // STABLEHLO_DIALECT_CHLO_ENUMS diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/dialect/ChloOps.cpp --- stablehlo/stablehlo/dialect/ChloOps.cpp +++ stablehlo/stablehlo/dialect/ChloOps.cpp @@ -2310,301 +421,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/Register.h b/stablehlo/stablehlo/dialec } // namespace stablehlo } // namespace mlir -diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ---- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -+++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -@@ -20,6 +20,7 @@ - #include - - #include "llvm/ADT/DenseSet.h" -+#include "llvm/ADT/STLExtras.h" - #include "llvm/ADT/SmallVector.h" - #include "llvm/ADT/StringRef.h" - #include "mlir/IR/Attributes.h" -@@ -31,57 +32,53 @@ - namespace mlir { - namespace stablehlo { - --static SmallVector> -+namespace { -+ -+struct ReindexedAxes { -+ SmallVector splitAxisSizes; -+ SmallVector groupedAxisIndices; -+}; -+ -+// Generates replica groups from the reshaped mesh axis sizes and the indices of -+// the communication axes using a Reshape-Transpose permutation. -+SmallVector> - flattenedReplicaGroupsFromTransposePermutation( -- const SmallVector& meshAxisNames, -- const SmallVector& commAxisNames, -- const llvm::DenseSet& commAxisSet, -- const SmallVector& axisSizes, -- const SmallVector& deviceIds, int64_t totalDevices) { -- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. -+ ArrayRef axisSizes, ArrayRef groupedAxisIndices, -+ ArrayRef deviceIds, int64_t totalDevices) { -+ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), -+ groupedAxisIndices.end()); - SmallVector transposeAxes; - // Non-grouped axes first -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (!commAxisSet.count(meshAxisNames[i])) { -+ for (size_t i = 0; i < axisSizes.size(); ++i) { -+ if (!groupedAxisSet.count(i)) { - transposeAxes.push_back(i); - } - } -- // Grouped axes -- for (const auto& name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- transposeAxes.push_back(i); -- break; -- } -- } -+ // Grouped axes in the specified order -+ for (int64_t idx : groupedAxisIndices) { -+ transposeAxes.push_back(idx); - } - -- SmallVector transposedSizes(meshAxisNames.size()); -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -+ SmallVector transposedSizes(axisSizes.size()); -+ for (size_t i = 0; i < axisSizes.size(); ++i) { - transposedSizes[i] = axisSizes[transposeAxes[i]]; - } - -- // Compute strides for original shape -- SmallVector originalStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ // Compute strides for reshaped shape -+ SmallVector originalStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; - } - - // Compute strides for transposed shape -- SmallVector transposedStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ SmallVector transposedStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; - } - -- // Generate chunks - int64_t numDevicesPerGroup = 1; -- for (auto name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- numDevicesPerGroup *= axisSizes[i]; -- break; -- } -- } -+ for (int64_t idx : groupedAxisIndices) { -+ numDevicesPerGroup *= axisSizes[idx]; - } - int64_t numGroups = totalDevices / numDevicesPerGroup; - -@@ -93,7 +90,7 @@ - for (int64_t j = 0; j < numDevicesPerGroup; ++j) { - int64_t linearTransposeIdx = i * numDevicesPerGroup + j; - int64_t originalIndex = 0; -- for (size_t k = 0; k < meshAxisNames.size(); ++k) { -+ for (size_t k = 0; k < axisSizes.size(); ++k) { - int64_t coord = - (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; - originalIndex += coord * originalStrides[transposeAxes[k]]; -@@ -102,10 +99,136 @@ - } - groups.push_back(std::move(group)); - } -- - return groups; - } - -+// Splits mesh axes based on sub-axis references and computes the corresponding -+// indices for the communication axes. -+FailureOr computeReindexedAxes(ArrayRef axesInMesh, -+ ArrayAttr commAxes, -+ Location loc) { -+ ReindexedAxes result; -+ -+ // Validate commAxes and verify that all mesh axes exist and have valid sizes. -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::dyn_cast(attr); -+ if (!shloAxisRef) { -+ return emitError(loc) << "expected AxisRefAttr in comm_axes"; -+ } -+ StringRef axisName = shloAxisRef.getName(); -+ bool found = false; -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ found = true; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ int64_t preSize = subAxisInfo.getPreSize(); -+ int64_t size = subAxisInfo.getSize(); -+ if (preSize < 1 || size < 1) { -+ return emitError(loc) -+ << "sub-axis pre_size and size must be at least 1"; -+ } -+ int64_t nextPreSize = preSize * size; -+ if (nextPreSize > meshAxis.getSize() || -+ meshAxis.getSize() % nextPreSize != 0) { -+ return emitError(loc) -+ << "sub-axis (pre_size * size) must divide mesh axis size"; -+ } -+ } -+ break; -+ } -+ } -+ if (!found) { -+ return emitError(loc) -+ << "axis '" << axisName << "' not found in mesh definition"; -+ } -+ } -+ -+ // Split each mesh axis according to the referenced subaxes. -+ struct SplitDim { -+ StringRef axisName; -+ int64_t preSize; -+ int64_t size; -+ int64_t dimIndex; -+ }; -+ SmallVector splitDims; -+ -+ for (auto meshAxis : axesInMesh) { -+ StringRef axisName = meshAxis.getName(); -+ int64_t axisSize = meshAxis.getSize(); -+ -+ if (axisSize == 1) { -+ int64_t dimIdx = result.splitAxisSizes.size(); -+ result.splitAxisSizes.push_back(1); -+ splitDims.push_back({axisName, /*preSize=*/1, /*size=*/1, dimIdx}); -+ continue; -+ } -+ -+ SmallVector preSizes = {1, axisSize}; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ if (shloAxisRef.getName() == axisName) { -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ preSizes.push_back(subAxisInfo.getPreSize()); -+ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); -+ } -+ } -+ } -+ -+ llvm::sort(preSizes); -+ preSizes.erase(llvm::unique(preSizes), preSizes.end()); -+ -+ for (size_t j = 0; j < preSizes.size() - 1; ++j) { -+ int64_t segPreSize = preSizes[j]; -+ int64_t segSize = preSizes[j + 1] / segPreSize; -+ int64_t dimIdx = result.splitAxisSizes.size(); -+ result.splitAxisSizes.push_back(segSize); -+ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); -+ } -+ } -+ -+ // Map each communication axis to its corresponding split dimension. -+ llvm::DenseSet groupedSet; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ StringRef axisName = shloAxisRef.getName(); -+ int64_t reqPreSize = 1; -+ int64_t reqSize = 0; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ reqPreSize = subAxisInfo.getPreSize(); -+ reqSize = subAxisInfo.getSize(); -+ } else { -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ reqSize = meshAxis.getSize(); -+ break; -+ } -+ } -+ } -+ -+ bool matched = false; -+ for (const auto& splitDim : splitDims) { -+ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && -+ splitDim.size == reqSize) { -+ if (!groupedSet.insert(splitDim.dimIndex).second) { -+ return emitError(loc) -+ << "Duplicate or overlapping communication axis: " << axisName; -+ } -+ result.groupedAxisIndices.push_back(splitDim.dimIndex); -+ matched = true; -+ break; -+ } -+ } -+ if (!matched) { -+ return emitError(loc) << "Invalid or overlapping communication axis on '" -+ << axisName << "'"; -+ } -+ } -+ -+ return result; -+} -+ -+} // namespace -+ - FailureOr>> flattenReplicaGroupMeshAxes( - Attribute meshAttr, ArrayAttr commAxes, Location loc) { - // Reference implementation for flattening a mesh-axes based replica group -@@ -120,34 +243,13 @@ - if (!mesh) - return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); - -- auto axesInMesh = mesh.getAxes(); -- -- // Identify which axes are communication axes. -- llvm::SmallVector commAxisNames; -- llvm::DenseSet commAxisSet; -- for (auto attr : commAxes) { -- auto shloAxisRef = llvm::dyn_cast(attr); -- if (!shloAxisRef) { -- return emitError(loc) << "expected AxisRefAttr in comm_axes"; -- } -- if (shloAxisRef.getSubAxisInfo()) { -- return emitError(loc) << "Subaxes are not supported in " -- "flattenReplicaGroupMeshAxes"; -- } -- commAxisNames.push_back(shloAxisRef.getName()); -- commAxisSet.insert(shloAxisRef.getName()); -- } -- -- // Calculate total devices and axis sizes -+ FailureOr reindexedAxes = -+ computeReindexedAxes(mesh.getAxes(), commAxes, loc); -+ if (failed(reindexedAxes)) return failure(); - - int64_t totalDevices = 1; -- SmallVector axisSizes; -- SmallVector meshAxisNames; -- for (auto meshAxis : axesInMesh) { -- auto typedMeshAxis = llvm::cast(meshAxis); -- axisSizes.push_back(typedMeshAxis.getSize()); -- meshAxisNames.push_back(typedMeshAxis.getName()); -- totalDevices *= typedMeshAxis.getSize(); -+ for (auto meshAxis : mesh.getAxes()) { -+ totalDevices *= llvm::cast(meshAxis).getSize(); - } - - SmallVector deviceIds; -@@ -160,8 +262,8 @@ - } - - return flattenedReplicaGroupsFromTransposePermutation( -- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, -- totalDevices); -+ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, -+ deviceIds, totalDevices); - } - - } // namespace stablehlo diff --ruN a/stablehlo/stablehlo/dialect/StablehloBytecode.h b/stablehlo/stablehlo/dialect/StablehloBytecode.h --- stablehlo/stablehlo/dialect/StablehloBytecode.h +++ stablehlo/stablehlo/dialect/StablehloBytecode.h @@ -2686,7 +502,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo def STABLEHLO_RNG_ALGORITHM_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>; def STABLEHLO_RNG_ALGORITHM_THREE_FRY : I32EnumAttrCase<"THREE_FRY", 1>; -@@ -220,6 +232,9 @@ +@@ -220,6 +232,8 @@ let cppNamespace = "::mlir::stablehlo"; } @@ -2696,7 +512,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo +} #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS -+ diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h --- stablehlo/stablehlo/dialect/VhloBytecode.h +++ stablehlo/stablehlo/dialect/VhloBytecode.h @@ -2883,315 +698,6 @@ diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp b/stable resultsVec.push_back(wrap(result)); } return mlirArrayAttrGet(mlirModuleGetContext(module), resultsVec.size(), -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel b/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel ---- stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel -+++ stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel -@@ -0,0 +1,281 @@ -+# Copyright 2025 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") -+load("@rules_cc//cc:cc_binary.bzl", "cc_binary") -+load("@rules_cc//cc:cc_library.bzl", "cc_library") -+load("@rules_cc//cc:cc_test.bzl", "cc_test") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+filegroup( -+ name = "stablehlo_builder_filegroup", -+ srcs = glob(["**"]), -+) -+ -+##### -+# TableGen -+##### -+ -+cc_binary( -+ name = "mlir_builder_tblgen", -+ srcs = ["MlirBuilderTblgen.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ "@llvm-project//llvm:Support", -+ "@llvm-project//llvm:TableGen", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//mlir:TableGen", -+ ], -+) -+ -+#### -+## MlirBuilder base class -+#### -+ -+cc_library( -+ name = "mlir_builder", -+ srcs = ["MlirBuilder.cpp"], -+ hdrs = ["MlirBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+#### -+## Attr / Type Builder Helpers -+#### -+ -+cc_library( -+ name = "attr_type_builder_util", -+ srcs = ["AttrTypeBuilderUtil.cpp"], -+ hdrs = ["AttrTypeBuilderUtil.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "attr_type_builder_util_test", -+ srcs = ["AttrTypeBuilderUtilTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+##### -+## Dialect-specific builders -+#### -+ -+gentbl_cc_library( -+ name = "chlo_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "ChloBuilder.h.inc": ["-gen-builder-decls"], -+ "ChloBuilder.cpp.inc": ["-gen-builder-defs"], -+ "ChloBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "//:stablehlo/dialect/ChloOps.td", -+ deps = [ -+ "//:chlo_ops_td_files", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "chlo_builder", -+ srcs = ["ChloBuilder.cpp"], -+ hdrs = ["ChloBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":chlo_builder_inc", -+ ":mlir_builder", -+ "//:chlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+gentbl_cc_library( -+ name = "func_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "FuncBuilder.h.inc": ["-gen-builder-decls"], -+ "FuncBuilder.cpp.inc": ["-gen-builder-defs"], -+ "FuncBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "@llvm-project//mlir:include/mlir/Dialect/Func/IR/FuncOps.td", -+ deps = [ -+ "@llvm-project//mlir:FuncTdFiles", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "func_builder", -+ srcs = ["FuncBuilder.cpp"], -+ hdrs = ["FuncBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":func_builder_inc", -+ ":mlir_builder", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+gentbl_cc_library( -+ name = "stablehlo_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "StablehloBuilder.h.inc": ["-gen-builder-decls"], -+ "StablehloBuilder.cpp.inc": ["-gen-builder-defs"], -+ "StablehloBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "//:stablehlo/dialect/StablehloOps.td", -+ deps = [ -+ "//:base_td_files", -+ "//:stablehlo_ops_td_filegroup", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:ShapeOpsTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "stablehlo_builder", -+ srcs = ["StablehloBuilder.cpp"], -+ hdrs = ["StablehloBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":mlir_builder", -+ ":stablehlo_builder_inc", -+ "//:stablehlo_broadcast_lowering", -+ "//:stablehlo_ops", -+ "//:stablehlo_type_inference", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:InferTypeOpInterface", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "stablehlo_builder_test", -+ srcs = ["StablehloBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_ops", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+cc_test( -+ name = "chlo_builder_test", -+ srcs = ["ChloBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":chlo_builder", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_broadcast_lowering", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+cc_test( -+ name = "mlir_builder_test", -+ srcs = ["MlirBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":chlo_builder", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//mlir:TosaDialect", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp b/stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp ---- stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp -+++ stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp -@@ -13,7 +13,7 @@ - limitations under the License. - ==============================================================================*/ - --#include "ChloBuilder.h" -+#include "stablehlo/integrations/cpp/builder/ChloBuilder.h" - - #include - -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp b/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp ---- stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp -+++ stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp -@@ -13,7 +13,7 @@ - limitations under the License. - ==============================================================================*/ - --#include "MlirBuilder.h" -+#include "stablehlo/integrations/cpp/builder/MlirBuilder.h" - - #include "mlir/IR/Location.h" - #include "mlir/IR/Value.h" diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp --- stablehlo/stablehlo/integrations/python/StablehloApi.cpp +++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp @@ -3927,267 +1433,6 @@ diff --ruN a/stablehlo/stablehlo/reference/Value.cpp b/stablehlo/stablehlo/refer if (isTensor()) getTensor().print(os); else if (isToken()) -diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/testdata/BUILD.bazel ---- stablehlo/stablehlo/testdata/BUILD.bazel -+++ stablehlo/stablehlo/testdata/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,46 +52,42 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_testdata_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob( -+lit_test_suite( -+ name = "stablehlo_testdata_small_tests", -+ srcs = glob( - ["**/*.mlir"], - exclude = ["**/*.large.mlir"], -- ) --] -+ ), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", -+ tags = [ -+ "stablehlo_testdata_small_tests", -+ "stablehlo_testdata_tests", -+ ], -+) - --# Give large tests more RAM / timeout --[ -- lit_test( -- name = "%s.test" % src, -- size = "large", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_testdata_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.large.mlir"]) --] -+lit_test_suite( -+ name = "stablehlo_testdata_large_tests", -+ srcs = glob(["**/*.large.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "large", -+ tags = [ -+ "stablehlo_testdata_large_tests", -+ "stablehlo_testdata_tests", -+ ], -+) - - test_suite( - name = "stablehlo_testdata_tests", -diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/BUILD.bazel ---- stablehlo/stablehlo/tests/BUILD.bazel -+++ stablehlo/stablehlo/tests/BUILD.bazel -@@ -11,16 +11,30 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") - load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") - load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - cc_library( - name = "check_ops", - srcs = [ -@@ -29,6 +43,7 @@ - hdrs = [ - "CheckOps.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":check_ops_inc_gen", -@@ -67,6 +82,7 @@ - - gentbl_cc_library( - name = "check_ops_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "CheckOps.cpp.inc": ["-gen-op-defs"], -@@ -84,6 +100,7 @@ - srcs = [ - "CheckOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "//:base_td_files", -@@ -98,6 +115,7 @@ - hdrs = [ - "TestUtils.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":test_utils_inc_gen", -@@ -120,11 +138,14 @@ - - gentbl_cc_library( - name = "test_utils_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", -- tbl_outs = {"TestUtils.h.inc": [ -- "-gen-pass-decls", -- "-name=HloTest", -- ]}, -+ tbl_outs = { -+ "TestUtils.h.inc": [ -+ "-gen-pass-decls", -+ "-name=HloTest", -+ ], -+ }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "TestUtils.td", - deps = [ -@@ -137,6 +158,7 @@ - srcs = [ - "TestUtils.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:PassBaseTdFiles", -@@ -163,31 +185,27 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- "@llvm-project//llvm:not", -- # Allow empty since not all `.mlir` tests have `.mlir.bc` test files. -- ] + glob( -- ["%s.bc" % src], -- allow_empty = True, -- ), -- tags = ["stablehlo_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_tests", -+ srcs = glob([ -+ "*.mlir", -+ "**/*.mlir", -+ ]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ "@llvm-project//llvm:not", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_tests"], - ) -+ -+test_suite( -+ name = "all_tests", # alias for `:stablehlo_tests` -+ tests = [ -+ "//stablehlo/tests:stablehlo_tests", -+ ], -+) diff --ruN a/stablehlo/stablehlo/tests/CheckOps.cpp b/stablehlo/stablehlo/tests/CheckOps.cpp --- stablehlo/stablehlo/tests/CheckOps.cpp +++ stablehlo/stablehlo/tests/CheckOps.cpp @@ -4348,196 +1593,6 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st // CHECK-LABEL: func.func @mulhi_i32( // CHECK-SAME: %[[ARG0:.*]]: tensor<4xi32>, // CHECK-SAME: %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> { -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir ---- stablehlo/stablehlo/tests/interpret/all_gather.mlir -+++ stablehlo/stablehlo/tests/interpret/all_gather.mlir -@@ -133,3 +133,41 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { -+ %result = "stablehlo.all_gather"(%operand) { -+ all_gather_dim = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<1xi64>) -> tensor<2xi64> -+ return %result : tensor<2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], -+ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -+ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> -+ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir ---- stablehlo/stablehlo/tests/interpret/all_reduce.mlir -+++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir -@@ -135,3 +135,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.all_reduce"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, -+ channel_handle = #stablehlo.channel_handle -+ } : (tensor<1xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], -+ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir ---- stablehlo/stablehlo/tests/interpret/all_to_all.mlir -+++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir -@@ -172,3 +172,43 @@ - func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { -+ %result = "stablehlo.all_to_all"(%operand) { -+ split_dimension = 0 : i64, -+ concat_dimension = 1 : i64, -+ split_count = 2 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2x1xi64>) -> tensor<1x2xi64> -+ return %result : tensor<1x2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> -+ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> -+ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> -+ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> -+ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> -+ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> -+ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> -+ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], -+ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] -+ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, -+ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> -+ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, -+ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) -+ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ---- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -+++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -@@ -90,3 +90,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.reduce_scatter"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ scatter_dimension = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> -+ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> -+ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> -+ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> -+ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> -+ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> -+ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> -+ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], -+ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] -+ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> -+ func.return -+ } -+} diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir --- stablehlo/stablehlo/tests/ops_chlo.mlir +++ stablehlo/stablehlo/tests/ops_chlo.mlir @@ -4626,39 +1681,13 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir --- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir -@@ -1,11 +1,11 @@ +@@ -1,4 +1,4 @@ -// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" -allow-unregistered-dialect | FileCheck %s +// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" | FileCheck %s module { - - // CHECK-LABEL: @all_reduce_rgv3 - func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -19,7 +19,7 @@ - - // CHECK-LABEL: @all_gather_rgv3 - func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_gather"(%arg0) { - all_gather_dim = 0 : i64, - replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -@@ -30,7 +30,7 @@ - - // CHECK-LABEL: @all_to_all_rgv3 - func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_to_all"(%arg0) { - concat_dimension = 0 : i64, - split_dimension = 0 : i64, -@@ -40,11 +40,11 @@ +@@ -40,7 +40,7 @@ return %0 : tensor<4xf32> } @@ -4667,12 +1696,7 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -55,11 +55,11 @@ +@@ -55,7 +55,7 @@ return %0 : tensor<4xf32> } @@ -4681,11 +1705,6 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh_dev func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor @@ -70,7 +70,7 @@ return %0 : tensor<4xf32> } @@ -4695,63 +1714,6 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh_max func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> { -@@ -84,4 +84,56 @@ - } : (tensor<4xf32>) -> tensor<4xf32> - return %0 : tensor<4xf32> - } -+ -+ // CHECK-LABEL: @all_reduce_subaxis -+ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_1 -+ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_2 -+ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_size_one_axis -+ func.func @all_reduce_size_one_axis(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } - } diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir --- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir +++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl index c8508914e189ec..8764fbd0a3f908 100644 --- a/third_party/xla/third_party/stablehlo/workspace.bzl +++ b/third_party/xla/third_party/stablehlo/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): # LINT.IfChange - STABLEHLO_COMMIT = "639335932274930617d9724fb25ec9f2f52c5e40" - STABLEHLO_SHA256 = "dd590949798e26838d9ac3d6016def9b508caa60d0d7a69535ed2f5454ad30b8" + STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019" + STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd" # LINT.ThenChange(Google-internal path) tf_http_archive( From f58f013223b2cac78fa4f148890ea4fe7e14cc24 Mon Sep 17 00:00:00 2001 From: David Dunleavy Date: Wed, 23 Sep 2026 13:54:46 -0700 Subject: [PATCH 36/42] Remove `bazel analyze-profile` from CI This was removed so it's failing on JAX builds. Just remove it for everything because we will likely bump soon and it provides little value PiperOrigin-RevId: 986984963 --- third_party/xla/build_tools/ci/build.py | 1 - .../xla/build_tools/ci/golden_commands.txt | 21 ------------------- 2 files changed, 22 deletions(-) diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py index c4b7150b1a2981..20ea7038739ac5 100755 --- a/third_party/xla/build_tools/ci/build.py +++ b/third_party/xla/build_tools/ci/build.py @@ -246,7 +246,6 @@ def commands(self) -> List[List[str]]: ) ) cmds.append(self.bazel_command(subcommand=self.subcommand)) - cmds.append(["bazel", "analyze-profile", "profile.json.gz"]) return cmds diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt index e65e3f069635d5..bffa12b8d1b742 100644 --- a/third_party/xla/build_tools/ci/golden_commands.txt +++ b/third_party/xla/build_tools/ci/golden_commands.txt @@ -1,105 +1,86 @@ # BEGIN BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test bazel test --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS # BEGIN BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda13 --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --@local_config_cuda//cuda:override_include_cuda_libs --nobuild -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test bazel test --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda13 --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --@local_config_cuda//cuda:override_include_cuda_libs -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS # BEGIN BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS bazel --output_base=$GITHUB_WORKSPACE\bazel_output_base build --build_tag_filters= --test_tag_filters= --config=rbe_windows_amd64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE\openxla\xla --override_module=xla=$GITHUB_WORKSPACE\openxla\xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_runfile_links=False -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only=False --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only=False --//xla/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS # BEGIN BuildType.XLA_MACOS_ARM64_CPU_KOKORO df -h bazel --version mkdir -p /tmpfs/bazel_output bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/tmpfs/bazel_output --test_size_filters=small,medium --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_MACOS_ARM64_CPU_KOKORO # BEGIN BuildType.XLA_MACOS_X86_CPU_KOKORO sudo wget --no-verbose -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64 @@ -107,9 +88,7 @@ chmod +x /usr/local/bin/bazel bazel --version mkdir -p /Volumes/BuildData/bazel_output bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_MACOS_X86_CPU_KOKORO # BEGIN BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS bazel --output_user_root=C:/x build --build_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_xla_windows_x86_cpu_2022 --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -//xla/tpu/... -//xla/backends/cpu/collectives:gloo_collectives_test -//xla/backends/cpu/collectives:mpi_collectives -//xla/backends/cpu/collectives:mpi_communicator -//xla/python/transfer/... -//xla/backends/profiler/subprocess:subprocess_profiling_session -//xla/backends/profiler/subprocess:subprocess_profiling_session_test -//xla/backends/profiler/subprocess:subprocess_registry -//xla/backends/profiler/subprocess:subprocess_registry_test -//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk -//xla/backends/cpu/runtime/ynnpack:ynn_interop -//xla/backends/cpu/runtime/ynnpack:ynn_threadpool -//xla/backends/gpu/... -//xla/codegen/emitters/tests/... -//xla/service/gpu/... -//xla/codegen/intrinsic/cpp:eigen_unary_test -bazel analyze-profile profile.json.gz # END BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS From 4470f2f1dd91ac973e4270bc47150a97d9644ff6 Mon Sep 17 00:00:00 2001 From: Yue Sheng Date: Wed, 23 Sep 2026 14:25:22 -0700 Subject: [PATCH 37/42] [Mosaic TPU] Add some fold rules for the transpose op. PiperOrigin-RevId: 987004821 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 22 +++++++++++++++++++ .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 1 + 2 files changed, 23 insertions(+) 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 89a0171789b02c..9c1eae7cc792a3 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -955,6 +955,28 @@ 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 084deae550db80..086a825b8b811c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -1878,6 +1878,7 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> { } }]; let hasVerifier = 1; + let hasFolder = 1; } def TPU_LogOp : TPU_Op<"log"> { From bf1d7cbd630ed028c94b6da8b6f63a44d8b4b16a Mon Sep 17 00:00:00 2001 From: Michael Kuperstein Date: Wed, 23 Sep 2026 14:36:35 -0700 Subject: [PATCH 38/42] [XLA] Disambiguate `kFusion` instructions in `HloCSE` `CseKey` hash based on the fused computation Previously, `CseKey::AbslHashValue` in `HloCSE` fell into the `default:` branch for `HloOpcode::kFusion`, only hashing the top-level operands, output dimensions, and the fused computation's root opcode. As a result, distinct `kFusion` instructions sharing operands, output shape, and root opcode collided in `HloCSE`'s `representatives` hash table and triggered expensive `HloFusionInstruction::IdenticalSlowPath` (`HloComputation::Equal`) comparisons. This change factors out `CseKey::HashInstruction` and applies it to the internal instructions of `instruction->fused_instructions_computation()` (in post-order from `root_instruction()`) when hashing a `kFusion` instruction, skipping operand `unique_id()`s for internal instructions so equivalent `kFusion` instructions continue to hash identically while distinct fused computations are disambiguated. PiperOrigin-RevId: 987011851 --- third_party/xla/xla/service/hlo_cse.cc | 80 +++++++++++++++------ third_party/xla/xla/service/hlo_cse_test.cc | 54 ++++++++++++++ 2 files changed, 111 insertions(+), 23 deletions(-) diff --git a/third_party/xla/xla/service/hlo_cse.cc b/third_party/xla/xla/service/hlo_cse.cc index 440eb16a92ad35..b67f565edf1a6f 100644 --- a/third_party/xla/xla/service/hlo_cse.cc +++ b/third_party/xla/xla/service/hlo_cse.cc @@ -112,12 +112,14 @@ absl::StatusOr CombineConstants( return combined > 0; } -// An instruction is considered to be equivalent to another only if they -// share the exact same set of operands. +// This differs from "normal" HLO Instruction hashing because it takes the +// operands of an instruction into account. For the purposes of CSEs, two +// instructions that have different operands are never equivalent, so +// disambiguating by operands is always useful - but not always sufficient. struct CseKey { template - friend H AbslHashValue(H h, const CseKey& key) { - auto instruction = key.hlo; + static H HashInstruction(H h, const HloInstruction* instruction, + bool hash_operands) { h = instruction->shape().IsArray() ? H::combine(std::move(h), instruction->opcode(), instruction->shape().dimensions()) @@ -144,26 +146,28 @@ struct CseKey { h = result_accuracy_hash(std::move(h), instruction->result_accuracy()); // Hash operands, ignoring operand order on commutative ops. - if (HloOpcodeIsBinaryCommutative(instruction->opcode())) { - CHECK_EQ(instruction->operand_count(), 2); - auto id0 = instruction->operand(0)->unique_id(); - if (instruction->operand(0)->opcode() == HloOpcode::kIota) { - id0 = 0; - } - auto id1 = instruction->operand(1)->unique_id(); - if (instruction->operand(1)->opcode() == HloOpcode::kIota) { - id1 = 0; - } - if (id0 > id1) { - std::swap(id0, id1); - } - h = H::combine(std::move(h), id0, id1); - } else { - for (auto operand : instruction->operands()) { - if (operand->opcode() == HloOpcode::kIota) { - continue; + if (hash_operands) { + if (HloOpcodeIsBinaryCommutative(instruction->opcode())) { + CHECK_EQ(instruction->operand_count(), 2); + auto id0 = instruction->operand(0)->unique_id(); + if (instruction->operand(0)->opcode() == HloOpcode::kIota) { + id0 = 0; + } + auto id1 = instruction->operand(1)->unique_id(); + if (instruction->operand(1)->opcode() == HloOpcode::kIota) { + id1 = 0; + } + if (id0 > id1) { + std::swap(id0, id1); + } + h = H::combine(std::move(h), id0, id1); + } else { + for (auto operand : instruction->operands()) { + if (operand->opcode() == HloOpcode::kIota) { + continue; + } + h = H::combine(std::move(h), operand->unique_id()); } - h = H::combine(std::move(h), operand->unique_id()); } } @@ -171,6 +175,31 @@ struct CseKey { h = H::combine(std::move(h), c->root_instruction()->opcode()); } switch (instruction->opcode()) { + case HloOpcode::kFusion: + // The return post-order for two equal computations will always be the + // same, since it's determined purely by operand order, so we should + // never expect to hash two equal computations to different values. Note + // that the reasons this works are somewhat subtle: (a) We never + // consider kFusion commutative, so a dependency on operand order is + // safe. (b) Control dependencies may make the post-order traversal + // order depend on insertion order, but that's not an issue for fusion + // computations. (c) The post-order returned by + // HloComputation::MakeInstructionPostOrder() can depend on insertion + // order if dead instructions are present, which is why we use + // MakeInstructionPostOrderFrom() on the root. + for (const HloInstruction* fused_instruction : + instruction->fused_instructions_computation() + ->MakeInstructionPostOrderFrom( + *instruction->fused_instructions_computation() + ->root_instruction())) { + // Internal instructions of equivalent fusions belong to different + // HloComputations (whose IDs are encoded into + // HloInstruction::unique_id()), so we can only hash operand + // unique_ids for top-level instructions. + h = HashInstruction(std::move(h), fused_instruction, + /*hash_operands=*/false); + } + return H::combine(std::move(h), instruction->fusion_kind()); case HloOpcode::kSlice: return H::combine(std::move(h), instruction->slice_starts(), instruction->slice_strides()); @@ -227,6 +256,11 @@ struct CseKey { return std::move(h); } } + + template + friend H AbslHashValue(H h, const CseKey& key) { + return HashInstruction(std::move(h), key.hlo, /*hash_operands=*/true); + } HloInstruction* hlo; }; diff --git a/third_party/xla/xla/service/hlo_cse_test.cc b/third_party/xla/xla/service/hlo_cse_test.cc index 68dad767e5f12c..9336949ec3e0cb 100644 --- a/third_party/xla/xla/service/hlo_cse_test.cc +++ b/third_party/xla/xla/service/hlo_cse_test.cc @@ -1089,6 +1089,60 @@ TEST_F(HloCseTest, ResultAccuracyCseKey) { EXPECT_EQ(root->operand(2), root->operand(3)); } +TEST_F(HloCseTest, FusionCseKey) { + constexpr absl::string_view kHlo = R"( + HloModule m + + fused_comp.0 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p0, p1) + ROOT add = f32[4] add(p0, mul) + } + + fused_comp.1 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p0, p1) + ROOT add = f32[4] add(p0, mul) + } + + fused_comp.2 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + sub = f32[4] subtract(p0, p1) + ROOT add = f32[4] add(p0, sub) + } + + fused_comp.3 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p1, p1) + ROOT add = f32[4] add(p0, mul) + } + + ENTRY main { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + f0 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.0 + f1 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.1 + f2 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.2 + f3 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.3 + ROOT t = tuple(f0, f1, f2, f3) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kHlo)); + HloCSE cse(/*is_layout_sensitive=*/false); + EXPECT_THAT(cse.Run(m.get()), absl_testing::IsOkAndHolds(true)); + const HloInstruction* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->operand_count(), 4); + EXPECT_EQ(root->operand(0), root->operand(1)); + EXPECT_NE(root->operand(0), root->operand(2)); + EXPECT_NE(root->operand(0), root->operand(3)); + EXPECT_NE(root->operand(2), root->operand(3)); +} + TEST_F(HloCseTest, ScalarCustomCallNoOperands) { constexpr absl::string_view kHlo = R"( HloModule main From 38393dfa728061ca57b52d39e7f84ee54523062c Mon Sep 17 00:00:00 2001 From: Saketh Reddy Pingili Date: Wed, 23 Sep 2026 15:22:21 -0700 Subject: [PATCH 39/42] PR #49149: [XLA:Evaluator] Include with _USE_MATH_DEFINES in hlo_evaluator.h 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/49149 ## [XLA:Evaluator] Include `` with `_USE_MATH_DEFINES` ### Problem `hlo_evaluator.h` defines `_USE_MATH_DEFINES` for MSVC but does not directly include ``. This makes the availability of constants such as `M_PI` dependent on indirect include order. ### Solution Include `` immediately after `_USE_MATH_DEFINES`, matching the existing pattern in `xla/fp_util.h`. ### Changes - `xla/hlo/evaluator/hlo_evaluator.h`: Add `#include ` after `_USE_MATH_DEFINES`. ### Testing - `git diff --check` — PASS - Unit tests not run locally; Bazel is unavailable. ### Risk Low risk. Header-only change with no API or runtime logic changes. Copybara import of the project: -- a8382c01fef2ee71002056b7bb91fe439175bcb3 by saketh reddy pingili : [XLA:Evaluator] Include with _USE_MATH_DEFINES in hlo_evaluator.h Merging this change closes #49149 PiperOrigin-RevId: 987040037 --- third_party/xla/xla/hlo/evaluator/hlo_evaluator.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h index acc6b37b43dfeb..42ba52fd318d12 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h @@ -18,6 +18,7 @@ limitations under the License. #define _USE_MATH_DEFINES +#include #include #include #include @@ -812,7 +813,6 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // Optional handler exercised when evaluating literals. EvalLiteralHandler eval_literal_handler_; - // Set by EvaluateInternal and opportunistically used by the HandleXXX // functions. When non-empty, the HandleXXX function may evaluate the // instruction at only the given shape index. From c96dc33008ee2d32f9496005dd0ba11f55826fe9 Mon Sep 17 00:00:00 2001 From: Naums Mogers Date: Wed, 23 Sep 2026 15:52:56 -0700 Subject: [PATCH 40/42] [Mosaic] Add n-D vector load with expanded semantics Add `tpu.vector_expand_load`, the logical counterpart of `tpu.vector_compress_store`. `result[i_0, i_1, ...]`, if not masked, is loaded from `base[j_0, j_1, ...]` where `j_n = indices_n + i_n` for `n != expand_dim`, and `j_n = indices_n + c` for `n == expand_dim`, with `c` the count of enabled mask values preceding position `i_n` in the slice. Elements at masked-off positions are undefined. PiperOrigin-RevId: 987058288 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 84 +++++--- .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 43 +++++ .../dialect/tpu/tpu_ops_verification_test.cc | 182 ++++++++++++++++++ 3 files changed, 279 insertions(+), 30 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 9c1eae7cc792a3..09432b11e5d080 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -1169,49 +1169,73 @@ LogicalResult CompressStoreVregOp::verify() { return success(); } -LogicalResult VectorCompressStoreOp::verify() { - MemRefType ref_ty = getBase().getType(); +template +LogicalResult verifyCompressStoreExpandLoad( + Op op, VectorType ty, int32_t compress_expand_dim, StringRef vector_name, + StringRef compress_expand_dim_name) { + MemRefType ref_ty = op.getBase().getType(); const int64_t rank = ref_ty.getRank(); - VectorType value_ty = getValueToStore().getType(); - if (value_ty.getRank() != rank) { - return emitOpError("Expected valueToStore to have the same rank as base (") - << rank << "). Got: " << value_ty.getRank() << "."; - } - if (llvm::size(getIndices()) != rank) { - return emitOpError("Expected ") << rank << " indices."; - } - const int32_t compress_dim = getCompressDim(); - if (compress_dim < 0 || compress_dim >= rank) { - return emitOpError("Expected compress_dim to be in [0, ") - << rank << "). Got: " << compress_dim << "."; - } + if (ty.getRank() != rank) { + return op.emitOpError("Expected ") + << vector_name << " to have the same rank as base (" << rank + << "). Got: " << ty.getRank() << "."; + } + if (llvm::size(op.getIndices()) != rank) { + return op.emitOpError("Expected ") << rank << " indices."; + } + if (compress_expand_dim < 0 || compress_expand_dim >= rank) { + return op.emitOpError("Expected ") + << compress_expand_dim_name << " to be in [0, " << rank + << "). Got: " << compress_expand_dim << "."; + } + const StringRef non_target_dim_adjective = + compress_expand_dim_name == "compress_dim" ? "Non-compressed" + : "Non-expanded"; for (int64_t i = 0; i < rank; ++i) { - if (i == compress_dim || ref_ty.isDynamicDim(i)) { + if (i == compress_expand_dim || ref_ty.isDynamicDim(i)) { continue; } - if (value_ty.getDimSize(i) > ref_ty.getDimSize(i)) { - return emitOpError("Non-compressed dimension ") - << i << " of valueToStore goes out of bounds of base (" - << value_ty.getDimSize(i) << " > " << ref_ty.getDimSize(i) << ")."; + if (ty.getDimSize(i) > ref_ty.getDimSize(i)) { + return op.emitOpError(non_target_dim_adjective) + << " dimension " << i << " of " << vector_name + << " goes out of bounds of base (" << ty.getDimSize(i) << " > " + << ref_ty.getDimSize(i) << ")."; } } if (ref_ty.getMemorySpace() && !HasMemorySpace(ref_ty, MemorySpace::kVmem)) { - return emitOpError("Expected base memref to be in VMEM."); + return op.emitOpError("Expected base memref to be in VMEM."); } - if (value_ty.getElementType() != ref_ty.getElementType()) { - return emitOpError("Expected base and valueToStore element type to match"); + if (ty.getElementType() != ref_ty.getElementType()) { + return op.emitOpError("Expected base and ") + << vector_name << " element type to match"; } - // Note: We deliberately do not go through verifyStoreOp, which rejects masked - // stores of non-32-bit element types. The mask here is mandatory and narrow - // element types are allowed. - if (value_ty.getShape() != getMask().getType().getShape()) { - return emitOpError("Expected mask shape to match value shape: (") - << value_ty.getShape() << "). Got: (" - << getMask().getType().getShape() << ")."; + // Note: We deliberately do not go through verifyStoreOp / verifyLoadOp, which + // reject masked stores/loads of non-32-bit element types. The mask here is + // mandatory and narrow element types are allowed. They are rejected later, in + // the vector layout passes. + if (ty.getShape() != op.getMask().getType().getShape()) { + return op.emitOpError("Expected mask shape to match ") + << (vector_name == "valueToStore" ? "value" : "result") + << " shape: (" << ty.getShape() << "). Got: (" + << op.getMask().getType().getShape() << ")."; } return success(); } +LogicalResult VectorCompressStoreOp::verify() { + return verifyCompressStoreExpandLoad( + *this, getValueToStore().getType(), getCompressDim(), + /*vector_name=*/"valueToStore", + /*compress_expand_dim_name=*/"compress_dim"); +} + +LogicalResult VectorExpandLoadOp::verify() { + return verifyCompressStoreExpandLoad( + *this, getResult().getType(), getExpandDim(), + /*vector_name=*/"result", + /*compress_expand_dim_name=*/"expand_dim"); +} + template LogicalResult verifyLoadOp(Op op) { MemRefType ref_ty = op.getBase().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 086a825b8b811c..1130919108db8c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -413,6 +413,49 @@ def TPU_ExpandLoadVregOp : TPU_Op<"expand_load_vreg", [DefaultMemRead]> { let hasVerifier = 1; } +def TPU_VectorExpandLoadOp : TPU_Op<"vector_expand_load", [DefaultMemRead]> { + let summary = "Performs a masked, expanded load of a vector."; + let description = [{ + In contrast to `tpu.vector_load`, which has non-expanded semantics + (masked-out elements leave gaps in the addresses that are read), masked-out + elements are skipped entirely when striding along `expand_dim`: + `result[i_0, i_1, ...]`, if not masked, is loaded from + `base[j_0, j_1, ...]` where: + - For `n != expand_dim`, `j_n = indices_n + i_n` (same as a regular load) + - For `n == expand_dim`, `j_n = indices_n + c`, where `c` is the count of + enabled (true) values in `mask[i_0, i_1, ..., 0:i_n, ...]` + + Elements at masked-off positions are undefined. + + For example, for a rank-1 `result` (and hence `expand_dim` = 0), loading + with mask [true, false, true, false] yields + [base[indices], , base[indices + 1], ]. + + base : Memref to load from. + indices : Scalar indices into `base`, identifying the position of the + first loaded element. Must be of the same rank as the base + memref shape. + mask : Mask of the same shape as `result`. Masking has a granularity + of 32 bits: elements that share a 32-bit word must have the + same mask value. + expand_dim: The dimension along which the enabled elements are read + consecutively. + result : The loaded vector. Must be of the same rank as the base memref + shape. + }]; + let arguments = (ins + AnyMemRef:$base, + Variadic:$indices, + VectorOfNonZeroRankOf<[I1]>:$mask, + I32Attr:$expand_dim + ); + let results = (outs AnyVectorOfNonZeroRank:$result); + let assemblyFormat = [{ + $base `[` $indices `]` `,` $mask attr-dict `:` type($base) `,` type($result) `,` type($mask) + }]; + let hasVerifier = 1; +} + // tpu.vector_load loads a vector from memory into a register. // // base : Memref to load from. diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc index 2fa8f8df4ecae9..0540083953a131 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc @@ -765,6 +765,188 @@ TEST_F(TpuOpsVerificationTest, VectorCompressStoreDynamicBaseDimAllowed) { ASSERT_OK(VerifyOp(cs)); } +TEST_F(TpuOpsVerificationTest, VectorExpandLoadValid) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +// Packed (sub-32-bit) types must pass verification. +TEST_F(TpuOpsVerificationTest, VectorExpandLoadValidPacked) { + auto c0 = Create(0); + Value memref = + Create( + GetMemRefType({4, 16}, builder().getBF16Type(), MemorySpace::kVmem)) + .getMemref(); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, builder().getBF16Type()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadInvalidMemorySpace) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kHbm); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected base memref to be in VMEM."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedRank) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{8}, + /*values=*/SmallVector(8, true)); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected result to have the same rank as " + "base (2). Got: 1."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedIndicesCount) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), StatusIs(_, HasSubstr("Expected 2 indices."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadExpandDimOutOfRange) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(2)); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, HasSubstr("Expected expand_dim to be in [0, 2). Got: 2."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadNonMinormostExpandDim) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(0)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedMaskShape) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 4}, + /*values=*/SmallVector(8, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, HasSubstr("Expected mask shape to match result shape"))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadNonExpandedDimOutOfBounds) { + auto c0 = Create(0); + Value memref = AllocaI32({2, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{4, 8}, + /*values=*/SmallVector(32, true)); + auto el = Create( + /*result=*/VectorType::get({4, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Non-expanded dimension 0 of result goes " + "out of bounds of base"))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadDynamicIndicesAllowed) { + // Non-expanded dimension fits (2 <= 4), but constant index offset plus size + // exceeds base dimension (3 + 2 = 5 > 4). All indices are treated as dynamic, + // so verification passes. + auto c0 = Create(0); + auto c3 = Create(3); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c3, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadDynamicBaseDimAllowed) { + auto c0 = Create(0); + Value memref = AllocaI32({ShapedType::kDynamic, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{8, 8}, + /*values=*/SmallVector(64, true)); + auto el = Create( + /*result=*/VectorType::get({8, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + TEST_F(TpuOpsVerificationTest, UnpackSubelementsValidIndex) { Value source = ConstantI8Vector(/*shape=*/{4, 8}, /*values=*/{1}); auto unpack = Create( From be9e67c0c69cc0d3456044dd17e313cbd570df3b Mon Sep 17 00:00:00 2001 From: Yanghui Ou Date: Wed, 23 Sep 2026 16:10:15 -0700 Subject: [PATCH 41/42] [HLO->MHLO] Support tracing async-start across tuples and optimization barriers PiperOrigin-RevId: 987069321 --- .../xla/xla/hlo/translate/hlo_to_mhlo/BUILD | 1 - .../translate/hlo_to_mhlo/async_importer.cc | 81 ++++++++++++++++++- .../hlo_to_mhlo/tests/import_async.hlo | 27 +++++++ 3 files changed, 105 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD index 15b4cc8d75fd10..cdf5463c660cee 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD @@ -63,7 +63,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/mlir_hlo", "//xla/mlir_hlo:unregistered_attributes", - "//xla/tsl/platform:errors", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc index c6e2739e5c4012..8349a5716dc71a 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -49,7 +50,6 @@ limitations under the License. #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h" #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" #include "xla/mlir_hlo/utils/unregistered_attributes.h" -#include "xla/tsl/platform/errors.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -201,6 +201,79 @@ absl::StatusOr ImportStablehloAsyncDone( return op.getOperation(); } +// Finds the defining MHLO AsyncStartOp by traversing backwards through +// transparent dataflow intermediaries (tuples, optimization barriers, and +// get-tuple-element). Maintains a tuple index stack to correctly navigate +// through multi-level nested tuples. +mlir::mhlo::AsyncStartOp FindMhloAsyncStart(mlir::Value value) { + llvm::SmallVector tuple_indices_stack; + while (value) { + mlir::Operation* defining_op = value.getDefiningOp(); + if (!defining_op) { + return nullptr; + } + + // Base case: Reached the target AsyncStartOp. + if (auto async_start = + llvm::dyn_cast(defining_op)) { + return async_start; + } + + // GetTupleElement (StableHLO / MHLO): Push the tuple index onto the + // stack so that when we encounter the matching TupleOp upstream, we select + // the correct operand branch. + if (auto gte = + llvm::dyn_cast(defining_op)) { + tuple_indices_stack.push_back(gte.getIndex()); + value = gte.getOperand(); + continue; + } + if (auto gte = llvm::dyn_cast(defining_op)) { + tuple_indices_stack.push_back(gte.getIndex()); + value = gte.getOperand(); + continue; + } + + // TupleOp (StableHLO / MHLO): Pop the expected tuple index from the + // stack to step into the corresponding tuple element operand. If no index + // was tracked, the tuple was consumed directly rather than by element, so + // it cannot be an async-start op. + if (llvm::isa(defining_op)) { + if (tuple_indices_stack.empty()) { + return nullptr; + } + int64_t idx = tuple_indices_stack.pop_back_val(); + if (idx >= defining_op->getNumOperands()) { + return nullptr; + } + value = defining_op->getOperand(idx); + continue; + } + + // OptimizationBarrierOp (StableHLO / MHLO): The barrier is transparent; + // trace the corresponding input operand matching this OpResult's result + // number. + if (llvm::isa(defining_op)) { + auto op_result = mlir::dyn_cast(value); + if (!op_result) { + return nullptr; + } + unsigned res_num = op_result.getResultNumber(); + if (res_num >= defining_op->getNumOperands()) { + return nullptr; + } + value = defining_op->getOperand(res_num); + continue; + } + + // Non-transparent defining op; async-start could not be found along this + // path. + return nullptr; + } + return nullptr; +} + absl::StatusOr ImportOldStyleAsyncDone( llvm::SmallVectorImpl& attributes, const llvm::SmallVectorImpl& operands, mlir::Location loc, @@ -208,7 +281,7 @@ absl::StatusOr ImportOldStyleAsyncDone( bool useBundleResult = false) { assert(operands.size() == 1 && "*-done ops must take only a single async_bundle operand"); - auto async_start = operands[0].getDefiningOp(); + auto async_start = FindMhloAsyncStart(operands[0]); if (!async_start) { return InvalidArgument("*-start requires *-done as input"); } @@ -229,7 +302,9 @@ absl::StatusOr ImportOldStyleAsyncDone( operands, attributes); return {op}; } - if (useBundleResult) result_type = async_bundle.getTypes()[1]; + if (useBundleResult) { + result_type = async_bundle.getTypes()[1]; + } auto op = mlir::mhlo::AsyncDoneOp::create(*builder, loc, Untuple(result_type), operands, attributes); return CreateTupleFromOpResults(builder, loc, op.getOperation(), result_type); diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo index 14672943c8e049..d3b603c33adf2a 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo @@ -295,6 +295,33 @@ ENTRY %test_all_gather_start_mesh { ROOT ag-done = f32[128,128] all-gather-done(ag-start) } +// ----- + +HloModule main + +// CHECK-LABEL: func.func private @copy_(%arg0: tensor<128x32xf32>) +// CHECK-NEXT: mhlo.copy + +// CHECK-LABEL: func.func @main +// CHECK-SAME: (%[[ARG0:.*]]: tensor<1024x1024xf32>, %[[ARG1:.*]]: tensor<128x32xf32>) +// CHECK: %[[BARRIER:.*]] = stablehlo.optimization_barrier %[[ARG0]] +// CHECK: %[[START:.*]] = "mhlo.async_start"(%[[ARG1]]) <{called_computation = @copy_, execution_thread = "main"}> +// CHECK: %[[DONE:.*]] = "mhlo.async_done"(%[[START]]) +// CHECK-SAME: called_computation = @copy_ +// CHECK-SAME: execution_thread = "main" +// CHECK: return %[[BARRIER]], %[[DONE]] +ENTRY %test_async_copy_opt_barrier (dot_general.1: f32[1024,1024], shard_map.5: f32[128,32]) -> (f32[1024,1024], f32[128,32]) { + %dot_general.1 = f32[1024,1024]{1,0} parameter(0) + %shard_map.5 = f32[128,32]{1,0} parameter(1) + %opt_bar = f32[1024,1024]{1,0} opt-barrier(%dot_general.1) + %copy-start = (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[]) copy-start(%shard_map.5) + %tuple.4 = (f32[1024,1024]{1,0}, (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[])) tuple(%opt_bar, %copy-start) + %get-tuple-element.0 = f32[1024,1024]{1,0} get-tuple-element(%tuple.4), index=0 + %get-tuple-element.1 = (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[]) get-tuple-element(%tuple.4), index=1 + %copy-done = f32[128,32]{1,0} copy-done(%get-tuple-element.1) + ROOT %root = (f32[1024,1024]{1,0}, f32[128,32]{1,0}) tuple(%get-tuple-element.0, %copy-done) +} + // BROKEN: b/TODO: support roundtrip of async custom calls? // HloModule main, entry_computation_layout={(f32[10]{0})->(f32[20]{0})} From 409704b2efab4bd56253cd8f479b217e4ca0667a Mon Sep 17 00:00:00 2001 From: Maxim Ermilov Date: Wed, 23 Sep 2026 16:24:54 -0700 Subject: [PATCH 42/42] implement BufferUses for AsyncThunks PiperOrigin-RevId: 987077956 --- .../xla/xla/backends/gpu/runtime/BUILD | 1 + .../xla/backends/gpu/runtime/async_thunk.h | 4 +- .../gpu/runtime/host_send_recv_thunk.cc | 65 +++++++++++++------ .../gpu/runtime/host_send_recv_thunk.h | 42 ++++++++---- .../xla/xla/backends/gpu/runtime/thunk.proto | 2 + .../xla/xla/service/gpu/thunk_emitter.cc | 12 +++- 6 files changed, 86 insertions(+), 40 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index bfe3b67572bec5..d96b6dfca36538 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -3500,6 +3500,7 @@ cc_library( "//xla/runtime:buffer_use", "//xla/runtime:device_id", "//xla/service:buffer_assignment", + "//xla/service:shaped_slice", "//xla/stream_executor:device_address", "//xla/stream_executor:event", "//xla/stream_executor:stream", diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h index df5c119602697a..419167ac764b5c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h @@ -62,8 +62,8 @@ class AsyncStartThunk : public Thunk { absl::Status Initialize(const InitializeParams& params) override; absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. + // AsyncStartThunk touches no buffers; this is explicit rather than relying on + // the base class default. BufferUses buffer_uses() const override { return {}; } ExecutionStreamId execution_stream_id() const { return execution_stream_id_; } diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc index 1ac64f4c193a20..d1c0ab247b1ce5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" +#include "xla/service/shaped_slice.h" #include "xla/shape.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/event.h" @@ -107,13 +108,12 @@ HostSendRecvAsyncEvents::Extract(se::StreamExecutor* executor, //===----------------------------------------------------------------------===// HostSendThunk::HostSendThunk( - ThunkInfo thunk_info, Shape shape, BufferAllocation::Slice buffer, - int64_t channel_id, std::shared_ptr events, + ThunkInfo thunk_info, ShapedSlice buffer, int64_t channel_id, + std::shared_ptr events, absl::flat_hash_map frontend_attrs, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostSend, thunk_info), - shape_(shape), - buffer_(buffer), + slice_(buffer), channel_id_(channel_id), events_(std::move(events)), frontend_attrs_(std::move(frontend_attrs)), @@ -123,8 +123,9 @@ absl::StatusOr HostSendThunk::ToProto() const { ThunkProto proto; *proto.mutable_thunk_info() = thunk_info().ToProto(); HostSendThunkProto& host_send_thunk_proto = *proto.mutable_host_send_thunk(); - *host_send_thunk_proto.mutable_shape() = shape_.ToProto(); - ABSL_ASSIGN_OR_RETURN(*host_send_thunk_proto.mutable_buffer(), buffer_.ToProto()); + *host_send_thunk_proto.mutable_shape() = slice_.shape.ToProto(); + ABSL_ASSIGN_OR_RETURN(*host_send_thunk_proto.mutable_buffer(), + slice_.slice.ToProto()); host_send_thunk_proto.set_channel_id(channel_id_); host_send_thunk_proto.mutable_frontend_attrs()->insert( frontend_attrs_.begin(), frontend_attrs_.end()); @@ -160,13 +161,13 @@ absl::StatusOr> HostSendThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); return std::make_unique( - thunk_info, std::move(shape), buffer, proto.channel_id(), + thunk_info, ShapedSlice{buffer, shape}, proto.channel_id(), async_event_it->second, std::move(frontend_attrs), device_constraint); } absl::Status HostSendThunk::ExecuteOnStream(const ExecuteParams& params) { VLOG(3) << "Send buffer: channel_id=" << channel_id_ - << "; shape=" << shape_.ToString(); + << "; shape=" << slice_.shape.ToString(); ABSL_ASSIGN_OR_RETURN(bool skip, ShouldSkip("sending buffer", params, device_constraint_)); @@ -186,13 +187,13 @@ absl::Status HostSendThunk::ExecuteOnStream(const ExecuteParams& params) { } se::DeviceAddressBase src = - params.buffer_allocations->GetDeviceAddress(buffer_); + params.buffer_allocations->GetDeviceAddress(slice_.slice); // Send buffer to a handler registered with the executable. if (auto* send = params.send_device_memory_function) { ABSL_ASSIGN_OR_RETURN( AsyncValueRef> done, - (*send)(channel_id_, stream, shape_, src, frontend_attrs_)); + (*send)(channel_id_, stream, slice_.shape, src, frontend_attrs_)); return events_->Emplace(stream->parent(), channel_id_, std::move(done)); } @@ -212,15 +213,15 @@ std::optional HostSendThunk::GetAsyncEventsUniqueId() //===----------------------------------------------------------------------===// // HostSendDoneThunk //===----------------------------------------------------------------------===// - HostSendDoneThunk::HostSendDoneThunk( - ThunkInfo thunk_info, int64_t channel_id, + ThunkInfo thunk_info, std::optional slice, int64_t channel_id, std::shared_ptr events, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostSendDone, thunk_info), channel_id_(channel_id), events_(std::move(events)), - device_constraint_(device_constraint) {} + device_constraint_(device_constraint), + slice_(slice) {} absl::StatusOr HostSendDoneThunk::ToProto() const { ThunkProto proto; @@ -239,6 +240,10 @@ absl::StatusOr HostSendDoneThunk::ToProto() const { } host_send_done_thunk_proto.set_async_events_unique_id( async_events_unique_id.value().value()); + if (slice_.has_value()) { + ABSL_ASSIGN_OR_RETURN(*host_send_done_thunk_proto.mutable_buffer(), + slice_->ToProto()); + } return proto; } @@ -255,9 +260,15 @@ absl::StatusOr> HostSendDoneThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); - return std::make_unique(thunk_info, proto.channel_id(), - async_event_it->second, - device_constraint); + std::optional slice; + if (proto.has_buffer()) { + ABSL_ASSIGN_OR_RETURN(slice, + ShapedSlice::FromProto(proto.buffer(), allocations)); + } + + return std::make_unique( + thunk_info, slice, proto.channel_id(), async_event_it->second, + device_constraint); } absl::Status HostSendDoneThunk::ExecuteOnStream(const ExecuteParams& params) { @@ -408,13 +419,14 @@ std::optional HostRecvThunk::GetAsyncEventsUniqueId() //===----------------------------------------------------------------------===// HostRecvDoneThunk::HostRecvDoneThunk( - ThunkInfo thunk_info, int64_t channel_id, + ThunkInfo thunk_info, std::optional slice, int64_t channel_id, std::shared_ptr events, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostRecvDone, thunk_info), channel_id_(channel_id), events_(std::move(events)), - device_constraint_(device_constraint) {} + device_constraint_(device_constraint), + slice_(slice) {} absl::StatusOr HostRecvDoneThunk::ToProto() const { ThunkProto proto; @@ -433,6 +445,11 @@ absl::StatusOr HostRecvDoneThunk::ToProto() const { } host_recv_done_thunk_proto.set_async_events_unique_id( async_events_unique_id.value().value()); + + if (slice_.has_value()) { + ABSL_ASSIGN_OR_RETURN(*host_recv_done_thunk_proto.mutable_buffer(), + slice_->ToProto()); + } return proto; } @@ -449,9 +466,15 @@ absl::StatusOr> HostRecvDoneThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); - return std::make_unique(thunk_info, proto.channel_id(), - async_event_it->second, - device_constraint); + std::optional slice; + if (proto.has_buffer()) { + ABSL_ASSIGN_OR_RETURN(slice, + ShapedSlice::FromProto(proto.buffer(), allocations)); + } + + return std::make_unique( + thunk_info, slice, proto.channel_id(), async_event_it->second, + device_constraint); } absl::Status HostRecvDoneThunk::ExecuteOnStream(const ExecuteParams& params) { diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h index a3517382828389..9a83f160943c35 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h @@ -34,6 +34,7 @@ limitations under the License. #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" +#include "xla/service/shaped_slice.h" #include "xla/shape.h" #include "xla/stream_executor/event.h" #include "xla/stream_executor/stream_executor.h" @@ -100,8 +101,7 @@ class HostSendThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostSendThunk(ThunkInfo thunk_info, Shape shape, - BufferAllocation::Slice buffer, int64_t channel_id, + HostSendThunk(ThunkInfo thunk_info, ShapedSlice buffer, int64_t channel_id, std::shared_ptr events, absl::flat_hash_map frontend_attrs, std::optional device_constraint); @@ -110,7 +110,7 @@ class HostSendThunk : public HostAsyncThunk { BufferUses buffer_uses() const override { return { - BufferUse::Read(buffer_, shape_), + BufferUse::Read(slice_.slice, slice_.shape), }; } @@ -121,8 +121,7 @@ class HostSendThunk : public HostAsyncThunk { bool IsAsyncStart() const override { return events_ != nullptr; } private: - Shape shape_; - BufferAllocation::Slice buffer_; + ShapedSlice slice_; int64_t channel_id_; @@ -142,15 +141,22 @@ class HostSendDoneThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostSendDoneThunk(ThunkInfo thunk_info, int64_t channel_id, + HostSendDoneThunk(ThunkInfo thunk_info, std::optional slice, + int64_t channel_id, std::shared_ptr events, std::optional device_constraint); absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. - BufferUses buffer_uses() const override { return {}; } + BufferUses buffer_uses() const override { + // TODO(b/527907619): make mandatory after 2027-1-1 + if (!slice_.has_value()) { + return {}; + } + return { + BufferUse::Read(slice_->slice, slice_->shape), + }; + } absl::StatusOr ToProto() const override; @@ -163,12 +169,12 @@ class HostSendDoneThunk : public HostAsyncThunk { std::shared_ptr events_; std::optional device_constraint_; + std::optional slice_; }; //===----------------------------------------------------------------------===// // HostRecvThunk //===----------------------------------------------------------------------===// - class HostRecvThunk : public HostAsyncThunk { public: static absl::StatusOr> FromProto( @@ -218,15 +224,22 @@ class HostRecvDoneThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostRecvDoneThunk(ThunkInfo thunk_info, int64_t channel_id, + HostRecvDoneThunk(ThunkInfo thunk_info, std::optional slice, + int64_t channel_id, std::shared_ptr events, std::optional device_constraint); absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. - BufferUses buffer_uses() const override { return {}; } + BufferUses buffer_uses() const override { + // TODO(b/527907619): make mandatory after 2027-1-1 + if (!slice_.has_value()) { + return {}; + } + return { + BufferUse::Write(slice_->slice, slice_->shape), + }; + } absl::StatusOr ToProto() const override; @@ -239,6 +252,7 @@ class HostRecvDoneThunk : public HostAsyncThunk { std::shared_ptr events_; std::optional device_constraint_; + std::optional slice_; }; } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.proto b/third_party/xla/xla/backends/gpu/runtime/thunk.proto index 24bbab38e2aed7..b32f8d74310d25 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.proto @@ -149,6 +149,7 @@ message HostSendDoneThunkProto { int64 channel_id = 1; optional int64 device_constraint = 2; uint64 async_events_unique_id = 3; + ShapedSliceProto buffer = 4; } message HostRecvThunkProto { @@ -164,6 +165,7 @@ message HostRecvDoneThunkProto { int64 channel_id = 1; optional int64 device_constraint = 2; uint64 async_events_unique_id = 3; + ShapedSliceProto buffer = 4; } message HostExecuteStartThunkProto { diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc index e39b70dfd60a3c..3713b9c37535a1 100644 --- a/third_party/xla/xla/service/gpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc @@ -2516,7 +2516,7 @@ absl::StatusOr ThunkEmitter::EmitHostSend( return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( instr, ir_emitter_context_->GetNextThunkId()), - src->shape(), slice.slice, *instr->channel_id(), send_recv_events_, + slice, *instr->channel_id(), send_recv_events_, ConvertFrontendAttributes(instr->frontend_attributes()), DeviceConstraint(instr)); } @@ -2550,10 +2550,13 @@ absl::StatusOr ThunkEmitter::EmitHostSendDone( "Unknown channel id in host transfer send done instruction"); } + const HloInstruction* src = host_transfer->operand(0); + ABSL_ASSIGN_OR_RETURN(ShapedSlice slice, GetShapedSliceForHlo(src, {})); + return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( done, ir_emitter_context_->GetNextThunkId()), - *host_transfer->channel_id(), send_recv_events_, + slice, *host_transfer->channel_id(), send_recv_events_, DeviceConstraint(host_transfer)); } @@ -2565,10 +2568,13 @@ absl::StatusOr ThunkEmitter::EmitHostRecvDone( "Unknown channel id in host transfer recv done instruction"); } + const HloInstruction* src = host_transfer->operand(0); + ABSL_ASSIGN_OR_RETURN(ShapedSlice slice, GetShapedSliceForHlo(src, {})); + return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( done, ir_emitter_context_->GetNextThunkId()), - *host_transfer->channel_id(), send_recv_events_, + slice, *host_transfer->channel_id(), send_recv_events_, DeviceConstraint(host_transfer)); }